From 417c53e0ce23e8c3f2843bc727c9493e35b1d5d0 Mon Sep 17 00:00:00 2001 From: Gergely Szilvasy Date: Sat, 2 May 2026 14:30:13 -0700 Subject: [PATCH] Add NoneSIMDLevel context manager for cross-level reference checks (#5158) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5158 Several DD-dispatched code paths produce integer outputs (codes, top-k indices) where bit-exact equivalence between SIMD levels is the right correctness bar, but the existing tests assert only level- independent properties (decode error, recall) that wouldn't catch silent dispatch drift between AVX2/AVX512/NEON/NONE implementations. Adds NoneSIMDLevel context manager in common_faiss_tests.py: codes = quantizer.compute_codes(xb) with NoneSIMDLevel(): codes_none = quantizer.compute_codes(xb) np.testing.assert_array_equal(codes, codes_none) The context manager pins SIMDConfig to NONE for the duration of the block, restores the previous level on exit (incl. exception), and raises SkipTest when NONE isn't compiled in (static-AVX2 lane). Adds test_codes_match_none / test_PQ4_search_matches_none methods to five existing for_all_simd_levels test classes: * TestScalarQuantizerEncodeDecode (test_scalar_quantizer_correctness) * TestProductQuantizer (test_product_quantizer) * TestRaBitQuantizerEncodeDecode (test_rabitq) * TestSearch (test_fast_scan; checks (D, I) output of PQ16x4fs search) * TestResidualQuantizer (test_residual_quantizer; replaces inlined set_level/restore from D102821184 with the new helper) Caveat documented in the docstring: SkipTest raised inside subTest is swallowed; tests using both must do an explicit availability check before the loop. Applied where needed. Reviewed By: mdouze Differential Revision: D102989453 fbshipit-source-id: 475ee504271e0505615841ef82263776245419e0 --- tests/common_faiss_tests.py | 45 ++++++++++++++++++ tests/test_fast_scan.py | 54 +++++++++++++++++++++- tests/test_product_quantizer.py | 19 +++++++- tests/test_rabitq.py | 18 +++++++- tests/test_residual_quantizer.py | 15 +----- tests/test_scalar_quantizer_correctness.py | 22 ++++++++- 6 files changed, 156 insertions(+), 17 deletions(-) diff --git a/tests/common_faiss_tests.py b/tests/common_faiss_tests.py index 4451c5468..39fe25e6a 100644 --- a/tests/common_faiss_tests.py +++ b/tests/common_faiss_tests.py @@ -8,6 +8,7 @@ from __future__ import absolute_import, division, print_function, unicode_literals import sys +import unittest import numpy as np import faiss @@ -140,6 +141,50 @@ _SIMD_LEVEL_NAMES = { } +class NoneSIMDLevel: + """Context manager that temporarily pins SIMDConfig to SIMDLevel.NONE. + + Use inside a @for_all_simd_levels test to compute a reference value at + NONE for a cross-level equivalence check:: + + codes = quantizer.compute_codes(xb) + with NoneSIMDLevel(): + codes_none = quantizer.compute_codes(xb) + np.testing.assert_array_equal(codes, codes_none) + + The previous level is restored on exit (including via exception). + + If SIMDLevel.NONE is not compiled into this build (static-AVX2 CI + lane), the context manager raises SkipTest from __enter__, so the + enclosing test_method is skipped cleanly without per-call guards. + + Caveat: SkipTest raised from inside a self.subTest(...) block is + swallowed by the subTest mechanism instead of skipping the parent + test. When iterating with subTest, do an up-front availability check + before the loop:: + + if not faiss.SIMDConfig.is_simd_level_available( + faiss.SIMDLevel_NONE): + self.skipTest("SIMDLevel.NONE not available") + for x in cases: + with self.subTest(x=x): + ... + with NoneSIMDLevel(): + ... + """ + + def __enter__(self): + if not faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_NONE): + raise unittest.SkipTest("SIMDLevel.NONE not available") + self._prev = faiss.SIMDConfig.get_level() + faiss.SIMDConfig.set_level(faiss.SIMDLevel_NONE) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + faiss.SIMDConfig.set_level(self._prev) + return False + + def for_all_simd_levels(cls): """Parameterize a test class to run at each available SIMD level. diff --git a/tests/test_fast_scan.py b/tests/test_fast_scan.py index 011ea087a..2dfeea457 100644 --- a/tests/test_fast_scan.py +++ b/tests/test_fast_scan.py @@ -12,7 +12,11 @@ import faiss from faiss.contrib import datasets -from common_faiss_tests import for_all_simd_levels +from common_faiss_tests import ( + compare_binary_result_lists, + for_all_simd_levels, + NoneSIMDLevel, +) # the tests tend to timeout in stress modes + dev otherwise faiss.omp_set_num_threads(4) @@ -37,6 +41,54 @@ class TestSearch(unittest.TestCase): recall_at_1 = (Iref[:, 0] == Ia[:, 0]).sum() / nq assert recall_at_1 > 0.6 + def test_PQ4_search_matches_none(self): + """Fast-scan PQ search dispatches via the 256/512-bit QBS path. + + The integer QBS kernel is deterministic across SIMD levels *only if* + the upstream float PQ distance table is bit-identical. On some + toolchains (cmake aarch64 NEON/SVE) the float distance-table + computation reduces in a different order at NEON/SVE vs scalar, + producing a non-bit-identical LUT and thus different integer + distances. We diagnose which case applies, then assert accordingly: + - LUT identical -> assert (D, I) matches exactly (tie-tolerantly). + - LUT diverges -> fall back to a recall@1 check (the integer + kernel can't be bit-exact if its input isn't). + """ + ds = datasets.SyntheticDataset(32, 2000, 5000, 200) + index = faiss.index_factory(32, 'PQ16x4fs') + index.train(ds.get_train()) + index.add(ds.get_database()) + + xq = ds.get_queries() + nq = xq.shape[0] + M, ksub = index.pq.M, index.pq.ksub + + tab = np.zeros((nq, M, ksub), dtype='float32') + index.pq.compute_distance_tables( + nq, faiss.swig_ptr(xq), faiss.swig_ptr(tab)) + D, I = index.search(xq, 10) + + with NoneSIMDLevel(): + tab_none = np.zeros((nq, M, ksub), dtype='float32') + index.pq.compute_distance_tables( + nq, faiss.swig_ptr(xq), faiss.swig_ptr(tab_none)) + D_none, I_none = index.search(xq, 10) + + lut_diffs = int((tab != tab_none).sum()) + if lut_diffs == 0: + # Float LUT is bit-identical -> integer kernel must be too + # (modulo tied-index ordering). + compare_binary_result_lists(D, I, D_none, I_none) + else: + # Float LUT diverges across SIMD levels (float reductions are + # not order-stable across vector widths). Integer distances + # cannot be expected to match; check result-list overlap. + recall_at_1 = (I[:, 0] == I_none[:, 0]).mean() + self.assertGreater( + recall_at_1, 0.95, + f"LUT diverges ({lut_diffs} entries differ); " + f"recall@1 vs NONE = {recall_at_1:.3f}") + # This is an experiment to see if we can catch performance # regressions. It runs 2 codes, one should be faster than the diff --git a/tests/test_product_quantizer.py b/tests/test_product_quantizer.py index 5a12d2c3b..7f4485c91 100644 --- a/tests/test_product_quantizer.py +++ b/tests/test_product_quantizer.py @@ -10,7 +10,7 @@ import numpy as np import faiss import unittest -from common_faiss_tests import for_all_simd_levels +from common_faiss_tests import for_all_simd_levels, NoneSIMDLevel @for_all_simd_levels @@ -74,6 +74,23 @@ class TestProductQuantizer(unittest.TestCase): for i in range(16): self.do_test_codec(i + 1) + def test_codes_match_none(self): + """PQ codes are integer; encode dispatch must produce bit-identical + output at every SIMD level.""" + if not faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_NONE): + self.skipTest("SIMDLevel.NONE not available") + d, cs = 64, 4 + np.random.seed(123) + x = np.random.random(size=(2000, d)).astype('float32') + for nbits in (4, 6, 8): + with self.subTest(nbits=nbits): + pq = faiss.ProductQuantizer(d, cs, nbits) + pq.train(x) + codes = pq.compute_codes(x) + with NoneSIMDLevel(): + codes_none = pq.compute_codes(x) + np.testing.assert_array_equal(codes, codes_none) + @for_all_simd_levels class TestPQTransposedCentroids(unittest.TestCase): diff --git a/tests/test_rabitq.py b/tests/test_rabitq.py index 8d50bb30f..3b66af3db 100644 --- a/tests/test_rabitq.py +++ b/tests/test_rabitq.py @@ -9,7 +9,7 @@ import numpy as np import faiss from faiss.contrib import datasets -from common_faiss_tests import for_all_simd_levels +from common_faiss_tests import for_all_simd_levels, NoneSIMDLevel def random_rotation(d, seed=123): @@ -539,6 +539,22 @@ class TestRaBitQuantizerEncodeDecode(unittest.TestCase): def test_encode_decode_IP(self): self.do_test_encode_decode(16, faiss.METRIC_INNER_PRODUCT) + def test_codes_match_none(self): + """RaBitQ codes are integer; encode must be bit-identical across + SIMD levels.""" + if not faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_NONE): + self.skipTest("SIMDLevel.NONE not available") + d = 64 + rs = np.random.RandomState(123) + vec = rs.standard_normal((100, d)).astype(np.float32) + for metric in (faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT): + with self.subTest(metric=metric): + quantizer = faiss.RaBitQuantizer(d, metric) + codes = quantizer.compute_codes(vec) + with NoneSIMDLevel(): + codes_none = quantizer.compute_codes(vec) + np.testing.assert_array_equal(codes, codes_none) + # ============================================================================== # Multi-bit RaBitQ Test Suite diff --git a/tests/test_residual_quantizer.py b/tests/test_residual_quantizer.py index 4ac932b83..a63c15bd4 100644 --- a/tests/test_residual_quantizer.py +++ b/tests/test_residual_quantizer.py @@ -8,7 +8,7 @@ import numpy as np import faiss import unittest -from common_faiss_tests import for_all_simd_levels +from common_faiss_tests import for_all_simd_levels, NoneSIMDLevel from faiss.contrib import datasets from faiss.contrib.inspect_tools import get_additive_quantizer_codebooks @@ -267,12 +267,7 @@ class TestResidualQuantizer(unittest.TestCase): """RQ encode is integer-valued, so codes must be bit-exact across SIMD levels. Catches dispatch drift in residual_quantizer_encode_steps.cpp:96,369 (with_simd_level_256bit). - Static-build CI lanes only have one SIMD level, so the comparison - is vacuous and skipped there. """ - if not faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_NONE): - self.skipTest("SIMDLevel.NONE not available in this build") - ds = datasets.SyntheticDataset(32, 1000, 200, 0) rq = faiss.ResidualQuantizer(ds.d, 4, 6) rq.train_type = faiss.ResidualQuantizer.Train_default @@ -281,14 +276,8 @@ class TestResidualQuantizer(unittest.TestCase): xb = ds.get_database() codes = rq.compute_codes(xb) - # Switch to NONE, compute the reference, and restore so the rest of - # this test instance (and tearDown) sees the decorator's level. - prev_level = faiss.SIMDConfig.get_level() - faiss.SIMDConfig.set_level(faiss.SIMDLevel_NONE) - try: + with NoneSIMDLevel(): codes_none = rq.compute_codes(xb) - finally: - faiss.SIMDConfig.set_level(prev_level) np.testing.assert_array_equal(codes, codes_none) def test_clipping(self): diff --git a/tests/test_scalar_quantizer_correctness.py b/tests/test_scalar_quantizer_correctness.py index 468a74a0e..3924a44f5 100644 --- a/tests/test_scalar_quantizer_correctness.py +++ b/tests/test_scalar_quantizer_correctness.py @@ -12,7 +12,7 @@ import unittest from faiss.contrib.datasets import SyntheticDataset -from common_faiss_tests import for_all_simd_levels +from common_faiss_tests import for_all_simd_levels, NoneSIMDLevel @for_all_simd_levels @@ -64,6 +64,26 @@ class TestScalarQuantizerEncodeDecode(unittest.TestCase): faiss.normalize_L2(self.xb) self.do_encode_decode(faiss.ScalarQuantizer.QT_8bit_tqmse, 0.1) + def test_codes_match_none(self): + """SQ codes are integer; encode dispatch (sq-dispatch.h) must + produce bit-identical output at every SIMD level. Catches drift in + the per-ISA scalar-quantizer encode kernels.""" + if not faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_NONE): + self.skipTest("SIMDLevel.NONE not available") + for qtype in ( + faiss.ScalarQuantizer.QT_8bit, + faiss.ScalarQuantizer.QT_4bit, + faiss.ScalarQuantizer.QT_8bit_uniform, + faiss.ScalarQuantizer.QT_4bit_uniform, + faiss.ScalarQuantizer.QT_fp16): + with self.subTest(qtype=qtype): + sq = faiss.ScalarQuantizer(self.d, qtype) + sq.train(self.xb) + codes = sq.compute_codes(self.xb) + with NoneSIMDLevel(): + codes_none = sq.compute_codes(self.xb) + np.testing.assert_array_equal(codes, codes_none) + @for_all_simd_levels class TestScalarQuantizerSearch(unittest.TestCase):