faiss 0.5.3 → 0.6.1

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

@@ -7,47 +7,25 @@
 
 #pragma once
 
-#ifdef __AVX512F__
-
 #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_8bit_single_impl<SIMDLevel::AVX512>(
+        size_t M,
+        const float* sim_table,
+        const uint8_t* code0) {
     float result0 = 0;
     constexpr size_t ksub = 1 << 8;
 
@@ -59,50 +37,38 @@ 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);
+        PQDecoder8 decoder0(code0 + m, 8);
         for (; m < M; m++) {
             result0 += tab[decoder0.decode()];
             tab += ksub;
@@ -112,56 +78,16 @@ 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_8bit_single_impl().
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_8bit_four_impl<SIMDLevel::AVX512>(
+        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,
@@ -180,59 +106,47 @@ 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);
-        PQDecoder8 decoder3(code3 + m, nbits);
+        PQDecoder8 decoder0(code0 + m, 8);
+        PQDecoder8 decoder1(code1 + m, 8);
+        PQDecoder8 decoder2(code2 + m, 8);
+        PQDecoder8 decoder3(code3 + m, 8);
         for (; m < M; m++) {
             result0 += tab[decoder0.decode()];
             result1 += tab[decoder1.decode()];
@@ -243,6 +157,5 @@ distance_four_codes_avx512(
     }
 }
 
+} // namespace pq_code_distance
 } // namespace faiss
-
-#endif

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

