faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/utils/distances.h

@@ -261,6 +261,9 @@ void inner_product_to_L2sqr(
  */
 void fvec_add(size_t d, const float* a, const float* b, float* c);
 
+template <SIMDLevel>
+void fvec_add(size_t d, const float* a, const float* b, float* c);
+
 /** compute c := a + b for a, c vectors and b a scalar
  *
  * c and a can overlap
@@ -270,6 +273,9 @@ void fvec_add(size_t d, const float* a, const float* b, float* c);
  */
 void fvec_add(size_t d, const float* a, float b, float* c);
 
+template <SIMDLevel>
+void fvec_add(size_t d, const float* a, float b, float* c);
+
 /** compute c := a - b for vectors
  *
  * c and a can overlap, c and b can overlap
@@ -280,6 +286,9 @@ void fvec_add(size_t d, const float* a, float b, float* c);
  */
 void fvec_sub(size_t d, const float* a, const float* b, float* c);
 
+template <SIMDLevel>
+void fvec_sub(size_t d, const float* a, const float* b, float* c);
+
 /***************************************************************************
  * Compute a subset of  distances
  ***************************************************************************/
@@ -360,7 +369,7 @@ void pairwise_indexed_inner_product(
  * KNN functions
  ***************************************************************************/
 
-// threshold on nx above which we switch to BLAS to compute distances
+// threshold on nx * d above which we switch to BLAS to compute distances
 FAISS_API extern int distance_compute_blas_threshold;
 
 // block sizes for BLAS distance computations
@@ -542,6 +551,16 @@ void compute_PQ_dis_tables_dsub2(
         bool is_inner_product,
         float* dis_tables);
 
+template <SIMDLevel>
+void compute_PQ_dis_tables_dsub2(
+        size_t d,
+        size_t ksub,
+        const float* centroids,
+        size_t nx,
+        const float* x,
+        bool is_inner_product,
+        float* dis_tables);
+
 /***************************************************************************
  * Templatized versions of distance functions
  ***************************************************************************/

data/vendor/faiss/faiss/utils/distances_dispatch.h

@@ -13,7 +13,7 @@
  *
  * This is a PRIVATE header. Do not include in public APIs or user code.
  *
- * These wrappers call DISPATCH_SIMDLevel to route to the correct SIMD
+ * These wrappers call with_simd_level to route to the correct SIMD
  * implementation. They are plain inline functions with a _dispatch suffix
  * (e.g. fvec_L2sqr_dispatch). Internal callers that want inlining include
  * this header and call the _dispatch variants directly.
@@ -24,30 +24,36 @@
 
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/extra_distances.h>
 
 namespace faiss {
 
 inline float fvec_L1_dispatch(const float* x, const float* y, size_t d) {
-    DISPATCH_SIMDLevel(fvec_L1, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_L1<SL>(x, y, d); });
 }
 
 inline float fvec_Linf_dispatch(const float* x, const float* y, size_t d) {
-    DISPATCH_SIMDLevel(fvec_Linf, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_Linf<SL>(x, y, d); });
 }
 
 inline float fvec_norm_L2sqr_dispatch(const float* x, size_t d) {
-    DISPATCH_SIMDLevel(fvec_norm_L2sqr, x, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_norm_L2sqr<SL>(x, d); });
 }
 
 inline float fvec_L2sqr_dispatch(const float* x, const float* y, size_t d) {
-    DISPATCH_SIMDLevel(fvec_L2sqr, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_L2sqr<SL>(x, y, d); });
 }
 
 inline float fvec_inner_product_dispatch(
         const float* x,
         const float* y,
         size_t d) {
-    DISPATCH_SIMDLevel(fvec_inner_product, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_inner_product<SL>(x, y, d); });
 }
 
 inline void fvec_inner_product_batch_4_dispatch(
@@ -61,18 +67,10 @@ inline void fvec_inner_product_batch_4_dispatch(
         float& dis1,
         float& dis2,
         float& dis3) {
-    DISPATCH_SIMDLevel(
-            fvec_inner_product_batch_4,
-            x,
-            y0,
-            y1,
-            y2,
-            y3,
-            d,
-            dis0,
-            dis1,
-            dis2,
-            dis3);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_inner_product_batch_4<SL>(
+                x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    });
 }
 
 inline void fvec_L2sqr_batch_4_dispatch(
@@ -86,8 +84,9 @@ inline void fvec_L2sqr_batch_4_dispatch(
         float& dis1,
         float& dis2,
         float& dis3) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_batch_4, x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_L2sqr_batch_4<SL>(x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    });
 }
 
 inline void fvec_L2sqr_ny_transposed_dispatch(
@@ -98,8 +97,9 @@ inline void fvec_L2sqr_ny_transposed_dispatch(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_ny_transposed, dis, x, y, y_sqlen, d, d_offset, ny);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_L2sqr_ny_transposed<SL>(dis, x, y, y_sqlen, d, d_offset, ny);
+    });
 }
 
 inline void fvec_inner_products_ny_dispatch(
@@ -108,7 +108,9 @@ inline void fvec_inner_products_ny_dispatch(
         const float* y,
         size_t d,
         size_t ny) {
-    DISPATCH_SIMDLevel(fvec_inner_products_ny, ip, x, y, d, ny);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_inner_products_ny<SL>(ip, x, y, d, ny);
+    });
 }
 
 inline void fvec_L2sqr_ny_dispatch(
@@ -117,7 +119,8 @@ inline void fvec_L2sqr_ny_dispatch(
         const float* y,
         size_t d,
         size_t ny) {
-    DISPATCH_SIMDLevel(fvec_L2sqr_ny, dis, x, y, d, ny);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { fvec_L2sqr_ny<SL>(dis, x, y, d, ny); });
 }
 
 inline size_t fvec_L2sqr_ny_nearest_dispatch(
@@ -126,8 +129,11 @@ inline size_t fvec_L2sqr_ny_nearest_dispatch(
         const float* y,
         size_t d,
         size_t ny) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_ny_nearest, distances_tmp_buffer, x, y, d, ny);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() {
+                return fvec_L2sqr_ny_nearest<SL>(
+                        distances_tmp_buffer, x, y, d, ny);
+            });
 }
 
 inline size_t fvec_L2sqr_ny_nearest_y_transposed_dispatch(
@@ -138,15 +144,11 @@ inline size_t fvec_L2sqr_ny_nearest_y_transposed_dispatch(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_ny_nearest_y_transposed,
-            distances_tmp_buffer,
-            x,
-            y,
-            y_sqlen,
-            d,
-            d_offset,
-            ny);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() {
+                return fvec_L2sqr_ny_nearest_y_transposed<SL>(
+                        distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
+            });
 }
 
 inline void fvec_madd_dispatch(
@@ -155,7 +157,8 @@ inline void fvec_madd_dispatch(
         float bf,
         const float* b,
         float* c) {
-    DISPATCH_SIMDLevel(fvec_madd, n, a, bf, b, c);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { fvec_madd<SL>(n, a, bf, b, c); });
 }
 
 inline int fvec_madd_and_argmin_dispatch(
@@ -164,7 +167,84 @@ inline int fvec_madd_and_argmin_dispatch(
         float bf,
         const float* b,
         float* c) {
-    DISPATCH_SIMDLevel(fvec_madd_and_argmin, n, a, bf, b, c);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() {
+                return fvec_madd_and_argmin<SL>(n, a, bf, b, c);
+            });
+}
+
+inline void fvec_sub_dispatch(
+        size_t d,
+        const float* a,
+        const float* b,
+        float* c) {
+    with_simd_level_256bit(
+            [&]<SIMDLevel level>() { fvec_sub<level>(d, a, b, c); });
+}
+
+inline void fvec_add_dispatch(
+        size_t d,
+        const float* a,
+        const float* b,
+        float* c) {
+    with_simd_level_256bit(
+            [&]<SIMDLevel level>() { fvec_add<level>(d, a, b, c); });
+}
+
+inline void fvec_add_scalar_dispatch(
+        size_t d,
+        const float* a,
+        float b,
+        float* c) {
+    with_simd_level_256bit(
+            [&]<SIMDLevel level>() { fvec_add<level>(d, a, b, c); });
+}
+
+inline void compute_PQ_dis_tables_dsub2_dispatch(
+        size_t d,
+        size_t ksub,
+        const float* centroids,
+        size_t nx,
+        const float* x,
+        bool is_inner_product,
+        float* dis_tables) {
+    with_simd_level_256bit([&]<SIMDLevel level>() {
+        compute_PQ_dis_tables_dsub2<level>(
+                d, ksub, centroids, nx, x, is_inner_product, dis_tables);
+    });
+}
+
+/***************************************************************************
+ * Dispatching function that takes a lambda directly.
+ * The lambda should be templated on VectorDistance, eg.:
+ *
+ *   auto result = with_VectorDistance(
+ *       metric, metric_arg, [&]<class VD>(VD vd) {
+ *           return vd(x, y);
+ *       });
+ **************************************************************************/
+
+template <typename LambdaType>
+auto with_VectorDistance(
+        size_t d,
+        MetricType metric,
+        float metric_arg,
+        LambdaType&& action) {
+    auto dispatch_metric = [&]<MetricType mt>() {
+        auto call = [&]<SIMDLevel level>() {
+            VectorDistance<mt, level> vd = {d, metric_arg};
+            return action(vd);
+        };
+
+        constexpr bool has_simd = mt == METRIC_INNER_PRODUCT ||
+                mt == METRIC_L2 || mt == METRIC_L1 || mt == METRIC_Linf;
+        if constexpr (!has_simd) {
+            return call.template operator()<SIMDLevel::NONE>();
+        } else {
+            return with_simd_level(call);
+        }
+    };
+    return with_metric_type(metric, dispatch_metric);
 }
 
 } // namespace faiss

