faiss 0.5.3 → 0.6.1

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,74 @@ 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_8bit_single_impl<SIMDLevel::ARM_SVE>(
+        size_t M,
         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;
+    constexpr size_t ksub = 1 << 8;
 
     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 +118,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 +134,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 +145,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 +158,16 @@ 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)
+// Combines 4 operations of pq_code_distance_8bit_single_impl().
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_8bit_four_impl<SIMDLevel::ARM_SVE>(
+        size_t M,
         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)
-        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 +188,10 @@ distance_four_codes_sve(
     }
 
     const float* tab = sim_table;
-
-    const size_t ksub = 1 << nbits;
+    constexpr size_t ksub = 1 << 8;
 
     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 +200,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 +344,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/pq_code_distance/rvv.cpp

@@ -0,0 +1,68 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_RISCV_RVV
+
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+// RISCV_RVV: no RVV-optimized PQ code distance exists yet. Use scalar.
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+float pq_code_distance_8bit_single_impl<SIMDLevel::RISCV_RVV>(
+        size_t M,
+        const float* sim_table,
+        const uint8_t* code) {
+    return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
+            M, 8, sim_table, code);
+}
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_8bit_four_impl<SIMDLevel::RISCV_RVV>(
+        size_t M,
+        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) {
+    PQCodeDistanceScalar<PQDecoder8>::distance_four_codes(
+            M,
+            8,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+
+} // namespace pq_code_distance
+} // namespace faiss
+
+#define THE_SIMD_LEVEL SIMDLevel::RISCV_RVV
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-rvv.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/PQDistanceComputer_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/IVFPQScanner_impl.h>
+
+#undef THE_SIMD_LEVEL
+
+#endif // COMPILE_SIMD_RISCV_RVV