faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-impl.h

@@ -7,6 +7,12 @@
 
 #pragma once
 
+#include <faiss/impl/scalar_quantizer/codecs.h>
+#include <faiss/impl/scalar_quantizer/distance_computers.h>
+#include <faiss/impl/scalar_quantizer/quantizers.h>
+#include <faiss/impl/scalar_quantizer/similarities.h>
+#include <faiss/impl/scalar_quantizer/sq8_coefficients.h>
+#include <faiss/impl/simdlib/simdlib_avx512.h>
 #include <cstring>
 
 namespace faiss {
@@ -14,11 +20,6 @@ namespace faiss {
 namespace scalar_quantizer {
 
 using simd16float32 = faiss::simd16float32_tpl<SIMDLevel::AVX512>;
-using simd512bit = faiss::simd512bit_tpl<SIMDLevel::AVX512>;
-
-/**********************************************************
- * TurboQuant bit-unpacking helpers
- **********************************************************/
 
 namespace {
 
@@ -100,7 +101,7 @@ FAISS_ALWAYS_INLINE __m512i unpack_16x4bit_to_u32(const uint8_t* code, int i) {
  **********************************************************/
 
 template <>
-struct Codec8bit<SIMDLevel::AVX512> : Codec8bit<SIMDLevel::NONE> {
+struct Codec8bit<THE_LEVEL_TO_DISPATCH> : Codec8bit<SIMDLevel::NONE> {
     static FAISS_ALWAYS_INLINE simd16float32
     decode_16_components(const uint8_t* code, size_t i) {
         const __m128i c16 = _mm_loadu_si128((__m128i*)(code + i));
@@ -113,7 +114,7 @@ struct Codec8bit<SIMDLevel::AVX512> : Codec8bit<SIMDLevel::NONE> {
 };
 
 template <>
-struct Codec4bit<SIMDLevel::AVX512> : Codec4bit<SIMDLevel::NONE> {
+struct Codec4bit<THE_LEVEL_TO_DISPATCH> : Codec4bit<SIMDLevel::NONE> {
     static FAISS_ALWAYS_INLINE simd16float32
     decode_16_components(const uint8_t* code, size_t i) {
         uint64_t c8 = *(uint64_t*)(code + (i >> 1));
@@ -135,7 +136,7 @@ struct Codec4bit<SIMDLevel::AVX512> : Codec4bit<SIMDLevel::NONE> {
 };
 
 template <>
-struct Codec6bit<SIMDLevel::AVX512> : Codec6bit<SIMDLevel::NONE> {
+struct Codec6bit<THE_LEVEL_TO_DISPATCH> : Codec6bit<SIMDLevel::NONE> {
     static FAISS_ALWAYS_INLINE simd16float32
     decode_16_components(const uint8_t* code, size_t i) {
         // pure AVX512 implementation (not necessarily the fastest).
@@ -192,7 +193,7 @@ template <class Codec>
 struct QuantizerTemplate<
         Codec,
         scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
-        SIMDLevel::AVX512>
+        THE_LEVEL_TO_DISPATCH>
         : QuantizerTemplate<
                   Codec,
                   scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
@@ -211,13 +212,19 @@ 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>
 struct QuantizerTemplate<
         Codec,
         scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
-        SIMDLevel::AVX512>
+        THE_LEVEL_TO_DISPATCH>
         : QuantizerTemplate<
                   Codec,
                   scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
@@ -244,36 +251,83 @@ struct QuantizerTemplate<
  * Lloyd-Max scalar quantizer
  **********************************************************/
 
-#define DEFINE_LLOYD_MAX_AVX512_SPECIALIZATION(NBITS, INDEX_EXPR)      \
-    template <>                                                        \
-    struct QuantizerLloydMax<NBITS, SIMDLevel::AVX512>                 \
-            : QuantizerLloydMax<NBITS, SIMDLevel::NONE> {              \
-        using Base = QuantizerLloydMax<NBITS, SIMDLevel::NONE>;        \
-                                                                       \
-        QuantizerLloydMax(size_t d, const std::vector<float>& trained) \
-                : Base(d, trained) {                                   \
-            assert(d % 16 == 0);                                       \
-        }                                                              \
-                                                                       \
-        FAISS_ALWAYS_INLINE simd16float32                              \
-        reconstruct_16_components(const uint8_t* code, int i) const {  \
-            const __m512i indices = (INDEX_EXPR);                      \
-            return simd16float32(_mm512_i32gather_ps(                  \
-                    indices, this->centroids, sizeof(float)));         \
-        }                                                              \
-    }
-
-DEFINE_LLOYD_MAX_AVX512_SPECIALIZATION(1, unpack_16x1bit_to_u32(code, i));
-DEFINE_LLOYD_MAX_AVX512_SPECIALIZATION(2, unpack_16x2bit_to_u32(code, i));
-DEFINE_LLOYD_MAX_AVX512_SPECIALIZATION(3, unpack_16x3bit_to_u32(code, i));
-DEFINE_LLOYD_MAX_AVX512_SPECIALIZATION(4, unpack_16x4bit_to_u32(code, i));
-
-#undef DEFINE_LLOYD_MAX_AVX512_SPECIALIZATION
+// 1-bit Lloyd-Max AVX512: 16 comparisons → 2 bytes via mask compare.
+// NLT_UQ is !(x < b), matching the scalar upper_bound tie-break.
+template <>
+struct QuantizerLloydMax<1, THE_LEVEL_TO_DISPATCH>
+        : QuantizerLloydMaxBase<1, THE_LEVEL_TO_DISPATCH> {
+    using Base = QuantizerLloydMaxBase<1, THE_LEVEL_TO_DISPATCH>;
+
+    QuantizerLloydMax(size_t d, const std::vector<float>& trained)
+            : Base(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        return simd16float32(_mm512_i32gather_ps(
+                unpack_16x1bit_to_u32(code, i),
+                this->centroids,
+                sizeof(float)));
+    }
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        __m512 boundary = _mm512_set1_ps(this->boundaries[0]);
+        for (size_t i = 0; i < this->d; i += 16) {
+            __m512 vals = _mm512_loadu_ps(x + i);
+            __mmask16 mask = _mm512_cmp_ps_mask(vals, boundary, _CMP_NLT_UQ);
+            uint16_t bits = _cvtmask16_u32(mask);
+            memcpy(code + i / 8, &bits, sizeof(uint16_t));
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < this->d; i += 16) {
+            simd16float32 xi =
+                    reconstruct_16_components(code, static_cast<int>(i));
+            _mm512_storeu_ps(x + i, xi.f);
+        }
+    }
+};
 
+// 2-4 bit Lloyd-Max AVX512: decode via gather, encode stays scalar.
+#define DEFINE_LLOYD_MAX_AVX512_MULTIBIT(NBITS, UNPACK_EXPR)                  \
+    template <>                                                               \
+    struct QuantizerLloydMax<NBITS, THE_LEVEL_TO_DISPATCH>                    \
+            : QuantizerLloydMaxBase<NBITS, THE_LEVEL_TO_DISPATCH> {           \
+        using Base = QuantizerLloydMaxBase<NBITS, THE_LEVEL_TO_DISPATCH>;     \
+                                                                              \
+        QuantizerLloydMax(size_t d, const std::vector<float>& trained)        \
+                : Base(d, trained) {                                          \
+            assert(d % 16 == 0);                                              \
+        }                                                                     \
+                                                                              \
+        FAISS_ALWAYS_INLINE simd16float32                                     \
+        reconstruct_16_components(const uint8_t* code, int i) const {         \
+            return simd16float32(_mm512_i32gather_ps(                         \
+                    (UNPACK_EXPR), this->centroids, sizeof(float)));          \
+        }                                                                     \
+                                                                              \
+        void decode_vector(const uint8_t* code, float* x) const final {       \
+            for (size_t i = 0; i < this->d; i += 16) {                        \
+                simd16float32 xi =                                            \
+                        reconstruct_16_components(code, static_cast<int>(i)); \
+                _mm512_storeu_ps(x + i, xi.f);                                \
+            }                                                                 \
+        }                                                                     \
+    }
+
+DEFINE_LLOYD_MAX_AVX512_MULTIBIT(2, unpack_16x2bit_to_u32(code, i));
+DEFINE_LLOYD_MAX_AVX512_MULTIBIT(3, unpack_16x3bit_to_u32(code, i));
+DEFINE_LLOYD_MAX_AVX512_MULTIBIT(4, unpack_16x4bit_to_u32(code, i));
+
+#undef DEFINE_LLOYD_MAX_AVX512_MULTIBIT
+
+// 8-bit Lloyd-Max AVX512
 template <>
-struct QuantizerLloydMax<8, SIMDLevel::AVX512>
-        : QuantizerLloydMax<8, SIMDLevel::NONE> {
-    using Base = QuantizerLloydMax<8, SIMDLevel::NONE>;
+struct QuantizerLloydMax<8, THE_LEVEL_TO_DISPATCH>
+        : QuantizerLloydMaxBase<8, THE_LEVEL_TO_DISPATCH> {
+    using Base = QuantizerLloydMaxBase<8, THE_LEVEL_TO_DISPATCH>;
 
     QuantizerLloydMax(size_t d, const std::vector<float>& trained)
             : Base(d, trained) {
@@ -288,6 +342,14 @@ struct QuantizerLloydMax<8, SIMDLevel::AVX512>
         return simd16float32(
                 _mm512_i32gather_ps(indices, this->centroids, sizeof(float)));
     }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < this->d; i += 16) {
+            simd16float32 xi =
+                    reconstruct_16_components(code, static_cast<int>(i));
+            _mm512_storeu_ps(x + i, xi.f);
+        }
+    }
 };
 
 /**********************************************************
@@ -295,17 +357,25 @@ struct QuantizerLloydMax<8, SIMDLevel::AVX512>
  **********************************************************/
 
 template <>
-struct QuantizerFP16<SIMDLevel::AVX512> : QuantizerFP16<SIMDLevel::NONE> {
-    QuantizerFP16(size_t d, const std::vector<float>& trained)
-            : QuantizerFP16<SIMDLevel::NONE>(d, trained) {
-        assert(d % 16 == 0);
-    }
+struct QuantizerFP16<THE_LEVEL_TO_DISPATCH>
+        : QuantizerFP16Base<THE_LEVEL_TO_DISPATCH> {
+    using QuantizerFP16Base<THE_LEVEL_TO_DISPATCH>::QuantizerFP16Base;
 
     FAISS_ALWAYS_INLINE simd16float32
     reconstruct_16_components(const uint8_t* code, int i) const {
         __m256i codei = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
         return simd16float32(_mm512_cvtph_ps(codei));
     }
+
+    /// Components i .. i + 15 selected by mask, the others 0. Reads no byte
+    /// past the selected components.
+    FAISS_ALWAYS_INLINE simd16float32 reconstruct_16_components_masked(
+            const uint8_t* code,
+            int i,
+            __mmask16 mask) const {
+        __m256i codei = _mm256_maskz_loadu_epi16(mask, code + 2 * i);
+        return simd16float32(_mm512_cvtph_ps(codei));
+    }
 };
 
 /**********************************************************
@@ -313,15 +383,26 @@ struct QuantizerFP16<SIMDLevel::AVX512> : QuantizerFP16<SIMDLevel::NONE> {
  **********************************************************/
 
 template <>
-struct QuantizerBF16<SIMDLevel::AVX512> : QuantizerBF16<SIMDLevel::NONE> {
-    QuantizerBF16(size_t d, const std::vector<float>& trained)
-            : QuantizerBF16<SIMDLevel::NONE>(d, trained) {
-        assert(d % 16 == 0);
-    }
+struct QuantizerBF16<THE_LEVEL_TO_DISPATCH>
+        : QuantizerBF16Base<THE_LEVEL_TO_DISPATCH> {
+    using QuantizerBF16Base<THE_LEVEL_TO_DISPATCH>::QuantizerBF16Base;
 
     FAISS_ALWAYS_INLINE simd16float32
     reconstruct_16_components(const uint8_t* code, int i) const {
         __m256i code_256i = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
+        return widen(code_256i);
+    }
+
+    /// Components i .. i + 15 selected by mask, the others 0. Reads no byte
+    /// past the selected components.
+    FAISS_ALWAYS_INLINE simd16float32 reconstruct_16_components_masked(
+            const uint8_t* code,
+            int i,
+            __mmask16 mask) const {
+        return widen(_mm256_maskz_loadu_epi16(mask, code + 2 * i));
+    }
+
+    static FAISS_ALWAYS_INLINE simd16float32 widen(__m256i code_256i) {
         __m512i code_512i = _mm512_cvtepu16_epi32(code_256i);
         code_512i = _mm512_slli_epi32(code_512i, 16);
         return simd16float32(_mm512_castsi512_ps(code_512i));
@@ -333,10 +414,10 @@ struct QuantizerBF16<SIMDLevel::AVX512> : QuantizerBF16<SIMDLevel::NONE> {
  **********************************************************/
 
 template <>
-struct Quantizer8bitDirect<SIMDLevel::AVX512>
-        : Quantizer8bitDirect<SIMDLevel::NONE> {
+struct Quantizer8bitDirect<THE_LEVEL_TO_DISPATCH>
+        : Quantizer8bitDirectBase<THE_LEVEL_TO_DISPATCH> {
     Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
-            : Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {
+            : Quantizer8bitDirectBase<THE_LEVEL_TO_DISPATCH>(d, trained) {
         assert(d % 16 == 0);
     }
 
@@ -353,10 +434,10 @@ struct Quantizer8bitDirect<SIMDLevel::AVX512>
  **********************************************************/
 
 template <>
-struct Quantizer8bitDirectSigned<SIMDLevel::AVX512>
-        : Quantizer8bitDirectSigned<SIMDLevel::NONE> {
+struct Quantizer8bitDirectSigned<THE_LEVEL_TO_DISPATCH>
+        : Quantizer8bitDirectSignedBase<THE_LEVEL_TO_DISPATCH> {
     Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
-            : Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {
+            : Quantizer8bitDirectSignedBase<THE_LEVEL_TO_DISPATCH>(d, trained) {
         assert(d % 16 == 0);
     }
 
@@ -375,9 +456,9 @@ struct Quantizer8bitDirectSigned<SIMDLevel::AVX512>
  **********************************************************/
 
 template <>
-struct SimilarityL2<SIMDLevel::AVX512> {
+struct SimilarityL2<THE_LEVEL_TO_DISPATCH> {
     static constexpr int simdwidth = 16;
-    static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
+    static constexpr SIMDLevel simd_level = THE_LEVEL_TO_DISPATCH;
     static constexpr MetricType metric_type = METRIC_L2;
 
     const float *y, *yi;
@@ -405,15 +486,38 @@ struct SimilarityL2<SIMDLevel::AVX512> {
         accu16 = accu16 + tmp * tmp;
     }
 
+    FAISS_ALWAYS_INLINE void fmadd_16_components_2(
+            simd16float32 x,
+            simd16float32 y_2) {
+        simd16float32 tmp = y_2 - x;
+        accu16 = fmadd(tmp, tmp, accu16);
+    }
+
     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 <>
-struct SimilarityIP<SIMDLevel::AVX512> {
+struct SimilarityIP<THE_LEVEL_TO_DISPATCH> {
     static constexpr int simdwidth = 16;
-    static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
+    static constexpr SIMDLevel simd_level = THE_LEVEL_TO_DISPATCH;
     static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
 
     const float *y, *yi;
@@ -439,32 +543,144 @@ struct SimilarityIP<SIMDLevel::AVX512> {
         accu16 = accu16 + x1 * x2;
     }
 
+    FAISS_ALWAYS_INLINE void fmadd_16_components_2(
+            simd16float32 x1,
+            simd16float32 x2) {
+        accu16 = fmadd(x1, x2, accu16);
+    }
+
     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;
+    }
 };
 
 /**********************************************************
  * Distance Computers
  **********************************************************/
 
-template <class Quantizer, class Similarity>
-struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
+namespace {
+
+/// 16 lanes to 4: lanes j, j + 4, j + 8 and j + 12 summed into lane j.
+FAISS_ALWAYS_INLINE __m128 reduce_16_to_4(__m512 a) {
+    const __m256 hi =
+            _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(a), 1));
+    const __m256 s8 = _mm256_add_ps(_mm512_castps512_ps256(a), hi);
+    return _mm_add_ps(_mm256_castps256_ps128(s8), _mm256_extractf128_ps(s8, 1));
+}
+
+} // namespace
+
+template <class Quantizer, class Similarity, bool four_acc>
+struct DCTemplate<Quantizer, Similarity, THE_LEVEL_TO_DISPATCH, four_acc>
         : SQDistanceComputer {
     using Sim = Similarity;
 
     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);
+        }
+    }
+
+    /// The quantizer can decode a partial block, so d need not be a
+    /// multiple of 16.
+    static constexpr bool has_masked_tail() {
+        return requires(const Quantizer& q, const uint8_t* c) {
+            { q.reconstruct_16_components_masked(c, 0, __mmask16(0)) };
+        };
+    }
+
+    /// Adds the partial last block, from component i, of a d that is not a
+    /// multiple of 16.
+    FAISS_ALWAYS_INLINE void add_tail(
+            Similarity& sim,
+            const float* x,
+            const uint8_t* code,
+            size_t i) const {
+        if constexpr (has_masked_tail()) {
+            if (i < quant.d) {
+                const __mmask16 mask = __mmask16((1u << (quant.d - i)) - 1);
+                sim.fmadd_16_components_2(
+                        quant.reconstruct_16_components_masked(
+                                code, static_cast<int>(i), mask),
+                        simd16float32(_mm512_maskz_loadu_ps(mask, x + i)));
+            }
+        }
+    }
+
+    /// Similarity terms of one code with four accumulators.
+    FAISS_ALWAYS_INLINE __m512
+    accumulate_4(const float* x, const uint8_t* code) const {
+        struct Acc : Similarity {
+            Acc() : Similarity(nullptr) {}
+        };
+        Acc s[4];
+        for (int j = 0; j < 4; j++) {
+            s[j].begin_16();
+        }
+        const size_t d = quant.d;
+        size_t i = 0;
+        for (; i + 64 <= d; i += 64) {
+            for (int j = 0; j < 4; j++) {
+                s[j].fmadd_16_components_2(
+                        quant.reconstruct_16_components(
+                                code, static_cast<int>(i + 16 * j)),
+                        simd16float32(x + i + 16 * j));
+            }
+        }
+        for (; i + 16 <= d; i += 16) {
+            s[0].fmadd_16_components_2(
+                    quant.reconstruct_16_components(code, static_cast<int>(i)),
+                    simd16float32(x + i));
+        }
+        add_tail(s[0], x, code, i);
+        return _mm512_add_ps(
+                _mm512_add_ps(s[0].accu16.f, s[1].accu16.f),
+                _mm512_add_ps(s[2].accu16.f, s[3].accu16.f));
+    }
 
     float compute_distance(const float* x, const uint8_t* code) const {
+        if constexpr (four_acc) {
+            return reduce_4(reduce_16_to_4(accumulate_4(x, code)));
+        }
         Similarity sim(x);
         sim.begin_16();
-        for (size_t i = 0; i < quant.d; i += 16) {
+        const size_t d16 = quant.d & ~size_t(15);
+        for (size_t i = 0; i < d16; i += 16) {
             simd16float32 xi = quant.reconstruct_16_components(code, i);
             sim.add_16_components(xi);
         }
+        add_tail(sim, x, code, d16);
         return sim.result_16();
     }
 
@@ -472,16 +688,46 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
             const {
         Similarity sim(nullptr);
         sim.begin_16();
-        for (size_t i = 0; i < quant.d; i += 16) {
+        const size_t d16 = quant.d & ~size_t(15);
+        for (size_t i = 0; i < d16; i += 16) {
             simd16float32 x1 = quant.reconstruct_16_components(code1, i);
             simd16float32 x2 = quant.reconstruct_16_components(code2, i);
             sim.add_16_components_2(x1, x2);
         }
+        if constexpr (has_masked_tail()) {
+            if (d16 < quant.d) {
+                const __mmask16 mask = __mmask16((1u << (quant.d - d16)) - 1);
+                const int i = static_cast<int>(d16);
+                sim.add_16_components_2(
+                        quant.reconstruct_16_components_masked(code1, i, mask),
+                        quant.reconstruct_16_components_masked(code2, i, mask));
+            }
+        }
         return sim.result_16();
     }
 
     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 {
@@ -490,64 +736,197 @@ 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);
+        }
     }
-};
 
-template <class Similarity>
-struct DistanceComputerByte<Similarity, SIMDLevel::AVX512>
-        : SQDistanceComputer {
-    using Sim = Similarity;
+    void query_to_codes_batch_4(
+            const uint8_t* code_0,
+            const uint8_t* code_1,
+            const uint8_t* code_2,
+            const uint8_t* code_3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) const final {
+        if constexpr (four_acc) {
+            alignas(16) float dis[4];
+            _mm_store_ps(
+                    dis,
+                    reduce_4x4(
+                            reduce_16_to_4(accumulate_4(q, code_0)),
+                            reduce_16_to_4(accumulate_4(q, code_1)),
+                            reduce_16_to_4(accumulate_4(q, code_2)),
+                            reduce_16_to_4(accumulate_4(q, code_3))));
+            dis0 = dis[0];
+            dis1 = dis[1];
+            dis2 = dis[2];
+            dis3 = dis[3];
+            return;
+        }
+        Similarity sim0(q);
+        Similarity sim1(q);
+        Similarity sim2(q);
+        Similarity sim3(q);
+
+        sim0.begin_16();
+        sim1.begin_16();
+        sim2.begin_16();
+        sim3.begin_16();
+
+        const size_t d16 = quant.d & ~size_t(15);
+        for (size_t i = 0; i < d16; i += 16) {
+            simd16float32 xi0 = quant.reconstruct_16_components(code_0, i);
+            simd16float32 xi1 = quant.reconstruct_16_components(code_1, i);
+            simd16float32 xi2 = quant.reconstruct_16_components(code_2, i);
+            simd16float32 xi3 = quant.reconstruct_16_components(code_3, i);
+            sim0.add_16_components(xi0);
+            sim1.add_16_components(xi1);
+            sim2.add_16_components(xi2);
+            sim3.add_16_components(xi3);
+        }
+        add_tail(sim0, q, code_0, d16);
+        add_tail(sim1, q, code_1, d16);
+        add_tail(sim2, q, code_2, d16);
+        add_tail(sim3, q, code_3, d16);
 
-    int d;
-    std::vector<uint8_t> tmp;
+        dis0 = sim0.result_16();
+        dis1 = sim1.result_16();
+        dis2 = sim2.result_16();
+        dis3 = sim3.result_16();
+    }
+};
 
-    DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
+// QT_8bit + IP: the precomputed form of the AVX2 specialization, 16 dims per
+// convert and FMA.
+template <bool four_acc>
+struct DCTemplate<
+        QuantizerTemplate<
+                Codec8bit<THE_LEVEL_TO_DISPATCH>,
+                QuantizerTemplateScaling::NON_UNIFORM,
+                THE_LEVEL_TO_DISPATCH>,
+        SimilarityIP<THE_LEVEL_TO_DISPATCH>,
+        THE_LEVEL_TO_DISPATCH,
+        four_acc> : SQDistanceComputer {
+    using Quantizer = QuantizerTemplate<
+            Codec8bit<THE_LEVEL_TO_DISPATCH>,
+            QuantizerTemplateScaling::NON_UNIFORM,
+            THE_LEVEL_TO_DISPATCH>;
+    using Sim = SimilarityIP<THE_LEVEL_TO_DISPATCH>;
 
-    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
-            const {
-        // compute 16 lanes of 32-bit products (16-bytes) at once for
-        // the supported metrics
-        __m512i accu = _mm512_setzero_si512();
-        constexpr int kLanes = 16;
-        for (int i = 0; i < d; i += kLanes) {
-            __m128i c1 = _mm_loadu_si128((__m128i*)(code1 + i));
-            __m128i c2 = _mm_loadu_si128((__m128i*)(code2 + i));
-            __m512i c1i = _mm512_cvtepu8_epi32(c1);
-            __m512i c2i = _mm512_cvtepu8_epi32(c2);
-
-            __m512i v;
-            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
-                v = _mm512_mullo_epi32(c1i, c2i);
-            } else {
-                __m512i diff = _mm512_sub_epi32(c1i, c2i);
-                v = _mm512_mullo_epi32(diff, diff);
-            }
-            accu = _mm512_add_epi32(accu, v);
-        }
-        return _mm512_reduce_add_epi32(accu);
+    Quantizer quant;
+    std::vector<float> a, rmin, qa;
+    float bias = 0;
+
+    DCTemplate(size_t d, const std::vector<float>& trained)
+            : quant(d, trained), a(d), rmin(d), qa(d) {
+        sq8::make_coefficients(trained.data(), d, a.data(), rmin.data());
     }
 
     void set_query(const float* x) final {
-        for (int i = 0; i < d; i++) {
-            tmp[i] = int(x[i]);
+        q = x;
+        bias = 0;
+        for (size_t i = 0; i < quant.d; i++) {
+            qa[i] = x[i] * a[i];
+            bias += x[i] * rmin[i];
         }
     }
 
-    int compute_distance(const float* x, const uint8_t* code) {
-        set_query(x);
-        return compute_code_distance(tmp.data(), code);
+    static FAISS_ALWAYS_INLINE __m512 codes_16(const uint8_t* code) {
+        return _mm512_cvtepi32_ps(
+                _mm512_cvtepu8_epi32(_mm_loadu_si128((const __m128i*)code)));
     }
 
-    float symmetric_dis(idx_t i, idx_t j) override {
-        return compute_code_distance(
-                codes + i * code_size, codes + j * code_size);
+    /// sum_i qa[i] * code[i] of N codes, in 4 lanes each. The N codes share
+    /// each qa load. With four_acc, dims i + 16 * j of each 64 go to
+    /// accumulator j of the code, so a code sums the same way for any N.
+    template <size_t N>
+    FAISS_ALWAYS_INLINE void accumulate(
+            const uint8_t* const* codes,
+            __m128* sums) const {
+        constexpr size_t n_acc = four_acc ? 4 : 1;
+        __m512 acc[N][n_acc];
+        for (size_t n = 0; n < N; n++) {
+            for (size_t j = 0; j < n_acc; j++) {
+                acc[n][j] = _mm512_setzero_ps();
+            }
+        }
+        const size_t d = quant.d;
+        size_t i = 0;
+        if constexpr (four_acc) {
+            for (; i + 64 <= d; i += 64) {
+                for (size_t j = 0; j < 4; j++) {
+                    const __m512 qa16 = _mm512_loadu_ps(qa.data() + i + 16 * j);
+                    for (size_t n = 0; n < N; n++) {
+                        acc[n][j] = _mm512_fmadd_ps(
+                                qa16,
+                                codes_16(codes[n] + i + 16 * j),
+                                acc[n][j]);
+                    }
+                }
+            }
+        }
+        for (; i < d; i += 16) {
+            const __m512 qa16 = _mm512_loadu_ps(qa.data() + i);
+            for (size_t n = 0; n < N; n++) {
+                acc[n][0] = _mm512_fmadd_ps(
+                        qa16, codes_16(codes[n] + i), acc[n][0]);
+            }
+        }
+        for (size_t n = 0; n < N; n++) {
+            __m512 sum = acc[n][0];
+            if constexpr (four_acc) {
+                sum = _mm512_add_ps(
+                        _mm512_add_ps(acc[n][0], acc[n][1]),
+                        _mm512_add_ps(acc[n][2], acc[n][3]));
+            }
+            sums[n] = reduce_16_to_4(sum);
+        }
     }
 
     float query_to_code(const uint8_t* code) const final {
-        return compute_code_distance(tmp.data(), code);
+        __m128 sum;
+        accumulate<1>(&code, &sum);
+        return bias + reduce_4(sum);
+    }
+
+    void query_to_codes_batch_4(
+            const uint8_t* code_0,
+            const uint8_t* code_1,
+            const uint8_t* code_2,
+            const uint8_t* code_3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) const final {
+        const uint8_t* codes[4] = {code_0, code_1, code_2, code_3};
+        __m128 sums[4];
+        accumulate<4>(codes, sums);
+        alignas(16) float dis[4];
+        _mm_store_ps(dis, reduce_4x4(sums[0], sums[1], sums[2], sums[3]));
+        dis0 = bias + dis[0];
+        dis1 = bias + dis[1];
+        dis2 = bias + dis[2];
+        dis3 = bias + dis[3];
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        const uint8_t* c1 = codes + i * code_size;
+        const uint8_t* c2 = codes + j * code_size;
+        Sim sim(nullptr);
+        sim.begin_16();
+        for (size_t k = 0; k < quant.d; k += 16) {
+            simd16float32 x1 = quant.reconstruct_16_components(c1, k);
+            simd16float32 x2 = quant.reconstruct_16_components(c2, k);
+            sim.add_16_components_2(x1, x2);
+        }
+        return sim.result_16();
     }
 };
 
 } // namespace scalar_quantizer
+
 } // namespace faiss