perf(vector): vectorize RaBitQ dist table quantization#7241
Conversation
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
| // would produce NaN/inf products on which the kernels' saturating | ||
| // narrows disagree. Either way the table carries no information at u8 | ||
| // resolution: emit a zeroed LUT that callers reconstruct from `min`. | ||
| if !range.is_finite() || !factor.is_finite() { |
There was a problem hiding this comment.
This fallback can still produce NaN in the caller. When qmax - qmin overflows to inf, we zero the LUT but return the original (qmin, qmax), and storage.rs later reconstructs with 0.0 * inf + sum_min.
Can we return an explicit fallback/dequantization state here and make the caller use exact distance when the reconstruction scale is non-finite? This needs an end-to-end distance_all_with_scratch test for both u8 and u16.
There was a problem hiding this comment.
Good catch — fixed in 4fd43e1.
quantize_dist_table_into / _u16_into now return an explicit DistTableDequant { Affine { qmin, qmax }, Exact }. When the caller's affine reconstruction would be non-finite, they return Exact and binary_distances_with_scratch / binary_distances_hacc_with_scratch compute exact per-row distances (compute_single_rq_distance) for every row instead of reconstructing 0 * inf.
The overflow predicate covers the whole reconstruction, not just the inf range you flagged: each row's reconstructed binary IP lies in [num_tables*qmin, num_tables*qmax] with a quantized term up to num_tables*(qmax-qmin), so sum_min = num_tables*qmin overflow is caught too. It's scale-independent, so the u8 and u16 paths share it.
Added an end-to-end distance_all test for both the u8 (Normal) and u16 (Accurate) paths that forces such a table and asserts the output is NaN-free and bit-identical to the always-exact per-row computation.
| unsafe fn quantize16_epi32(src: *const f32, min: __m512, factor: __m512) -> __m512i { | ||
| // SAFETY: the caller guarantees 16 floats are readable at `src`. | ||
| let v = unsafe { _mm512_loadu_ps(src) }; | ||
| _mm512_cvtps_epi32(_mm512_mul_ps(_mm512_sub_ps(v, min), factor)) |
There was a problem hiding this comment.
This relies on the current MXCSR rounding mode. The scalar path uses fixed round_ties_even(), but _mm512_cvtps_epi32 / _mm256_cvtps_epi32 follow thread-local MXCSR, so SIMD can diverge if native code changes rounding mode.
Can we make rounding explicit before conversion and add an x86 test that changes MXCSR?
There was a problem hiding this comment.
Fixed. All quantize paths now round with fixed-mode rounding, independent of MXCSR:
- AVX-512:
_mm512_cvt_roundps_epi32::<{ _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC }>(embedded static rounding). - AVX2 has no embedded-rounding convert, so it rounds with
_mm256_round_ps::<{ _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC }>(vroundps, ignores MXCSR) beforecvtps, which then sees an integral value.
Validating this on an AVX-512 host surfaced a second, related issue worth flagging: f32::round_ties_even() itself lowers to an MXCSR-honoring instruction on x86 (it truncated under round-toward-zero in the test). Since the SIMD kernels delegate their scalar tails (len % step) to the scalar quantizer, the kernel tails — and the scalar fallback — were still MXCSR-dependent even after the bodies were fixed. So the scalar path on x86 now builds nearest-even from f32::floor (always fixed-mode) instead of round_ties_even; aarch64 keeps round_ties_even (frintn, already fixed-mode and autovectorizable).
Added test_quantize_rounding_ignores_mxcsr: runs every kernel (scalar, AVX2, AVX-512) under a forced round-toward-zero MXCSR and requires bit-exact agreement with the nearest-even reference. The length (511) is deliberately not a multiple of the SIMD step so the scalar tails are exercised. Verified on a c4 (AVX-512) host.
…SR-independent SIMD rounding Addresses review feedback on #7241: - quantize_dist_table_into / _u16_into now return an explicit DistTableDequant { Affine { qmin, qmax }, Exact }. When the affine reconstruction num_tables * {qmin, qmax, qmax - qmin} would overflow f32, they signal Exact and the caller computes exact per-row distances instead of reconstructing 0 * inf = NaN. Adds an end-to-end distance_all test for both the u8 (Normal) and u16 (Accurate) paths. - The AVX-512/AVX2 quantizers now round with static round-to-nearest-even (_mm512_cvt_roundps_epi32 with embedded rounding; _mm256_round_ps before the convert) so the SIMD result no longer depends on the dynamic MXCSR rounding mode. Adds an x86 test that runs the kernels under a truncating MXCSR and requires they still match the scalar round_ties_even reference. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
f32::round_ties_even can itself lower to an MXCSR-honoring instruction on x86, so the scalar fallback is not a valid invariant under a forced rounding mode. The test now covers only the AVX2/AVX-512 kernels, whose static rounding is the property under test. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The AVX2/AVX-512 kernels delegate their scalar tails (len % step) to quantize_u8_scalar/_u16_scalar, which rounded with f32::round_ties_even. On x86 that can lower to an MXCSR-honoring instruction, so the kernel tails (and the scalar fallback) drifted from nearest-even under a non-default rounding mode while the SIMD bodies did not. Round via f32::floor (a fixed-mode op on every target) instead, so the scalar path, the SIMD tails, and the SIMD bodies all round half-to-even independent of MXCSR. The MXCSR test now covers all kernels (scalar included) and uses a length whose tail is non-empty. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Only x86 needs the floor-based fixed-mode rounding (round_ties_even can lower to an MXCSR-honoring instruction there). On aarch64 round_ties_even is frintn, already fixed-mode and autovectorizable, so the floor-based form would only cost vectorization with no benefit. Gate it to x86_64. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2 participants
What
quantize_dist_table_into/quantize_dist_table_u16_intoquantize the per-(query, partition)dim * 4-entry f32 FastScan distance table into u8 (fast/normal approx modes) or u16 (accuratemode) LUT entries. Both were scalar:
itertools::minmax_by(total_cmp)for the min/max pass(branchy pairwise compares that never vectorize;
minmax_implalone is 6.0-6.6% of query time incpu-clock profiles of IVF_RQ on dbpedia-openai-1M, dim=1536, num_bits=5, nprobes=24) plus a
scalar quantize-write loop.
This moves them into a dedicated module
vector/bq/dist_table_quant.rswith the sameruntime-dispatch treatment as the ex-dot kernels (#7205):
fold that LLVM auto-vectorizes (NEON is part of the aarch64 baseline).
(d - qmin) * factor,cvtps_epi32(nearest-even, MXCSR default), thenunsigned-saturating narrows (
vpmovusdb/vpmovusdwon AVX-512;packus+ lane-restorepermutes on AVX2); scalar fallback uses
round_ties_even.Numeric semantics
f32::round(half away from zero) to half-to-even so the scalarfallback and all SIMD paths are bit-exact with each other. Relative to the old code this can
move a LUT entry by 1 only on exact .5 ties, within the table's inherent quantization error.
total_cmpdiffers only on NaN (inputs are finite sums of rotated-querycomponents) and the sign of zero, which callers cannot observe (
d - qminand theqmin == qmaxearly-out are unchanged either way).255/f32::MAX, which overflowsfactorto inf, orabove
f32::MAX) now collapse to the zeroed-LUT early-out. Previously these produced garbage(NaN -> 0 casts); the SIMD narrows would additionally disagree across kernels on the NaN
products, so the early-out keeps every path deterministic and bit-exact.
Benchmarks
GCP c4-standard-16 (AVX-512),
taskset -c 4, criterion baseline = main (b6a99cd),cargo bench -p lance-index --bench rq:The untouched
RQ bulk ex kernel loopcontrol group moved within +-3% between runs, so thebinary-only wins are far above the machine noise floor. Full-path results for num_bits=5/9 are
within that noise envelope (the binary stage is a small share of those).
On Apple M-series (NEON: portable min/max fold + autovectorized scalar quantize), the same
paired criterion run improves binary-only distance_all by -24% / -13% / -8% (num_bits=3/5/9);
the Mac's run-to-run noise is larger than the pinned GCP runs but all three move well past it.
Tests
against a straightforward
total_cmp+round_ties_evenreference and require bit-exactagreement: random inputs (lengths 1..6160 including non-multiples of 16/32, scales
1e-3..1e4), exact .5 ties (integer tables constructed so
factor == 0.5), all-equal inputs,signed-zero mixes, degenerate ranges, and scratch-buffer reuse.
cargo test -p lance-index --lib vector::bq: 207 pass on aarch64 (scalar + NEON dispatch)and on AVX-512/AVX2 hardware (GCP c4, both debug and release).
cargo fmt --allandcargo clippy --all --tests --benches -- -D warningsare clean.🤖 Generated with Claude Code