Stop allocating per (query, probe) on OpenMP workers in IVF RaBitQ fast-scan (#5716)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5716

`IndexIVFRaBitQFastScan` heap-allocated up to three times per (query, probe) inside the OpenMP parallel regions of `compute_LUT` / `compute_LUT_uint8`:
- `compute_residual_LUT` built a fresh `std::vector<uint8_t> rotated_qq(d)` on every call.
- For multi-bit indexes, it copied `rotated_q` (d floats) into a temporary `QueryFactorsData`.
- `compute_LUT_uint8` then copied that temporary into `context.query_factors[ij]` instead of moving it.

Under jemalloc heap profiling on aarch64, which Tupperware enables by default in task 0 of every job, each sampled allocation on an OpenMP worker unwinds through libomp's `__kmp_invoke_microtask`. That frame has no CFI, so libunwind's fallback maps and unmaps a whole ELF image per sample, serialised on `mmap_lock` (same mechanism as D123990523; details in that diff). After D123990523, a 74-scenario ServiceLab sweep with profiling on still showed IVF RaBitQ fast-scan at up to 30k (1-bit) and 57k (multi-bit) minor faults per pass, with up to 58 of 72 threads blocked.

gdb stack snapshots of a profiled `hr_ivf1024_rabitqfs4` run (348 threads caught inside `je_prof_backtrace`) attributed 149 samples to the `rotated_q` copy and 31 to `rotated_qq` inside `compute_residual_LUT`. The second copy, in `compute_LUT_uint8`, is visible in the code; its frames were not symbolised in the snapshot.

This diff:
- makes `rotated_qq` a per-thread buffer, like `rotated_q` and `centroid_buf` already are;
- stores multi-bit rotated queries in one block allocated once per search in `search_preassigned`, outside the parallel region and uninitialised. Each (query, probe) writes its slice and the handler reads it by storage index, via a new `FastScanDistancePostProcessing::rotated_q` that is offset per thread slice like `query_factors`;
- moves rather than copies into `context.query_factors[ij]`;
- keeps the capacity of a reused `QueryFactorsData` in `compute_residual_LUT`, so the single-query `IVFRaBitQFastScanScanner` no longer reallocates `rotated_q` on every `set_list`.

Memory is unchanged: the same n × nprobe × d floats as before, now contiguous instead of in n × nprobe separate vectors.

**Effect** (devbig334 Grace, Tupperware-matched environment, heap profiling on as in production, nq=10000, single runs on a busy shared host; recall unchanged):

| hr_ivf1024 | nprobe | QPS before → after | minor faults per pass before → after |
|---|---|---|---|
| rabitqfs4 (multi-bit) | 16 | 27,032 → 36,069 | 4,006 → 1,911 |
| rabitqfs4 (multi-bit) | 64 | 10,168 → 23,295 | 14,932 → 5,822 |
| rabitqfs4 (multi-bit) | 256 | 2,883 → 6,438 | 49,841 → 20,371 |
| rabitqfs (1-bit) | 256 | 3,179 → 6,854 | 23,935 → 20,654 |

Profiling off, rabitqfs4 at nprobe=256: 14,337 → 13,151 QPS, within this host's noise.

**What this does not fix.** With these allocations gone, 1-bit and multi-bit converge on about 20k faults per pass at nprobe=256. gdb attributes 163 of the 206 remaining samples to `posix_memalign` in `IndexIVFFastScan::search_implem_12` on worker threads: the LUT tables (`dis_tables`, n × nprobe × LUT size) that each thread slice allocates. Those allocations are needed, and jemalloc samples them by bytes. The remedy for that remainder is unwind info for libomp's aarch64 `__kmp_invoke_microtask`, being raised separately, rather than more FAISS changes.

Reviewed By: alibeklfc

Differential Revision: D124204736

fbshipit-source-id: 28be93c1cabbcb679bd5328808207fd8e24828b8

faiss/IndexIVFFastScan.cpp

@@ -663,6 +663,9 @@ void IndexIVFFastScan::search_dispatch_implem(
                     if (thread_context.query_factors != nullptr) {
                         thread_context.query_factors += i0 * cur_nprobe;
                     }
+                    if (thread_context.rotated_q != nullptr) {
+                        thread_context.rotated_q += i0 * cur_nprobe * d;
+                    }
 
                     auto scanner = make_knn_scanner(
                             is_max,

faiss/IndexIVFRaBitQFastScan.cpp

@@ -402,12 +402,14 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
         uint8_t qb_param,
         bool centered_param,
         std::vector<float>& rotated_q,
-        std::vector<float>& centroid_buf) const {
+        std::vector<float>& centroid_buf,
+        std::vector<uint8_t>& rotated_qq,
+        float* rotated_q_out) const {
     const size_t d_val = static_cast<size_t>(d);
     FAISS_THROW_IF_NOT(d_val > 0);
     rotated_q.resize(d_val);
     centroid_buf.resize(d_val);
-    std::vector<uint8_t> rotated_qq(d_val);
+    rotated_qq.resize(d_val);
 
     // Compute residual
     quantizer->reconstruct(centroid_id, centroid_buf.data());
@@ -415,7 +417,9 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
         rotated_q[i] = query[i] - centroid_buf[i];
     }
 
-    // Compute query factors using shared utility
+    // Keep the caller's rotated_q capacity across the assignment below, so a
+    // reused QueryFactorsData does not reallocate on every call.
+    std::vector<float> rotated_q_keep = std::move(query_factors.rotated_q);
     query_factors = rabitq_utils::compute_query_factors(
             rotated_q.data(),
             d_val,
@@ -433,7 +437,12 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
     }
 
     if (rabitq.nb_bits > 1) {
-        query_factors.rotated_q = rotated_q;
+        if (rotated_q_out != nullptr) {
+            std::copy(rotated_q.begin(), rotated_q.end(), rotated_q_out);
+        } else {
+            rotated_q_keep.assign(rotated_q.begin(), rotated_q.end());
+            query_factors.rotated_q = std::move(rotated_q_keep);
+        }
     }
 
     // Build LUT using branchless subset-sum construction
@@ -529,8 +538,15 @@ void IndexIVFRaBitQFastScan::search_preassigned(
     }
 
     std::vector<QueryFactorsData> query_factors_storage(n * cur_nprobe);
+    // One block for the multi-bit rotated queries, allocated here rather than
+    // per (query, probe) on OpenMP workers: under jemalloc heap profiling on
+    // aarch64, worker-thread allocations are very expensive to sample.
+    // Left uninitialised; every slot read is written by compute_LUT first.
+    std::unique_ptr<float[]> rotated_q_storage(
+            rabitq.nb_bits > 1 ? new float[n * cur_nprobe * d] : nullptr);
     FastScanDistancePostProcessing context;
     context.query_factors = query_factors_storage.data();
+    context.rotated_q = rotated_q_storage.get();
     context.nprobe = cur_nprobe;
     context.qb = used_qb;
     context.centered = used_centered;
@@ -571,6 +587,7 @@ void IndexIVFRaBitQFastScan::compute_LUT(
     {
         std::vector<float> rotated_q(d);
         std::vector<float> centroid_buf(d);
+        std::vector<uint8_t> rotated_qq(d);
 
 #pragma omp for
         for (idx_t ij = 0; ij < static_cast<idx_t>(n * cq_nprobe); ij++) {
@@ -588,7 +605,10 @@ void IndexIVFRaBitQFastScan::compute_LUT(
                         used_qb,
                         used_centered,
                         rotated_q,
-                        centroid_buf);
+                        centroid_buf,
+                        rotated_qq,
+                        context.rotated_q ? context.rotated_q + ij * d
+                                          : nullptr);
 
                 if (context.query_factors != nullptr) {
                     context.query_factors[ij] = std::move(query_factors_data);
@@ -629,6 +649,7 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8(
         AlignedTable<float> lut_float(cur_nprobe * dim12);
         std::vector<float> rotated_q(d);
         std::vector<float> centroid_buf(d);
+        std::vector<uint8_t> rotated_qq(d);
         std::vector<float> all_mins(cur_nprobe * M);
         std::vector<float> probe_b(cur_nprobe);
         std::vector<float> probe_span(cur_nprobe);
@@ -652,10 +673,13 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8(
                             used_qb,
                             used_centered,
                             rotated_q,
-                            centroid_buf);
+                            centroid_buf,
+                            rotated_qq,
+                            context.rotated_q ? context.rotated_q + ij * d
+                                              : nullptr);
 
                     if (context.query_factors != nullptr) {
-                        context.query_factors[ij] = qf;
+                        context.query_factors[ij] = std::move(qf);
                     }
                 } else {
                     memset(lut_float.get() + j * dim12,
@@ -879,6 +903,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
     AlignedTable<float> lut_float;
     std::vector<float> rotated_q;
     std::vector<float> centroid_buf;
+    std::vector<uint8_t> rotated_qq;
     QueryFactorsData query_factors;
     FastScanDistancePostProcessing context;
     std::vector<int> probe_map;
@@ -900,6 +925,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
               lut_float(16 * index_in.M),
               rotated_q(index_in.d),
               centroid_buf(index_in.d),
+              rotated_qq(index_in.d),
               probe_map({0}),
               mins_buf(index_in.M) {
         this->keep_max = is_similarity_metric(index_in.metric_type);
@@ -931,7 +957,9 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
                 qb,
                 centered,
                 rotated_q,
-                centroid_buf);
+                centroid_buf,
+                rotated_qq,
+                nullptr);
 
         // Single-probe quantization (simplified inline, no OMP, no 3D)
         const size_t M = index.M;

faiss/IndexIVFRaBitQFastScan.h

@@ -131,7 +131,9 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
             uint8_t qb_param,
             bool centered_param,
             std::vector<float>& rotated_q,
-            std::vector<float>& centroid_buf) const;
+            std::vector<float>& centroid_buf,
+            std::vector<uint8_t>& rotated_qq,
+            float* rotated_q_out) const;
 
    private:
     /// Decode FastScan code to RaBitQ residual vector with explicit
@@ -399,7 +401,8 @@ float IVFRaBitQHeapHandler<C, SL>::compute_full_multibit_distance(
             unpack_buf.data(),
             ex_code,
             ex_fac,
-            query_factors.rotated_q.data(),
+            context->rotated_q ? context->rotated_q + storage_idx_val * dim
+                               : query_factors.rotated_q.data(),
             is_similarity ? query_factors.q_dot_c : query_factors.qr_to_c_L2sqr,
             dim,
             ex_bits,

faiss/impl/fast_scan/FastScanDistancePostProcessing.h

@@ -29,6 +29,10 @@ struct FastScanDistancePostProcessing {
     /// QueryFactorsData subset for this context.
     rabitq_utils::QueryFactorsData* query_factors = nullptr;
 
+    /// Multi-bit IVF RaBitQ: rotated query per (query, probe), d floats each,
+    /// indexed like query_factors (nullptr if unused).
+    float* rotated_q = nullptr;
+
     /// The nprobe value used when allocating query_factors storage.
     /// This is needed because the allocation size (n * nprobe) may use a
     /// different nprobe than index->nprobe if search params override it.