Eliminate per-code denormalization in uniform SQ distance computation (#5166)

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.** ``` <q, x_hat> = sum_i q[i] * (vmin + vdiff * n[i]) = vmin * sum_i q[i] + vdiff * sum_i q[i] * n[i] = bias + scale * <q, n> ``` 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
2026-10-12 02:13:36 +00:00 · 2026-05-29 16:52:27 -07:00
parent 34eb98936f
commit 1cb7601829
4 changed files with 242 additions and 8 deletions
@@ -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)
@@ -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 <class Codec>
@@ -399,6 +405,22 @@ struct SimilarityL2<SIMDLevel::AVX2> {
        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<SIMDLevel::AVX2> {
        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<Quantizer, Similarity, SIMDLevel::AVX2> : SQDistanceComputer {

    Quantizer quant;

+    // Pre-adjusted query buffer for uniform quantizers
+    std::vector<float> 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<float>& 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<Quantizer, Similarity, SIMDLevel::AVX2> : 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<int>(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<Quantizer, Similarity, SIMDLevel::AVX2> : 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(
@@ -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 <class Codec>
@@ -411,6 +417,22 @@ struct SimilarityL2<SIMDLevel::AVX512> {
    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<SIMDLevel::AVX512> {
    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, Similarity, SIMDLevel::AVX512>

    Quantizer quant;

+    // Pre-adjusted query buffer for uniform quantizers
+    std::vector<float> 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<float>& 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<Quantizer, Similarity, SIMDLevel::AVX512>

    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<Quantizer, Similarity, SIMDLevel::AVX512>
    }

    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(
@@ -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 <class Codec>
@@ -397,6 +403,22 @@ struct SimilarityL2<SIMDLevel::ARM_NEON> {
    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<SIMDLevel::ARM_NEON> {
    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, Similarity, SIMDLevel::ARM_NEON>

    Quantizer quant;

+    // Pre-adjusted query buffer for uniform quantizers
+    std::vector<float> 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<float>& 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<Quantizer, Similarity, SIMDLevel::ARM_NEON>

    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<Quantizer, Similarity, SIMDLevel::ARM_NEON>
    }

    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(