From 7a4e7975d1aa88a783e1304932346e3dce7214d1 Mon Sep 17 00:00:00 2001 From: lyang24 Date: Fri, 31 Jul 2026 07:49:57 -0700 Subject: [PATCH] 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: https://github.com/facebookresearch/faiss/pull/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 --- faiss/IndexIVFRaBitQFastScan.h | 9 ++-- faiss/IndexRaBitQFastScan.h | 6 ++- tests/CMakeLists.txt | 1 + tests/test_rabitq_fastscan.cpp | 79 ++++++++++++++++++++++++++++++++++ tests/test_rabitq_fastscan.py | 34 +++++++++++++++ 5 files changed, 124 insertions(+), 5 deletions(-) create mode 100644 tests/test_rabitq_fastscan.cpp diff --git a/faiss/IndexIVFRaBitQFastScan.h b/faiss/IndexIVFRaBitQFastScan.h index b07839460..51d6761c9 100644 --- a/faiss/IndexIVFRaBitQFastScan.h +++ b/faiss/IndexIVFRaBitQFastScan.h @@ -267,11 +267,12 @@ void IVFRaBitQHeapHandler::handle( } const size_t max_positions = std::min(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. diff --git a/faiss/IndexRaBitQFastScan.h b/faiss/IndexRaBitQFastScan.h index 5c9bf0833..8a1891cd6 100644 --- a/faiss/IndexRaBitQFastScan.h +++ b/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; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a72172f78..c157676ef 100644 --- a/tests/CMakeLists.txt +++ b/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 ) diff --git a/tests/test_rabitq_fastscan.cpp b/tests/test_rabitq_fastscan.cpp new file mode 100644 index 000000000..f58adbb0d --- /dev/null +++ b/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 + +#include + +#include +#include +#include + +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> 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; + 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); +} diff --git a/tests/test_rabitq_fastscan.py b/tests/test_rabitq_fastscan.py index b4be13876..b06045aee 100644 --- a/tests/test_rabitq_fastscan.py +++ b/tests/test_rabitq_fastscan.py @@ -510,6 +510,40 @@ class TestRaBitQFastScan(unittest.TestCase): 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):