reuse unpack buffer in multibit IVFRaBitQFastScan refinement (#5003)

Summary:
This removes two pieces of avoidable overhead from multibit `IndexIVFRaBitQFastScan` refinement.

  In `compute_full_multibit_distance()`, we were allocating a temporary unpack buffer and reacquiring the current list codes for every refined candidate. The handler already has the list pointer, so this patch just reuses that and keeps a reusable unpack buffer on
  the handler.

  I only measured this on multibit workloads. On my AVX2 machine (`d=128`, `METRIC_L2`), the gain is small but consistent:
  - `127us -> 126us`
  - `1335us -> 1292us`
  - `6414us -> 6084us`
  - `17955us -> 16800us`

  No change to 1-bit logic or search semantics.

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

Reviewed By: mnorris11

Differential Revision: D98806563

Pulled By: alibeklfc

fbshipit-source-id: 8ae878d0f027bf2e534fed78e2fc33aa1f69e64c

faiss/IndexIVFRaBitQFastScan.h

@@ -199,7 +199,8 @@ IVFRaBitQHeapHandler<C, SL>::IVFRaBitQHeapHandler(
           storage_size(idx->compute_per_vector_storage_size()),
           packed_block_size(((idx->M2 + 1) / 2) * idx->bbs),
           full_block_size(idx->get_block_stride()),
-          packer(idx->get_CodePacker()) {
+          packer(idx->get_CodePacker()),
+          unpack_buf(idx->code_size) {
     current_list_no = 0;
     probe_indices.clear();
     for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
@@ -352,7 +353,7 @@ float IVFRaBitQHeapHandler<C, SL>::compute_full_multibit_distance(
         size_t /*db_idx*/,
         size_t local_q,
         size_t global_q,
-        size_t local_offset) const {
+        size_t local_offset) {
     const size_t ex_bits = index->rabitq.nb_bits - 1;
     const size_t dim = index->d;
 
@@ -374,12 +375,13 @@ float IVFRaBitQHeapHandler<C, SL>::compute_full_multibit_distance(
     size_t storage_idx_val = global_q * nprobe_val + probe_rank;
     const auto& query_factors = context->query_factors[storage_idx_val];
 
-    InvertedLists::ScopedCodes list_codes(index->invlists, current_list_no);
-    std::vector<uint8_t> unpacked_code(index->code_size);
-    packer->unpack_1(list_codes.get(), local_offset, unpacked_code.data());
+    // Use list_codes_ptr (already set by set_list_context) and the
+    // pre-allocated unpack_buf to avoid per-refinement ScopedCodes
+    // re-acquisition and heap allocation.
+    packer->unpack_1(list_codes_ptr, local_offset, unpack_buf.data());
 
     return rabitq_utils::compute_full_multibit_distance(
-            unpacked_code.data(),
+            unpack_buf.data(),
             ex_code,
             ex_fac,
             query_factors.rotated_q.data(),

faiss/IndexRaBitQFastScan.h

@@ -148,6 +148,9 @@ struct RaBitQHeapHandler
     const size_t packed_block_size;
     const size_t full_block_size;
     std::unique_ptr<CodePacker> packer; // cached for unpack in hot path
+    // Handler-local scratch reused across refinements. This assumes a handler
+    // instance is confined to one search slice and not entered concurrently.
+    std::vector<uint8_t> unpack_buf; // reusable buffer for unpack_1
 
     // Use float-based comparator for heap operations
     using Cfloat = typename std::conditional<
@@ -175,7 +178,8 @@ struct RaBitQHeapHandler
               storage_size(index->compute_per_vector_storage_size()),
               packed_block_size(((index->M2 + 1) / 2) * index->bbs),
               full_block_size(index->get_block_stride()),
-              packer(index->get_CodePacker()) {
+              packer(index->get_CodePacker()),
+              unpack_buf(index->code_size) {
 #pragma omp parallel for if (nq > 100)
         for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
             float* heap_dis = heap_distances + q * k;
@@ -293,7 +297,7 @@ struct RaBitQHeapHandler
     }
 
    private:
-    float compute_full_multibit_distance(size_t db_idx, size_t q) const {
+    float compute_full_multibit_distance(size_t db_idx, size_t q) {
         const size_t ex_bits = rabitq_index->rabitq.nb_bits - 1;
         const size_t dim = rabitq_index->d;
 
@@ -315,10 +319,10 @@ struct RaBitQHeapHandler
         const rabitq_utils::QueryFactorsData& query_factors =
                 context->query_factors[q];
 
-        std::vector<uint8_t> unpacked_code(rabitq_index->code_size);
-        packer->unpack_1(
-                rabitq_index->codes.get(), db_idx, unpacked_code.data());
-        const uint8_t* sign_bits = unpacked_code.data();
+        // Reuse pre-allocated unpack_buf to avoid per-refinement heap
+        // allocation.
+        packer->unpack_1(rabitq_index->codes.get(), db_idx, unpack_buf.data());
+        const uint8_t* sign_bits = unpack_buf.data();
 
         return rabitq_utils::compute_full_multibit_distance(
                 sign_bits,

faiss/impl/fast_scan/rabitq_result_handler.h

@@ -60,6 +60,9 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare<C, true, SL> {
     const size_t packed_block_size;
     const size_t full_block_size;
     std::unique_ptr<CodePacker> packer; // cached for unpack in hot path
+    // Handler-local scratch reused across refinements. This assumes a handler
+    // instance is confined to one search slice and not entered concurrently.
+    std::vector<uint8_t> unpack_buf; // reusable buffer for unpack_1
 
     // Use float-based comparator for heap operations
     using Cfloat = typename std::conditional<
@@ -97,7 +100,7 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare<C, true, SL> {
             size_t db_idx,
             size_t local_q,
             size_t global_q,
-            size_t local_offset) const;
+            size_t local_offset);
 };
 
 } // namespace simd_result_handlers