faiss 0.5.3 → 0.6.0

data/vendor/faiss/faiss/utils/Heap.h

@@ -642,6 +642,27 @@ void merge_knn_results(
         typename C::T* distances,
         idx_t* labels);
 
+/** Reduces k_base pairs (base_labels, base_distances) into k pairs
+ * (labels, distances). The function is used for the refining process.
+ *
+ * @param n              number of vectors to process
+ * @param k              number of output nearest neighbors per vector
+ * @param labels         output labels, size (n, k)
+ * @param distances      output distances, size (n, k)
+ * @param k_base         number of input nearest neighbors per vector
+ * @param base_labels    input labels, size (n, k_base)
+ * @param base_distances input distances, size (n, k_base)
+ */
+template <class C>
+void reorder_2_heaps(
+        int64_t n,
+        int64_t k,
+        typename C::TI* __restrict labels,
+        float* __restrict distances,
+        int64_t k_base,
+        const typename C::TI* __restrict base_labels,
+        const float* __restrict base_distances);
+
 } // namespace faiss
 
 #endif /* FAISS_Heap_h */

data/vendor/faiss/faiss/utils/NeuralNet.cpp

@@ -12,6 +12,7 @@
 #include <cstring>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 
 /* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */
@@ -265,14 +266,16 @@ nn::Int32Tensor2D QINCoStep::encode(
         const float* db = zqs_r.data() + i * K * d;
         float dis_min = HUGE_VALF;
         int64_t idx = -1;
-        for (size_t j = 0; j < K; j++) {
-            float dis = fvec_L2sqr(q, db, d);
-            if (dis < dis_min) {
-                dis_min = dis;
-                idx = j;
+        with_simd_level([&]<SIMDLevel SL>() {
+            for (size_t j = 0; j < K; j++) {
+                float dis = fvec_L2sqr<SL>(q, db, d);
+                if (dis < dis_min) {
+                    dis_min = dis;
+                    idx = j;
+                }
+                db += d;
             }
-            db += d;
-        }
+        });
         codes.v[i] = idx;
         if (res) {
             const float* xhat_row = xhat.data() + i * d;

data/vendor/faiss/faiss/utils/distances.cpp

@@ -27,6 +27,7 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
 
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/distances_fused/distances_fused.h>
 
 #ifndef FINTEGER
@@ -55,6 +56,122 @@ int sgemm_(
 
 namespace faiss {
 
+/***************************************************************************
+ * Public API dispatch wrappers
+ ***************************************************************************/
+
+float fvec_L1(const float* x, const float* y, size_t d) {
+    return fvec_L1_dispatch(x, y, d);
+}
+
+float fvec_Linf(const float* x, const float* y, size_t d) {
+    return fvec_Linf_dispatch(x, y, d);
+}
+
+float fvec_norm_L2sqr(const float* x, size_t d) {
+    return fvec_norm_L2sqr_dispatch(x, d);
+}
+
+float fvec_L2sqr(const float* x, const float* y, size_t d) {
+    return fvec_L2sqr_dispatch(x, y, d);
+}
+
+float fvec_inner_product(const float* x, const float* y, size_t d) {
+    return fvec_inner_product_dispatch(x, y, d);
+}
+
+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) {
+    fvec_inner_product_batch_4_dispatch(
+            x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+}
+
+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) {
+    fvec_L2sqr_batch_4_dispatch(x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+}
+
+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) {
+    fvec_L2sqr_ny_transposed_dispatch(dis, x, y, y_sqlen, d, d_offset, ny);
+}
+
+void fvec_inner_products_ny(
+        float* ip,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+    fvec_inner_products_ny_dispatch(ip, x, y, d, ny);
+}
+
+void fvec_L2sqr_ny(
+        float* dis,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+    fvec_L2sqr_ny_dispatch(dis, x, y, d, ny);
+}
+
+size_t fvec_L2sqr_ny_nearest(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+    return fvec_L2sqr_ny_nearest_dispatch(distances_tmp_buffer, x, y, d, ny);
+}
+
+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) {
+    return fvec_L2sqr_ny_nearest_y_transposed_dispatch(
+            distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
+}
+
+void fvec_madd(size_t n, const float* a, float bf, const float* b, float* c) {
+    fvec_madd_dispatch(n, a, bf, b, c);
+}
+
+int fvec_madd_and_argmin(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* c) {
+    return fvec_madd_and_argmin_dispatch(n, a, bf, b, c);
+}
+
 /***************************************************************************
  * Matrix/vector ops
  ***************************************************************************/
@@ -67,7 +184,7 @@ void fvec_norms_L2(
         size_t nx) {
 #pragma omp parallel for if (nx > 10000)
     for (int64_t i = 0; i < nx; i++) {
-        nr[i] = sqrtf(fvec_norm_L2sqr(x + i * d, d));
+        nr[i] = sqrtf(fvec_norm_L2sqr_dispatch(x + i * d, d));
     }
 }
 
@@ -78,7 +195,7 @@ void fvec_norms_L2sqr(
         size_t nx) {
 #pragma omp parallel for if (nx > 10000)
     for (int64_t i = 0; i < nx; i++) {
-        nr[i] = fvec_norm_L2sqr(x + i * d, d);
+        nr[i] = fvec_norm_L2sqr_dispatch(x + i * d, d);
     }
 }
 
@@ -93,16 +210,16 @@ void fvec_norms_L2sqr(
 // The workaround below is explicitly branching
 // off to a codepath without omp.
 
-#define FVEC_RENORM_L2_IMPL                   \
-    float* __restrict xi = x + i * d;         \
-                                              \
-    float nr = fvec_norm_L2sqr(xi, d);        \
-                                              \
-    if (nr > 0) {                             \
-        size_t j;                             \
-        const float inv_nr = 1.0 / sqrtf(nr); \
-        for (j = 0; j < d; j++)               \
-            xi[j] *= inv_nr;                  \
+#define FVEC_RENORM_L2_IMPL                     \
+    float* __restrict xi = x + i * d;           \
+                                                \
+    float nr = fvec_norm_L2sqr_dispatch(xi, d); \
+                                                \
+    if (nr > 0) {                               \
+        size_t j;                               \
+        const float inv_nr = 1.0 / sqrtf(nr);   \
+        for (j = 0; j < d; j++)                 \
+            xi[j] *= inv_nr;                    \
     }
 
 void fvec_renorm_L2_noomp(size_t d, size_t nx, float* __restrict x) {
@@ -159,7 +276,7 @@ void exhaustive_inner_product_seq(
                 if (!res.is_in_selection(j)) {
                     continue;
                 }
-                float ip = fvec_inner_product(x_i, y_j, d);
+                float ip = fvec_inner_product_dispatch(x_i, y_j, d);
                 resi.add_result(ip, j);
             }
             resi.end();
@@ -191,7 +308,7 @@ void exhaustive_L2sqr_seq(
                 if (!res.is_in_selection(j)) {
                     continue;
                 }
-                float disij = fvec_L2sqr(x_i, y_j, d);
+                float disij = fvec_L2sqr_dispatch(x_i, y_j, d);
                 resi.add_result(disij, j);
             }
             resi.end();
@@ -998,7 +1115,7 @@ void fvec_inner_products_by_idx(
             if (idsj[i] < 0) {
                 ipj[i] = -INFINITY;
             } else {
-                ipj[i] = fvec_inner_product(xj, y + d * idsj[i], d);
+                ipj[i] = fvec_inner_product_dispatch(xj, y + d * idsj[i], d);
             }
         }
     }
@@ -1023,7 +1140,7 @@ void fvec_L2sqr_by_idx(
             if (idsj[i] < 0) {
                 disj[i] = INFINITY;
             } else {
-                disj[i] = fvec_L2sqr(xj, y + d * idsj[i], d);
+                disj[i] = fvec_L2sqr_dispatch(xj, y + d * idsj[i], d);
             }
         }
     }
@@ -1040,7 +1157,7 @@ void pairwise_indexed_L2sqr(
 #pragma omp parallel for if (n > 1)
     for (int64_t j = 0; j < n; j++) {
         if (ix[j] >= 0 && iy[j] >= 0) {
-            dis[j] = fvec_L2sqr(x + d * ix[j], y + d * iy[j], d);
+            dis[j] = fvec_L2sqr_dispatch(x + d * ix[j], y + d * iy[j], d);
         } else {
             dis[j] = INFINITY;
         }
@@ -1058,7 +1175,8 @@ void pairwise_indexed_inner_product(
 #pragma omp parallel for if (n > 1)
     for (int64_t j = 0; j < n; j++) {
         if (ix[j] >= 0 && iy[j] >= 0) {
-            dis[j] = fvec_inner_product(x + d * ix[j], y + d * iy[j], d);
+            dis[j] = fvec_inner_product_dispatch(
+                    x + d * ix[j], y + d * iy[j], d);
         } else {
             dis[j] = -INFINITY;
         }
@@ -1096,7 +1214,7 @@ void knn_inner_products_by_idx(
             if (idsi[j] < 0 || idsi[j] >= ny) {
                 break;
             }
-            float ip = fvec_inner_product(x_, y + d * idsi[j], d);
+            float ip = fvec_inner_product_dispatch(x_, y + d * idsi[j], d);
 
             if (ip > simi[0]) {
                 minheap_replace_top(k, simi, idxi, ip, idsi[j]);
@@ -1132,7 +1250,7 @@ void knn_L2sqr_by_idx(
             if (idsi[j] < 0 || idsi[j] >= ny) {
                 break;
             }
-            float disij = fvec_L2sqr(x_, y + d * idsi[j], d);
+            float disij = fvec_L2sqr_dispatch(x_, y + d * idsi[j], d);
 
             if (disij < simi[0]) {
                 maxheap_replace_top(k, simi, idxi, disij, idsi[j]);
@@ -1170,19 +1288,19 @@ void pairwise_L2sqr(
 
 #pragma omp parallel for if (nb > 1)
     for (int64_t i = 0; i < nb; i++) {
-        b_norms[i] = fvec_norm_L2sqr(xb + i * ldb, d);
+        b_norms[i] = fvec_norm_L2sqr_dispatch(xb + i * ldb, d);
     }
 
 #pragma omp parallel for
     for (int64_t i = 1; i < nq; i++) {
-        float q_norm = fvec_norm_L2sqr(xq + i * ldq, d);
+        float q_norm = fvec_norm_L2sqr_dispatch(xq + i * ldq, d);
         for (int64_t j = 0; j < nb; j++) {
             dis[i * ldd + j] = q_norm + b_norms[j];
         }
     }
 
     {
-        float q_norm = fvec_norm_L2sqr(xq, d);
+        float q_norm = fvec_norm_L2sqr_dispatch(xq, d);
         for (int64_t j = 0; j < nb; j++) {
             dis[j] += q_norm;
         }

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
@@ -473,6 +559,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 +573,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,170 @@
+/*
+ * 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 DISPATCH_SIMDLevel 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>
+
+namespace faiss {
+
+inline float fvec_L1_dispatch(const float* x, const float* y, size_t d) {
+    DISPATCH_SIMDLevel(fvec_L1, x, y, d);
+}
+
+inline float fvec_Linf_dispatch(const float* x, const float* y, size_t d) {
+    DISPATCH_SIMDLevel(fvec_Linf, x, y, d);
+}
+
+inline float fvec_norm_L2sqr_dispatch(const float* x, size_t d) {
+    DISPATCH_SIMDLevel(fvec_norm_L2sqr, x, d);
+}
+
+inline float fvec_L2sqr_dispatch(const float* x, const float* y, size_t d) {
+    DISPATCH_SIMDLevel(fvec_L2sqr, 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);
+}
+
+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) {
+    DISPATCH_SIMDLevel(
+            fvec_inner_product_batch_4,
+            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) {
+    DISPATCH_SIMDLevel(
+            fvec_L2sqr_batch_4, 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) {
+    DISPATCH_SIMDLevel(
+            fvec_L2sqr_ny_transposed, 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) {
+    DISPATCH_SIMDLevel(fvec_inner_products_ny, 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) {
+    DISPATCH_SIMDLevel(fvec_L2sqr_ny, 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) {
+    DISPATCH_SIMDLevel(
+            fvec_L2sqr_ny_nearest, 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) {
+    DISPATCH_SIMDLevel(
+            fvec_L2sqr_ny_nearest_y_transposed,
+            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) {
+    DISPATCH_SIMDLevel(fvec_madd, 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) {
+    DISPATCH_SIMDLevel(fvec_madd_and_argmin, n, a, bf, b, c);
+}
+
+} // namespace faiss