faiss 0.3.1 → 0.3.3

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.
@@ -18,6 +18,8 @@
 
 #ifdef __AVX2__
 #include <immintrin.h>
+#elif defined(__ARM_FEATURE_SVE)
+#include <arm_sve.h>
 #endif
 
 #include <faiss/impl/AuxIndexStructures.h>
@@ -130,20 +132,17 @@ void fvec_renorm_L2(size_t d, size_t nx, float* __restrict x) {
 namespace {
 
 /* Find the nearest neighbors for nx queries in a set of ny vectors */
-template <class BlockResultHandler, bool use_sel = false>
+template <class BlockResultHandler>
 void exhaustive_inner_product_seq(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        BlockResultHandler& res,
-        const IDSelector* sel = nullptr) {
+        BlockResultHandler& res) {
     using SingleResultHandler =
             typename BlockResultHandler::SingleResultHandler;
-    int nt = std::min(int(nx), omp_get_max_threads());
-
-    FAISS_ASSERT(use_sel == (sel != nullptr));
+    [[maybe_unused]] int nt = std::min(int(nx), omp_get_max_threads());
 
 #pragma omp parallel num_threads(nt)
     {
@@ -156,7 +155,7 @@ void exhaustive_inner_product_seq(
             resi.begin(i);
 
             for (size_t j = 0; j < ny; j++, y_j += d) {
-                if (use_sel && !sel->is_member(j)) {
+                if (!res.is_in_selection(j)) {
                     continue;
                 }
                 float ip = fvec_inner_product(x_i, y_j, d);
@@ -167,20 +166,17 @@ void exhaustive_inner_product_seq(
     }
 }
 
-template <class BlockResultHandler, bool use_sel = false>
+template <class BlockResultHandler>
 void exhaustive_L2sqr_seq(
         const float* x,
         const float* y,
         size_t d,
         size_t nx,
         size_t ny,
-        BlockResultHandler& res,
-        const IDSelector* sel = nullptr) {
+        BlockResultHandler& res) {
     using SingleResultHandler =
             typename BlockResultHandler::SingleResultHandler;
-    int nt = std::min(int(nx), omp_get_max_threads());
-
-    FAISS_ASSERT(use_sel == (sel != nullptr));
+    [[maybe_unused]] int nt = std::min(int(nx), omp_get_max_threads());
 
 #pragma omp parallel num_threads(nt)
     {
@@ -191,7 +187,7 @@ void exhaustive_L2sqr_seq(
             const float* y_j = y;
             resi.begin(i);
             for (size_t j = 0; j < ny; j++, y_j += d) {
-                if (use_sel && !sel->is_member(j)) {
+                if (!res.is_in_selection(j)) {
                     continue;
                 }
                 float disij = fvec_L2sqr(x_i, y_j, d);
@@ -326,6 +322,9 @@ void exhaustive_L2sqr_blas_default_impl(
                     float ip = *ip_line;
                     float dis = x_norms[i] + y_norms[j] - 2 * ip;
 
+                    if (!res.is_in_selection(j)) {
+                        dis = HUGE_VALF;
+                    }
                     // negative values can occur for identical vectors
                     // due to roundoff errors
                     if (dis < 0)
@@ -560,6 +559,183 @@ void exhaustive_L2sqr_blas_cmax_avx2(
         InterruptCallback::check();
     }
 }
+#elif defined(__ARM_FEATURE_SVE)
+void exhaustive_L2sqr_blas_cmax_sve(
+        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) {
+    // BLAS does not like empty matrices
+    if (nx == 0 || ny == 0)
+        return;
+
+    /* block sizes */
+    const size_t bs_x = distance_compute_blas_query_bs;
+    const size_t bs_y = distance_compute_blas_database_bs;
+    // const size_t bs_x = 16, bs_y = 16;
+    std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]);
+    std::unique_ptr<float[]> x_norms(new float[nx]);
+    std::unique_ptr<float[]> del2;
+
+    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
+
+    const size_t lanes = svcntw();
+
+    if (!y_norms) {
+        float* y_norms2 = new float[ny];
+        del2.reset(y_norms2);
+        fvec_norms_L2sqr(y_norms2, y, d, ny);
+        y_norms = y_norms2;
+    }
+
+    for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
+        size_t i1 = i0 + bs_x;
+        if (i1 > nx)
+            i1 = nx;
+
+        res.begin_multiple(i0, i1);
+
+        for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
+            size_t j1 = j0 + bs_y;
+            if (j1 > ny)
+                j1 = ny;
+            /* compute the actual dot products */
+            {
+                float one = 1, zero = 0;
+                FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
+                sgemm_("Transpose",
+                       "Not transpose",
+                       &nyi,
+                       &nxi,
+                       &di,
+                       &one,
+                       y + j0 * d,
+                       &di,
+                       x + i0 * d,
+                       &di,
+                       &zero,
+                       ip_block.get(),
+                       &nyi);
+            }
+#pragma omp parallel for
+            for (int64_t i = i0; i < i1; i++) {
+                const size_t count = j1 - j0;
+                float* ip_line = ip_block.get() + (i - i0) * count;
+
+                svprfw(svwhilelt_b32_u64(0, count), ip_line, SV_PLDL1KEEP);
+                svprfw(svwhilelt_b32_u64(lanes, count),
+                       ip_line + lanes,
+                       SV_PLDL1KEEP);
+
+                // Track lanes min distances + lanes min indices.
+                // All the distances tracked do not take x_norms[i]
+                //   into account in order to get rid of extra
+                //   vaddq_f32(x_norms[i], ...) instructions
+                //   is distance computations.
+                auto min_distances = svdup_n_f32(res.dis_tab[i] - x_norms[i]);
+
+                // these indices are local and are relative to j0.
+                // so, value 0 means j0.
+                auto min_indices = svdup_n_u32(0u);
+
+                auto current_indices = svindex_u32(0u, 1u);
+
+                // process lanes * 2 elements per loop
+                for (size_t idx_j = 0; idx_j < count;
+                     idx_j += lanes * 2, ip_line += lanes * 2) {
+                    svprfw(svwhilelt_b32_u64(idx_j + lanes * 2, count),
+                           ip_line + lanes * 2,
+                           SV_PLDL1KEEP);
+                    svprfw(svwhilelt_b32_u64(idx_j + lanes * 3, count),
+                           ip_line + lanes * 3,
+                           SV_PLDL1KEEP);
+
+                    // mask
+                    const auto mask_0 = svwhilelt_b32_u64(idx_j, count);
+                    const auto mask_1 = svwhilelt_b32_u64(idx_j + lanes, count);
+
+                    // load values for norms
+                    const auto y_norm_0 =
+                            svld1_f32(mask_0, y_norms + idx_j + j0 + 0);
+                    const auto y_norm_1 =
+                            svld1_f32(mask_1, y_norms + idx_j + j0 + lanes);
+
+                    // load values for dot products
+                    const auto ip_0 = svld1_f32(mask_0, ip_line + 0);
+                    const auto ip_1 = svld1_f32(mask_1, ip_line + lanes);
+
+                    // compute dis = y_norm[j] - 2 * dot(x_norm[i], y_norm[j]).
+                    // x_norm[i] was dropped off because it is a constant for a
+                    // given i. We'll deal with it later.
+                    const auto distances_0 =
+                            svmla_n_f32_z(mask_0, y_norm_0, ip_0, -2.f);
+                    const auto distances_1 =
+                            svmla_n_f32_z(mask_1, y_norm_1, ip_1, -2.f);
+
+                    // compare the new distances to the min distances
+                    // for each of the first group of 4 ARM SIMD components.
+                    auto comparison =
+                            svcmpgt_f32(mask_0, min_distances, distances_0);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances =
+                            svsel_f32(comparison, distances_0, min_distances);
+                    min_indices =
+                            svsel_u32(comparison, current_indices, min_indices);
+                    current_indices = svadd_n_u32_x(
+                            mask_0,
+                            current_indices,
+                            static_cast<uint32_t>(lanes));
+
+                    // compare the new distances to the min distances
+                    // for each of the second group of 4 ARM SIMD components.
+                    comparison =
+                            svcmpgt_f32(mask_1, min_distances, distances_1);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances =
+                            svsel_f32(comparison, distances_1, min_distances);
+                    min_indices =
+                            svsel_u32(comparison, current_indices, min_indices);
+                    current_indices = svadd_n_u32_x(
+                            mask_1,
+                            current_indices,
+                            static_cast<uint32_t>(lanes));
+                }
+
+                // add missing x_norms[i]
+                // negative values can occur for identical vectors
+                //    due to roundoff errors.
+                auto mask = svwhilelt_b32_u64(0, count);
+                min_distances = svadd_n_f32_z(
+                        svcmpge_n_f32(mask, min_distances, -x_norms[i]),
+                        min_distances,
+                        x_norms[i]);
+                min_indices = svadd_n_u32_x(
+                        mask, min_indices, static_cast<uint32_t>(j0));
+                mask = svcmple_n_f32(mask, min_distances, res.dis_tab[i]);
+                if (svcntp_b32(svptrue_b32(), mask) == 0)
+                    res.add_result(i, res.dis_tab[i], res.ids_tab[i]);
+                else {
+                    const auto min_distance = svminv_f32(mask, min_distances);
+                    const auto min_index = svminv_u32(
+                            svcmpeq_n_f32(mask, min_distances, min_distance),
+                            min_indices);
+                    res.add_result(i, min_distance, min_index);
+                }
+            }
+        }
+        // Does nothing for SingleBestResultHandler, but
+        // keeping the call for the consistency.
+        res.end_multiple();
+        InterruptCallback::check();
+    }
+}
 #endif
 
 // an override if only a single closest point is needed
@@ -582,6 +758,16 @@ void exhaustive_L2sqr_blas<Top1BlockResultHandler<CMax<float, int64_t>>>(
     // run the specialized AVX2 implementation
     exhaustive_L2sqr_blas_cmax_avx2(x, y, d, nx, ny, res, y_norms);
 
+#elif defined(__ARM_FEATURE_SVE)
+    // use a faster fused kernel if available
+    if (exhaustive_L2sqr_fused_cmax(x, y, d, nx, ny, res, y_norms)) {
+        // the kernel is available and it is complete, we're done.
+        return;
+    }
+
+    // run the specialized SVE implementation
+    exhaustive_L2sqr_blas_cmax_sve(x, y, d, nx, ny, res, y_norms);
+
 #elif defined(__aarch64__)
     // use a faster fused kernel if available
     if (exhaustive_L2sqr_fused_cmax(x, y, d, nx, ny, res, y_norms)) {
@@ -601,44 +787,40 @@ void exhaustive_L2sqr_blas<Top1BlockResultHandler<CMax<float, int64_t>>>(
 #endif
 }
 
-template <class BlockResultHandler>
-void knn_L2sqr_select(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        BlockResultHandler& res,
-        const float* y_norm2,
-        const IDSelector* sel) {
-    if (sel) {
-        exhaustive_L2sqr_seq<BlockResultHandler, true>(
-                x, y, d, nx, ny, res, sel);
-    } else if (nx < distance_compute_blas_threshold) {
-        exhaustive_L2sqr_seq(x, y, d, nx, ny, res);
-    } else {
-        exhaustive_L2sqr_blas(x, y, d, nx, ny, res, y_norm2);
+struct Run_search_inner_product {
+    using T = void;
+    template <class BlockResultHandler>
+    void f(BlockResultHandler& res,
+           const float* x,
+           const float* y,
+           size_t d,
+           size_t nx,
+           size_t ny) {
+        if (res.sel || nx < distance_compute_blas_threshold) {
+            exhaustive_inner_product_seq(x, y, d, nx, ny, res);
+        } else {
+            exhaustive_inner_product_blas(x, y, d, nx, ny, res);
+        }
     }
-}
-
-template <class BlockResultHandler>
-void knn_inner_product_select(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        BlockResultHandler& res,
-        const IDSelector* sel) {
-    if (sel) {
-        exhaustive_inner_product_seq<BlockResultHandler, true>(
-                x, y, d, nx, ny, res, sel);
-    } else if (nx < distance_compute_blas_threshold) {
-        exhaustive_inner_product_seq(x, y, d, nx, ny, res);
-    } else {
-        exhaustive_inner_product_blas(x, y, d, nx, ny, res);
+};
+
+struct Run_search_L2sqr {
+    using T = void;
+    template <class BlockResultHandler>
+    void f(BlockResultHandler& res,
+           const float* x,
+           const float* y,
+           size_t d,
+           size_t nx,
+           size_t ny,
+           const float* y_norm2) {
+        if (res.sel || nx < distance_compute_blas_threshold) {
+            exhaustive_L2sqr_seq(x, y, d, nx, ny, res);
+        } else {
+            exhaustive_L2sqr_blas(x, y, d, nx, ny, res, y_norm2);
+        }
     }
-}
+};
 
 } // anonymous namespace
 
@@ -675,16 +857,9 @@ void knn_inner_product(
         return;
     }
 
-    if (k == 1) {
-        Top1BlockResultHandler<CMin<float, int64_t>> res(nx, vals, ids);
-        knn_inner_product_select(x, y, d, nx, ny, res, sel);
-    } else if (k < distance_compute_min_k_reservoir) {
-        HeapBlockResultHandler<CMin<float, int64_t>> res(nx, vals, ids, k);
-        knn_inner_product_select(x, y, d, nx, ny, res, sel);
-    } else {
-        ReservoirBlockResultHandler<CMin<float, int64_t>> res(nx, vals, ids, k);
-        knn_inner_product_select(x, y, d, nx, ny, res, sel);
-    }
+    Run_search_inner_product r;
+    dispatch_knn_ResultHandler(
+            nx, vals, ids, k, METRIC_INNER_PRODUCT, sel, r, x, y, d, nx, ny);
 
     if (imin != 0) {
         for (size_t i = 0; i < nx * k; i++) {
@@ -730,16 +905,11 @@ void knn_L2sqr(
         knn_L2sqr_by_idx(x, y, sela->ids, d, nx, ny, sela->n, k, vals, ids, 0);
         return;
     }
-    if (k == 1) {
-        Top1BlockResultHandler<CMax<float, int64_t>> res(nx, vals, ids);
-        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2, sel);
-    } else if (k < distance_compute_min_k_reservoir) {
-        HeapBlockResultHandler<CMax<float, int64_t>> res(nx, vals, ids, k);
-        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2, sel);
-    } else {
-        ReservoirBlockResultHandler<CMax<float, int64_t>> res(nx, vals, ids, k);
-        knn_L2sqr_select(x, y, d, nx, ny, res, y_norm2, sel);
-    }
+
+    Run_search_L2sqr r;
+    dispatch_knn_ResultHandler(
+            nx, vals, ids, k, METRIC_L2, sel, r, x, y, d, nx, ny, y_norm2);
+
     if (imin != 0) {
         for (size_t i = 0; i < nx * k; i++) {
             if (ids[i] >= 0) {
@@ -766,6 +936,7 @@ void knn_L2sqr(
  * Range search
  ***************************************************************************/
 
+// TODO accept a y_norm2 as well
 void range_search_L2sqr(
         const float* x,
         const float* y,
@@ -775,15 +946,9 @@ void range_search_L2sqr(
         float radius,
         RangeSearchResult* res,
         const IDSelector* sel) {
-    using RH = RangeSearchBlockResultHandler<CMax<float, int64_t>>;
-    RH resh(res, radius);
-    if (sel) {
-        exhaustive_L2sqr_seq<RH, true>(x, y, d, nx, ny, resh, sel);
-    } else if (nx < distance_compute_blas_threshold) {
-        exhaustive_L2sqr_seq(x, y, d, nx, ny, resh, sel);
-    } else {
-        exhaustive_L2sqr_blas(x, y, d, nx, ny, resh);
-    }
+    Run_search_L2sqr r;
+    dispatch_range_ResultHandler(
+            res, radius, METRIC_L2, sel, r, x, y, d, nx, ny, nullptr);
 }
 
 void range_search_inner_product(
@@ -795,15 +960,9 @@ void range_search_inner_product(
         float radius,
         RangeSearchResult* res,
         const IDSelector* sel) {
-    using RH = RangeSearchBlockResultHandler<CMin<float, int64_t>>;
-    RH resh(res, radius);
-    if (sel) {
-        exhaustive_inner_product_seq<RH, true>(x, y, d, nx, ny, resh, sel);
-    } else if (nx < distance_compute_blas_threshold) {
-        exhaustive_inner_product_seq(x, y, d, nx, ny, resh);
-    } else {
-        exhaustive_inner_product_blas(x, y, d, nx, ny, resh);
-    }
+    Run_search_inner_product r;
+    dispatch_range_ResultHandler(
+            res, radius, METRIC_INNER_PRODUCT, sel, r, x, y, d, nx, ny);
 }
 
 /***************************************************************************

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.
@@ -285,7 +285,7 @@ FAISS_API extern int distance_compute_blas_database_bs;
 // rather than a heap
 FAISS_API extern int distance_compute_min_k_reservoir;
 
-/** Return the k nearest neighors of each of the nx vectors x among the ny
+/** Return the k nearest neighbors of each of the nx vectors x among the ny
  *  vector y, w.r.t to max inner product.
  *
  * @param x    query vectors, size nx * d
@@ -301,7 +301,7 @@ void knn_inner_product(
         float_minheap_array_t* res,
         const IDSelector* sel = nullptr);
 
-/**  Return the k nearest neighors of each of the nx vectors x among the ny
+/**  Return the k nearest neighbors of each of the nx vectors x among the ny
  *  vector y, for the inner product metric.
  *
  * @param x    query vectors, size nx * d
@@ -320,7 +320,7 @@ void knn_inner_product(
         int64_t* indexes,
         const IDSelector* sel = nullptr);
 
-/** Return the k nearest neighors of each of the nx vectors x among the ny
+/** Return the k nearest neighbors of each of the nx vectors x among the ny
  *  vector y, for the L2 distance
  * @param x    query vectors, size nx * d
  * @param y    database vectors, size ny * d
@@ -338,7 +338,7 @@ void knn_L2sqr(
         const float* y_norm2 = nullptr,
         const IDSelector* sel = nullptr);
 
-/**  Return the k nearest neighors of each of the nx vectors x among the ny
+/**  Return the k nearest neighbors of each of the nx vectors x among the ny
  *  vector y, for the L2 distance
  *
  * @param x    query vectors, size nx * d
@@ -413,7 +413,7 @@ void knn_L2sqr_by_idx(
 /// Forward declaration, see AuxIndexStructures.h
 struct RangeSearchResult;
 
-/** Return the k nearest neighors of each of the nx vectors x among the ny
+/** Return the k nearest neighbors of each of the nx vectors x among the ny
  *  vector y, w.r.t to max inner product
  *
  * @param x      query vectors, size nx * d
@@ -469,7 +469,7 @@ void compute_PQ_dis_tables_dsub2(
  * @param n   size of the tables
  * @param a   size n
  * @param b   size n
- * @param c   restult table, size n
+ * @param c   result table, size n
  */
 void fvec_madd(size_t n, const float* a, float bf, const float* b, float* c);
 

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.