From 1cb760182917697677fe100f8a1cc2b4d0b9f751 Mon Sep 17 00:00:00 2001 From: Mulugeta Mammo Date: Fri, 29 May 2026 16:52:27 -0700 Subject: [PATCH] Eliminate per-code denormalization in uniform SQ distance computation (#5166) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: This PR removes per-code denormalization from the inner loop of L2 and inner product distance computations for uniform scalar quantizers (`QT_8bit_uniform` and `QT_4bit_uniform`), yielding a speedup of up to **1.39x**. For uniform integer scalar quantizers, `vmin` and `vdiff` are scalars shared across all dimensions. The reconstructed value for each component is therefore a function of a per-code decode `n` that depends only on the byte code: ``` x_hat = vmin + vdiff * n ``` This structure lets us factor `vmin` and `vdiff` out of the per-database-vector inner loop entirely, instead of recomputing the transform on every code on every distance evaluation. **L2 distance.** ``` ||q - x_hat||^2 = ||q - (vmin + vdiff * n)||^2 = vdiff^2 * ||(q - vmin) / vdiff - n||^2 ``` We pre-adjust the query once in `set_query()` to `q_adj = (q - vmin) / vdiff` and precompute `scale = vdiff^2`. The hot loop then compares codes directly against `q_adj` in the codec's native decode space, applying `scale` exactly once at the end. **Inner product.** ``` = sum_i q[i] * (vmin + vdiff * n[i]) = vmin * sum_i q[i] + vdiff * sum_i q[i] * n[i] = bias + scale * ``` We compute `bias = vmin * sum_i q[i]` and `scale = vdiff` once in `set_query()`. The hot loop accumulates the dot product against the raw decode `n` and applies `bias + scale` once at the end. In both cases this removes one FMA (the `vmin + vdiff * n` denormalization) from the inner loop per 8 components per database vector. Beyond the raw FLOP reduction, shortening the dependency chain lets SIMD pipelines better overlap the decode of the next lane with the accumulator update of the previous lane. **The Change** ``` Current This PR -------- --------- set_query(q): store q q_adj = (q - vmin) / vdiff [once] per code (×N): raw = decode(bytes) raw = decode(bytes) x = vmin + raw * vdiff ← gone diff = q_adj - raw diff = q - x accu += diff^2 accu += diff^2 ``` The optimization is gated on a C++20 `requires` check for a new `decode_8_raw()` method, defined only on the uniform `QuantizerTemplate` specializations. All other quantizer types fall through to the original `compute_distance` path unchanged. **Speedup** Below are the results from running `benchs/bench_scalar_quantizer.py` for a `dd` build on SPR, compared to the existing implementation. Similar results were observed for `avx-2` as well. ``` | | QT_4bit_uniform | QT_8bit_uniform | |--------------|-----------------|-----------------| | RS_minmax | 0.99x | 1.05x | | RS_minmax | 1.07x | 1.03x | | RS_minmax | 0.83x | 1.03x | | RS_minmax | 0.96x | 1.05x | | RS_minmax | 0.89x | 1.03x | | RS_minmax | 1.14x | 1.03x | | RS_minmax | 0.99x | 1.05x | | RS_meanstd | 1.28x | 1.09x | | RS_meanstd | 1.10x | 1.01x | | RS_meanstd | 1.14x | 1.06x | | RS_meanstd | 1.18x | 1.08x | | RS_meanstd | 1.11x | 1.05x | | RS_meanstd | 1.18x | 1.06x | | RS_meanstd | 1.16x | 1.08x | | RS_quantiles | 1.39x | 1.07x | | RS_quantiles | 1.08x | 1.01x | | RS_quantiles | 1.21x | 0.99x | | RS_quantiles | 1.25x | 1.10x | | RS_optim | 1.03x | 1.03x | ``` The raw performance results are available here: https://gist.github.com/mulugetam/7db50f89279bb270a1fe336206730d60 Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5166 Test Plan: Ran `python benchs/bench_scalar_quantizer.py` on SPR to validate performance results (see summary table). Ran `pytest tests/test_index_accuracy.py` to verify correctness is preserved for uniform SQ types (`QT_8bit_uniform`, `QT_4bit_uniform`). Reviewed By: mdouze Differential Revision: D106148760 Pulled By: mnorris11 fbshipit-source-id: 65769cfe9cf8d2d1ef5de05b6bdecf9dbee96d5c --- benchs/bench_scalar_quantizer.py | 4 +- faiss/impl/scalar_quantizer/sq-avx2.cpp | 82 ++++++++++++++++++++++- faiss/impl/scalar_quantizer/sq-avx512.cpp | 82 ++++++++++++++++++++++- faiss/impl/scalar_quantizer/sq-neon.cpp | 82 ++++++++++++++++++++++- 4 files changed, 242 insertions(+), 8 deletions(-) diff --git a/benchs/bench_scalar_quantizer.py b/benchs/bench_scalar_quantizer.py index d6830ab9a..e4a4e6a7e 100644 --- a/benchs/bench_scalar_quantizer.py +++ b/benchs/bench_scalar_quantizer.py @@ -17,7 +17,7 @@ ncent = 256 variants = [(name, getattr(faiss.ScalarQuantizer, name)) for name in dir(faiss.ScalarQuantizer) - if name.startswith('QT_')] + if name.startswith('QT_') and name != 'QT_count'] quantizer = faiss.IndexFlatL2(d) # quantizer.add(np.zeros((1, d), dtype='float32')) @@ -67,7 +67,7 @@ if True: index.sq.rangestat = getattr(faiss.ScalarQuantizer, rsname) - index.rangestat_arg = val + index.sq.rangestat_arg = val index.train(xt) index.add(xb) diff --git a/faiss/impl/scalar_quantizer/sq-avx2.cpp b/faiss/impl/scalar_quantizer/sq-avx2.cpp index d3d7cec6b..01dc262f8 100644 --- a/faiss/impl/scalar_quantizer/sq-avx2.cpp +++ b/faiss/impl/scalar_quantizer/sq-avx2.cpp @@ -196,6 +196,12 @@ struct QuantizerTemplate< return simd8float32(_mm256_fmadd_ps( xi, _mm256_set1_ps(this->vdiff), _mm256_set1_ps(this->vmin))); } + + /// Raw codec decode without denormalization + FAISS_ALWAYS_INLINE simd8float32 + decode_8_raw(const uint8_t* code, int i) const { + return Codec::decode_8_components(code, i); + } }; template @@ -399,6 +405,22 @@ struct SimilarityL2 { const __m128 v3 = _mm_add_ps(v1, v2); return _mm_cvtss_f32(v3); } + + static void adjust_query_for_raw_decode( + const float* x, + float* q_adj, + size_t d, + float vmin, + float vdiff, + float& scale_factor, + float& bias) { + float inv_vdiff = (vdiff != 0) ? 1.0f / vdiff : 0.0f; + for (size_t i = 0; i < d; i++) { + q_adj[i] = (x[i] - vmin) * inv_vdiff; + } + scale_factor = vdiff * vdiff; + bias = 0; + } }; template <> @@ -442,6 +464,23 @@ struct SimilarityIP { const __m128 v3 = _mm_add_ps(v1, v2); return _mm_cvtss_f32(v3); } + + static void adjust_query_for_raw_decode( + const float* x, + float* q_adj, + size_t d, + float vmin, + float vdiff, + float& scale_factor, + float& bias) { + float sum_q = 0; + for (size_t i = 0; i < d; i++) { + q_adj[i] = x[i]; + sum_q += x[i]; + } + scale_factor = vdiff; + bias = vmin * sum_q; + } }; /********************************************************** @@ -454,8 +493,23 @@ struct DCTemplate : SQDistanceComputer { Quantizer quant; + // Pre-adjusted query buffer for uniform quantizers + std::vector q_adj; + float scale_factor = 0; + float bias = 0; + + static constexpr bool has_decode_raw() { + return requires(const Quantizer& q, const uint8_t* c, int i) { + { q.decode_8_raw(c, i) }; + }; + } + DCTemplate(size_t d, const std::vector& trained) - : quant(d, trained) {} + : quant(d, trained) { + if constexpr (has_decode_raw()) { + q_adj.resize(d); + } + } float compute_distance(const float* x, const uint8_t* code) const { Similarity sim(x); @@ -484,6 +538,26 @@ struct DCTemplate : SQDistanceComputer { void set_query(const float* x) final { q = x; + if constexpr (has_decode_raw()) { + Sim::adjust_query_for_raw_decode( + x, + q_adj.data(), + quant.d, + quant.vmin, + quant.vdiff, + scale_factor, + bias); + } + } + + float query_to_code_predecoded(const uint8_t* code) const { + Similarity sim(q_adj.data()); + sim.begin_8(); + for (size_t i = 0; i < quant.d; i += 8) { + simd8float32 xi = quant.decode_8_raw(code, static_cast(i)); + sim.add_8_components(xi); + } + return bias + scale_factor * sim.result_8(); } float symmetric_dis(idx_t i, idx_t j) override { @@ -492,7 +566,11 @@ struct DCTemplate : SQDistanceComputer { } float query_to_code(const uint8_t* code) const final { - return compute_distance(q, code); + if constexpr (has_decode_raw()) { + return query_to_code_predecoded(code); + } else { + return compute_distance(q, code); + } } void query_to_codes_batch_4( diff --git a/faiss/impl/scalar_quantizer/sq-avx512.cpp b/faiss/impl/scalar_quantizer/sq-avx512.cpp index 3459ae4b3..5550fc8d4 100644 --- a/faiss/impl/scalar_quantizer/sq-avx512.cpp +++ b/faiss/impl/scalar_quantizer/sq-avx512.cpp @@ -214,6 +214,12 @@ struct QuantizerTemplate< return simd16float32(_mm512_fmadd_ps( xi, _mm512_set1_ps(this->vdiff), _mm512_set1_ps(this->vmin))); } + + /// Raw codec decode without denormalization + FAISS_ALWAYS_INLINE simd16float32 + decode_16_raw(const uint8_t* code, int i) const { + return Codec::decode_16_components(code, i); + } }; template @@ -411,6 +417,22 @@ struct SimilarityL2 { FAISS_ALWAYS_INLINE float result_16() { return horizontal_add(accu16); } + + static void adjust_query_for_raw_decode( + const float* x, + float* q_adj, + size_t d, + float vmin, + float vdiff, + float& scale_factor, + float& bias) { + float inv_vdiff = (vdiff != 0) ? 1.0f / vdiff : 0.0f; + for (size_t i = 0; i < d; i++) { + q_adj[i] = (x[i] - vmin) * inv_vdiff; + } + scale_factor = vdiff * vdiff; + bias = 0; + } }; template <> @@ -445,6 +467,23 @@ struct SimilarityIP { FAISS_ALWAYS_INLINE float result_16() { return horizontal_add(accu16); } + + static void adjust_query_for_raw_decode( + const float* x, + float* q_adj, + size_t d, + float vmin, + float vdiff, + float& scale_factor, + float& bias) { + float sum_q = 0; + for (size_t i = 0; i < d; i++) { + q_adj[i] = x[i]; + sum_q += x[i]; + } + scale_factor = vdiff; + bias = vmin * sum_q; + } }; /********************************************************** @@ -458,8 +497,23 @@ struct DCTemplate Quantizer quant; + // Pre-adjusted query buffer for uniform quantizers + std::vector q_adj; + float scale_factor = 0; + float bias = 0; + + static constexpr bool has_decode_raw() { + return requires(const Quantizer& q, const uint8_t* c, int i) { + { q.decode_16_raw(c, i) }; + }; + } + DCTemplate(size_t d, const std::vector& trained) - : quant(d, trained) {} + : quant(d, trained) { + if constexpr (has_decode_raw()) { + q_adj.resize(d); + } + } float compute_distance(const float* x, const uint8_t* code) const { Similarity sim(x); @@ -485,6 +539,26 @@ struct DCTemplate void set_query(const float* x) final { q = x; + if constexpr (has_decode_raw()) { + Sim::adjust_query_for_raw_decode( + x, + q_adj.data(), + quant.d, + quant.vmin, + quant.vdiff, + scale_factor, + bias); + } + } + + float query_to_code_predecoded(const uint8_t* code) const { + Similarity sim(q_adj.data()); + sim.begin_16(); + for (size_t i = 0; i < quant.d; i += 16) { + simd16float32 xi = quant.decode_16_raw(code, i); + sim.add_16_components(xi); + } + return bias + scale_factor * sim.result_16(); } float symmetric_dis(idx_t i, idx_t j) override { @@ -493,7 +567,11 @@ struct DCTemplate } float query_to_code(const uint8_t* code) const final { - return compute_distance(q, code); + if constexpr (has_decode_raw()) { + return query_to_code_predecoded(code); + } else { + return compute_distance(q, code); + } } void query_to_codes_batch_4( diff --git a/faiss/impl/scalar_quantizer/sq-neon.cpp b/faiss/impl/scalar_quantizer/sq-neon.cpp index 7d07f6c4e..543adf846 100644 --- a/faiss/impl/scalar_quantizer/sq-neon.cpp +++ b/faiss/impl/scalar_quantizer/sq-neon.cpp @@ -180,6 +180,12 @@ struct QuantizerTemplate< xi.data.val[1], this->vdiff)}); } + + /// Raw codec decode without denormalization (for pre-decode opt) + FAISS_ALWAYS_INLINE simd8float32 + decode_8_raw(const uint8_t* code, int i) const { + return Codec::decode_8_components(code, i); + } }; template @@ -397,6 +403,22 @@ struct SimilarityL2 { FAISS_ALWAYS_INLINE float result_8() { return horizontal_add(accu8); } + + static void adjust_query_for_raw_decode( + const float* x, + float* q_adj, + size_t d, + float vmin, + float vdiff, + float& scale_factor, + float& bias) { + float inv_vdiff = (vdiff != 0) ? 1.0f / vdiff : 0.0f; + for (size_t i = 0; i < d; i++) { + q_adj[i] = (x[i] - vmin) * inv_vdiff; + } + scale_factor = vdiff * vdiff; + bias = 0; + } }; template <> @@ -431,6 +453,23 @@ struct SimilarityIP { FAISS_ALWAYS_INLINE float result_8() { return horizontal_add(accu8); } + + static void adjust_query_for_raw_decode( + const float* x, + float* q_adj, + size_t d, + float vmin, + float vdiff, + float& scale_factor, + float& bias) { + float sum_q = 0; + for (size_t i = 0; i < d; i++) { + q_adj[i] = x[i]; + sum_q += x[i]; + } + scale_factor = vdiff; + bias = vmin * sum_q; + } }; /********************************************************** @@ -444,8 +483,23 @@ struct DCTemplate Quantizer quant; + // Pre-adjusted query buffer for uniform quantizers + std::vector q_adj; + float scale_factor = 0; + float bias = 0; + + static constexpr bool has_decode_raw() { + return requires(const Quantizer& q, const uint8_t* c, int i) { + { q.decode_8_raw(c, i) }; + }; + } + DCTemplate(size_t d, const std::vector& trained) - : quant(d, trained) {} + : quant(d, trained) { + if constexpr (has_decode_raw()) { + q_adj.resize(d); + } + } float compute_distance(const float* x, const uint8_t* code) const { Similarity sim(x); @@ -471,6 +525,26 @@ struct DCTemplate void set_query(const float* x) final { q = x; + if constexpr (has_decode_raw()) { + Sim::adjust_query_for_raw_decode( + x, + q_adj.data(), + quant.d, + quant.vmin, + quant.vdiff, + scale_factor, + bias); + } + } + + float query_to_code_predecoded(const uint8_t* code) const { + Similarity sim(q_adj.data()); + sim.begin_8(); + for (size_t i = 0; i < quant.d; i += 8) { + simd8float32 xi = quant.decode_8_raw(code, i); + sim.add_8_components(xi); + } + return bias + scale_factor * sim.result_8(); } float symmetric_dis(idx_t i, idx_t j) override { @@ -479,7 +553,11 @@ struct DCTemplate } float query_to_code(const uint8_t* code) const final { - return compute_distance(q, code); + if constexpr (has_decode_raw()) { + return query_to_code_predecoded(code); + } else { + return compute_distance(q, code); + } } void query_to_codes_batch_4(