data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp

@@ -7,9 +7,9 @@
 
 // -*- c++ -*-
 
-#include <faiss/utils/distances_fused/avx512.h>
+#include <faiss/utils/distances_fused/distances_fused.h>
 
-#ifdef __AVX512F__
+#ifdef COMPILE_SIMD_AVX512
 
 #include <immintrin.h>
 
@@ -78,7 +78,7 @@ void kernel(
     const float* const __restrict xd_0 = x + i * DIM;
 
     // prefetch the next point
-    _mm_prefetch(xd_0 + DIM * sizeof(float), _MM_HINT_NTA);
+    _mm_prefetch((char*)(xd_0 + DIM * sizeof(float)), _MM_HINT_NTA);
 
     // load a single point from x
     // load -2 * value
@@ -262,10 +262,10 @@ void exhaustive_L2sqr_fused_cmax(
         }
     }
 
-    const size_t nx_p = (nx / NX_POINTS_PER_LOOP) * NX_POINTS_PER_LOOP;
+    const idx_t nx_p = (nx / NX_POINTS_PER_LOOP) * NX_POINTS_PER_LOOP;
     // the main loop.
 #pragma omp parallel for schedule(dynamic)
-    for (size_t i = 0; i < nx_p; i += NX_POINTS_PER_LOOP) {
+    for (idx_t i = 0; i < nx_p; i += NX_POINTS_PER_LOOP) {
         kernel<DIM, NX_POINTS_PER_LOOP, NY_POINTS_PER_LOOP>(
                 x, y, y_transposed.data(), ny, res, y_norms, i);
     }
@@ -283,7 +283,8 @@ void exhaustive_L2sqr_fused_cmax(
 
 } // namespace
 
-bool exhaustive_L2sqr_fused_cmax_AVX512(
+template <>
+bool exhaustive_L2sqr_fused_cmax<SIMDLevel::AVX512>(
         const float* x,
         const float* y,
         size_t d,
@@ -343,4 +344,4 @@ bool exhaustive_L2sqr_fused_cmax_AVX512(
 
 } // namespace faiss
 
-#endif
+#endif // COMPILE_SIMD_AVX512

data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp

@@ -7,13 +7,37 @@
 
 #include <faiss/utils/distances_fused/distances_fused.h>
 
-#include <faiss/impl/platform_macros.h> // NOLINT
-
-#include <faiss/utils/distances_fused/avx512.h> // NOLINT
-#include <faiss/utils/distances_fused/simdlib_based.h>
+#include <faiss/impl/simd_dispatch.h>
 
 namespace faiss {
 
+// Scalar fallback: no fused kernel available.
+template <>
+bool exhaustive_L2sqr_fused_cmax<SIMDLevel::NONE>(
+        const float*,
+        const float*,
+        size_t,
+        size_t,
+        size_t,
+        Top1BlockResultHandler<CMax<float, int64_t>>&,
+        const float*) {
+    return false;
+}
+
+#ifdef COMPILE_SIMD_RISCV_RVV
+template <>
+bool exhaustive_L2sqr_fused_cmax<SIMDLevel::RISCV_RVV>(
+        const float*,
+        const float*,
+        size_t,
+        size_t,
+        size_t,
+        Top1BlockResultHandler<CMax<float, int64_t>>&,
+        const float*) {
+    return false;
+}
+#endif // COMPILE_SIMD_RISCV_RVV
+
 bool exhaustive_L2sqr_fused_cmax(
         const float* x,
         const float* y,
@@ -27,16 +51,11 @@ bool exhaustive_L2sqr_fused_cmax(
         return true;
     }
 
-#ifdef __AVX512F__
-    // avx512 kernel
-    return exhaustive_L2sqr_fused_cmax_AVX512(x, y, d, nx, ny, res, y_norms);
-#elif defined(__AVX2__) || defined(__aarch64__)
-    // avx2 or arm neon kernel
-    return exhaustive_L2sqr_fused_cmax_simdlib(x, y, d, nx, ny, res, y_norms);
-#else
-    // not supported, please use a general-purpose kernel
-    return false;
-#endif
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+            [&]<SIMDLevel SL>() {
+                return exhaustive_L2sqr_fused_cmax<SL>(
+                        x, y, d, nx, ny, res, y_norms);
+            });
 }
 
 } // namespace faiss

data/vendor/faiss/faiss/utils/distances_fused/distances_fused.h

@@ -21,8 +21,8 @@
 #pragma once
 
 #include <faiss/impl/ResultHandler.h>
-
 #include <faiss/utils/Heap.h>
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
@@ -37,4 +37,15 @@ bool exhaustive_L2sqr_fused_cmax(
         Top1BlockResultHandler<CMax<float, int64_t>>& res,
         const float* y_norms);
 
+// Per-SIMD-level implementation (defined in per-SIMD TUs).
+template <SIMDLevel>
+bool exhaustive_L2sqr_fused_cmax(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms);
+
 } // namespace faiss