faiss 0.3.2 → 0.3.3

data/vendor/faiss/faiss/impl/PolysemousTraining.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.
@@ -11,7 +11,6 @@
 #include <cstdio>
 #include <cstring>
 #include <memory>
-#include <random>
 
 #include <algorithm>
 
@@ -20,7 +19,6 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
 
 extern "C" {
 

data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ProductQuantizer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/Quantizer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.
@@ -492,40 +492,6 @@ void ResidualQuantizer::refine_beam(
  * Functions using the dot products between codebook entries
  *******************************************************************/
 
-void ResidualQuantizer::compute_codebook_tables() {
-    cent_norms.resize(total_codebook_size);
-    fvec_norms_L2sqr(
-            cent_norms.data(), codebooks.data(), d, total_codebook_size);
-    size_t cross_table_size = 0;
-    for (int m = 0; m < M; m++) {
-        size_t K = (size_t)1 << nbits[m];
-        cross_table_size += K * codebook_offsets[m];
-    }
-    codebook_cross_products.resize(cross_table_size);
-    size_t ofs = 0;
-    for (int m = 1; m < M; m++) {
-        FINTEGER ki = (size_t)1 << nbits[m];
-        FINTEGER kk = codebook_offsets[m];
-        FINTEGER di = d;
-        float zero = 0, one = 1;
-        assert(ofs + ki * kk <= cross_table_size);
-        sgemm_("Transposed",
-               "Not transposed",
-               &ki,
-               &kk,
-               &di,
-               &one,
-               codebooks.data() + d * kk,
-               &di,
-               codebooks.data(),
-               &di,
-               &zero,
-               codebook_cross_products.data() + ofs,
-               &ki);
-        ofs += ki * kk;
-    }
-}
-
 void ResidualQuantizer::refine_beam_LUT(
         size_t n,
         const float* query_norms, // size n

data/vendor/faiss/faiss/impl/ResidualQuantizer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.
@@ -112,7 +112,7 @@ struct ResidualQuantizer : AdditiveQuantizer {
 
     /** lower-level encode function
      *
-     * @param n              number of vectors to hanlde
+     * @param n              number of vectors to handle
      * @param residuals      vectors to encode, size (n, beam_size, d)
      * @param beam_size      input beam size
      * @param new_beam_size  output beam size (should be <= K * beam_size)
@@ -143,16 +143,6 @@ struct ResidualQuantizer : AdditiveQuantizer {
      * @param beam_size  if != -1, override the beam size
      */
     size_t memory_per_point(int beam_size = -1) const;
-
-    /** Cross products used in codebook tables used for beam_LUT = 1
-     */
-    void compute_codebook_tables();
-
-    /// dot products of all codebook entries with the previous codebooks
-    /// size sum(codebook_offsets[m] * 2^nbits[m], m=0..M-1)
-    std::vector<float> codebook_cross_products;
-    /// norms of all codebook entries (size total_codebook_size)
-    std::vector<float> cent_norms;
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/impl/ResultHandler.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ScalarQuantizer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/ThreadedIndex.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/code_distance/code_distance-avx2.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/code_distance/code_distance-avx512.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/code_distance/code_distance-generic.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

data/vendor/faiss/faiss/impl/code_distance/code_distance-sve.h

@@ -0,0 +1,440 @@
+/*
+ * 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
+
+#ifdef __ARM_FEATURE_SVE
+
+#include <arm_sve.h>
+
+#include <tuple>
+#include <type_traits>
+
+#include <faiss/impl/code_distance/code_distance-generic.h>
+
+namespace faiss {
+
+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);
+}
+
+static 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 float distance_single_code_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 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);
+    mm1 = svadd_u32_x(pg, mm1, offsets_0);
+    const auto collected0 = svld1_gather_u32index_f32(pg, sim_table, mm1);
+    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
+        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) {
+    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;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                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;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+        }
+
+        {
+            const auto mm1hi = svunpkhi_u16(mm1);
+            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;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                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;
+            }
+
+            m += quad_lanes;
+        }
+    }
+
+    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 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));
+
+    const auto quad_lanes = svcntw();
+
+    // 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)
+        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) {
+    if (M <= svcntw()) {
+        distance_four_codes_sve_for_small_m(
+                M,
+                sim_table,
+                code0,
+                code1,
+                code2,
+                code3,
+                result0,
+                result1,
+                result2,
+                result3);
+        return;
+    }
+
+    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);
+    auto partialSum3 = 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 mm10 = svld1_u8(pg, code0 + m);
+        const auto mm11 = svld1_u8(pg, code1 + m);
+        const auto mm12 = svld1_u8(pg, code2 + m);
+        const auto mm13 = svld1_u8(pg, code3 + m);
+        {
+            const auto mm10lo = svunpklo_u16(mm10);
+            const auto mm11lo = svunpklo_u16(mm11);
+            const auto mm12lo = svunpklo_u16(mm12);
+            const auto mm13lo = svunpklo_u16(mm13);
+            const auto pglo = svunpklo_b(pg);
+
+            {
+                const auto pglolo = svunpklo_b(pglo);
+                {
+                    const auto idx1 = svunpklo_u32(mm10lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm11lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm12lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm13lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+
+            {
+                const auto pglohi = svunpkhi_b(pglo);
+                {
+                    const auto idx1 = svunpkhi_u32(mm10lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm11lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm12lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm13lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+        }
+
+        {
+            const auto mm10hi = svunpkhi_u16(mm10);
+            const auto mm11hi = svunpkhi_u16(mm11);
+            const auto mm12hi = svunpkhi_u16(mm12);
+            const auto mm13hi = svunpkhi_u16(mm13);
+            const auto pghi = svunpkhi_b(pg);
+
+            {
+                const auto pghilo = svunpklo_b(pghi);
+                {
+                    const auto idx1 = svunpklo_u32(mm10hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm11hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm12hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm13hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+
+            {
+                const auto pghihi = svunpkhi_b(pghi);
+                {
+                    const auto idx1 = svunpkhi_u32(mm10hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm11hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm12hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm13hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+        }
+    }
+
+    result0 = svaddv_f32(svptrue_b32(), partialSum0);
+    result1 = svaddv_f32(svptrue_b32(), partialSum1);
+    result2 = svaddv_f32(svptrue_b32(), partialSum2);
+    result3 = svaddv_f32(svptrue_b32(), partialSum3);
+}
+
+} // namespace faiss
+
+#endif