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

@@ -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.
 

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

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):

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

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):

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):