facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Hoist per-list allocations out of the IVF fast-scan search loops (#5717)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5717

`IndexIVFFastScan::search_implem_12` and `search_implem_14` allocate a fresh `AlignedTable<uint8_t> LUT`, `q_map` and `lut_entries` for every inverted list they visit, and call `get_block_stride()`, which heap-allocates a `CodePacker` on each call; `search_implem_10` makes the same `get_block_stride()` call per list. These loops run on OpenMP worker threads, so with jemalloc heap profiling on (the production default for task 0 of every job) on aarch64 every sampled allocation pays libunwind's slow fallback at the `__kmp_invoke_microtask` frame.

gdb thread snapshots of the D124204736 build under profiling (`hr_ivf1024_rabitqfs4`, Grace): of 206 threads caught in jemalloc's `prof_backtrace`, 163 were in `posix_memalign` called from `search_implem_12` (the per-list LUT), 6 in its per-list vectors and 4 in `get_block_stride()`.

This allocates those buffers once per call / per thread, sized for the largest batch (`qbs2` queries x `dim12`), and computes the block stride once. Each iteration still uses only the first `nc` entries, so the search is unchanged.

Reviewed By: junjieqi

Differential Revision: D124387842

fbshipit-source-id: 6bcfe66270fcc0b208b0c88aaf578fea952f6174
This commit is contained in:
Weixiang Hong
2026-10-09 14:58:46 -07:00
committed by meta-codesync[bot]
parent 39079668a8
commit db6aa06d96
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/IndexIVFFastScan.cpp
+10 -7
View File
@@ -1063,6 +1063,7 @@ void IndexIVFFastScan::search_implem_10(
// Allocate probe_map once and reuse it
std::vector<int> probe_map;
probe_map.reserve(1);
const size_t block_stride = get_block_stride();
for (idx_t i = 0; i < n; i++) {
const uint8_t* LUT = nullptr;
@@ -1137,7 +1138,7 @@ void IndexIVFFastScan::search_implem_10(
codes.get(),
LUT,
context.pq2x4_scale,
get_block_stride());
block_stride);
ndis += ls;
nlist_visited++;
@@ -1229,6 +1230,9 @@ void IndexIVFFastScan::search_implem_12(
// Allocate vectors once and reuse them
std::vector<int> probe_map;
probe_map.reserve(actual_qbs2);
std::vector<int> q_map(actual_qbs2), lut_entries(actual_qbs2);
AlignedTable<uint8_t> LUT(actual_qbs2 * dim12);
const size_t block_stride = get_block_stride();
size_t i0 = 0;
uint64_t t_copy_pack = 0, t_scan = 0;
@@ -1255,8 +1259,6 @@ void IndexIVFFastScan::search_implem_12(
// re-organize LUTs and biases into the right order
int nc = static_cast<int>(i1 - i0);
std::vector<int> q_map(nc), lut_entries(nc);
AlignedTable<uint8_t> LUT(nc * dim12);
memset(LUT.get(), -1, nc * dim12);
int qbs_for_list = pq4_preferred_qbs(nc);
@@ -1306,7 +1308,7 @@ void IndexIVFFastScan::search_implem_12(
codes.get(),
LUT.get(),
context.pq2x4_scale,
get_block_stride());
block_stride);
// prepare for next loop
i0 = i1;
}
@@ -1464,6 +1466,9 @@ void IndexIVFFastScan::search_implem_14(
// Allocate probe_map once per thread and reuse it
std::vector<int> probe_map;
probe_map.reserve(actual_qbs2);
std::vector<int> q_map(actual_qbs2), lut_entries(actual_qbs2);
AlignedTable<uint8_t> LUT(actual_qbs2 * dim12);
const size_t block_stride = get_block_stride();
#pragma omp for schedule(dynamic)
for (idx_t cluster = 0; cluster < static_cast<idx_t>(ses.size());
@@ -1477,8 +1482,6 @@ void IndexIVFFastScan::search_implem_14(
// re-organize LUTs and biases into the right order
int nc = static_cast<int>(i1 - i0);
std::vector<int> q_map(nc), lut_entries(nc);
AlignedTable<uint8_t> LUT(nc * dim12);
memset(LUT.get(), -1, nc * dim12);
int qbs_for_list = pq4_preferred_qbs(nc);
@@ -1529,7 +1532,7 @@ void IndexIVFFastScan::search_implem_14(
codes.get(),
LUT.get(),
context.pq2x4_scale,
get_block_stride());
block_stride);
}
// labels is in-place for HeapHC