Fix RaBitQ fast scan aux offsets for large bbs (#5421)

Summary:
Fix RaBitQ FastScan aux-factor lookup for bbs > 32 so SIMD sub-blocks read factors from the correct vector offsets.

Pull Request resolved: #5421

Test Plan:
Added regression coverage for the bbs>32 aux-factor offset fix.

Non-IVF path (imported C++ test, now wired into the internal Buck build via a new `test_rabitq_fastscan_cpp` target):
  buck test fbcode//faiss/tests:test_rabitq_fastscan_cpp  =>  Pass 1

IVF path (new end-to-end regression test `test_ivf_large_bbs_aux_offsets`; `nlist=1` + `bbs=64` so a single list spans multiple bbs blocks and the `b>=1` SIMD sub-blocks exercise `idx_base % bbs != 0`):
  buck test fbcode//faiss/tests:test_rabitq_fastscan -- test_ivf_large_bbs_aux_offsets  =>  Pass 1

Red/green verified: the IVF test FAILS on the pre-fix offset (`aux_base + j * storage_size`) and PASSES with the fix (`aux_base + ((idx_base % bbs) + j) * storage_size`).

Reviewed By: mnorris11

Differential Revision: D112829433

Pulled By: alibeklfc

fbshipit-source-id: 0271fddc459a6af6675ebfb9078c44e9476733b3

Author: lyang24

faiss/IndexIVFRaBitQFastScan.h

@@ -267,11 +267,12 @@ void IVFRaBitQHeapHandler<C, SL>::handle(
     }
     const size_t max_positions = std::min<size_t>(32, this->ntotal - idx_base);
 
-    // Hoist aux pointer base out of loop: all 32 elements in this block share
-    // the same block base. Only the per-element offset (j * storage_size)
-    // varies.
+    // Hoist aux pointer base out of loop: it points at this 32-lane sub-block's
+    // factors, i.e. the bbs block base plus the loop-invariant intra-block
+    // offset ((idx_base % bbs)). Only the per-element j term varies below.
     const uint8_t* aux_base = this->list_codes_ptr +
-            (idx_base / index->bbs) * full_block_size + packed_block_size;
+            (idx_base / index->bbs) * full_block_size + packed_block_size +
+            (idx_base % index->bbs) * storage_size;
 
     // Cache index fields used in the inner loop.
     // Use overridden qb/centered from context if provided, else index defaults.

faiss/IndexRaBitQFastScan.h

@@ -209,8 +209,12 @@ struct RaBitQHeapHandler
                 : 0;
 
         const size_t block_idx = base_db_idx / rabitq_index->bbs;
+        // aux_base points at this 32-lane sub-block's factors: the bbs block
+        // base plus the loop-invariant intra-block offset
+        // ((base_db_idx % bbs)). Only the per-element i term varies below.
         const uint8_t* aux_base = rabitq_index->codes.get() +
-                block_idx * full_block_size + packed_block_size;
+                block_idx * full_block_size + packed_block_size +
+                (base_db_idx % rabitq_index->bbs) * storage_size;
 
         for (size_t i = 0; i < max_vectors; i++) {
             const size_t db_idx = base_db_idx + i;

tests/CMakeLists.txt

@@ -44,6 +44,7 @@ set(FAISS_TEST_SRC
   test_distances_simd.cpp
   test_simd_levels.cpp
   test_fast_scan_distance_to_code.cpp
+  test_rabitq_fastscan.cpp
   test_super_kmeans_foundations.cpp
 )
 

tests/test_rabitq_fastscan.cpp

@@ -0,0 +1,79 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <cstring>
+
+#include <faiss/IndexRaBitQFastScan.h>
+#include <faiss/impl/RaBitQUtils.h>
+#include <faiss/utils/Heap.h>
+
+namespace {
+
+void set_aux_factor(
+        faiss::IndexRaBitQFastScan& index,
+        size_t vec_pos,
+        float distance) {
+    const size_t packed_block_size = ((index.M2 + 1) / 2) * index.bbs;
+    auto* ptr = faiss::rabitq_utils::get_block_aux_ptr(
+            index.codes.get(),
+            vec_pos,
+            index.bbs,
+            packed_block_size,
+            index.get_block_stride(),
+            index.compute_per_vector_storage_size());
+
+    faiss::SignBitFactors factors;
+    factors.or_minus_c_l2sqr = distance;
+    factors.dp_multiplier = 0.0f;
+    memcpy(ptr, &factors, sizeof(factors));
+}
+
+} // namespace
+
+TEST(RaBitQFastScan, HeapHandlerUsesBbsLocalAuxOffset) {
+    faiss::IndexRaBitQFastScan index(4, faiss::METRIC_L2, 64, 1);
+    index.is_trained = true;
+    index.ntotal = 64;
+    index.ntotal2 = 64;
+    index.codes.resize(index.get_block_stride());
+    memset(index.codes.get(), 0, index.codes.size());
+
+    for (size_t i = 0; i < 64; i++) {
+        set_aux_factor(index, i, 100.0f);
+    }
+
+    // If b=1 incorrectly reads aux offsets 0..31, lane 1 wins with label 33.
+    // The correct bbs-local offsets 32..63 make lane 0 win with label 32.
+    set_aux_factor(index, 0, 1000.0f);
+    set_aux_factor(index, 1, 0.0f);
+    set_aux_factor(index, 32, 1.0f);
+
+    float distances[1];
+    int64_t labels[1];
+    faiss::FastScanDistancePostProcessing context;
+    faiss::RaBitQHeapHandler<faiss::CMax<float, int64_t>> handler(
+            &index,
+            /*nq_val=*/1,
+            /*k_val=*/1,
+            distances,
+            labels,
+            /*sel_in=*/nullptr,
+            &context,
+            /*multi_bit=*/false);
+
+    handler.set_block_origin(/*i0_in=*/0, /*j0_in=*/0);
+
+    using Simd16 = faiss::simd16uint16_tpl<faiss::SINGLE_SIMD_LEVEL_256>;
+    Simd16 zero(0);
+    handler.handle(/*q=*/0, /*b=*/1, zero, zero);
+    handler.end();
+
+    EXPECT_EQ(labels[0], 32);
+    EXPECT_FLOAT_EQ(distances[0], 1.0f);
+}

tests/test_rabitq_fastscan.py

@@ -510,6 +510,40 @@ def test_search_with_parameters(self):
         recall = faiss.eval_intersection(I, I_gt) / (nq * k)
         self.assertGreater(recall, 0.4)
 
+    def test_ivf_large_bbs_aux_offsets(self):
+        """Regression for aux-factor offsets when bbs > 32 (IVF path).
+
+        A bbs block holds bbs/32 SIMD sub-blocks. With nlist=1 every vector
+        lands in a single list that spans many bbs=64 blocks, so the b>=1
+        sub-blocks exercise ``idx_base % bbs != 0`` in IVFRaBitQHeapHandler.
+        If the per-element aux offset drops the ``idx_base % bbs`` term, roughly
+        half the vectors read another sub-block's factors and recall collapses,
+        so bbs=64 must match the bbs=32 baseline, which is always block-aligned.
+        """
+        d, nlist, nprobe, k = 64, 1, 1, 10
+        ds = datasets.SyntheticDataset(d, 2000, 2000, 100)
+        I_gt = ds.get_groundtruth(k)
+
+        def recall_for_bbs(bbs):
+            quantizer = faiss.IndexFlat(d, faiss.METRIC_L2)
+            index = faiss.IndexIVFRaBitQFastScan(
+                quantizer, d, nlist, faiss.METRIC_L2, bbs, True, 1
+            )
+            index.qb = 8
+            index.nprobe = nprobe
+            index.train(ds.get_train())
+            index.add(ds.get_database())
+            _, I = index.search(ds.get_queries(), k)
+            return faiss.eval_intersection(I[:, :k], I_gt[:, :k]) / (ds.nq * k)
+
+        recall_bbs32 = recall_for_bbs(32)
+        recall_bbs64 = recall_for_bbs(64)
+        np.testing.assert_(
+            abs(recall_bbs32 - recall_bbs64) < 0.02,
+            f"bbs=64 recall ({recall_bbs64:.3f}) diverged from bbs=32 "
+            f"({recall_bbs32:.3f}); aux offsets likely wrong for bbs>32",
+        )
+
 
 @for_all_simd_levels
 class TestIVFRaBitQFastScanFiltering(unittest.TestCase):