mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Enable byte-LUT gather path for QT_8bit_eden distance computation (#5498)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5498 `EDENQuantizer::supports_byte_lut` excludes `nb_bits == 8`, forcing the highest-precision EDEN scalar-quantizer variant onto a fully scalar, per-dimension bit-unpacking distance path via `eden_utils::get_distance_computer` (`EDENOptimizedDistanceComputer`), the computer used by `IndexEDEN` and `IndexIVFEDEN`. ## Root cause `build_code_dot_lut` gates on `supports_byte_lut(nb_bits)` before constructing a per-query LUT. For `nb_bits == 8`, the check returned false, so `lut_kind` stayed `CodeDotLUTKind::None` and every call to `distance_to_code` fell back to `compute_code_dot_reference`, a per-dimension loop through `BitstringReader::read` (bit-shift, mask, buffer-position bookkeeping per dimension). `QT_1bit_eden`, `QT_2bit_eden`, and `QT_4bit_eden` already have this LUT gated through the same predicate. The generic "Byte" LUT branch in `build_code_dot_lut` (lines ~191-210) was already correct for `nb_bits == 8` (`values_per_byte == 1`, one 256-entry row per dimension) but unreachable because of the early return. ## Fix Add `nb_bits == 8` to `supports_byte_lut`'s predicate. This is the only change required to `EDENQuantizer.cpp`: no new code paths, no signature changes. `distance_to_code` now uses `compute_code_dot_lut` (direct `table[code[i]]` array reads, partially unrolled) in place of the bit-unpacking reference loop. Note this fix activates the scalar table-lookup path reachable via `distance_to_code` (used by both `IndexEDEN`'s brute-force scan and `IndexIVFEDEN`'s per-list scan). It does not activate the AVX2/AVX512 *batch*-gather kernels (`compute_code_dot_lut_batch_8_avx2`, `compute_code_dot_lut_batch_8_avx512`/`_batch_16_avx512`) for `IndexEDEN::search` specifically: `IndexEDEN.cpp`'s `use_eden_batch_scan` gate was not updated by this diff and still returns `false` for `nb_bits == 8`, so those kernels remain unreachable for flat EDEN search. That gap is logged as a separate follow-up (see Notes) rather than bundled into this diff, since it touches a different file/function and needs its own benchmark. ## Benchmark `IndexEDEN` built via `index_factory(d, "EDEN8")`, `SyntheticDataset`, k=10, best-of-5 search latency (`mode/opt`): | d | threads | before | after | speedup | |---|---|---|---|---| | 128 | 1 | 27,194.11 ms (36.8 QPS) | 8,629.36 ms (115.9 QPS) | 3.15x | | 128 | 56 | 1,197.53 ms (835.1 QPS) | 392.58 ms (2,547.3 QPS) | 3.05x | | 768 | 1 | 161,997.57 ms (6.2 QPS) | 64,010.54 ms (15.6 QPS) | 2.53x | | 768 | 56 | 5,250.05 ms (190.5 QPS) | 2,886.25 ms (346.5 QPS) | 1.82x | nb=100,000, nq=1,000 in all four configurations. The measured speedup comes from replacing the scalar bit-unpacking loop with a scalar table lookup (both single-threaded per-code cost, no new SIMD instructions on this reachable path) -- it exceeds the 10% impact-gate threshold by a wide margin regardless of mechanism. ## Tests Added `bits=8` to three existing parity/regression loops in `tests/test_eden.py` (`test_distance_computer_matches_reference_l2`, `test_high_dimensional_search_matches_distance_computer`, `test_search_matches_none_simd_level`) -- these previously covered only `[1, 2, 4]` and never exercised the code path this diff activates (`test_scalar_quantizer.cpp`'s existing `QT_8bit_eden` coverage exercises a separate implementation, `ScalarQuantizer`'s generic `QuantizerLloydMax<8, SL>` dispatch in `quantizers.h`, not `EDENQuantizer.cpp`'s `eden_utils::get_distance_computer`). The new assertions directly validate the new LUT math against a NumPy reference and against `SIMDLevel::NONE` output; all pass, confirming bit-exact parity with the pre-fix scalar reference computation. `test_scalar_quantizer` (16 tests), `test_eden` (8 tests, includes the 3 newly-added `bits=8` assertions) pass. Full Monday quantization rotation suite (`test_pq_encoding`, `test_scalar_quantizer_correctness`, `test_residual_quantizer`, `test_local_search_quantizer`, `test_standalone_codec`, `test_product_quantizer`, `test_rabitq`, `test_rabitq_fastscan`) passes: 245 pass, 0 fail (2 skips are a pre-existing devserver limitation, `SIMDLevel.NONE` unavailable, unrelated to this change). `buck2 build fbcode//faiss:faiss` clean, 0 warnings. ## Notes for reviewer Two related follow-ups identified during self-review, logged to the FAISS improver findings backlog rather than bundled into this diff: 1. `IndexEDEN.cpp`'s `use_eden_batch_scan` doesn't include `nb_bits == 8`, so the AVX2/AVX512 batch-gather kernels remain unreachable for flat EDEN search even after this fix. Mechanical follow-up, needs its own benchmark to quantify incremental gain. 2. `EDENOptimizedDistanceComputer::set_query` (and hence `IVFEDENDistanceComputer::operator()`, reachable via the generic `Index::get_distance_computer()` API used by `IndexRefine`, HNSW-on-flat-storage, NNDescent, NSG) rebuilds a full LUT on every single-code query. This is a pre-existing pattern already present for `nb_bits` 4 at `d < 256` (unbounded byte-LUT path, same as this diff's default for `nb_bits == 8`), not introduced by this diff, but worth a dedicated look given `nb_bits == 8` has no half-byte-style size cap the way 1/2/4-bit do. ___ Differential Revision: D114590565 fbshipit-source-id: eacae7855324c5f593b3de7ea8df6c84c171080e
This commit is contained in:
committed by
meta-codesync[bot]
parent
6187d940e3
commit
4332ceb994
@@ -59,7 +59,7 @@ float compute_code_dot_reference(
|
||||
}
|
||||
|
||||
bool supports_byte_lut(size_t nb_bits) {
|
||||
return nb_bits == 1 || nb_bits == 2 || nb_bits == 4;
|
||||
return nb_bits == 1 || nb_bits == 2 || nb_bits == 4 || nb_bits == 8;
|
||||
}
|
||||
|
||||
bool use_half_byte_lut(size_t nb_bits, size_t num_bytes) {
|
||||
|
||||
+3
-3
@@ -156,7 +156,7 @@ class TestEDENScalarQuantizer(unittest.TestCase):
|
||||
xb = rs.randn(40, 32).astype("float32")
|
||||
xq = rs.randn(1, 32).astype("float32")
|
||||
|
||||
for bits in [1, 2, 4]:
|
||||
for bits in [1, 2, 4, 8]:
|
||||
with self.subTest(bits=bits):
|
||||
index = faiss.IndexEDEN(32, faiss.METRIC_L2, bits)
|
||||
index.train(xt)
|
||||
@@ -227,7 +227,7 @@ class TestEDENScalarQuantizer(unittest.TestCase):
|
||||
rs = np.random.RandomState(2468)
|
||||
previous = faiss.SIMDConfig.get_level()
|
||||
try:
|
||||
for bits, d in [(1, 256), (2, 64), (4, 256)]:
|
||||
for bits, d in [(1, 256), (2, 64), (4, 256), (8, 128)]:
|
||||
xt = rs.randn(400, d).astype("float32")
|
||||
xb = rs.randn(128, d).astype("float32")
|
||||
xq = rs.randn(5, d).astype("float32")
|
||||
@@ -256,7 +256,7 @@ class TestEDENScalarQuantizer(unittest.TestCase):
|
||||
xb = rs.randn(nb, d).astype("float32")
|
||||
xq = rs.randn(nq, d).astype("float32")
|
||||
|
||||
for bits in [1, 2, 4]:
|
||||
for bits in [1, 2, 4, 8]:
|
||||
with self.subTest(bits=bits):
|
||||
index = faiss.IndexEDEN(d, faiss.METRIC_L2, bits)
|
||||
index.train(xt)
|
||||
|
||||
Reference in New Issue
Block a user