faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/impl/{code_distance/code_distance-sve.h → pq_code_distance/pq_code_distance-sve.cpp}

@@ -5,64 +5,34 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#pragma once
-
-#ifdef __ARM_FEATURE_SVE
+#ifdef COMPILE_SIMD_ARM_SVE
 
 #include <arm_sve.h>
 
-#include <tuple>
-#include <type_traits>
-
-#include <faiss/impl/ProductQuantizer.h>
-#include <faiss/impl/code_distance/code_distance-generic.h>
-
-namespace faiss {
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
 
-template <typename PQDecoderT>
-std::enable_if_t<!std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_single_code_sve(
-        // the product quantizer
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        const uint8_t* code) {
-    // default implementation
-    return distance_single_code_generic<PQDecoderT>(M, nbits, sim_table, code);
-}
+namespace {
 
-static inline void distance_codes_kernel(
+inline void distance_codes_kernel(
         svbool_t pg,
         svuint32_t idx1,
         svuint32_t offsets_0,
         const float* tab,
         svfloat32_t& partialSum) {
-    // add offset
     const auto indices_to_read_from = svadd_u32_x(pg, idx1, offsets_0);
-
-    // gather values, similar to some operations of tab[index]
     const auto collected =
             svld1_gather_u32index_f32(pg, tab, indices_to_read_from);
-
-    // collect partial sum
     partialSum = svadd_f32_m(pg, partialSum, collected);
 }
 
-static inline float distance_single_code_sve_for_small_m(
-        // the product quantizer
+inline float distance_single_code_sve_for_small_m(
         const size_t M,
-        // precomputed distances, layout (M, ksub)
         const float* sim_table,
-        // codes
         const uint8_t* __restrict code) {
     constexpr size_t nbits = 8u;
-
     const size_t ksub = 1 << nbits;
 
     const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
-
-    // loop
     const auto pg = svwhilelt_b32_u64(0, M);
 
     auto mm1 = svld1ub_u32(pg, code);
@@ -71,44 +41,75 @@ static inline float distance_single_code_sve_for_small_m(
     return svaddv_f32(pg, collected0);
 }
 
-template <typename PQDecoderT>
-std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_single_code_sve(
-        // the product quantizer
+inline void distance_four_codes_sve_for_small_m(
         const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    constexpr size_t nbits = 8u;
+    const size_t ksub = 1 << nbits;
+
+    const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
+    const auto pg = svwhilelt_b32_u64(0, M);
+
+    auto mm10 = svld1ub_u32(pg, code0);
+    auto mm11 = svld1ub_u32(pg, code1);
+    auto mm12 = svld1ub_u32(pg, code2);
+    auto mm13 = svld1ub_u32(pg, code3);
+    mm10 = svadd_u32_x(pg, mm10, offsets_0);
+    mm11 = svadd_u32_x(pg, mm11, offsets_0);
+    mm12 = svadd_u32_x(pg, mm12, offsets_0);
+    mm13 = svadd_u32_x(pg, mm13, offsets_0);
+    const auto collected0 = svld1_gather_u32index_f32(pg, sim_table, mm10);
+    const auto collected1 = svld1_gather_u32index_f32(pg, sim_table, mm11);
+    const auto collected2 = svld1_gather_u32index_f32(pg, sim_table, mm12);
+    const auto collected3 = svld1_gather_u32index_f32(pg, sim_table, mm13);
+    result0 = svaddv_f32(pg, collected0);
+    result1 = svaddv_f32(pg, collected1);
+    result2 = svaddv_f32(pg, collected2);
+    result3 = svaddv_f32(pg, collected3);
+}
+
+} // namespace
+
+namespace faiss {
+namespace pq_code_distance {
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+float pq_code_distance_single_impl<SIMDLevel::ARM_SVE>(
+        size_t M,
+        size_t nbits,
         const float* sim_table,
         const uint8_t* code) {
     if (M <= svcntw())
         return distance_single_code_sve_for_small_m(M, sim_table, code);
 
     const float* tab = sim_table;
-
     const size_t ksub = 1 << nbits;
 
     const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
-
-    // accumulators of partial sums
     auto partialSum = svdup_n_f32(0.f);
 
     const auto lanes = svcntb();
     const auto quad_lanes = lanes / 4;
 
-    // loop
     for (std::size_t m = 0; m < M;) {
         const auto pg = svwhilelt_b8_u64(m, M);
-
         const auto mm1 = svld1_u8(pg, code + m);
         {
             const auto mm1lo = svunpklo_u16(mm1);
             const auto pglo = svunpklo_b(pg);
 
             {
-                // convert uint8 values to uint32 values
                 const auto idx1 = svunpklo_u32(mm1lo);
                 const auto pglolo = svunpklo_b(pglo);
-
                 distance_codes_kernel(pglolo, idx1, offsets_0, tab, partialSum);
                 tab += ksub * quad_lanes;
             }
@@ -118,10 +119,8 @@ std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_
                 break;
 
             {
-                // convert uint8 values to uint32 values
                 const auto idx1 = svunpkhi_u32(mm1lo);
                 const auto pglohi = svunpkhi_b(pglo);
-
                 distance_codes_kernel(pglohi, idx1, offsets_0, tab, partialSum);
                 tab += ksub * quad_lanes;
             }
@@ -136,10 +135,8 @@ std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_
             const auto pghi = svunpkhi_b(pg);
 
             {
-                // convert uint8 values to uint32 values
                 const auto idx1 = svunpklo_u32(mm1hi);
                 const auto pghilo = svunpklo_b(pghi);
-
                 distance_codes_kernel(pghilo, idx1, offsets_0, tab, partialSum);
                 tab += ksub * quad_lanes;
             }
@@ -149,10 +146,8 @@ std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_
                 break;
 
             {
-                // convert uint8 values to uint32 values
                 const auto idx1 = svunpkhi_u32(mm1hi);
                 const auto pghihi = svunpkhi_b(pghi);
-
                 distance_codes_kernel(pghihi, idx1, offsets_0, tab, partialSum);
                 tab += ksub * quad_lanes;
             }
@@ -164,97 +159,17 @@ std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_
     return svaddv_f32(svptrue_b32(), partialSum);
 }
 
-template <typename PQDecoderT>
-std::enable_if_t<!std::is_same_v<PQDecoderT, PQDecoder8>, void>
-distance_four_codes_sve(
-        // the product quantizer
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // codes
-        const uint8_t* __restrict code0,
-        const uint8_t* __restrict code1,
-        const uint8_t* __restrict code2,
-        const uint8_t* __restrict code3,
-        // computed distances
-        float& result0,
-        float& result1,
-        float& result2,
-        float& result3) {
-    distance_four_codes_generic<PQDecoderT>(
-            M,
-            nbits,
-            sim_table,
-            code0,
-            code1,
-            code2,
-            code3,
-            result0,
-            result1,
-            result2,
-            result3);
-}
-
-static inline void distance_four_codes_sve_for_small_m(
-        // the product quantizer
-        const size_t M,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // codes
-        const uint8_t* __restrict code0,
-        const uint8_t* __restrict code1,
-        const uint8_t* __restrict code2,
-        const uint8_t* __restrict code3,
-        // computed distances
-        float& result0,
-        float& result1,
-        float& result2,
-        float& result3) {
-    constexpr size_t nbits = 8u;
-
-    const size_t ksub = 1 << nbits;
-
-    const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
-
-    // loop
-    const auto pg = svwhilelt_b32_u64(0, M);
-
-    auto mm10 = svld1ub_u32(pg, code0);
-    auto mm11 = svld1ub_u32(pg, code1);
-    auto mm12 = svld1ub_u32(pg, code2);
-    auto mm13 = svld1ub_u32(pg, code3);
-    mm10 = svadd_u32_x(pg, mm10, offsets_0);
-    mm11 = svadd_u32_x(pg, mm11, offsets_0);
-    mm12 = svadd_u32_x(pg, mm12, offsets_0);
-    mm13 = svadd_u32_x(pg, mm13, offsets_0);
-    const auto collected0 = svld1_gather_u32index_f32(pg, sim_table, mm10);
-    const auto collected1 = svld1_gather_u32index_f32(pg, sim_table, mm11);
-    const auto collected2 = svld1_gather_u32index_f32(pg, sim_table, mm12);
-    const auto collected3 = svld1_gather_u32index_f32(pg, sim_table, mm13);
-    result0 = svaddv_f32(pg, collected0);
-    result1 = svaddv_f32(pg, collected1);
-    result2 = svaddv_f32(pg, collected2);
-    result3 = svaddv_f32(pg, collected3);
-}
-
-// Combines 4 operations of distance_single_code()
-template <typename PQDecoderT>
-std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, void>
-distance_four_codes_sve(
-        // the product quantizer
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
+// Combines 4 operations of pq_code_distance_single_impl().
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_four_impl<SIMDLevel::ARM_SVE>(
+        size_t M,
+        size_t nbits,
         const float* sim_table,
-        // codes
         const uint8_t* __restrict code0,
         const uint8_t* __restrict code1,
         const uint8_t* __restrict code2,
         const uint8_t* __restrict code3,
-        // computed distances
         float& result0,
         float& result1,
         float& result2,
@@ -275,12 +190,10 @@ distance_four_codes_sve(
     }
 
     const float* tab = sim_table;
-
     const size_t ksub = 1 << nbits;
 
     const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
 
-    // accumulators of partial sums
     auto partialSum0 = svdup_n_f32(0.f);
     auto partialSum1 = svdup_n_f32(0.f);
     auto partialSum2 = svdup_n_f32(0.f);
@@ -289,7 +202,6 @@ distance_four_codes_sve(
     const auto lanes = svcntb();
     const auto quad_lanes = lanes / 4;
 
-    // loop
     for (std::size_t m = 0; m < M;) {
         const auto pg = svwhilelt_b8_u64(m, M);
 
@@ -434,6 +346,7 @@ distance_four_codes_sve(
     result3 = svaddv_f32(svptrue_b32(), partialSum3);
 }
 
+} // namespace pq_code_distance
 } // namespace faiss
 
-#endif
+#endif // COMPILE_SIMD_ARM_SVE

data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp

@@ -10,6 +10,7 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResidualQuantizer.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/simdlib.h>
@@ -877,12 +878,14 @@ void compute_codes_add_centroids_mp_lut0(
         pool.norms.resize(n);
         // recover the norms of reconstruction as
         // || original_vector - residual ||^2
-        for (size_t i = 0; i < n; i++) {
-            pool.norms[i] = fvec_L2sqr(
-                    x + i * rq.d,
-                    pool.residuals.data() + i * rq.max_beam_size * rq.d,
-                    rq.d);
-        }
+        with_simd_level([&]<SIMDLevel SL>() {
+            for (size_t i = 0; i < n; i++) {
+                pool.norms[i] = fvec_L2sqr<SL>(
+                        x + i * rq.d,
+                        pool.residuals.data() + i * rq.max_beam_size * rq.d,
+                        rq.d);
+            }
+        });
     }
 
     // pack only the first code of the beam

data/vendor/faiss/faiss/impl/scalar_quantizer/codecs.h

@@ -0,0 +1,121 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/utils/simd_levels.h>
+#include <faiss/utils/simdlib.h>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+/*******************************************************************
+ * Codec: converts between values in [0, 1] and an index in a code
+ * array. The "i" parameter is the vector component index (not byte
+ * index).
+ */
+
+template <SIMDLevel SL>
+struct Codec8bit {};
+
+template <>
+struct Codec8bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE void encode_component(
+            float x,
+            uint8_t* code,
+            size_t i) {
+        code[i] = (int)(255 * x);
+    }
+
+    static FAISS_ALWAYS_INLINE float decode_component(
+            const uint8_t* code,
+            size_t i) {
+        return (code[i] + 0.5f) / 255.0f;
+    }
+};
+
+template <SIMDLevel SL>
+struct Codec4bit {};
+
+template <>
+struct Codec4bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE void encode_component(
+            float x,
+            uint8_t* code,
+            size_t i) {
+        code[i / 2] |= (int)(x * 15.0) << ((i & 1) << 2);
+    }
+
+    static FAISS_ALWAYS_INLINE float decode_component(
+            const uint8_t* code,
+            size_t i) {
+        return (((code[i / 2] >> ((i & 1) << 2)) & 0xf) + 0.5f) / 15.0f;
+    }
+};
+
+template <SIMDLevel SL>
+struct Codec6bit {};
+
+template <>
+struct Codec6bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE void encode_component(
+            float x,
+            uint8_t* code,
+            size_t i) {
+        int bits = (int)(x * 63.0);
+        code += (i >> 2) * 3;
+        switch (i & 3) {
+            case 0:
+                code[0] |= bits;
+                break;
+            case 1:
+                code[0] |= bits << 6;
+                code[1] |= bits >> 2;
+                break;
+            case 2:
+                code[1] |= bits << 4;
+                code[2] |= bits >> 4;
+                break;
+            case 3:
+                code[2] |= bits << 2;
+                break;
+            default:
+                break;
+        }
+    }
+
+    static FAISS_ALWAYS_INLINE float decode_component(
+            const uint8_t* code,
+            size_t i) {
+        uint8_t bits = 0;
+        code += (i >> 2) * 3;
+        switch (i & 3) {
+            case 0:
+                bits = code[0] & 0x3f;
+                break;
+            case 1:
+                bits = code[0] >> 6;
+                bits |= (code[1] & 0xf) << 2;
+                break;
+            case 2:
+                bits = code[1] >> 4;
+                bits |= (code[2] & 3) << 4;
+                break;
+            case 3:
+                bits = code[2] >> 2;
+                break;
+            default:
+                break;
+        }
+        return (bits + 0.5f) / 63.0f;
+    }
+};
+
+} // namespace scalar_quantizer
+} // namespace faiss

data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h

@@ -0,0 +1,136 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/utils/simd_levels.h>
+#include <faiss/utils/simdlib.h>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+using SQDistanceComputer = ScalarQuantizer::SQDistanceComputer;
+
+/*******************************************************************
+ * DistanceComputer: combines a similarity and a quantizer to do
+ * code-to-vector or code-to-code comparisons
+ *******************************************************************/
+
+template <class Quantizer, class Similarity, SIMDLevel SL>
+struct DCTemplate : SQDistanceComputer {};
+
+template <class Quantizer, class Similarity>
+struct DCTemplate<Quantizer, Similarity, SIMDLevel::NONE> : SQDistanceComputer {
+    using Sim = Similarity;
+
+    Quantizer quant;
+
+    DCTemplate(size_t d, const std::vector<float>& trained)
+            : quant(d, trained) {}
+
+    float compute_distance(const float* x, const uint8_t* code) const {
+        Similarity sim(x);
+        sim.begin();
+        for (size_t i = 0; i < quant.d; i++) {
+            float xi = quant.reconstruct_component(code, i);
+            sim.add_component(xi);
+        }
+        return sim.result();
+    }
+
+    float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        Similarity sim(nullptr);
+        sim.begin();
+        for (size_t i = 0; i < quant.d; i++) {
+            float x1 = quant.reconstruct_component(code1, i);
+            float x2 = quant.reconstruct_component(code2, i);
+            sim.add_component_2(x1, x2);
+        }
+        return sim.result();
+    }
+
+    void set_query(const float* x) final {
+        q = x;
+    }
+
+    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_distance(q, code);
+    }
+};
+
+/*******************************************************************
+ * DistanceComputerByte: computes distances in the integer domain
+ *******************************************************************/
+
+template <class Similarity, SIMDLevel SL>
+struct DistanceComputerByte : SQDistanceComputer {};
+
+template <class Similarity>
+struct DistanceComputerByte<Similarity, SIMDLevel::NONE> : SQDistanceComputer {
+    using Sim = Similarity;
+
+    int d;
+    std::vector<uint8_t> tmp;
+
+    DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
+
+    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        int accu = 0;
+        for (int i = 0; i < d; i++) {
+            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
+                accu += int(code1[i]) * code2[i];
+            } else {
+                int diff = int(code1[i]) - code2[i];
+                accu += diff * diff;
+            }
+        }
+        return accu;
+    }
+
+    void set_query(const float* x) final {
+        for (int i = 0; i < d; i++) {
+            tmp[i] = int(x[i]);
+        }
+    }
+
+    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);
+    }
+};
+
+/*******************************************************************
+ * Selection function
+ *******************************************************************/
+
+template <SIMDLevel SL>
+SQDistanceComputer* sq_select_distance_computer(
+        MetricType metric,
+        ScalarQuantizer::QuantizerType qtype,
+        size_t d,
+        const std::vector<float>& trained);
+
+} // namespace scalar_quantizer
+} // namespace faiss