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:
Alibek Zhakubayev
2026-08-03 07:51:52 -07:00
committed by meta-codesync[bot]
parent 6187d940e3
commit 4332ceb994
2 changed files with 4 additions and 4 deletions
@@ -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
View File
@@ -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)