Skip to content

RaBitQ: add batch-four estimates and 9-bit SIMD scoring - #5647

Open
CLiqing wants to merge 2 commits into
facebookresearch:mainfrom
CLiqing:feature/rabitq-batch-estimates
Open

CLiqing wants to merge 2 commits into
facebookresearch:mainfrom
CLiqing:feature/rabitq-batch-estimates

Conversation

@CLiqing

@CLiqing CLiqing commented Sep 16, 2026

Copy link
Copy Markdown

Summary

Extend RaBitQ's distance-computation API with four independent one-bit estimates and add optimized kernels for standard RaBitQ codes.

  • Add RaBitQDistanceComputer::distance_to_code_1bit_batch_4, with a default implementation that delegates to four single-code evaluations.
  • For non-centered, four-bit query quantization, reuse query bit planes across four database codes. The optimized kernel uses the existing AVX512_VPOPCNT dispatch level; other SIMD levels and query modes retain fallbacks.
  • Add AVX2 and AVX-512 full-distance kernels for 9-bit database codes, exploiting byte-aligned extra codes rather than scalar extraction. The scalar tail reads only the required extra bytes.
  • Isolate BMI2 bit-plane helpers behind function target attributes and runtime feature checks on GNU-compatible x86-64 builds, while retaining the existing compile-time path for other supported toolchains.
  • Use local fixed-size memcpy for byte-aligned scalar bit-plane loads.

API contract and scope

Call set_query() before evaluating a batch. The four row pointers may be non-contiguous or repeated. Results are returned in input order. The method computes estimates only: it does not refine candidates, update thresholds, or increment search counters.

Existing query modes and database precisions remain supported. No changes to HNSW traversal, filtering, graph construction, defaults, or serialized code layout are included. This is an opt-in batch API; existing search call sites are unchanged.

Tests and benchmark

New standalone tests cover batch/single agreement for database bits 1–9, query bits 0–8, both centering modes, L2/IP, available SIMD overrides, repeated/non-contiguous row pointers, query reuse, and unchanged statistics. Independent references check packed multi-bit values and unaligned/tail buffers, including exactly sized RBQ9 extra codes. A round-trip test checks stored codes and distances.

bench_rabitq_batch_estimates compares four single-code estimates with the batch API at dimensions 128, 768, 960, and 1536. It prepares queries outside timing, shuffles code pointers, warms both paths, alternates measurement order, and prints per-code times and checksums. It is a scorer microbenchmark, not an end-to-end search benchmark.

Local validation (GCC 12, Release dynamic-dispatch build, Intel Xeon Platinum 8375C):

  • Built faiss_test and bench_rabitq_batch_estimates.
  • 61 tests passed with --gtest_filter='TestLowLevelIVF.*:*RaBitQ*:*HNSW*', including all five new tests.
  • 47 individual-process tests passed with ctest -R 'RaBitQ|HNSW'.
  • All five new tests passed with the changed RaBitQ translation units and test file instrumented by ASan/UBSan, including leak detection. Unchanged dependencies were linked from the Release build; this was not a fully instrumented library build.
  • Checked NONE, AVX2, AVX512 and AVX512_VPOPCNT on the local host. Clang 15 syntax checks passed for the three affected SIMD units, with baseline BMI2 disabled for the AVX2/AVX512 checks. ARM, non-BMI2 hardware, and native SPR execution were not performed locally.
  • No on-disk format changes; code and distance round-trip checks passed.

Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant