faiss 0.6.4 → 0.6.5

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

@@ -15,6 +15,7 @@
 #include <faiss/impl/scalar_quantizer/similarities.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/rabitq_simd.h>
+#include <faiss/utils/simd_impl/masked_sum.h>
 #include <limits>
 
 #ifndef THE_LEVEL_TO_DISPATCH
@@ -36,35 +37,6 @@ template <SIMDLevel SL0>
 inline constexpr SIMDLevel rabitq_popcount_level =
         SL0 == SIMDLevel::AVX512_SPR ? SIMDLevel::AVX512_VPOPCNT : SL0;
 
-/*******************************************************************
- * TurboQuant SIMD kernel: masked_sum
- * Compute sum of arr[j] where bit j of the bitmask is set.
- * NONE specialization is inline; AVX2/AVX512/NEON specializations
- * live in sq-avx2.cpp / sq-avx512.cpp / sq-neon.cpp.
- *******************************************************************/
-
-template <SIMDLevel SL0>
-float turboq_masked_sum(const float* arr, const uint8_t* bits, size_t d);
-
-template <>
-inline float turboq_masked_sum<SIMDLevel::NONE>(
-        const float* arr,
-        const uint8_t* bits,
-        size_t d) {
-    float result = 0;
-    for (size_t byte_idx = 0; byte_idx < (d + 7) / 8; byte_idx++) {
-        uint8_t b = bits[byte_idx];
-        size_t base = byte_idx * 8;
-        size_t end = std::min(base + 8, d);
-        for (size_t j = base; j < end; j++) {
-            if (b & (1 << (j - base))) {
-                result += arr[j];
-            }
-        }
-    }
-    return result;
-}
-
 /*******************************************************************
  * Full TurboQuant DC — lives here because it needs both
  * quantizers.h (QuantizerTurboQuantFull, SQTurboQFactors) and
@@ -73,8 +45,10 @@ inline float turboq_masked_sum<SIMDLevel::NONE>(
  *******************************************************************/
 template <int NBits, class Similarity, SIMDLevel SL2>
 struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer {
+    // The AVX-512 levels share the AVX512 masked_sum specialization.
+    static constexpr SIMDLevel SUM_LEVEL = simd512_level_selector<SL2>::value;
     using Sim = Similarity;
-    QuantizerTurboQuantFull<NBits, SIMDLevel::NONE> quant;
+    QuantizerTurboQuantFull<NBits, SL2> quant;
     std::vector<float> query;
     std::vector<float> query_proj;
     float q_norm_sq = 0;
@@ -277,19 +251,19 @@ struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer {
                         mse_popcnt_scale * static_cast<float>(pop);
             } else {
                 // Float path: masked accumulation
-                float pos_sum = turboq_masked_sum<SL2>(query.data(), code, d);
+                float pos_sum = masked_sum<SUM_LEVEL>(query.data(), code, d);
                 mse_dot = scaled_c0 * total_q_sum + delta_centroid * pos_sum;
             }
         } else if constexpr (NBits == 3) {
             // 2-bit MSE: decompose into 3 masked sums over bit-planes.
             size_t pb = quant.mse_plane_bytes;
-            float s0 = turboq_masked_sum<SL2>(query.data(), code, d);
-            float s1 = turboq_masked_sum<SL2>(query.data(), code + pb, d);
+            float s0 = masked_sum<SUM_LEVEL>(query.data(), code, d);
+            float s1 = masked_sum<SUM_LEVEL>(query.data(), code + pb, d);
             for (size_t i = 0; i < pb; i++) {
                 scratch_and[i] = code[i] & code[pb + i];
             }
             float s01 =
-                    turboq_masked_sum<SL2>(query.data(), scratch_and.data(), d);
+                    masked_sum<SUM_LEVEL>(query.data(), scratch_and.data(), d);
             mse_dot = mse_multi_base + mse_coeff_s0 * s0 + mse_coeff_s1 * s1 +
                     mse_coeff_s01 * s01;
         } else {
@@ -336,7 +310,7 @@ struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer {
             qjl_dot = qjl_coeff * gamma * (2.0f * pos_sum - total_qproj_sum);
         } else {
             float pos_sum =
-                    turboq_masked_sum<SL2>(query_proj.data(), qjl_code, d);
+                    masked_sum<SUM_LEVEL>(query_proj.data(), qjl_code, d);
             qjl_dot = qjl_coeff * gamma * (2.0f * pos_sum - total_qproj_sum);
         }
 
@@ -373,6 +347,50 @@ constexpr bool is_dimension_compatible(size_t d) {
     }
 }
 