@@ -0,0 +1,59 @@
+/*
+ * 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 non-templated PQ code distance dispatch wrappers
+// (pq_code_distance_8bit_single, pq_code_distance_8bit_four) declared
+// in pq_code_distance-inl.h. These use with_simd_level to route to the
+// best available SIMD implementation via pq_code_distance_8bit_*_impl
+// function template specializations.
+//
+// The NONE and ARM_NEON _impl specializations are defined inline in
+// pq_code_distance-generic.h (included transitively). The AVX2, AVX512,
+// and ARM_SVE specializations are in their respective per-SIMD files.
+
+#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+float pq_code_distance_8bit_single(
+        size_t M,
+        const float* sim_table,
+        const uint8_t* code) {
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return pq_code_distance_8bit_single_impl<SL>(M, sim_table, code);
+    });
+}
+
+void pq_code_distance_8bit_four(
+        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) {
+    with_simd_level([&]<SIMDLevel SL>() {
+        pq_code_distance_8bit_four_impl<SL>(
+                M,
+                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-generic.h

@@ -0,0 +1,96 @@
+/*
+ * 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/pq_code_distance/pq_code_distance-inl.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+// NONE: use scalar directly.
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+inline float pq_code_distance_8bit_single_impl<SIMDLevel::NONE>(
+        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 <>
+inline void pq_code_distance_8bit_four_impl<SIMDLevel::NONE>(
+        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);
+}
+
+#ifdef COMPILE_SIMD_ARM_NEON
+// ARM_NEON: No NEON-optimized PQ code distance exists. Use scalar.
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+inline float pq_code_distance_8bit_single_impl<SIMDLevel::ARM_NEON>(
+        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 <>
+inline void pq_code_distance_8bit_four_impl<SIMDLevel::ARM_NEON>(
+        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);
+}
+#endif // COMPILE_SIMD_ARM_NEON
+
+} // namespace pq_code_distance
+} // namespace faiss

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

@@ -0,0 +1,256 @@
+/*
+ * 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 PQ code distance SIMD-dispatched 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.
+ */
+
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include <faiss/impl/ProductQuantizer.h>
+#include <faiss/impl/platform_macros.h>
+#include <faiss/impl/simd_dispatch.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+/*********************************************************************
+ * PQCodeDistance — SIMD-dispatched PQ code distance
+ *
+ * Computes the distance from a PQ-encoded vector to a query vector,
+ * given a precomputed table of sub-distances (one per subquantizer
+ * per centroid). Originally extracted from IndexIVFPQ.cpp.
+ *
+ * DESIGN:
+ *
+ * PQCodeDistance<PQDecoderT, SL> computes PQ code distances at a given
+ * SIMD level. The dispatch site (IndexIVFPQ.cpp, IndexPQ.cpp) uses
+ * with_simd_level to select SL at runtime, which instantiates
+ * PQCodeDistance for ALL decoder types (PQDecoder8, PQDecoder16,
+ * PQDecoderGeneric) at the chosen level.
+ *
+ * Only PQDecoder8 has SIMD-optimized implementations (AVX2, AVX512,
+ * ARM_SVE). The other decoders always use scalar code — their decode()
+ * method is inherently sequential, so SIMD doesn't help.
+ *
+ * The primary template is always complete (no forward declarations
+ * needed). For PQDecoder8, it delegates to _impl dispatch bridge
+ * functions whose specializations are defined in per-SIMD .cpp files
+ * and resolved at link time. For other decoders, it uses scalar.
+ *
+ * ADDING A NEW SIMD LEVEL:
+ *
+ *   1. Add the level to SIMDLevel enum (simd_levels.h)
+ *   2. Add dispatch_config entry (simd_dispatch.bzl)
+ *   3. Define pq_code_distance_8bit_single_impl<NEW_LEVEL> and
+ *      pq_code_distance_8bit_four_impl<NEW_LEVEL> specializations in a
+ *      new .cpp file compiled with appropriate SIMD flags
+ *   4. Add the .cpp to the build (CMakeLists.txt, xplat.bzl)
+ *********************************************************************/
+
+/// Scalar PQ code distance implementation.
+/// Templated only on decoder type, independent of SIMD level.
+/// Used directly by non-PQDecoder8 decoders (PQDecoder16,
+/// PQDecoderGeneric) and as fallback for PQDecoder8 at NONE/NEON.
+template <typename PQDecoderT>
+struct PQCodeDistanceScalar {
+    using PQDecoder = PQDecoderT;
+
+    static float distance_single_code(
+            // number of subquantizers
+            size_t M,
+            size_t nbits,
+            // precomputed distances, layout (M, ksub)
+            const float* sim_table,
+            const uint8_t* code) {
+        PQDecoderT decoder(code, nbits);
+        const size_t ksub = 1 << nbits;
+
+        const float* tab = sim_table;
+        float result = 0;
+
+        for (size_t m = 0; m < M; m++) {
+            result += tab[decoder.decode()];
+            tab += ksub;
+        }
+
+        return result;
+    }
+
+    static void distance_four_codes(
+            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) {
+        PQDecoderT decoder0(code0, nbits);
+        PQDecoderT decoder1(code1, nbits);
+        PQDecoderT decoder2(code2, nbits);
+        PQDecoderT decoder3(code3, nbits);
+        const size_t ksub = 1 << nbits;
+
+        const float* tab = sim_table;
+        result0 = 0;
+        result1 = 0;
+        result2 = 0;
+        result3 = 0;
+
+        for (size_t m = 0; m < M; m++) {
+            result0 += tab[decoder0.decode()];
+            result1 += tab[decoder1.decode()];
+            result2 += tab[decoder2.decode()];
+            result3 += tab[decoder3.decode()];
+            tab += ksub;
+        }
+    }
+};
+
+/*********************************************************************
+ * Dispatch bridge — function templates for PQDecoder8 SIMD dispatch.
+ *
+ * Primary declarations only; specializations are defined in per-SIMD
+ * .cpp files (AVX2, AVX512, ARM_SVE) and pq_code_distance-generic.cpp
+ * (NONE, ARM_NEON). Same pattern as fvec_L2sqr et al. in distances.h.
+ *********************************************************************/
+
+template <SIMDLevel SL>
+float pq_code_distance_8bit_single_impl(
+        size_t M,
+        const float* sim_table,
+        const uint8_t* code);
+
+template <SIMDLevel SL>
+void pq_code_distance_8bit_four_impl(
+        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);
+
+/// Primary template — always complete.
+/// For PQDecoder8, delegates to _impl dispatch bridges (resolved at
+/// link time to per-SIMD implementations). For other decoders, uses
+/// scalar — their sequential decode() methods don't benefit from SIMD.
+template <typename PQDecoderT, SIMDLevel SL>
+struct PQCodeDistance {
+    using PQDecoder = PQDecoderT;
+    static constexpr SIMDLevel simd_level = SL;
+
+    static float distance_single_code(
+            size_t M,
+            size_t nbits,
+            const float* sim_table,
+            const uint8_t* code) {
+        if constexpr (std::is_same_v<PQDecoderT, PQDecoder8>) {
+            return pq_code_distance_8bit_single_impl<SL>(M, sim_table, code);
+        } else {
+            return PQCodeDistanceScalar<PQDecoderT>::distance_single_code(
+                    M, nbits, sim_table, code);
+        }
+    }
+
+    static void distance_four_codes(
+            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) {
+        if constexpr (std::is_same_v<PQDecoderT, PQDecoder8>) {
+            pq_code_distance_8bit_four_impl<SL>(
+                    M,
+                    sim_table,
+                    code0,
+                    code1,
+                    code2,
+                    code3,
+                    result0,
+                    result1,
+                    result2,
+                    result3);
+        } else {
+            PQCodeDistanceScalar<PQDecoderT>::distance_four_codes(
+                    M,
+                    nbits,
+                    sim_table,
+                    code0,
+                    code1,
+                    code2,
+                    code3,
+                    result0,
+                    result1,
+                    result2,
+                    result3);
+        }
+    }
+};
+
+/*********************************************************************
+ * Non-templated PQ code distance dispatch (PQDecoder8 only).
+ *
+ * These follow the same pattern as distances.h: the caller does not
+ * name a SIMDLevel. Internally they dispatch via with_simd_level
+ * to the best available SIMD implementation (DD: runtime detection,
+ * static: compile-time selection). Definitions are in
+ * pq_code_distance-generic.cpp.
+ *********************************************************************/
+
+/// Compute PQ distance for a single code, dispatching to the best
+/// available SIMD level.
+FAISS_API float pq_code_distance_8bit_single(
+        size_t M,
+        const float* sim_table,
+        const uint8_t* code);
+
+/// Compute PQ distances for four codes simultaneously, dispatching
+/// to the best available SIMD level.
+FAISS_API void pq_code_distance_8bit_four(
+        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);
+
+} // namespace pq_code_distance
+
+// Re-export public API into namespace faiss for convenience
+using pq_code_distance::pq_code_distance_8bit_four;
+using pq_code_distance::pq_code_distance_8bit_single;
+using pq_code_distance::PQCodeDistance;
+using pq_code_distance::PQCodeDistanceScalar;
+
+} // namespace faiss