faiss 0.6.0 → 0.6.2

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

@@ -5,114 +5,52 @@
  * 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.
+// 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-inl.h>
+#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
+
+#define THE_SIMD_LEVEL SIMDLevel::NONE
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/pq_scan_impl.h>
+#undef THE_SIMD_LEVEL
 
 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>(
+void pq_scan_8bit(
         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);
+        const float* dis_table,
+        const uint8_t* codes,
+        size_t ncodes,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids,
+        bool max_heap) {
+    with_simd_level([&]<SIMDLevel SL>() {
+        pq_scan_8bit_impl<SL>(
+                M, dis_table, codes, ncodes, k, heap_dis, heap_ids, max_heap);
+    });
 }
-#endif // COMPILE_SIMD_ARM_NEON
 
-float pq_code_distance_single(
+float pq_code_distance_8bit_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);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return pq_code_distance_8bit_single_impl<SL>(M, sim_table, code);
+    });
 }
 
-void pq_code_distance_four(
+void pq_code_distance_8bit_four(
         size_t M,
-        size_t nbits,
         const float* sim_table,
         const uint8_t* __restrict code0,
         const uint8_t* __restrict code1,
@@ -122,19 +60,19 @@ void pq_code_distance_four(
         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);
+    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

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

@@ -9,15 +9,284 @@
 
 /**
  * @file pq_code_distance-inl.h
- * @brief Private header for PQ code distance SIMD implementations.
+ * @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.
- *
- * This header re-exports the public API (pq_code_distance.h) plus the
- * simd_dispatch.h machinery needed by the implementation files.
  */
 
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include <faiss/impl/ProductQuantizer.h>
+#include <faiss/impl/platform_macros.h>
 #include <faiss/impl/simd_dispatch.h>
-#include <faiss/utils/pq_code_distance.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);
+
+/*********************************************************************
+ * Standalone PQ scan — SIMD-dispatched full-index scan.
+ *
+ * Scans all ncodes PQ codes against a precomputed distance table,
+ * maintaining a k-nearest-neighbor heap. Uses the SIMD PQ distance
+ * kernels (AVX2 gathers, etc.) for the inner loop, with the SIMD
+ * gathers inlined into the scan loop in each per-SIMD TU.
+ *
+ * Definitions are in pq_scan_impl.h (per-SIMD TUs) and
+ * pq_code_distance-generic.cpp (dispatch wrapper).
+ *********************************************************************/
+
+template <SIMDLevel SL>
+void pq_scan_8bit_impl(
+        size_t M,
+        const float* dis_table,
+        const uint8_t* codes,
+        size_t ncodes,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids,
+        bool max_heap);
+
+/// Scan all ncodes 8-bit PQ codes, dispatching to the best SIMD level.
+/// max_heap=true for L2 (CMax), false for IP (CMin).
+FAISS_API void pq_scan_8bit(
+        size_t M,
+        const float* dis_table,
+        const uint8_t* codes,
+        size_t ncodes,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids,
+        bool max_heap);
+
+} // 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::pq_scan_8bit;
+using pq_code_distance::PQCodeDistance;
+using pq_code_distance::PQCodeDistanceScalar;
+
+} // namespace faiss

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

@@ -83,16 +83,15 @@ namespace pq_code_distance {
 
 // NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
 template <>
-float pq_code_distance_single_impl<SIMDLevel::ARM_SVE>(
+float pq_code_distance_8bit_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;
+    constexpr size_t ksub = 1 << 8;
 
     const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
     auto partialSum = svdup_n_f32(0.f);
@@ -159,12 +158,11 @@ float pq_code_distance_single_impl<SIMDLevel::ARM_SVE>(
     return svaddv_f32(svptrue_b32(), partialSum);
 }
 
-// Combines 4 operations of pq_code_distance_single_impl().
+// Combines 4 operations of pq_code_distance_8bit_single_impl().
 // NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
 template <>
-void pq_code_distance_four_impl<SIMDLevel::ARM_SVE>(
+void pq_code_distance_8bit_four_impl<SIMDLevel::ARM_SVE>(
         size_t M,
-        size_t nbits,
         const float* sim_table,
         const uint8_t* __restrict code0,
         const uint8_t* __restrict code1,
@@ -190,7 +188,7 @@ void pq_code_distance_four_impl<SIMDLevel::ARM_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));
 
@@ -349,4 +347,9 @@ void pq_code_distance_four_impl<SIMDLevel::ARM_SVE>(
 } // namespace pq_code_distance
 } // namespace faiss
 
+#define THE_SIMD_LEVEL SIMDLevel::ARM_SVE
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/pq_scan_impl.h>
+#undef THE_SIMD_LEVEL
+
 #endif // COMPILE_SIMD_ARM_SVE