+// Distance computers and scanners. At AVX512 the 16-bit float codecs decode
+// a partial last block, so they take any d.
+template <SIMDLevel SL2>
+constexpr bool is_dimension_compatible(
+        size_t d,
+        ScalarQuantizer::QuantizerType qtype) {
+    return (SL2 == SIMDLevel::AVX512 &&
+            (qtype == ScalarQuantizer::QT_fp16 ||
+             qtype == ScalarQuantizer::QT_bf16)) ||
+            is_dimension_compatible<SL2>(d);
+}
+
+/// Calls f with the DCTemplate type for (Quantizer, Sim, SL2) and d. The
+/// 16-bit float quantizers and the QT_8bit inner product use it. Four
+/// accumulators per code hide the FMA latency at a large d. At a small d their
+/// reduction costs more than it saves. At AVX512 they are faster only on AMD
+/// (SIMDConfig::sq_avx512_4_accumulators).
+template <class Quantizer, class Sim, SIMDLevel SL2, class F>
+auto with_dc_template(size_t d, F f) {
+    if constexpr (SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::AVX512) {
+        if (d >= 384 &&
+            (SL2 == SIMDLevel::AVX2 || SIMDConfig::sq_avx512_4_accumulators)) {
+            return f.template
+            operator()<DCTemplate<Quantizer, Sim, SL2, true>>();
+        }
+    }
+    return f.template operator()<DCTemplate<Quantizer, Sim, SL2>>();
+}
+
+/// with_dc_template for QT_8bit. Only its inner product has a four-accumulator
+/// loop.
+template <class Sim, SIMDLevel SL2, class F>
+auto with_8bit_dc_template(size_t d, F f) {
+    using Quantizer = QuantizerTemplate<
+            Codec8bit<SL2>,
+            QuantizerTemplateScaling::NON_UNIFORM,
+            SL2>;
+    if constexpr (Sim::metric_type == METRIC_INNER_PRODUCT) {
+        return with_dc_template<Quantizer, Sim, SL2>(d, f);
+    } else {
+        return f.template operator()<DCTemplate<Quantizer, Sim, SL2>>();
+    }
+}
+
 /*******************************************************************
  * sq_select_quantizer: the big switch returning SQuantizer*
  *******************************************************************/
