faiss 0.5.2 → 0.6.0

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

@@ -5,49 +5,28 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#pragma once
-
-#ifdef __AVX512F__
+#ifdef COMPILE_SIMD_AVX512
 
 #include <immintrin.h>
 
-#include <type_traits>
-
-#include <faiss/impl/ProductQuantizer.h>
-#include <faiss/impl/code_distance/code_distance-generic.h>
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
 
 namespace faiss {
+namespace pq_code_distance {
 
 // According to experiments, the AVX-512 version may be SLOWER than
-//   the AVX2 version, which is somewhat unexpected.
-// This version is not used for now, but it may be used later.
+// the AVX2 version, which is somewhat unexpected.
+// This version is kept for completeness.
 //
 // TODO: test for AMD CPUs.
 
-template <typename PQDecoderT>
-typename std::enable_if<!std::is_same<PQDecoderT, PQDecoder8>::value, float>::
-        type inline distance_single_code_avx512(
-                // number of subquantizers
-                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);
-}
-
-template <typename PQDecoderT>
-typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, float>::
-        type inline distance_single_code_avx512(
-                // number of subquantizers
-                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* code0) {
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+float pq_code_distance_single_impl<SIMDLevel::AVX512>(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code0) {
     float result0 = 0;
     constexpr size_t ksub = 1 << 8;
 
@@ -59,49 +38,37 @@ typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, float>::
     const float* tab = sim_table;
 
     if (pqM16 > 0) {
-        // process 16 values per loop
         const __m512i vksub = _mm512_set1_epi32(ksub);
         __m512i offsets_0 = _mm512_setr_epi32(
                 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
         offsets_0 = _mm512_mullo_epi32(offsets_0, vksub);
 
-        // accumulators of partial sums
         __m512 partialSums[N];
         for (intptr_t j = 0; j < N; j++) {
             partialSums[j] = _mm512_setzero_ps();
         }
 
-        // loop
+        // Process 16 values per loop iteration.
         for (m = 0; m < pqM16 * 16; m += 16) {
-            // load 16 uint8 values
             __m128i mm1[N];
             mm1[0] = _mm_loadu_si128((const __m128i_u*)(code0 + m));
 
-            // process first 8 codes
             for (intptr_t j = 0; j < N; j++) {
                 const __m512i idx1 = _mm512_cvtepu8_epi32(mm1[j]);
-
-                // add offsets
                 const __m512i indices_to_read_from =
                         _mm512_add_epi32(idx1, offsets_0);
-
-                // gather 16 values, similar to 16 operations of tab[idx]
                 __m512 collected = _mm512_i32gather_ps(
                         indices_to_read_from, tab, sizeof(float));
-
-                // collect partial sums
                 partialSums[j] = _mm512_add_ps(partialSums[j], collected);
             }
             tab += ksub * 16;
         }
 
-        // horizontal sum for partialSum
         result0 += _mm512_reduce_add_ps(partialSums[0]);
     }
 
-    //
+    // Process leftovers.
     if (m < M) {
-        // process leftovers
         PQDecoder8 decoder0(code0 + m, nbits);
         for (; m < M; m++) {
             result0 += tab[decoder0.decode()];
@@ -112,56 +79,17 @@ typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, float>::
     return result0;
 }
 
-template <typename PQDecoderT>
-typename std::enable_if<!std::is_same<PQDecoderT, PQDecoder8>::value, void>::
-        type
-        distance_four_codes_avx512(
-                // number of subquantizers
-                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);
-}
-
-// Combines 4 operations of distance_single_code()
-template <typename PQDecoderT>
-typename std::enable_if<std::is_same<PQDecoderT, PQDecoder8>::value, void>::type
-distance_four_codes_avx512(
-        // number of subquantizers
-        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::AVX512>(
+        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,
@@ -180,55 +108,43 @@ distance_four_codes_avx512(
     const float* tab = sim_table;
 
     if (pqM16 > 0) {
-        // process 16 values per loop
         const __m512i vksub = _mm512_set1_epi32(ksub);
         __m512i offsets_0 = _mm512_setr_epi32(
                 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
         offsets_0 = _mm512_mullo_epi32(offsets_0, vksub);
 
-        // accumulators of partial sums
         __m512 partialSums[N];
         for (intptr_t j = 0; j < N; j++) {
             partialSums[j] = _mm512_setzero_ps();
         }
 
-        // loop
+        // Process 16 values per loop iteration.
         for (m = 0; m < pqM16 * 16; m += 16) {
-            // load 16 uint8 values
             __m128i mm1[N];
             mm1[0] = _mm_loadu_si128((const __m128i_u*)(code0 + m));
             mm1[1] = _mm_loadu_si128((const __m128i_u*)(code1 + m));
             mm1[2] = _mm_loadu_si128((const __m128i_u*)(code2 + m));
             mm1[3] = _mm_loadu_si128((const __m128i_u*)(code3 + m));
 
-            // process first 8 codes
             for (intptr_t j = 0; j < N; j++) {
                 const __m512i idx1 = _mm512_cvtepu8_epi32(mm1[j]);
-
-                // add offsets
                 const __m512i indices_to_read_from =
                         _mm512_add_epi32(idx1, offsets_0);
-
-                // gather 16 values, similar to 16 operations of tab[idx]
                 __m512 collected = _mm512_i32gather_ps(
                         indices_to_read_from, tab, sizeof(float));
-
-                // collect partial sums
                 partialSums[j] = _mm512_add_ps(partialSums[j], collected);
             }
             tab += ksub * 16;
         }
 
-        // horizontal sum for partialSum
         result0 += _mm512_reduce_add_ps(partialSums[0]);
         result1 += _mm512_reduce_add_ps(partialSums[1]);
         result2 += _mm512_reduce_add_ps(partialSums[2]);
         result3 += _mm512_reduce_add_ps(partialSums[3]);
     }
 
-    //
+    // Process leftovers.
     if (m < M) {
-        // process leftovers
         PQDecoder8 decoder0(code0 + m, nbits);
         PQDecoder8 decoder1(code1 + m, nbits);
         PQDecoder8 decoder2(code2 + m, nbits);
@@ -243,6 +159,51 @@ distance_four_codes_avx512(
     }
 }
 
+#ifdef COMPILE_SIMD_AVX512_SPR
+// AVX512_SPR: Sapphire Rapids is a superset of AVX512. Reuse the
+// AVX512 implementation until a dedicated SPR specialization is written.
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+float pq_code_distance_single_impl<SIMDLevel::AVX512_SPR>(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code) {
+    return pq_code_distance_single_impl<SIMDLevel::AVX512>(
+            M, nbits, sim_table, code);
+}
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_four_impl<SIMDLevel::AVX512_SPR>(
+        size_t M,
+        size_t nbits,
+        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) {
+    pq_code_distance_four_impl<SIMDLevel::AVX512>(
+            M,
+            nbits,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+#endif // COMPILE_SIMD_AVX512_SPR
+
+} // namespace pq_code_distance
 } // namespace faiss
 
-#endif
+#endif // COMPILE_SIMD_AVX512

data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp

@@ -0,0 +1,141 @@
+/*
+ * 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.
+ */
+
+// This TU provides:
+// 1. _impl specializations for NONE (and ARM_NEON), using scalar code.
+// 2. Non-templated PQ code distance dispatch wrappers
+//    (pq_code_distance_single, pq_code_distance_four) declared in
+//    pq_code_distance.h. These use DISPATCH_SIMDLevel to route to the
+//    best available SIMD implementation via pq_code_distance_*_impl
+//    function template specializations defined in the per-SIMD .cpp files.
+
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+// NONE: use scalar directly.
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+float pq_code_distance_single_impl<SIMDLevel::NONE>(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code) {
+    return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
+            M, nbits, sim_table, code);
+}
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_four_impl<SIMDLevel::NONE>(
+        size_t M,
+        size_t nbits,
+        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,
+            nbits,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+
+#ifdef COMPILE_SIMD_ARM_NEON
+// ARM_NEON: No NEON-optimized PQ code distance exists. Use scalar.
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+float pq_code_distance_single_impl<SIMDLevel::ARM_NEON>(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code) {
+    return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
+            M, nbits, sim_table, code);
+}
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_four_impl<SIMDLevel::ARM_NEON>(
+        size_t M,
+        size_t nbits,
+        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,
+            nbits,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+#endif // COMPILE_SIMD_ARM_NEON
+
+float pq_code_distance_single(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code) {
+    DISPATCH_SIMDLevel(pq_code_distance_single_impl, M, nbits, sim_table, code);
+}
+
+void pq_code_distance_four(
+        size_t M,
+        size_t nbits,
+        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) {
+    DISPATCH_SIMDLevel(
+            pq_code_distance_four_impl,
+            M,
+            nbits,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+
+} // namespace pq_code_distance
+} // namespace faiss

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

@@ -0,0 +1,23 @@
+/*
+ * 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
+
+/**
+ * @file pq_code_distance-inl.h
+ * @brief Private header for PQ code distance SIMD implementations.
+ *
+ * This is a PRIVATE header — do not include in public APIs or user code.
+ * Only faiss internal .cpp files (the per-SIMD implementation files and
+ * pq_code_distance-generic.cpp) should include this header.
+ *
+ * This header re-exports the public API (pq_code_distance.h) plus the
+ * simd_dispatch.h machinery needed by the implementation files.
+ */
+
+#include <faiss/impl/simd_dispatch.h>
+#include <faiss/utils/pq_code_distance.h>