faiss 0.6.1 → 0.6.3

data/vendor/faiss/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 <class Codec>
@@ -244,41 +250,87 @@ struct QuantizerTemplate<
 };
 
 /**********************************************************
- * TurboQuant MSE quantizer
+ * Lloyd-Max scalar quantizer
  **********************************************************/
 
-#define DEFINE_TQMSE_AVX512_SPECIALIZATION(NBITS, INDEX_EXPR)               \
-    template <>                                                             \
-    struct QuantizerTurboQuantMSE<NBITS, SIMDLevel::AVX512>                 \
-            : QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE> {              \
-        using Base = QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE>;        \
-                                                                            \
-        QuantizerTurboQuantMSE(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_TQMSE_AVX512_SPECIALIZATION(1, unpack_16x1bit_to_u32(code, i));
-DEFINE_TQMSE_AVX512_SPECIALIZATION(2, unpack_16x2bit_to_u32(code, i));
-DEFINE_TQMSE_AVX512_SPECIALIZATION(3, unpack_16x3bit_to_u32(code, i));
-DEFINE_TQMSE_AVX512_SPECIALIZATION(4, unpack_16x4bit_to_u32(code, i));
-
-#undef DEFINE_TQMSE_AVX512_SPECIALIZATION
+// 1-bit Lloyd-Max AVX512: 16 comparisons → 2 bytes via mask compare.
+template <>
+struct QuantizerLloydMax<1, SIMDLevel::AVX512>
+        : QuantizerLloydMax<1, SIMDLevel::NONE> {
+    using Base = QuantizerLloydMax<1, 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 {
+        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_GT_OQ);
+            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, 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 {         \
+            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 QuantizerTurboQuantMSE<8, SIMDLevel::AVX512>
-        : QuantizerTurboQuantMSE<8, SIMDLevel::NONE> {
-    using Base = QuantizerTurboQuantMSE<8, SIMDLevel::NONE>;
+struct QuantizerLloydMax<8, SIMDLevel::AVX512>
+        : QuantizerLloydMax<8, SIMDLevel::NONE> {
+    using Base = QuantizerLloydMax<8, SIMDLevel::NONE>;
 
-    QuantizerTurboQuantMSE(size_t d, const std::vector<float>& trained)
+    QuantizerLloydMax(size_t d, const std::vector<float>& trained)
             : Base(d, trained) {
         assert(d % 16 == 0);
     }
@@ -291,6 +343,14 @@ struct QuantizerTurboQuantMSE<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);
+        }
+    }
 };
 
 /**********************************************************
@@ -411,6 +471,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 +521,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 +551,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 +593,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 +621,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(
@@ -545,24 +677,36 @@ struct DistanceComputerByte<Similarity, SIMDLevel::AVX512>
 
     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
+        // compute 32 lanes of 16-bit products (32-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;
+        constexpr int kLanes = 32;
+        int i = 0;
+        for (; i + kLanes <= d; i += kLanes) {
+            __m256i c1 = _mm256_loadu_epi8(code1 + i);
+            __m256i c2 = _mm256_loadu_epi8(code2 + i);
+            __m512i c1i16 = _mm512_cvtepu8_epi16(c1);
+            __m512i c2i16 = _mm512_cvtepu8_epi16(c2);
             if (Sim::metric_type == METRIC_INNER_PRODUCT) {
-                v = _mm512_mullo_epi32(c1i, c2i);
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(c1i16, c2i16));
             } else {
-                __m512i diff = _mm512_sub_epi32(c1i, c2i);
-                v = _mm512_mullo_epi32(diff, diff);
+                __m512i diff = _mm512_sub_epi16(c1i16, c2i16);
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(diff, diff));
+            }
+        }
+        // tail handling for dimensions not divisible by 32
+        if (i < d) {
+            __mmask32 mask = (__mmask32)((1ULL << (d - i)) - 1ULL);
+            __m256i c1 = _mm256_maskz_loadu_epi8(mask, code1 + i);
+            __m256i c2 = _mm256_maskz_loadu_epi8(mask, code2 + i);
+            __m512i c1i16 = _mm512_cvtepu8_epi16(c1);
+            __m512i c2i16 = _mm512_cvtepu8_epi16(c2);
+            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(c1i16, c2i16));
+            } else {
+                __m512i diff = _mm512_sub_epi16(c1i16, c2i16);
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(diff, diff));
             }
-            accu = _mm512_add_epi32(accu, v);
         }
         return _mm512_reduce_add_epi32(accu);
     }
@@ -588,10 +732,105 @@ struct DistanceComputerByte<Similarity, SIMDLevel::AVX512>
     }
 };
 
+template <class Similarity>
+struct DistanceComputerByteSigned<Similarity, SIMDLevel::AVX512>
+        : SQDistanceComputer {
+    using Sim = Similarity;
+
+    int d;
+    std::vector<uint8_t> tmp;
+
+    DistanceComputerByteSigned(int d, const std::vector<float>&)
+            : d(d), tmp(d) {}
+
+    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        // codes store value + 128. madd_epi16 is signed, so IP unbiases the
+        // bytes before multiplying; for L2 the +128 cancels in the difference.
+        // Only dispatched for d % 32 == 0, so the loop needs no tail.
+        __m512i accu = _mm512_setzero_si512();
+        constexpr int kLanes = 32;
+        for (int i = 0; i + kLanes <= d; i += kLanes) {
+            __m512i c1 = _mm512_cvtepu8_epi16(_mm256_loadu_epi8(code1 + i));
+            __m512i c2 = _mm512_cvtepu8_epi16(_mm256_loadu_epi8(code2 + i));
+            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
+                const __m512i bias = _mm512_set1_epi16(128);
+                c1 = _mm512_sub_epi16(c1, bias);
+                c2 = _mm512_sub_epi16(c2, bias);
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(c1, c2));
+            } else {
+                __m512i diff = _mm512_sub_epi16(c1, c2);
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(diff, diff));
+            }
+        }
+        return _mm512_reduce_add_epi32(accu);
+    }
+
+    void set_query(const float* x) final {
+        for (int i = 0; i < d; i++) {
+            tmp[i] = uint8_t(int(x[i]) + 128);
+        }
+    }
+
+    int compute_distance(const float* x, const uint8_t* code) {
+        set_query(x);
+        return compute_code_distance(tmp.data(), code);
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return compute_code_distance(
+                codes + i * code_size, codes + j * code_size);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return compute_code_distance(tmp.data(), code);
+    }
+};
+
+/**********************************************************
+ * TurboQuant masked_sum AVX512 specialization
+ **********************************************************/
+
+template <SIMDLevel SL0>
+float turboq_masked_sum(const float* arr, const uint8_t* bits, size_t d);
+
+template <>
+float turboq_masked_sum<SIMDLevel::AVX512>(
+        const float* arr,
+        const uint8_t* bits,
+        size_t d) {
+    __m512 acc = _mm512_setzero_ps();
+    size_t i = 0;
+    size_t full_16 = (d / 16) * 16;
+    for (; i < full_16; i += 16) {
+        uint16_t mask16;
+        memcpy(&mask16, bits + i / 8, sizeof(mask16));
+        __mmask16 k = _cvtu32_mask16(mask16);
+        __m512 vals = _mm512_loadu_ps(arr + i);
+        acc = _mm512_mask_add_ps(acc, k, acc, vals);
+    }
+    float result = _mm512_reduce_add_ps(acc);
+    if (i < d) {
+        size_t remaining = d - i;
+        __mmask16 tail_mask = _cvtu32_mask16((1u << remaining) - 1);
+        __m512 tail_vals = _mm512_maskz_loadu_ps(tail_mask, arr + i);
+        uint16_t bits_tail = 0;
+        size_t bytes_remaining = (remaining + 7) / 8;
+        memcpy(&bits_tail, bits + i / 8, bytes_remaining);
+        __mmask16 bits_k = _cvtu32_mask16(bits_tail);
+        __mmask16 combined = _kand_mask16(tail_mask, bits_k);
+        __m512 masked_tail = _mm512_maskz_mov_ps(combined, tail_vals);
+        result += _mm512_reduce_add_ps(masked_tail);
+    }
+    return result;
+}
+
 } // namespace scalar_quantizer
 } // namespace faiss
 
+#ifndef SQ_AVX512_SKIP_DISPATCH
 #define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX512
 #include <faiss/impl/scalar_quantizer/sq-dispatch.h>
+#endif
 
 #endif // COMPILE_SIMD_AVX512