@@ -487,10 +505,13 @@ SQDistanceComputer* select_distance_computer_body(
         const std::vector<float>& trained) {
     // Return nullptr for incompatible dimensions in SIMD cases
     if constexpr (SL2 != SIMDLevel::NONE) {
-        if (!is_dimension_compatible<SL2>(d)) {
+        if (!is_dimension_compatible<SL2>(d, qtype)) {
             return nullptr;
         }
     }
+    auto new_dc = [&]<class DC>() -> SQDistanceComputer* {
+        return new DC(d, trained);
+    };
     switch (qtype) {
         case ScalarQuantizer::QT_8bit_uniform:
             return new DCTemplate<
@@ -511,13 +532,7 @@ SQDistanceComputer* select_distance_computer_body(
                     SL2>(d, trained);
 
         case ScalarQuantizer::QT_8bit:
-            return new DCTemplate<
-                    QuantizerTemplate<
-                            Codec8bit<SL2>,
-                            QuantizerTemplateScaling::NON_UNIFORM,
-                            SL2>,
-                    Sim,
-                    SL2>(d, trained);
+            return with_8bit_dc_template<Sim, SL2>(d, new_dc);
 
         case ScalarQuantizer::QT_6bit:
             return new DCTemplate<
@@ -538,10 +553,10 @@ SQDistanceComputer* select_distance_computer_body(
                     SL2>(d, trained);
 
         case ScalarQuantizer::QT_fp16:
-            return new DCTemplate<QuantizerFP16<SL2>, Sim, SL2>(d, trained);
+            return with_dc_template<QuantizerFP16<SL2>, Sim, SL2>(d, new_dc);
 
         case ScalarQuantizer::QT_bf16:
-            return new DCTemplate<QuantizerBF16<SL2>, Sim, SL2>(d, trained);
+            return with_dc_template<QuantizerBF16<SL2>, Sim, SL2>(d, new_dc);
 
         case ScalarQuantizer::QT_8bit_direct:
             if constexpr (is_avx512_family<SL2>) {
@@ -705,7 +720,7 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
             [&]<SIMDLevel SL2, class Similarity>() -> InvertedListScanner* {
         // Return nullptr for incompatible dimensions in SIMD cases
         if constexpr (SL2 != SIMDLevel::NONE) {
-            if (!is_dimension_compatible<SL2>(d)) {
+            if (!is_dimension_compatible<SL2>(d, qtype)) {
                 return nullptr;
             }
         }
@@ -727,13 +742,7 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                         Similarity,
                         SL2>>();
             case ScalarQuantizer::QT_8bit:
-                return scan.template operator()<DCTemplate<
-                        QuantizerTemplate<
-                                Codec8bit<SL2>,
-                                QuantizerTemplateScaling::NON_UNIFORM,
-                                SL2>,
-                        Similarity,
-                        SL2>>();
+                return with_8bit_dc_template<Similarity, SL2>(d, scan);
             case ScalarQuantizer::QT_4bit:
                 return scan.template operator()<DCTemplate<
                         QuantizerTemplate<
@@ -751,11 +760,11 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                         Similarity,
                         SL2>>();
             case ScalarQuantizer::QT_fp16:
-                return scan.template
-                operator()<DCTemplate<QuantizerFP16<SL2>, Similarity, SL2>>();
+                return with_dc_template<QuantizerFP16<SL2>, Similarity, SL2>(
+                        d, scan);
             case ScalarQuantizer::QT_bf16:
-                return scan.template
-                operator()<DCTemplate<QuantizerBF16<SL2>, Similarity, SL2>>();
+                return with_dc_template<QuantizerBF16<SL2>, Similarity, SL2>(
+                        d, scan);
             case ScalarQuantizer::QT_8bit_direct:
                 if constexpr (is_avx512_family<SL2>) {
                     if (d % 32 == 0) {

data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp

@@ -5,11 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#ifdef COMPILE_SIMD_ARM_NEON
-
 #include <faiss/impl/simdlib/simdlib_neon.h>
 
-#include <algorithm>
 #include <cstring>
 
 #include <faiss/impl/scalar_quantizer/codecs.h>
@@ -17,6 +14,7 @@
 #include <faiss/impl/scalar_quantizer/quantizers.h>
 #include <faiss/impl/scalar_quantizer/scanners.h>
 #include <faiss/impl/scalar_quantizer/similarities.h>
+#include <faiss/impl/scalar_quantizer/sq8_coefficients.h>
 
 namespace faiss {
 
@@ -105,14 +103,15 @@ template <>
 struct Codec8bit<SIMDLevel::ARM_NEON> : Codec8bit<SIMDLevel::NONE> {
     static FAISS_ALWAYS_INLINE simd8float32
     decode_8_components(const uint8_t* code, size_t i) {
-        float32_t result[8] = {};
-        for (size_t j = 0; j < 8; j++) {
-            result[j] =
-                    Codec8bit<SIMDLevel::NONE>::decode_component(code, i + j);
-        }
-        float32x4_t res1 = vld1q_f32(result);
-        float32x4_t res2 = vld1q_f32(result + 4);
-        return simd8float32(float32x4x2_t{res1, res2});
+        const uint16x8_t c16 = vmovl_u8(vld1_u8(code + i));
+        const float32x4_t lo = vcvtq_f32_u32(vmovl_u16(vget_low_u16(c16)));
+        const float32x4_t hi = vcvtq_f32_u32(vmovl_u16(vget_high_u16(c16)));
+        const float32x4_t one_255 = vdupq_n_f32(1.f / 255.f);
+        const float32x4_t half_one_255 = vdupq_n_f32(0.5f / 255.f);
+        return simd8float32(
+                float32x4x2_t{
+                        vfmaq_f32(half_one_255, lo, one_255),
+                        vfmaq_f32(half_one_255, hi, one_255)});
     }
 };
 
@@ -618,6 +617,97 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
     }
 };
 
+// QT_8bit + L2. With a, rmin from sq8_coefficients.h and e = q - rmin, each
+// term is (e_i - a_i * c_i)^2: one fmls and one fmla per 4 dims, instead of
+// separate codec, vmin/vdiff, sub, mul and add steps.
+template <>
+struct DCTemplate<
+        QuantizerTemplate<
+                Codec8bit<SIMDLevel::ARM_NEON>,
+                QuantizerTemplateScaling::NON_UNIFORM,
+                SIMDLevel::ARM_NEON>,
+        SimilarityL2<SIMDLevel::ARM_NEON>,
+        SIMDLevel::ARM_NEON> : SQDistanceComputer {
+    using Sim = SimilarityL2<SIMDLevel::ARM_NEON>;
+
+    size_t d;
+    std::vector<float> a, rmin, e;
+
+    DCTemplate(size_t d_in, const std::vector<float>& trained)
+            : d(d_in), a(d_in), rmin(d_in), e(d_in) {
+        assert(d % 8 == 0);
+        sq8::make_coefficients(trained.data(), d, a.data(), rmin.data());
+    }
+
+    void set_query(const float* x) final {
+        q = x;
+        for (size_t i = 0; i < d; i++) {
+            e[i] = x[i] - rmin[i];
+        }
+    }
+
+    // Each code has its own accumulators, so the batched results equal the
+    // single-code ones bit for bit.
+    template <size_t N>
+    FAISS_ALWAYS_INLINE void l2(
+            const uint8_t* const (&codes)[N],
+            float (&dis)[N]) const {
+        float32x4_t acc[N][2] = {};
+        for (size_t i = 0; i < d; i += 8) {
+            const float32x4_t a0 = vld1q_f32(a.data() + i);
+            const float32x4_t a1 = vld1q_f32(a.data() + i + 4);
+            const float32x4_t e0 = vld1q_f32(e.data() + i);
+            const float32x4_t e1 = vld1q_f32(e.data() + i + 4);
+            for (size_t n = 0; n < N; n++) {
+                const uint16x8_t c16 = vmovl_u8(vld1_u8(codes[n] + i));
+                const float32x4_t t0 = vfmsq_f32(
+                        e0, a0, vcvtq_f32_u32(vmovl_u16(vget_low_u16(c16))));
+                const float32x4_t t1 =
+                        vfmsq_f32(e1, a1, vcvtq_f32_u32(vmovl_high_u16(c16)));
+                acc[n][0] = vfmaq_f32(acc[n][0], t0, t0);
+                acc[n][1] = vfmaq_f32(acc[n][1], t1, t1);
+            }
+        }
+        for (size_t n = 0; n < N; n++) {
+            dis[n] = vaddvq_f32(vaddq_f32(acc[n][0], acc[n][1]));
+        }
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        float dis[1];
+        l2<1>({code}, dis);
+        return dis[0];
+    }
+
+    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 {
+        float dis[4];
+        l2<4>({code_0, code_1, code_2, code_3}, dis);
+        dis0 = dis[0];
+        dis1 = dis[1];
+        dis2 = dis[2];
+        dis3 = 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;
+        float accu = 0;
+        for (size_t k = 0; k < d; k++) {
+            const float diff = a[k] * float(int(c1[k]) - int(c2[k]));
+            accu += diff * diff;
+        }
+        return accu;
+    }
+};
+
 // Byte-domain kernels for QT_8bit_direct{,_signed}. The dispatch only
 // selects them when d % 16 == 0, so no loop needs a tail.
 
@@ -762,93 +852,8 @@ struct DistanceComputerByteSigned<Similarity, SIMDLevel::ARM_NEON>
     }
 };
 
-/**********************************************************
- * TurboQuant masked_sum NEON specialization (scalar fallback)
- **********************************************************/
-
-template <SIMDLevel SL0>
-float turboq_masked_sum(const float* arr, const uint8_t* bits, size_t d);
-
-template <>
-float turboq_masked_sum<SIMDLevel::ARM_NEON>(
-        const float* arr,
-        const uint8_t* bits,
-        size_t d) {
-    float result = 0;
-    for (size_t byte_idx = 0; byte_idx < (d + 7) / 8; byte_idx++) {
-        uint8_t b = bits[byte_idx];
-        size_t base = byte_idx * 8;
-        size_t end = std::min(base + 8, d);
-        for (size_t j = base; j < end; j++) {
-            if (b & (1 << (j - base))) {
-                result += arr[j];
-            }
-        }
-    }
-    return result;
-}
-
 } // namespace scalar_quantizer
 } // namespace faiss
 
 #define THE_LEVEL_TO_DISPATCH SIMDLevel::ARM_NEON
 #include <faiss/impl/scalar_quantizer/sq-dispatch.h>
-
-#ifdef COMPILE_SIMD_ARM_SVE
-
-// ARM_SVE: SVE is a superset of NEON. Forward to the NEON implementation
-// until a dedicated SVE specialization is written.
-
-namespace faiss {
-namespace scalar_quantizer {
-
-// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
-template <>
-ScalarQuantizer::SQuantizer* sq_select_quantizer<SIMDLevel::ARM_SVE>(
-        QuantizerType qtype,
-        size_t d,
-        const std::vector<float>& trained) {
-    return sq_select_quantizer<SIMDLevel::ARM_NEON>(qtype, d, trained);
-}
-
-// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
-template <>
-SQDistanceComputer* sq_select_distance_computer<SIMDLevel::ARM_SVE>(
-        MetricType metric,
-        ScalarQuantizer::QuantizerType qtype,
-        size_t d,
-        const std::vector<float>& trained) {
-    return sq_select_distance_computer<SIMDLevel::ARM_NEON>(
-            metric, qtype, d, trained);
-}
-
-// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
-template <>
-InvertedListScanner* sq_select_InvertedListScanner<SIMDLevel::ARM_SVE>(
-        QuantizerType qtype,
-        MetricType mt,
-        size_t d,
-        size_t code_size,
-        const std::vector<float>& trained,
-        const Index* quantizer,
-        bool store_pairs,
-        const IDSelector* sel,
-        bool by_residual) {
-    return sq_select_InvertedListScanner<SIMDLevel::ARM_NEON>(
-            qtype,
-            mt,
-            d,
-            code_size,
-            trained,
-            quantizer,
-            store_pairs,
-            sel,
-            by_residual);
-}
-
-} // namespace scalar_quantizer
-} // namespace faiss
-
-#endif // COMPILE_SIMD_ARM_SVE
-
-#endif // COMPILE_SIMD_ARM_NEON