faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/utils/random.cpp

@@ -100,7 +100,7 @@ void float_rand(float* x, size_t n, int64_t seed) {
     int a0 = rng0.rand_int(), b0 = rng0.rand_int();
 
 #pragma omp parallel for
-    for (int64_t j = 0; j < nblock; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nblock); j++) {
         RandomGenerator rng(a0 + j * b0);
 
         const size_t istart = j * n / nblock;
@@ -120,7 +120,7 @@ void float_randn(float* x, size_t n, int64_t seed) {
     int a0 = rng0.rand_int(), b0 = rng0.rand_int();
 
 #pragma omp parallel for
-    for (int64_t j = 0; j < nblock; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nblock); j++) {
         RandomGenerator rng(a0 + j * b0);
 
         double a = 0, b = 0, s = 0;
@@ -155,7 +155,7 @@ void int64_rand(int64_t* x, size_t n, int64_t seed) {
     int a0 = rng0.rand_int(), b0 = rng0.rand_int();
 
 #pragma omp parallel for
-    for (int64_t j = 0; j < nblock; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nblock); j++) {
         RandomGenerator rng(a0 + j * b0);
 
         const size_t istart = j * n / nblock;
@@ -174,7 +174,7 @@ void int64_rand_max(int64_t* x, size_t n, uint64_t max, int64_t seed) {
     int a0 = rng0.rand_int(), b0 = rng0.rand_int();
 
 #pragma omp parallel for
-    for (int64_t j = 0; j < nblock; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nblock); j++) {
         RandomGenerator rng(a0 + j * b0);
 
         const size_t istart = j * n / nblock;
@@ -219,7 +219,7 @@ void byte_rand(uint8_t* x, size_t n, int64_t seed) {
     int a0 = rng0.rand_int(), b0 = rng0.rand_int();
 
 #pragma omp parallel for
-    for (int64_t j = 0; j < nblock; j++) {
+    for (int64_t j = 0; j < static_cast<int64_t>(nblock); j++) {
         RandomGenerator rng(a0 + j * b0);
 
         const size_t istart = j * n / nblock;
@@ -261,7 +261,7 @@ void rand_smooth_vectors(size_t n, size_t d, float* x, int64_t seed) {
     float_rand(scales.data(), d, seed + 2);
 
 #pragma omp parallel for if (n * d > 10000)
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
         for (size_t j = 0; j < d; j++) {
             x[i * d + j] = sinf(x[i * d + j] * (scales[j] * 4 + 0.1));
         }

data/vendor/faiss/faiss/utils/simd_impl/IVFFlatScanner-inl.h

@@ -0,0 +1,51 @@
+/*
+ * 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 <faiss/IndexIVFFlat.h>
+#include <faiss/impl/expanded_scanners.h>
+
+#ifndef THE_SIMD_LEVEL
+#error "THE_SIMD_LEVEL not defined"
+#endif
+
+namespace faiss {
+
+constexpr faiss::SIMDLevel THE_SL = THE_SIMD_LEVEL;
+
+#define DEFINE_IVFFLAT_SCANNER_METHODS(mt)                                     \
+    template <>                                                                \
+    float IVFFlatScanner<VectorDistance<mt, THE_SL>>::distance_to_code(        \
+            const uint8_t* code) const {                                       \
+        const float* yj = (float*)code;                                        \
+        return vd(xi, yj);                                                     \
+    }                                                                          \
+    template <>                                                                \
+    size_t IVFFlatScanner<VectorDistance<mt, THE_SL>>::scan_codes(             \
+            size_t list_size,                                                  \
+            const uint8_t* codes,                                              \
+            const idx_t* ids,                                                  \
+            ResultHandler& handler) const {                                    \
+        return run_scan_codes_fix_C<C>(*this, list_size, codes, ids, handler); \
+    }
+
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_L2)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_INNER_PRODUCT)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_L1)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Linf)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Lp)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Canberra)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_BrayCurtis)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_JensenShannon)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Jaccard)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_NaNEuclidean)
+DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_GOWER)
+
+#undef DEFINE_IVFFLAT_SCANNER_METHODS
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/distances_aarch64.cpp

@@ -12,8 +12,12 @@
 #include <arm_neon.h>
 #include <limits>
 
-#define AUTOVEC_LEVEL SIMDLevel::ARM_NEON
+#define THE_SIMD_LEVEL SIMDLevel::ARM_NEON
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
+
+#include <faiss/utils/simd_impl/distances_simdlib256.h>
 
 namespace faiss {
 

data/vendor/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp

@@ -9,18 +9,47 @@
 
 #include <faiss/utils/distances.h>
 
-#define AUTOVEC_LEVEL SIMDLevel::ARM_SVE
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances_fused/distances_fused.h>
+#include <faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h>
+
+#ifndef FINTEGER
+#define FINTEGER long
+#endif
+
+extern "C" {
+
+int sgemm_(
+        const char* transa,
+        const char* transb,
+        FINTEGER* m,
+        FINTEGER* n,
+        FINTEGER* k,
+        const float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* b,
+        FINTEGER* ldb,
+        float* beta,
+        float* c,
+        FINTEGER* ldc);
+}
+
+#define THE_SIMD_LEVEL SIMDLevel::ARM_SVE
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
 
 namespace faiss {
 
 template <>
 void fvec_madd<SIMDLevel::ARM_SVE>(
         const size_t n,
-        const float* __restrict a,
+        const float* a,
         const float bf,
-        const float* __restrict b,
-        float* __restrict c) {
+        const float* b,
+        float* c) {
     const size_t lanes = static_cast<size_t>(svcntw());
     const size_t lanes2 = lanes * 2;
     const size_t lanes3 = lanes * 3;
@@ -565,4 +594,184 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
     return nearest_idx;
 }
 
+template <>
+void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_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 schedule(static) if ((i1 - i0) >= 16)
+            for (int64_t i = static_cast<int64_t>(i0);
+                 i < static_cast<int64_t>(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();
+    }
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/distances_autovec-inl.h

@@ -7,14 +7,26 @@
 
 #pragma once
 
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/extra_distances.h>
+
+#ifndef THE_SIMD_LEVEL
+#error "THE_SIMD_LEVEL not defined"
+#endif
 
 namespace faiss {
 
+constexpr faiss::SIMDLevel SL = THE_SIMD_LEVEL;
+/******************************************************************
+ * These functions are simple enough that the compile will do a good job
+ * vectorizing them given the appropriate flags.
+ ******************************************************************/
+
 FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
-float fvec_norm_L2sqr<AUTOVEC_LEVEL>(const float* x, size_t d) {
+float fvec_norm_L2sqr<SL>(const float* x, size_t d) {
     // the double in the _ref is suspected to be a typo. Some of the manual
     // implementations this replaces used float.
     float res = 0;
@@ -29,7 +41,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
 FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
-float fvec_L2sqr<AUTOVEC_LEVEL>(const float* x, const float* y, size_t d) {
+float fvec_L2sqr<SL>(const float* x, const float* y, size_t d) {
     size_t i;
     float res = 0;
     FAISS_PRAGMA_IMPRECISE_LOOP
@@ -43,10 +55,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
 FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
-float fvec_inner_product<AUTOVEC_LEVEL>(
-        const float* x,
-        const float* y,
-        size_t d) {
+float fvec_inner_product<SL>(const float* x, const float* y, size_t d) {
     float res = 0.F;
     FAISS_PRAGMA_IMPRECISE_LOOP
     for (size_t i = 0; i != d; ++i) {
@@ -58,7 +67,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
 FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
-float fvec_L1<AUTOVEC_LEVEL>(const float* x, const float* y, size_t d) {
+float fvec_L1<SL>(const float* x, const float* y, size_t d) {
     size_t i;
     float res = 0;
     FAISS_PRAGMA_IMPRECISE_LOOP
@@ -72,7 +81,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
 FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
-float fvec_Linf<AUTOVEC_LEVEL>(const float* x, const float* y, size_t d) {
+float fvec_Linf<SL>(const float* x, const float* y, size_t d) {
     float res = 0;
     FAISS_PRAGMA_IMPRECISE_LOOP
     for (size_t i = 0; i < d; i++) {
@@ -84,7 +93,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
 FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
-void fvec_inner_product_batch_4<AUTOVEC_LEVEL>(
+void fvec_inner_product_batch_4<SL>(
         const float* x,
         const float* y0,
         const float* y1,
@@ -116,7 +125,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
 FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
-void fvec_L2sqr_batch_4<AUTOVEC_LEVEL>(
+void fvec_L2sqr_batch_4<SL>(
         const float* x,
         const float* y0,
         const float* y1,
@@ -150,4 +159,148 @@ void fvec_L2sqr_batch_4<AUTOVEC_LEVEL>(
 }
 FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
+/******************************************************************
+ * VectorDistance::operator() specializations — defined out-of-class
+ * so that SIMD compilation units produce externally-linkable symbols.
+ ******************************************************************/
+
+template <>
+float VectorDistance<METRIC_L2, SL>::operator()(const float* x, const float* y)
+        const {
+    return fvec_L2sqr<SL>(x, y, this->d);
+}
+
+template <>
+float VectorDistance<METRIC_INNER_PRODUCT, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    return fvec_inner_product<SL>(x, y, this->d);
+}
+
+template <>
+float VectorDistance<METRIC_L1, SL>::operator()(const float* x, const float* y)
+        const {
+    return fvec_L1<SL>(x, y, this->d);
+}
+
+template <>
+float VectorDistance<METRIC_Linf, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    return fvec_Linf<SL>(x, y, this->d);
+}
+
+template <>
+float VectorDistance<METRIC_Lp, SL>::operator()(const float* x, const float* y)
+        const {
+    float accu = 0;
+    for (size_t i = 0; i < this->d; i++) {
+        float diff = fabs(x[i] - y[i]);
+        accu += powf(diff, this->metric_arg);
+    }
+    return accu;
+}
+
+template <>
+float VectorDistance<METRIC_Canberra, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    float accu = 0;
+    for (size_t i = 0; i < this->d; i++) {
+        float xi = x[i], yi = y[i];
+        accu += fabs(xi - yi) / (fabs(xi) + fabs(yi));
+    }
+    return accu;
+}
+
+template <>
+float VectorDistance<METRIC_BrayCurtis, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    float accu_num = 0, accu_den = 0;
+    for (size_t i = 0; i < this->d; i++) {
+        float xi = x[i], yi = y[i];
+        accu_num += fabs(xi - yi);
+        accu_den += fabs(xi + yi);
+    }
+    return accu_num / accu_den;
+}
+
+template <>
+float VectorDistance<METRIC_JensenShannon, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    float accu = 0;
+    for (size_t i = 0; i < this->d; i++) {
+        float xi = x[i], yi = y[i];
+        float mi = 0.5 * (xi + yi);
+        float kl1 = -xi * log(mi / xi);
+        float kl2 = -yi * log(mi / yi);
+        accu += kl1 + kl2;
+    }
+    return 0.5 * accu;
+}
+
+template <>
+float VectorDistance<METRIC_Jaccard, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    float accu_num = 0, accu_den = 0;
+    for (size_t i = 0; i < this->d; i++) {
+        accu_num += fmin(x[i], y[i]);
+        accu_den += fmax(x[i], y[i]);
+    }
+    return accu_num / accu_den;
+}
+
+template <>
+float VectorDistance<METRIC_NaNEuclidean, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    float accu = 0;
+    size_t present = 0;
+    for (size_t i = 0; i < this->d; i++) {
+        if (!std::isnan(x[i]) && !std::isnan(y[i])) {
+            float diff = x[i] - y[i];
+            accu += diff * diff;
+            present++;
+        }
+    }
+    if (present == 0) {
+        return NAN;
+    }
+    return float(this->d) / float(present) * accu;
+}
+
+template <>
+float VectorDistance<METRIC_GOWER, SL>::operator()(
+        const float* x,
+        const float* y) const {
+    float accu = 0;
+    size_t valid_dims = 0;
+
+    for (size_t i = 0; i < this->d; i++) {
+        if (std::isnan(x[i]) || std::isnan(y[i])) {
+            continue;
+        }
+
+        if (x[i] >= 0 && y[i] >= 0) {
+            if (x[i] > 1 || y[i] > 1) {
+                return std::numeric_limits<float>::quiet_NaN();
+            }
+            accu += fabs(x[i] - y[i]);
+        } else if (x[i] < 0 && y[i] < 0) {
+            accu += float(int(x[i] != y[i]));
+        } else {
+            return std::numeric_limits<float>::quiet_NaN();
+        }
+        valid_dims++;
+    }
+
+    if (valid_dims == 0) {
+        return std::numeric_limits<float>::quiet_NaN();
+    }
+    return accu / valid_dims;
+}
+
 } // namespace faiss