facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
faiss/tests/test_scalar_quantizer_correctness.py
T
Michael Norris b5a14632a6 Run the eight SIMD kernel tests that no build executed (#5620) 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5620

**TL;DR:** Eight SIMD kernel tests never ran, and two byte-domain kernels trust their caller. This change makes three fixes:

1. **Six tests never compiled.** `test_rabitq_simd.cpp` hides them behind `COMPILE_SIMD_*` macros that the faiss target does not export.
2. **Two tests skipped themselves in a static build.** `TestSQByteDirectAcrossSIMDLevels` needs `SIMDLevel::NONE` as a reference, and a static build does not report that level.
3. **The byte kernels never test their dimension requirement.** `DistanceComputerByte` and `DistanceComputerByteSigned` read 16 bytes per loop step.

## 1. Six tests never compiled

`test_rabitq_simd.cpp` wraps two blocks in `#ifdef COMPILE_SIMD_AVX2` and `#ifdef COMPILE_SIMD_AVX512`. The faiss target sets those macros as `preprocessor_flags`, and it does not export them. A test translation unit never sees them. So the preprocessor deleted both blocks. The file holds nine tests and ran three.

The AVX2 block loses its guard, because every x86 faiss build holds the AVX2 kernels. Its three tests also lose an `#ifdef FAISS_ENABLE_DD` around a runtime skip, which the same problem disabled. `is_simd_level_available()` reports the right answer in a static build too.

The AVX512 block moves to a dynamic dispatch target. A static x86 build compiles the AVX2 kernels only, so the test fails to link there rather than skips. `test_rabitq_simd_vpopcnt.cpp` becomes `test_rabitq_simd_dd.cpp`, because it now holds every RaBitQ test that a static build cannot link. A new `test_rabitq_simd_util.h` holds the scalar reference that both files call.

## 2. Two tests skipped themselves in a static build

`TestSQByteDirectAcrossSIMDLevels` measured the index a second time at `SIMDLevel::NONE`, then compared the two runs. A static build reports only the level it compiles, and omits `NONE`, so both tests skipped.

The tests now compute the distances from the definition. The distances are integers, so the reference is exact. The tests check absolute values now, instead of the agreement of two faiss paths.

## 3. The byte kernels trust the caller

`DistanceComputerByte` and `DistanceComputerByteSigned` read 16 bytes per loop step, so they require `d % 16 == 0`. `sq-dispatch.h` enforces that, and no other caller builds them.

A caller that ignores the requirement reads up to 15 bytes past each code and returns a wrong distance. Each constructor now tests the requirement and throws. The AVX2 and the ARM_NEON specializations both get the test.

Reviewed By: alibeklfc

Differential Revision: D119545379

fbshipit-source-id: ddfdcc30d3662f469b2ec3dbb9d5b9f4dafa21aa
2026-09-10 20:05:56 -07:00

684 lines
25 KiB
Python
Raw Permalink Blame History

# 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.
from __future__ import absolute_import, division, print_function
import os
import tempfile
import numpy as np
import faiss
import unittest
from faiss.contrib.datasets import SyntheticDataset
from common_faiss_tests import for_all_simd_levels, NoneSIMDLevel
@for_all_simd_levels
class TestScalarQuantizerEncodeDecode(unittest.TestCase):
def setUp(self):
self.d = 32
self.ds = SyntheticDataset(d=self.d, nt=0, nb=1000, nq=0, seed=42)
self.xb = self.ds.get_database()
def do_encode_decode(self, qtype, max_err):
sq = faiss.ScalarQuantizer(self.d, qtype)
sq.train(self.xb)
codes = sq.compute_codes(self.xb)
xb_decoded = sq.decode(codes)
self.assertEqual(xb_decoded.shape, self.xb.shape)
err = np.abs(self.xb - xb_decoded).max()
self.assertLess(err, max_err)
def test_4bit(self):
self.do_encode_decode(faiss.ScalarQuantizer.QT_4bit, 0.1)
def test_fp16(self):
self.do_encode_decode(faiss.ScalarQuantizer.QT_fp16, 0.001)
def test_8bit_uniform(self):
self.do_encode_decode(faiss.ScalarQuantizer.QT_8bit_uniform, 0.01)
def test_4bit_uniform(self):
self.do_encode_decode(faiss.ScalarQuantizer.QT_4bit_uniform, 0.1)
def test_tqmse_1bit(self):
faiss.normalize_L2(self.xb)
self.do_encode_decode(faiss.ScalarQuantizer.QT_1bit_tqmse, 1.0)
def test_tqmse_2bit(self):
faiss.normalize_L2(self.xb)
self.do_encode_decode(faiss.ScalarQuantizer.QT_2bit_tqmse, 0.8)
def test_tqmse_3bit(self):
faiss.normalize_L2(self.xb)
self.do_encode_decode(faiss.ScalarQuantizer.QT_3bit_tqmse, 0.6)
def test_tqmse_4bit(self):
faiss.normalize_L2(self.xb)
self.do_encode_decode(faiss.ScalarQuantizer.QT_4bit_tqmse, 0.2)
def test_tqmse_8bit(self):
faiss.normalize_L2(self.xb)
self.do_encode_decode(faiss.ScalarQuantizer.QT_8bit_tqmse, 0.1)
def test_tqmse_preserves_norm(self):
"""tqmse encodes unit-norm vectors with no per-vector scale factor in
the code, so reconstructions must keep unit norm. The Lloyd-Max table
is optimal for unit *variance*, so it only lines up with unit-norm
input once scaled by 1/sqrt(d) -- without that, every component lands
in the innermost cells and the codebook collapses toward a sign
quantizer, inflating norms by 15% at d=64 and 258% at d=768."""
for qtype, nbits in (
(faiss.ScalarQuantizer.QT_1bit_tqmse, 1),
(faiss.ScalarQuantizer.QT_2bit_tqmse, 2),
(faiss.ScalarQuantizer.QT_3bit_tqmse, 3),
(faiss.ScalarQuantizer.QT_4bit_tqmse, 4),
(faiss.ScalarQuantizer.QT_8bit_tqmse, 8),
):
for d in (32, 64, 768):
with self.subTest(nbits=nbits, d=d):
rs = np.random.RandomState(123)
x = rs.randn(200, d).astype("float32")
faiss.normalize_L2(x)
sq = faiss.ScalarQuantizer(d, qtype)
sq.train(x)
decoded = sq.decode(sq.compute_codes(x))
norms = np.linalg.norm(decoded, axis=1)
# 1-bit cannot resolve magnitude at all (every component
# reconstructs to +/-c), so it gets a looser band.
tol = 0.25 if nbits == 1 else 0.1
self.assertLess(abs(np.median(norms) - 1.0), tol)
# Guard the collapse directly: a codebook matched to the
# data reconstructs most of its 2^nbits distinct levels.
# This is what catches 8-bit, whose norm stays near 1 even
# when collapsed (31 of 256 levels used at d=768).
used = len(np.unique(decoded))
self.assertGreaterEqual(used, 2**nbits // 2)
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):
def setUp(self):
self.d = 32
self.ds = SyntheticDataset(d=self.d, nt=0, nb=10000, nq=100, seed=42)
self.xb = self.ds.get_database()
self.xq = self.ds.get_queries()
self.xb_unit = self.xb.copy()
self.xq_unit = self.xq.copy()
faiss.normalize_L2(self.xb_unit)
faiss.normalize_L2(self.xq_unit)
def do_search(self, factory_str, min_recall):
index_gt = faiss.IndexFlatL2(self.d)
index_gt.add(self.xb)
_, I_gt = index_gt.search(self.xq, 10)
index = faiss.index_factory(self.d, factory_str)
index.train(self.xb)
index.add(self.xb)
_, I = index.search(self.xq, 10)
recall = (I_gt[:, 0] == I[:, 0]).mean()
self.assertGreater(recall, min_recall)
def test_SQ4(self):
self.do_search("SQ4", 0.5)
def test_SQfp16(self):
self.do_search("SQfp16", 0.99)
def test_tqmse_search(self):
index_gt = faiss.IndexFlatL2(self.d)
index_gt.add(self.xb_unit)
_, I_gt = index_gt.search(self.xq_unit, 10)
recalls = {}
for factory_str in ["L2norm,RR,SQtqmse4", "L2norm,RR,SQtqmse8"]:
index = faiss.index_factory(self.d, factory_str)
index.train(self.xb)
index.add(self.xb)
_, I = index.search(self.xq, 10)
recalls[factory_str] = (I_gt[:, 0] == I[:, 0]).mean()
self.assertGreater(recalls["L2norm,RR,SQtqmse4"], 0.2)
self.assertGreaterEqual(
recalls["L2norm,RR,SQtqmse8"], recalls["L2norm,RR,SQtqmse4"]
)
@for_all_simd_levels
class TestScalarQuantizerDistances(unittest.TestCase):
def test_distance_matches_reconstruct(self):
d = 32
ds = SyntheticDataset(d=d, nt=0, nb=100, nq=1, seed=42)
x = ds.get_queries()
xb = ds.get_database()
index = faiss.index_factory(d, "SQ8")
index.train(xb)
index.add(xb)
D, I = index.search(x, 10)
for i in range(10):
xr = index.reconstruct(int(I[0, i]))
dist = ((x[0] - xr) ** 2).sum()
self.assertAlmostEqual(D[0, i], dist, places=4)
def test_tqmse_distance_matches_reconstruct(self):
d = 32
ds = SyntheticDataset(d=d, nt=0, nb=100, nq=1, seed=42)
x = ds.get_queries()
xb = ds.get_database()
faiss.normalize_L2(x)
faiss.normalize_L2(xb)
index = faiss.index_factory(d, "SQtqmse8")
index.train(xb)
index.add(xb)
D, I = index.search(x, 10)
for i in range(10):
xr = index.reconstruct(int(I[0, i]))
dist = ((x[0] - xr) ** 2).sum()
self.assertAlmostEqual(D[0, i], dist, places=4)
@for_all_simd_levels
class TestScalarQuantizerEdgeCases(unittest.TestCase):
def test_zero_vectors(self):
d = 32
xb = np.zeros((100, d), dtype="float32")
index = faiss.index_factory(d, "SQ8")
index.train(xb)
index.add(xb)
D, _ = index.search(np.zeros((1, d), dtype="float32"), 10)
self.assertTrue(np.allclose(D, 0, atol=1e-5))
def test_constant_vectors(self):
d = 32
xb = np.ones((100, d), dtype="float32") * 0.5
index = faiss.index_factory(d, "SQ8")
index.train(xb)
index.add(xb)
D, _ = index.search(np.ones((1, d), dtype="float32") * 0.5, 10)
self.assertTrue(np.allclose(D, 0, atol=1e-3))
def test_extreme_dims(self):
"""Test extreme dimension values (very small and very large)."""
for d in [1, 1024]:
with self.subTest(d=d):
ds = SyntheticDataset(d=d, nt=0, nb=100, nq=10, seed=42)
xb = ds.get_database()
xq = ds.get_queries()
index = faiss.index_factory(d, "SQ8")
index.train(xb)
index.add(xb)
D, I = index.search(xq, 10)
self.assertEqual(D.shape, (10, 10))
self.assertEqual(I.shape, (10, 10))
def test_non_simd_dims(self):
"""Test dimensions not aligned to SIMD width (8/16)."""
for d in [7, 9, 15, 17, 31, 33, 63, 65]:
with self.subTest(d=d):
ds = SyntheticDataset(d=d, nt=0, nb=100, nq=5, seed=42)
xb = ds.get_database()
xq = ds.get_queries()
index = faiss.index_factory(d, "SQ8")
index.train(xb)
index.add(xb)
D, I = index.search(xq, 10)
self.assertEqual(D.shape, (5, 10))
self.assertEqual(I.shape, (5, 10))
def test_tqmse_non_simd_dims(self):
factory_strings = [
"SQtqmse1",
"SQtqmse2",
"SQtqmse3",
"SQtqmse4",
"SQtqmse8",
]
for d in [7, 9, 15, 17, 31, 33, 63, 65]:
for factory_str in factory_strings:
with self.subTest(d=d, factory_str=factory_str):
ds = SyntheticDataset(d=d, nt=0, nb=100, nq=5, seed=42)
xb = ds.get_database()
xq = ds.get_queries()
faiss.normalize_L2(xb)
faiss.normalize_L2(xq)
index = faiss.index_factory(d, factory_str)
index.train(xb)
index.add(xb)
D, I = index.search(xq, 10)
self.assertEqual(D.shape, (5, 10))
self.assertEqual(I.shape, (5, 10))
@for_all_simd_levels
class TestScalarQuantizerIP(unittest.TestCase):
def test_inner_product(self):
d = 32
ds = SyntheticDataset(d=d, nt=0, nb=1000, nq=10, seed=42)
xb = ds.get_database().copy()
xq = ds.get_queries().copy()
faiss.normalize_L2(xb)
faiss.normalize_L2(xq)
index = faiss.index_factory(d, "SQ8", faiss.METRIC_INNER_PRODUCT)
index.train(xb)
index.add(xb)
D, _ = index.search(xq, 10)
# normalized vectors: max IP should be close to 1
self.assertTrue(np.all(D[:, 0] <= 1.1))
self.assertTrue(np.all(D[:, 0] >= 0.5))
class TestTurboQUnbiasedness(unittest.TestCase):
"""Verify TurboQuant reconstruction preserves inner products."""
def test_reconstruction_quality(self):
rng = np.random.RandomState(42)
print("\n=== TurboQuant reconstruction quality ===")
print(f"{'d':>6} {'qt':>8} {'self_ip':>10} {'cos_sim':>10}")
print("-" * 42)
for d in [128, 256, 384, 512, 768]:
n = 500
xb = rng.randn(n, d).astype("float32")
faiss.normalize_L2(xb)
for qt_name, qt in [
("SQtq2", faiss.ScalarQuantizer.QT_2bit_tq),
("SQtq3", faiss.ScalarQuantizer.QT_3bit_tq),
]:
index = faiss.IndexScalarQuantizer(d, qt)
index.train(xb)
index.add(xb)
xr = np.vstack([index.reconstruct(i) for i in range(n)])
# Self inner product: <x, x_hat> for each vector
self_ip = np.array([xb[i] @ xr[i] for i in range(n)])
mean_self_ip = self_ip.mean()
# Cosine similarity
xr_norms = np.linalg.norm(xr, axis=1)
cos_sim = np.array(
[xb[i] @ xr[i] / (xr_norms[i] + 1e-30) for i in range(n)]
).mean()
print(
f"{d:>6} {qt_name:>8} "
f"{mean_self_ip:>10.4f} {cos_sim:>10.4f}"
)
# Reconstruction should point in the right direction
self.assertGreater(
mean_self_ip,
0.1,
msg=f"d={d} {qt_name}: self IP {mean_self_ip:.4f} too low",
)
class TestTurboQFullFactory(unittest.TestCase):
"""Test factory string construction for TurboQ full types."""
def test_factory_sqtq(self):
cases = [
("SQtq2", faiss.ScalarQuantizer.QT_2bit_tq),
("SQtq3", faiss.ScalarQuantizer.QT_3bit_tq),
("SQtq4", faiss.ScalarQuantizer.QT_4bit_tq),
("SQtq5", faiss.ScalarQuantizer.QT_5bit_tq),
]
for factory_str, qtype in cases:
with self.subTest(factory_str=factory_str):
index = faiss.index_factory(32, factory_str)
self.assertIsInstance(index, faiss.IndexScalarQuantizer)
self.assertEqual(index.sq.qtype, qtype)
def test_factory_ivf_sqtq(self):
d = 64
for sq_str in ["SQtq2", "SQtq3", "SQtq4"]:
with self.subTest(sq_str=sq_str):
factory_str = f"IVF16,{sq_str}"
index = faiss.index_factory(d, factory_str)
self.assertIsInstance(index, faiss.IndexIVFScalarQuantizer)
ivf = faiss.downcast_index(index)
self.assertFalse(ivf.by_residual)
class TestTurboQFullEncodeDecode(unittest.TestCase):
"""Encode/decode roundtrip for TurboQ full types."""
def test_encode_decode_roundtrip(self):
for d in [32, 64, 128]:
for qtype, max_err in [
(faiss.ScalarQuantizer.QT_2bit_tq, 1.5),
(faiss.ScalarQuantizer.QT_3bit_tq, 1.2),
(faiss.ScalarQuantizer.QT_4bit_tq, 0.8),
]:
with self.subTest(d=d, qtype=qtype):
rng = np.random.RandomState(42)
xb = rng.randn(200, d).astype("float32")
faiss.normalize_L2(xb)
sq = faiss.ScalarQuantizer(d, qtype)
sq.train(xb)
codes = sq.compute_codes(xb)
xr = sq.decode(codes)
self.assertEqual(xr.shape, xb.shape)
err = np.abs(xb - xr).max()
self.assertLess(err, max_err)
def test_non_simd_aligned_dims(self):
for d in [7, 9, 15, 17, 33, 65]:
with self.subTest(d=d):
rng = np.random.RandomState(42)
xb = rng.randn(100, d).astype("float32")
faiss.normalize_L2(xb)
sq = faiss.ScalarQuantizer(d, faiss.ScalarQuantizer.QT_2bit_tq)
sq.train(xb)
codes = sq.compute_codes(xb)
xr = sq.decode(codes)
self.assertEqual(xr.shape, xb.shape)
class TestTurboQFullSearch(unittest.TestCase):
"""Search recall tests for TurboQ full types."""
def setUp(self):
self.d = 64
self.rng = np.random.RandomState(42)
self.nb = 5000
self.nq = 50
self.xb = self.rng.randn(self.nb, self.d).astype("float32")
self.xq = self.rng.randn(self.nq, self.d).astype("float32")
faiss.normalize_L2(self.xb)
faiss.normalize_L2(self.xq)
gt_index = faiss.IndexFlatL2(self.d)
gt_index.add(self.xb)
_, self.I_gt = gt_index.search(self.xq, 10)
def _recall_at_1(self, index):
_, I = index.search(self.xq, 10)
return (self.I_gt[:, 0] == I[:, 0]).mean()
def test_flat_sqtq_search(self):
for sq_str in ["SQtq2", "SQtq3", "SQtq4"]:
with self.subTest(sq_str=sq_str):
index = faiss.index_factory(self.d, sq_str)
index.train(self.xb)
index.add(self.xb)
recall = self._recall_at_1(index)
self.assertGreater(recall, 0.05)
def test_recall_increases_with_bits(self):
recalls = {}
for sq_str in ["SQtq2", "SQtq3", "SQtq4"]:
index = faiss.index_factory(self.d, sq_str)
index.train(self.xb)
index.add(self.xb)
recalls[sq_str] = self._recall_at_1(index)
self.assertGreaterEqual(recalls["SQtq3"], recalls["SQtq2"] - 0.05)
self.assertGreaterEqual(recalls["SQtq4"], recalls["SQtq3"] - 0.05)
def test_ivf_sqtq_search(self):
for sq_str in ["SQtq2", "SQtq3"]:
with self.subTest(sq_str=sq_str):
factory_str = f"IVF32,{sq_str}"
index = faiss.index_factory(self.d, factory_str)
index.train(self.xb)
index.add(self.xb)
index.nprobe = 8
_, I = index.search(self.xq, 10)
self.assertEqual(I.shape, (self.nq, 10))
self.assertTrue(np.all(I[:, 0] >= 0))
class TestTurboQFullSerialization(unittest.TestCase):
"""Serialization roundtrip for TurboQ full types."""
def _roundtrip(self, index):
fd, fname = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, fname)
index2 = faiss.read_index(fname)
return index2
finally:
os.unlink(fname)
def test_flat_serialize_roundtrip(self):
d = 64
rng = np.random.RandomState(42)
xb = rng.randn(500, d).astype("float32")
xq = rng.randn(10, d).astype("float32")
faiss.normalize_L2(xb)
faiss.normalize_L2(xq)
for qtype in [
faiss.ScalarQuantizer.QT_2bit_tq,
faiss.ScalarQuantizer.QT_3bit_tq,
faiss.ScalarQuantizer.QT_4bit_tq,
faiss.ScalarQuantizer.QT_5bit_tq,
]:
with self.subTest(qtype=qtype):
index = faiss.IndexScalarQuantizer(d, qtype)
index.train(xb)
index.add(xb)
D1, I1 = index.search(xq, 5)
index2 = self._roundtrip(index)
D2, I2 = index2.search(xq, 5)
np.testing.assert_array_equal(I1, I2)
np.testing.assert_allclose(D1, D2, atol=1e-6)
def test_ivf_serialize_roundtrip(self):
d = 64
rng = np.random.RandomState(42)
xb = rng.randn(1000, d).astype("float32")
xq = rng.randn(10, d).astype("float32")
faiss.normalize_L2(xb)
faiss.normalize_L2(xq)
for sq_str in ["SQtq2", "SQtq3"]:
with self.subTest(sq_str=sq_str):
index = faiss.index_factory(d, f"IVF16,{sq_str}")
index.train(xb)
index.add(xb)
index.nprobe = 4
D1, I1 = index.search(xq, 5)
index2 = self._roundtrip(index)
index2.nprobe = 4
D2, I2 = index2.search(xq, 5)
np.testing.assert_array_equal(I1, I2)
np.testing.assert_allclose(D1, D2, atol=1e-6)
def test_serialize_preserves_qtype(self):
d = 32
rng = np.random.RandomState(42)
xb = rng.randn(100, d).astype("float32")
faiss.normalize_L2(xb)
for qtype in [
faiss.ScalarQuantizer.QT_2bit_tq,
faiss.ScalarQuantizer.QT_3bit_tq,
]:
with self.subTest(qtype=qtype):
index = faiss.IndexScalarQuantizer(d, qtype)
index.train(xb)
index.add(xb)
index2 = self._roundtrip(index)
self.assertEqual(index2.sq.qtype, qtype)
self.assertEqual(index2.ntotal, index.ntotal)
self.assertEqual(index2.d, d)
class TestTurboQFullDistances(unittest.TestCase):
"""Verify distance computation is consistent for TurboQ."""
def test_search_distances_are_finite(self):
d = 64
rng = np.random.RandomState(42)
xb = rng.randn(500, d).astype("float32")
xq = rng.randn(10, d).astype("float32")
faiss.normalize_L2(xb)
faiss.normalize_L2(xq)
for qtype in [
faiss.ScalarQuantizer.QT_2bit_tq,
faiss.ScalarQuantizer.QT_3bit_tq,
]:
with self.subTest(qtype=qtype):
index = faiss.IndexScalarQuantizer(d, qtype)
index.train(xb)
index.add(xb)
D, I = index.search(xq, 5)
self.assertTrue(np.all(np.isfinite(D)))
self.assertTrue(np.all(D[:, 0] >= 0))
self.assertTrue(np.all(I >= 0))
def test_search_distances_sorted(self):
d = 64
rng = np.random.RandomState(42)
xb = rng.randn(500, d).astype("float32")
xq = rng.randn(10, d).astype("float32")
faiss.normalize_L2(xb)
faiss.normalize_L2(xq)
for qtype in [
faiss.ScalarQuantizer.QT_2bit_tq,
faiss.ScalarQuantizer.QT_3bit_tq,
]:
with self.subTest(qtype=qtype):
index = faiss.IndexScalarQuantizer(d, qtype)
index.train(xb)
index.add(xb)
D, _ = index.search(xq, 10)
for q in range(len(xq)):
for k in range(1, 10):
self.assertLessEqual(D[q, k - 1], D[q, k])
@for_all_simd_levels
class TestSQByteDirectAcrossSIMDLevels(unittest.TestCase):
"""Every SIMD level must return the exact QT_8bit_direct{,_signed}
distances.
The distances are integers, so this class computes them from the
definition and compares. An earlier version measured a second time at
SIMDLevel.NONE and compared the two runs. A static build reports only the
one level it compiles, and omits NONE, so that version skipped itself in
the build shape that ships.
d = 16 and d = 32 pass the d % 16 == 0 test in sq-dispatch.h. That test
routes the query to DistanceComputerByte and DistanceComputerByteSigned.
Two dispatch chains reach those kernels, and this class covers both.
search() uses sq_select_InvertedListScanner. get_distance_computer() uses
sq_select_distance_computer, which is the chain Refine(SQ8) follows.
"""
def reference(self, xq, xb, metric):
"""Returns the exact distances. The values are small integers. Every
product and sum stays below 2**24, so the float32 result loses
nothing."""
a = xq.astype("int64")
b = xb.astype("int64")
if metric == faiss.METRIC_L2:
return ((a[:, None, :] - b[None, :, :]) ** 2).sum(-1)
return a @ b.T
def do_test(self, measure, expect):
rng = np.random.RandomState(1234)
for qtype in (
faiss.ScalarQuantizer.QT_8bit_direct,
faiss.ScalarQuantizer.QT_8bit_direct_signed,
):
lo, hi = (
(-128, 128)
if qtype == faiss.ScalarQuantizer.QT_8bit_direct_signed
else (0, 256)
)
for metric in (faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT):
for d in (16, 32):
with self.subTest(qtype=qtype, metric=metric, d=d):
xb = rng.randint(lo, hi, (500, d)).astype("float32")
xq = rng.randint(lo, hi, (10, d)).astype("float32")
index = faiss.IndexScalarQuantizer(d, qtype, metric)
index.add(xb)
got = measure(index, xq)
truth = self.reference(xq, xb, metric)
np.testing.assert_array_equal(got, expect(truth, metric))
def test_scanner(self):
k = 10
def measure(index, xq):
return index.search(xq, k)[0]
def expect(truth, metric):
# search returns L2 ascending and inner product descending.
ordered = np.sort(truth, axis=1)
if metric == faiss.METRIC_INNER_PRODUCT:
ordered = ordered[:, ::-1]
return ordered[:, :k]
self.do_test(measure, expect)
def test_distance_computer(self):
def measure(index, xq):
dc = index.get_distance_computer()
out = np.empty((len(xq), index.ntotal), dtype="float32")
for q in range(len(xq)):
dc.set_query(faiss.swig_ptr(xq[q]))
for i in range(index.ntotal):
out[q, i] = dc(int(i))
return out
self.do_test(measure, lambda truth, metric: truth)
Reference in New Issue View Git Blame Copy Permalink