faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/utils/partitioning.h

@@ -11,6 +11,8 @@
 #include <stdio.h>
 
 #include <faiss/impl/platform_macros.h>
+#include <faiss/utils/ordered_key_value.h>
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
@@ -58,6 +60,35 @@ void simd_histogram_16(
         int shift,
         int* hist);
 
+/** SIMD-dispatched partition for aligned uint16_t arrays.
+ *  Specializations live in per-ISA TUs (partitioning_avx2.cpp, etc.). */
+template <SIMDLevel SL, class C>
+typename C::T partition_fuzzy_simd(
+        uint16_t* vals,
+        typename C::TI* ids,
+        size_t n,
+        size_t q_min,
+        size_t q_max,
+        size_t* q_out);
+
+/** SIMD-dispatched histogram functions.
+ *  Specializations live in per-ISA TUs. */
+template <SIMDLevel SL>
+void simd_histogram_8(
+        const uint16_t* data,
+        int n,
+        uint16_t min,
+        int shift,
+        int* hist);
+
+template <SIMDLevel SL>
+void simd_histogram_16(
+        const uint16_t* data,
+        int n,
+        uint16_t min,
+        int shift,
+        int* hist);
+
 struct PartitionStats {
     uint64_t bisect_cycles;
     uint64_t compress_cycles;

data/vendor/faiss/faiss/utils/popcount.h

@@ -0,0 +1,29 @@
+/*
+ * 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 <cstdint>
+
+// popcount = count number of 1s in the binary representation of a number
+
+#ifdef _MSC_VER
+#include <intrin.h>
+inline int popcount32(uint32_t x) {
+    return __popcnt(x);
+}
+inline int popcount64(uint64_t x) {
+    return static_cast<int>(__popcnt64(x));
+}
+#else
+inline int popcount32(uint32_t x) {
+    return __builtin_popcount(x);
+}
+inline int popcount64(uint64_t x) {
+    return __builtin_popcountll(x);
+}
+#endif

data/vendor/faiss/faiss/utils/pq_code_distance.h

@@ -0,0 +1,251 @@
+/*
+ * 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 <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include <faiss/impl/ProductQuantizer.h>
+#include <faiss/impl/platform_macros.h>
+#include <faiss/utils/simd_levels.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_single_impl<NEW_LEVEL> and
+ *      pq_code_distance_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_single_impl(
+        size_t M,
+        size_t nbits,
+        const float* sim_table,
+        const uint8_t* code);
+
+template <SIMDLevel SL>
+void pq_code_distance_four_impl(
+        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);
+
+/// 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 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_single_impl<SL>(M, nbits, 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_four_impl<SL>(
+                    M,
+                    nbits,
+                    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_single(
+        size_t M,
+        size_t nbits,
+        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_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);
+
+} // namespace pq_code_distance
+
+// Re-export public API into namespace faiss for convenience
+using pq_code_distance::pq_code_distance_four;
+using pq_code_distance::pq_code_distance_single;
+using pq_code_distance::PQCodeDistance;
+using pq_code_distance::PQCodeDistanceScalar;
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/prefetch.h

@@ -9,9 +9,9 @@
 
 // prefetches
 
-#ifdef __AVX__
+#if defined(__x86_64__) || defined(_M_X64)
 
-// AVX
+// x86_64
 
 #include <xmmintrin.h>
 

data/vendor/faiss/faiss/utils/quantize_lut.cpp

@@ -27,7 +27,7 @@ namespace {
 // there can be NaNs in tables, they should be ignored
 float tab_min(const float* tab, size_t n) {
     float min = HUGE_VAL;
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         if (tab[i] < min) {
             min = tab[i];
         }
@@ -37,7 +37,7 @@ float tab_min(const float* tab, size_t n) {
 
 float tab_max(const float* tab, size_t n) {
     float max = -HUGE_VAL;
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         if (tab[i] > max) {
             max = tab[i];
         }
@@ -46,14 +46,14 @@ float tab_max(const float* tab, size_t n) {
 }
 
 void round_tab(float* tab, size_t n, float a, float bi) {
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         tab[i] = floorf((tab[i] - bi) * a + 0.5);
     }
 }
 
 template <typename T>
 void round_tab(const float* tab, size_t n, float a, float bi, T* tab_out) {
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         tab_out[i] = (T)floorf((tab[i] - bi) * a + 0.5);
     }
 }
@@ -68,7 +68,7 @@ void round_uint8_per_column(
         float* b_out) {
     float max_span = 0;
     std::vector<float> mins(n);
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         mins[i] = tab_min(tab + i * d, d);
         float span = tab_max(tab + i * d, d) - mins[i];
         if (span > max_span) {
@@ -77,7 +77,7 @@ void round_uint8_per_column(
     }
     float a = 255 / max_span;
     float b = 0;
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         b += mins[i];
         round_tab(tab + i * d, d, a, mins[i]);
     }
@@ -98,10 +98,10 @@ void round_uint8_per_column_multi(
         float* b_out) {
     float max_span = 0;
     std::vector<float> mins(n);
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         float min_i = HUGE_VAL;
         float max_i = -HUGE_VAL;
-        for (int j = 0; j < m; j++) {
+        for (size_t j = 0; j < m; j++) {
             min_i = std::min(min_i, tab_min(tab + (j * n + i) * d, d));
             max_i = std::max(max_i, tab_max(tab + (j * n + i) * d, d));
         }
@@ -113,9 +113,9 @@ void round_uint8_per_column_multi(
     }
     float a = 255 / max_span;
     float b = 0;
-    for (int i = 0; i < n; i++) {
+    for (size_t i = 0; i < n; i++) {
         b += mins[i];
-        for (int j = 0; j < m; j++) {
+        for (size_t j = 0; j < m; j++) {
             round_tab(tab + (j * n + i) * d, d, a, mins[i]);
         }
     }
@@ -147,7 +147,7 @@ void quantize_LUT_and_bias(
         std::vector<float> mins(M);
         float max_span_LUT = -HUGE_VAL, max_span_dis = 0;
         b = 0;
-        for (int i = 0; i < M; i++) {
+        for (size_t i = 0; i < M; i++) {
             mins[i] = tab_min(LUT + i * ksub, ksub);
             float span = tab_max(LUT + i * ksub, ksub) - mins[i];
             max_span_LUT = std::max(max_span_LUT, span);
@@ -156,7 +156,7 @@ void quantize_LUT_and_bias(
         }
         a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
 
-        for (int i = 0; i < M; i++) {
+        for (size_t i = 0; i < M; i++) {
             round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
         }
         memset(LUTq + M * ksub, 0, ksub * (M2 - M));
@@ -167,7 +167,7 @@ void quantize_LUT_and_bias(
         float bias_max = tab_max(bias, nprobe);
         max_span_dis = bias_max - bias_min;
         b = 0;
-        for (int i = 0; i < M; i++) {
+        for (size_t i = 0; i < M; i++) {
             mins[i] = tab_min(LUT + i * ksub, ksub);
             float span = tab_max(LUT + i * ksub, ksub) - mins[i];
             max_span_LUT = std::max(max_span_LUT, span);
@@ -177,7 +177,7 @@ void quantize_LUT_and_bias(
         a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
         b += bias_min;
 
-        for (int i = 0; i < M; i++) {
+        for (size_t i = 0; i < M; i++) {
             round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
         }
         memset(LUTq + M * ksub, 0, ksub * (M2 - M));
@@ -192,10 +192,10 @@ void quantize_LUT_and_bias(
 
         b = HUGE_VAL;
         size_t ij = 0;
-        for (int j = 0; j < nprobe; j++) {
+        for (size_t j = 0; j < nprobe; j++) {
             float max_span_dis_j = bias[j] - bias_min;
             float b2j = bias[j];
-            for (int i = 0; i < M; i++) {
+            for (size_t i = 0; i < M; i++) {
                 mins[ij] = tab_min(LUT + ij * ksub, ksub);
                 float span = tab_max(LUT + ij * ksub, ksub) - mins[ij];
                 max_span_LUT = std::max(max_span_LUT, span);
@@ -212,8 +212,8 @@ void quantize_LUT_and_bias(
 
         ij = 0;
         size_t ij_2 = 0;
-        for (int j = 0; j < nprobe; j++) {
-            for (int i = 0; i < M; i++) {
+        for (size_t j = 0; j < nprobe; j++) {
+            for (size_t i = 0; i < M; i++) {
                 round_tab(
                         LUT + ij * ksub, ksub, a, mins[ij], LUTq + ij_2 * ksub);
                 ij++;
@@ -230,10 +230,10 @@ void quantize_LUT_and_bias(
         std::vector<float> LUT2_storage(nprobe * M * ksub);
         float* LUT2 = LUT2_storage.data();
         size_t ijc = 0;
-        for (int j = 0; j < nprobe; j++) {
+        for (size_t j = 0; j < nprobe; j++) {
             float bias_j = bias[j] / M;
-            for (int i = 0; i < M; i++) {
-                for (int c = 0; c < ksub; c++) {
+            for (size_t i = 0; i < M; i++) {
+                for (size_t c = 0; c < ksub; c++) {
                     LUT2[ijc] = LUT[ijc] + bias_j;
                     ijc++;
                 }
@@ -241,8 +241,8 @@ void quantize_LUT_and_bias(
         }
         std::vector<float> mins(M, HUGE_VAL), maxs(M, -HUGE_VAL);
         size_t ij = 0;
-        for (int j = 0; j < nprobe; j++) {
-            for (int i = 0; i < M; i++) {
+        for (size_t j = 0; j < nprobe; j++) {
+            for (size_t i = 0; i < M; i++) {
                 mins[i] = std::min(mins[i], tab_min(LUT2 + ij * ksub, ksub));
                 maxs[i] = std::max(maxs[i], tab_max(LUT2 + ij * ksub, ksub));
                 ij++;
@@ -251,7 +251,7 @@ void quantize_LUT_and_bias(
 
         float max_span = -HUGE_VAL;
         b = 0;
-        for (int i = 0; i < M; i++) {
+        for (size_t i = 0; i < M; i++) {
             float span = maxs[i] - mins[i];
             max_span = std::max(max_span, span);
             b += mins[i];
@@ -259,8 +259,8 @@ void quantize_LUT_and_bias(
         a = 255 / max_span;
         ij = 0;
         size_t ij_2 = 0;
-        for (int j = 0; j < nprobe; j++) {
-            for (int i = 0; i < M; i++) {
+        for (size_t j = 0; j < nprobe; j++) {
+            for (size_t i = 0; i < M; i++) {
                 round_tab(
                         LUT2 + ij * ksub, ksub, a, mins[i], LUTq + ij_2 * ksub);
                 ij++;
@@ -298,7 +298,7 @@ void aq_quantize_LUT_and_bias(
     float bias_max = tab_max(bias, nprobe);
     max_span_dis = bias_max - bias_min;
     b = 0;
-    for (int i = 0; i < M; i++) {
+    for (size_t i = 0; i < M; i++) {
         mins[i] = tab_min(LUT + i * ksub, ksub);
         float span = tab_max(LUT + i * ksub, ksub) - mins[i];
         max_span_LUT = std::max(max_span_LUT, span);
@@ -308,7 +308,7 @@ void aq_quantize_LUT_and_bias(
     a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
     b += bias_min;
 
-    for (int i = 0; i < M; i++) {
+    for (size_t i = 0; i < M; i++) {
         round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
     }
     memset(LUTq + M * ksub, 0, ksub * (M2 - M));
@@ -324,14 +324,14 @@ float aq_estimate_norm_scale(
         size_t M_norm,
         const float* LUT) {
     float max_span_LUT = -HUGE_VAL;
-    for (int i = 0; i < M - M_norm; i++) {
+    for (size_t i = 0; i < M - M_norm; i++) {
         float min = tab_min(LUT + i * ksub, ksub);
         float span = tab_max(LUT + i * ksub, ksub) - min;
         max_span_LUT = std::max(max_span_LUT, span);
     }
 
     float max_span_LUT_norm = -HUGE_VAL;
-    for (int i = M - M_norm; i < M; i++) {
+    for (size_t i = M - M_norm; i < M; i++) {
         float min = tab_min(LUT + i * ksub, ksub);
         float span = tab_max(LUT + i * ksub, ksub) - min;
         max_span_LUT_norm = std::max(max_span_LUT_norm, span);

data/vendor/faiss/faiss/utils/quantize_lut.h

@@ -50,7 +50,7 @@ void round_uint8_per_column_multi(
 
 /** LUT quantization to uint8 and bias to uint16.
  *
- * (nprobe, M, ksub, lut_is_3d) determine the size of the the LUT
+ * (nprobe, M, ksub, lut_is_3d) determine the size of the LUT
  *
  *  LUT input:
  *  - 2D size (M, ksub): single matrix per probe (lut_is_3d=false)