faiss 0.5.3 → 0.6.1

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

@@ -15,6 +15,7 @@
 
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
@@ -27,15 +28,27 @@ struct IDSelector;
 /// Squared L2 distance between two vectors
 float fvec_L2sqr(const float* x, const float* y, size_t d);
 
+template <SIMDLevel>
+float fvec_L2sqr(const float* x, const float* y, size_t d);
+
 /// inner product
 float fvec_inner_product(const float* x, const float* y, size_t d);
 
+template <SIMDLevel>
+float fvec_inner_product(const float* x, const float* y, size_t d);
+
 /// L1 distance
 float fvec_L1(const float* x, const float* y, size_t d);
 
+template <SIMDLevel>
+float fvec_L1(const float* x, const float* y, size_t d);
+
 /// infinity distance
 float fvec_Linf(const float* x, const float* y, size_t d);
 
+template <SIMDLevel>
+float fvec_Linf(const float* x, const float* y, size_t d);
+
 /// Special version of inner product that computes 4 distances
 /// between x and yi, which is performance oriented.
 void fvec_inner_product_batch_4(
@@ -50,6 +63,19 @@ void fvec_inner_product_batch_4(
         float& dis2,
         float& dis3);
 
+template <SIMDLevel>
+void fvec_inner_product_batch_4(
+        const float* x,
+        const float* y0,
+        const float* y1,
+        const float* y2,
+        const float* y3,
+        const size_t d,
+        float& dis0,
+        float& dis1,
+        float& dis2,
+        float& dis3);
+
 /// Special version of L2sqr that computes 4 distances
 /// between x and yi, which is performance oriented.
 void fvec_L2sqr_batch_4(
@@ -64,6 +90,19 @@ void fvec_L2sqr_batch_4(
         float& dis2,
         float& dis3);
 
+template <SIMDLevel>
+void fvec_L2sqr_batch_4(
+        const float* x,
+        const float* y0,
+        const float* y1,
+        const float* y2,
+        const float* y3,
+        const size_t d,
+        float& dis0,
+        float& dis1,
+        float& dis2,
+        float& dis3);
+
 /** Compute pairwise distances between sets of vectors
  *
  * @param d     dimension of the vectors
@@ -93,6 +132,14 @@ void fvec_inner_products_ny(
         size_t d,
         size_t ny);
 
+template <SIMDLevel>
+void fvec_inner_products_ny(
+        float* ip, /* output inner product */
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny);
+
 /* compute ny square L2 distance between x and a set of contiguous y vectors */
 void fvec_L2sqr_ny(
         float* dis,
@@ -101,6 +148,14 @@ void fvec_L2sqr_ny(
         size_t d,
         size_t ny);
 
+template <SIMDLevel>
+void fvec_L2sqr_ny(
+        float* dis,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny);
+
 /* compute ny square L2 distance between x and a set of transposed contiguous
    y vectors. squared lengths of y should be provided as well */
 void fvec_L2sqr_ny_transposed(
@@ -112,6 +167,16 @@ void fvec_L2sqr_ny_transposed(
         size_t d_offset,
         size_t ny);
 
+template <SIMDLevel>
+void fvec_L2sqr_ny_transposed(
+        float* dis,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t ny);
+
 /* compute ny square L2 distance between x and a set of contiguous y vectors
    and return the index of the nearest vector.
    return 0 if ny == 0. */
@@ -122,6 +187,14 @@ size_t fvec_L2sqr_ny_nearest(
         size_t d,
         size_t ny);
 
+template <SIMDLevel>
+size_t fvec_L2sqr_ny_nearest(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny);
+
 /* compute ny square L2 distance between x and a set of transposed contiguous
    y vectors and return the index of the nearest vector.
    squared lengths of y should be provided as well
@@ -135,9 +208,22 @@ size_t fvec_L2sqr_ny_nearest_y_transposed(
         size_t d_offset,
         size_t ny);
 
+template <SIMDLevel>
+size_t fvec_L2sqr_ny_nearest_y_transposed(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t ny);
+
 /** squared norm of a vector */
 float fvec_norm_L2sqr(const float* x, size_t d);
 
+template <SIMDLevel>
+float fvec_norm_L2sqr(const float* x, size_t d);
+
 /** compute the L2 norms for a set of vectors
  *
  * @param  norms    output norms, size nx
@@ -175,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
@@ -184,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
@@ -194,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
  ***************************************************************************/
@@ -274,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
@@ -456,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
  ***************************************************************************/
@@ -473,6 +578,10 @@ void compute_PQ_dis_tables_dsub2(
  */
 void fvec_madd(size_t n, const float* a, float bf, const float* b, float* c);
 
+/* same statically */
+template <SIMDLevel>
+void fvec_madd(size_t n, const float* a, float bf, const float* b, float* c);
+
 /** same as fvec_madd, also return index of the min of the result table
  * @return    index of the min of table c
  */
@@ -483,4 +592,12 @@ int fvec_madd_and_argmin(
         const float* b,
         float* c);
 
+template <SIMDLevel>
+int fvec_madd_and_argmin(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* c);
+
 } // namespace faiss

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

@@ -0,0 +1,250 @@
+/*
+ * 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 distances_dispatch.h
+ * @brief Inlineable dispatch wrappers for distance functions.
+ *
+ * This is a PRIVATE header. Do not include in public APIs or user code.
+ *
+ * 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.
+ *
+ * The public API functions (fvec_L2sqr, etc.) are defined as regular extern
+ * functions in distances.cpp and simply delegate to these _dispatch variants.
+ */
+
+#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) {
+    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) {
+    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) {
+    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) {
+    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) {
+    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(
+        const float* x,
+        const float* y0,
+        const float* y1,
+        const float* y2,
+        const float* y3,
+        const size_t d,
+        float& dis0,
+        float& dis1,
+        float& dis2,
+        float& 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(
+        const float* x,
+        const float* y0,
+        const float* y1,
+        const float* y2,
+        const float* y3,
+        const size_t d,
+        float& dis0,
+        float& dis1,
+        float& dis2,
+        float& 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(
+        float* dis,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t 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(
+        float* ip,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t 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(
+        float* dis,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t 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(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t 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(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t 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(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* 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(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* 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