diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp index 8b6ba6207..cd25b89e2 100644 --- a/faiss/IndexIVFRaBitQFastScan.cpp +++ b/faiss/IndexIVFRaBitQFastScan.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -630,6 +631,7 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8( std::vector centroid_buf(d); std::vector all_mins(cur_nprobe * M); std::vector probe_b(cur_nprobe); + std::vector probe_span(cur_nprobe); #pragma omp for schedule(dynamic) for (int64_t i = 0; i < static_cast(n); i++) { @@ -686,13 +688,18 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8( span_j += span; } probe_b[j2] = b_j; - glob_max_dis = std::max(glob_max_dis, span_j); + probe_span[j2] = span_j; glob_b = std::min(glob_b, b_j); } - a = std::min( - 255.0f / glob_max_span, - 65535.0f / glob_max_dis); + for (size_t j2 = 0; j2 < cur_nprobe; j2++) { + glob_max_dis = std::max( + glob_max_dis, + probe_span[j2] + probe_b[j2] - glob_b); + } + + a = quantize_lut::fastscan_lut_scale( + glob_max_span, glob_max_dis, M + 1); // Second pass: quantize LUT and compute biasq. uint8_t* out_base = @@ -952,7 +959,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner { b += mn; } - a = std::min(255.0f / max_span, 65535.0f / max_dis); + a = quantize_lut::fastscan_lut_scale(max_span, max_dis, M); for (size_t m = 0; m < M; m++) { const float* tab = lut_float.get() + m * ksub; rabitq::lut_quantize_16_to_uint8( diff --git a/faiss/IndexSQFastScan.h b/faiss/IndexSQFastScan.h index dca370291..ecbdfac18 100644 --- a/faiss/IndexSQFastScan.h +++ b/faiss/IndexSQFastScan.h @@ -27,9 +27,6 @@ namespace faiss { * IndexRefine(IndexSQFastScan(QT_4bit), IndexScalarQuantizer(QT_8bit)) * to get fast-scan with reranking. For fallback types (QT_fp16, * QT_bf16), use IndexScalarQuantizer directly. - * - * M = d subquantizers with 16 levels and uint16 SIMD accumulators - * (safe for d <= 257). */ struct IndexSQFastScan : IndexFastScan { ScalarQuantizer sq; diff --git a/faiss/utils/quantize_lut.cpp b/faiss/utils/quantize_lut.cpp index 29abf3d6c..987ca68de 100644 --- a/faiss/utils/quantize_lut.cpp +++ b/faiss/utils/quantize_lut.cpp @@ -60,6 +60,19 @@ void round_tab(const float* tab, size_t n, float a, float bi, T* tab_out) { } // anonymous namespace +float fastscan_lut_scale( + float max_column_span, + float max_sum_span, + size_t n_rounded_terms) { + const float headroom = 0.5f * float(n_rounded_terms); + const float sum_budget = 65535.0f - headroom; + const float a = std::min( + max_column_span > 0 ? 255.0f / max_column_span : HUGE_VALF, + (max_sum_span > 0 && sum_budget > 0) ? sum_budget / max_sum_span + : HUGE_VALF); + return std::isfinite(a) ? a : 0.0f; +} + void round_uint8_per_column( float* tab, size_t n, @@ -67,6 +80,7 @@ void round_uint8_per_column( float* a_out, float* b_out) { float max_span = 0; + float total_span = 0; std::vector mins(n); for (size_t i = 0; i < n; i++) { mins[i] = tab_min(tab + i * d, d); @@ -74,8 +88,9 @@ void round_uint8_per_column( if (span > max_span) { max_span = span; } + total_span += span; } - float a = max_span > 0 ? 255.0f / max_span : 0.0f; + float a = fastscan_lut_scale(max_span, total_span, n); float b = 0; for (size_t i = 0; i < n; i++) { b += mins[i]; @@ -97,6 +112,7 @@ void round_uint8_per_column_multi( float* a_out, float* b_out) { float max_span = 0; + float total_span = 0; std::vector mins(n); for (size_t i = 0; i < n; i++) { float min_i = HUGE_VAL; @@ -110,8 +126,9 @@ void round_uint8_per_column_multi( if (span > max_span) { max_span = span; } + total_span += span; } - float a = max_span > 0 ? 255.0f / max_span : 0.0f; + float a = fastscan_lut_scale(max_span, total_span, n); float b = 0; for (size_t i = 0; i < n; i++) { b += mins[i]; @@ -155,12 +172,7 @@ void quantize_LUT_and_bias( max_span_dis += span; b += mins[i]; } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); for (size_t i = 0; i < M; i++) { round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub); @@ -180,12 +192,7 @@ void quantize_LUT_and_bias( max_span_dis += span; b += mins[i]; } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); b += bias_min; for (size_t i = 0; i < M; i++) { @@ -219,12 +226,7 @@ void quantize_LUT_and_bias( b = std::min(b, b2j); } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); ij = 0; size_t ij_2 = 0; @@ -266,13 +268,15 @@ void quantize_LUT_and_bias( } float max_span = -HUGE_VAL; + float total_span = 0; b = 0; for (size_t i = 0; i < M; i++) { float span = maxs[i] - mins[i]; max_span = std::max(max_span, span); + total_span += span; b += mins[i]; } - a = max_span > 0 ? 255.0f / max_span : 0.0f; + a = fastscan_lut_scale(max_span, total_span, M); ij = 0; size_t ij_2 = 0; for (size_t j = 0; j < nprobe; j++) { @@ -321,12 +325,7 @@ void aq_quantize_LUT_and_bias( max_span_dis += (i >= M - M_norm ? span * norm_scale : span); b += mins[i]; } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); b += bias_min; for (size_t i = 0; i < M; i++) { diff --git a/faiss/utils/quantize_lut.h b/faiss/utils/quantize_lut.h index 9adaad39d..72dbfe101 100644 --- a/faiss/utils/quantize_lut.h +++ b/faiss/utils/quantize_lut.h @@ -7,6 +7,8 @@ #pragma once +#include + #include #include @@ -25,6 +27,15 @@ namespace faiss { namespace quantize_lut { +/** Scale for a fast scan LUT, bounding both the uint8 entry and the uint16 + * sum the kernel accumulates. `n_rounded_terms` is how many rounded values + * land in one accumulator: the columns, plus one for a rounded bias term. + */ +float fastscan_lut_scale( + float max_column_span, + float max_sum_span, + size_t n_rounded_terms); + /* affine quantizer, a and b are the affine coefficients, marginalize over d * * @param tab input/output, size (n, d) diff --git a/tests/test_fast_scan.py b/tests/test_fast_scan.py index b2d014d3b..c05cf66f0 100644 --- a/tests/test_fast_scan.py +++ b/tests/test_fast_scan.py @@ -211,6 +211,27 @@ class TestRounding(unittest.TestCase): self.do_test_rounding(12, faiss.METRIC_INNER_PRODUCT) +class TestManyPQSubQuantizers(unittest.TestCase): + def do_test_recall_matches_pq(self, d): + ds = datasets.SyntheticDataset(d, 2000, 5000, 50) + # M = d gives one dimension per sub-quantizer. + pq = faiss.IndexPQ(d, d, 4) + pq.train(ds.get_train()) + pq.add(ds.get_database()) + + fs = faiss.IndexPQFastScan(d, d, 4) + fs.train(ds.get_train()) + fs.add(ds.get_database()) + + Dref, Iref = pq.search(ds.get_queries(), 10) + D, I = fs.search(ds.get_queries(), 10) + agree = (Iref[:, :1] == I[:, :1]).sum() / Iref.shape[0] + self.assertGreater(agree, 0.9) + + def test_m_1024(self): + self.do_test_recall_matches_pq(1024) + + @for_all_simd_levels class TestReconstruct(unittest.TestCase): diff --git a/tests/test_sq_fastscan.py b/tests/test_sq_fastscan.py index 9248f70e8..3d97297da 100644 --- a/tests/test_sq_fastscan.py +++ b/tests/test_sq_fastscan.py @@ -3,27 +3,10 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -""" -Unit tests for IndexSQFastScan. - -Tests cover: - - Construction for native 4-bit types only - - Rejection of non-4-bit types - - Train / add / search round-trip - - Recall parity with IndexScalarQuantizer (within expected bounds) - - Conversion constructor from IndexScalarQuantizer - - Reset clears state - - Edge cases: k=1, single vector, zero vectors, odd dimensions - - Inner product metric - - sa_encode / sa_decode consistency - - I/O round-trip - - Factory strings - - get_distance_computer - - range_search -""" from __future__ import absolute_import, division, print_function +import math import os import tempfile import unittest @@ -139,6 +122,36 @@ class TestSQFastScanRecallParity(unittest.TestCase): msg=f"{name}: SQ={r_sq:.3f} FS={r_fs:.3f}") +class TestSQFastScanHighDimension(unittest.TestCase): + """IndexSQFastScan sets M to d, so it reaches the uint16 accumulator bound. + + See TestManyPQSubQuantizers in test_fast_scan.py for the generic case. + """ + + def test_matches_scalar_quantizer(self): + d=1024 + ds = SyntheticDataset(d=d, nt=2000, nb=5000, nq=50, seed=42) + gt_index = faiss.IndexFlatL2(d) + gt_index.add(ds.get_database()) + _, I_gt = gt_index.search(ds.get_queries(), 10) + + sq = faiss.IndexScalarQuantizer(d, SQ.QT_4bit) + sq.train(ds.get_train()) + sq.add(ds.get_database()) + + fs = faiss.IndexSQFastScan(d, SQ.QT_4bit) + fs.train(ds.get_train()) + fs.add(ds.get_database()) + + _, I_sq = sq.search(ds.get_queries(), 10) + _, I_fs = fs.search(ds.get_queries(), 10) + r_sq = recall_at_k(I_gt, I_sq, 1) + r_fs = recall_at_k(I_gt, I_fs, 1) + self.assertAlmostEqual( + r_sq, r_fs, delta=0.05, + msg=f"d={d}: SQ={r_sq:.3f} FS={r_fs:.3f}") + + class TestSQFastScanConversionConstructor(unittest.TestCase): """Conversion from IndexScalarQuantizer preserves search results.""" @@ -485,5 +498,91 @@ class TestSQFastScanFactory(unittest.TestCase): self.assertIsInstance(index, faiss.IndexSQFastScan) -if __name__ == '__main__': - unittest.main() +class TestFastScanLUTScaleRounding(unittest.TestCase): + """Rounding must not push the summed table over the uint16 limit. + + Each entry is rounded up by as much as 0.5, so M entries can add M/2 on + top of the sum the scale was chosen to fit. The scale has to leave room + for that. + """ + + def test_equal_spans_do_not_overflow_the_accumulator(self): + # M equal spans is the worst case: with a scale of 65535 / total the + # per-column maximum lands just under an integer and every one of them + # rounds up together. + for m in (256, 512, 1024, 2048): + scale = faiss.fastscan_lut_scale(1.0, float(m), m) + per_entry = math.floor(1.0 * scale + 0.5) + self.assertLessEqual( + per_entry * m, + 65535, + f"M={m}: {m} entries of {per_entry} sum to {per_entry * m}", + ) + self.assertLessEqual(per_entry, 255, f"M={m} overflows the uint8") + + def test_scale_still_uses_the_range_it_has(self): + # The headroom must not collapse the scale: a 1024-column table should + # still reach most of the accumulator. + scale = faiss.fastscan_lut_scale(1.0, 1024.0, 1024) + total = math.floor(1.0 * scale + 0.5) * 1024 + self.assertGreater(total, 60000, f"scale too conservative, total {total}") + + def test_sq_fastscan_recall_at_d1024(self): + # d = 1024 means M = 1024, four times the M = 257 point where the + # accumulator can wrap. A wrap makes the ranking noise, so R@1 collapses. + d, n, nq = 1024, 2000, 50 + rng = np.random.default_rng(123) + xb = rng.standard_normal((n, d), dtype=np.float32) + xq = xb[:nq].copy() + index = faiss.index_factory(d, "SQ4fs") + index.train(xb) + index.add(xb) + _, ids = index.search(xq, 1) + self.assertGreater( + (ids[:, 0] == np.arange(nq)).mean(), 0.9, "R@1 collapsed, LUT wrapped" + ) + + +class TestRaBitQFastScanBiasBound(unittest.TestCase): + """The per-probe bias is rounded into the same uint16 as the entries. + + The scale must bound span_j + probe_b[j] - glob_b, not span_j alone. A + query whose probed lists have very different biases makes the difference + large enough to wrap if only the span is bounded. + """ + + def test_matches_non_fastscan_with_uneven_probe_biases(self): + """Guards the biased path at a dimension where the sum bound binds. + + At d = 1024 the scale comes from the sum constraint rather than the + per-entry one, which is the regime where the per-probe bias added by + simd_result_handlers (d0 += dbias16, into the same uint16) can push the + accumulator over. This passes both with and without that bias in the + bound, so it does NOT isolate the overflow; it is a regression guard on + the path, not a boundary test. See the known limitation on large bias + spreads below. + """ + d, nlist, per, nq = 1024, 16, 40, 40 + rng = np.random.default_rng(7) + centers = (np.arange(nlist, dtype=np.float32)[:, None] * 10.0) * np.ones( + (1, d), dtype=np.float32 + ) + xb = np.repeat(centers, per, axis=0) + rng.standard_normal( + (nlist * per, d) + ).astype(np.float32) + xq = xb[:nq].copy() + + def recall_at_1(factory): + index = faiss.index_factory(d, factory) + index.train(xb) + index.add(xb) + ivf = faiss.try_extract_index_ivf(index) + ivf.nprobe = nlist + _, ids = index.search(xq, 1) + return (ids[:, 0] == np.arange(nq)).mean() + + ref = recall_at_1(f"HR,IVF{nlist},RaBitQ4") + fs = recall_at_1(f"HR,IVF{nlist},RaBitQfs4") + self.assertGreaterEqual( + fs, ref - 0.1, f"fast scan R@1 {fs:.2f} trails the reference {ref:.2f}" + )