faiss 0.6.2 → 0.6.4

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

@@ -129,6 +129,22 @@ struct ElementOpIP {
     }
 };
 
+struct ElementOpL2 {
+    static svfloat32_t op(svbool_t pg, svfloat32_t x, svfloat32_t y) {
+        const svfloat32_t diff = svsub_f32_x(pg, x, y);
+        return svmul_f32_x(pg, diff, diff);
+    }
+
+    static svfloat32_t merge(
+            svbool_t pg,
+            svfloat32_t z,
+            svfloat32_t x,
+            svfloat32_t y) {
+        const svfloat32_t diff = svsub_f32_x(pg, x, y);
+        return svmla_f32_x(pg, z, diff, diff);
+    }
+};
+
 template <typename ElementOp>
 void fvec_op_ny_sve_d1(float* dis, const float* x, const float* y, size_t ny) {
     const size_t lanes = svcntw();
@@ -513,13 +529,87 @@ void fvec_L2sqr_ny<SIMDLevel::ARM_SVE>(
         const float* y,
         size_t d,
         size_t ny) {
-    // Use autovectorized L2sqr in a loop
-    for (size_t i = 0; i < ny; i++) {
-        dis[i] = fvec_L2sqr<SIMDLevel::ARM_SVE>(x, y, d);
-        y += d;
+    const size_t lanes = static_cast<size_t>(svcntw());
+
+    switch (d) {
+        case 1:
+            fvec_op_ny_sve_d1<ElementOpL2>(dis, x, y, ny);
+            break;
+
+        case 2:
+            fvec_op_ny_sve_d2<ElementOpL2>(dis, x, y, ny);
+            break;
+
+        case 4:
+            fvec_op_ny_sve_d4<ElementOpL2>(dis, x, y, ny);
+            break;
+
+        case 8:
+            fvec_op_ny_sve_d8<ElementOpL2>(dis, x, y, ny);
+            break;
+
+        default:
+            if (d == lanes)
+                fvec_op_ny_sve_lanes1<ElementOpL2>(dis, x, y, ny);
+            else if (d == lanes * 2)
+                fvec_op_ny_sve_lanes2<ElementOpL2>(dis, x, y, ny);
+            else if (d == lanes * 3)
+                fvec_op_ny_sve_lanes3<ElementOpL2>(dis, x, y, ny);
+            else if (d == lanes * 4)
+                fvec_op_ny_sve_lanes4<ElementOpL2>(dis, x, y, ny);
+            else {
+                // Fallback: use autovectorized L2sqr
+                for (size_t i = 0; i < ny; i++) {
+                    dis[i] = fvec_L2sqr<SIMDLevel::ARM_SVE>(x, y, d);
+                    y += d;
+                }
+            }
+            break;
+    }
+}
+
+namespace {
+
+/// Low-dimensional L2sqr nearest (D in {2,4,8}). The data is row-major
+/// (centroid c is y[c*D .. c*D+D]). Each SVE lane tracks one centroid across
+/// batches, keeping the lane-local minimum and index in registers. The scratch
+/// buffer is not written, matching the AVX2/AVX512 D2/D4/D8 implementations.
+template <int D>
+size_t fvec_L2sqr_ny_nearest_lowdim(
+        float* /*distances_tmp_buffer*/,
+        const float* x,
+        const float* y,
+        size_t ny) {
+    const size_t lanes = svcntw();
+    svfloat32_t global_mins = svdup_n_f32(HUGE_VALF);
+    svuint32_t global_ids = svdup_n_u32(0);
+    svuint32_t current_ids = svindex_u32(0, 1);
+
+    for (size_t c = 0; c < ny; c += lanes) {
+        const svbool_t pg = svwhilelt_b32_u64(c, ny);
+        svfloat32_t distances = svdup_n_f32(0.0f);
+        for (uint32_t j = 0; j < D; ++j) {
+            const svuint32_t offsets = svindex_u32(j, D);
+            const svfloat32_t yv =
+                    svld1_gather_u32index_f32(pg, y + c * D, offsets);
+            const svfloat32_t diff = svsub_n_f32_x(pg, yv, x[j]);
+            distances = svmla_f32_m(pg, distances, diff, diff);
+        }
+
+        const svbool_t closer = svcmplt_f32(pg, distances, global_mins);
+        global_mins = svsel_f32(closer, distances, global_mins);
+        global_ids = svsel_u32(closer, current_ids, global_ids);
+        current_ids =
+                svadd_n_u32_x(pg, current_ids, static_cast<uint32_t>(lanes));
     }
+
+    const svbool_t all = svptrue_b32();
+    const float global_min = svminv_f32(all, global_mins);
+    return svminv_u32(svcmpeq_n_f32(all, global_mins, global_min), global_ids);
 }
 
+} // namespace
+
 template <>
 size_t fvec_L2sqr_ny_nearest<SIMDLevel::ARM_SVE>(
         float* distances_tmp_buffer,
@@ -527,14 +617,52 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::ARM_SVE>(
         const float* y,
         size_t d,
         size_t ny) {
-    fvec_L2sqr_ny<SIMDLevel::ARM_SVE>(distances_tmp_buffer, x, y, d, ny);
+    switch (d) {
+        case 2:
+            return fvec_L2sqr_ny_nearest_lowdim<2>(
+                    distances_tmp_buffer, x, y, ny);
+        case 4:
+            return fvec_L2sqr_ny_nearest_lowdim<4>(
+                    distances_tmp_buffer, x, y, ny);
+        case 8:
+            return fvec_L2sqr_ny_nearest_lowdim<8>(
+                    distances_tmp_buffer, x, y, ny);
+    }
+
+    const size_t lanes = static_cast<size_t>(svcntw());
 
     size_t nearest_idx = 0;
     float min_dis = HUGE_VALF;
 
-    for (size_t i = 0; i < ny; i++) {
-        if (distances_tmp_buffer[i] < min_dis) {
-            min_dis = distances_tmp_buffer[i];
+    for (size_t i = 0; i < ny; ++i) {
+        const float* yi = y + i * d;
+        size_t j = 0;
+
+        svfloat32_t accv = svdup_n_f32(0.0f);
+
+        for (; j + lanes <= d; j += lanes) {
+            const svbool_t pg = svptrue_b32();
+            const svfloat32_t xv = svld1_f32(pg, x + j);
+            const svfloat32_t yv = svld1_f32(pg, yi + j);
+            const svfloat32_t diff = svsub_f32_x(pg, xv, yv);
+            accv = svmla_f32_x(pg, accv, diff, diff);
+        }
+
+        if (j < d) {
+            const svbool_t pg = svwhilelt_b32_u64(j, d);
+            const svfloat32_t xv = svld1_f32(pg, x + j);
+            const svfloat32_t yv = svld1_f32(pg, yi + j);
+            const svfloat32_t diff = svsub_f32_x(pg, xv, yv);
+            // Merging predication: lanes outside the tail keep the partial
+            // sums accumulated above, which the reduction below still adds in.
+            // `_x` would leave them unspecified.
+            accv = svmla_f32_m(pg, accv, diff, diff);
+        }
+
+        const float dist = svaddv_f32(svptrue_b32(), accv);
+        distances_tmp_buffer[i] = dist;
+        if (dist < min_dis) {
+            min_dis = dist;
             nearest_idx = i;
         }
     }
@@ -542,7 +670,6 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::ARM_SVE>(
     return nearest_idx;
 }
 
-FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
 template <>
 void fvec_L2sqr_ny_transposed<SIMDLevel::ARM_SVE>(
         float* dis,
@@ -552,22 +679,26 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::ARM_SVE>(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    float x_sqlen = 0;
-    FAISS_PRAGMA_IMPRECISE_LOOP
-    for (size_t j = 0; j < d; j++) {
-        x_sqlen += x[j] * x[j];
-    }
+    const size_t lanes = static_cast<size_t>(svcntw());
+    const float x_sq = fvec_norm_L2sqr(x, d);
+
+    for (size_t k = 0; k < ny; k += lanes) {
+        svbool_t pg = svwhilelt_b32_u64(k, ny);
+        svfloat32_t acc = svdup_n_f32(0.0f);
 
-    for (size_t i = 0; i < ny; i++) {
-        float dp = 0;
-        FAISS_PRAGMA_IMPRECISE_LOOP
-        for (size_t j = 0; j < d; j++) {
-            dp += x[j] * y[i + j * d_offset];
+        for (size_t j = 0; j < d; ++j) {
+            svfloat32_t ychunk = svld1_f32(pg, y + j * d_offset + k);
+            svfloat32_t xj = svdup_n_f32(x[j]);
+            acc = svmla_f32_x(pg, acc, xj, ychunk);
         }
-        dis[i] = x_sqlen + y_sqlen[i] - 2 * dp;
+
+        svfloat32_t ysq = svld1_f32(pg, y_sqlen + k);
+        svfloat32_t two_acc = svmul_f32_x(pg, acc, svdup_n_f32(2.0f));
+        svfloat32_t sum = svadd_f32_x(pg, svdup_n_f32(x_sq), ysq);
+        svfloat32_t res = svsub_f32_x(pg, sum, two_acc);
+        svst1_f32(pg, dis + k, res);
     }
 }
-FAISS_PRAGMA_IMPRECISE_FUNCTION_END
 
 template <>
 size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
@@ -578,20 +709,53 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    fvec_L2sqr_ny_transposed<SIMDLevel::ARM_SVE>(
-            distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
+    const size_t lanes = svcntw();
+    const float x_sq = fvec_norm_L2sqr(x, d);
 
-    size_t nearest_idx = 0;
-    float min_dis = HUGE_VALF;
+    size_t current_min_idx = 0;
+    svfloat32_t current_min_v = svdup_n_f32(HUGE_VALF);
 
-    for (size_t i = 0; i < ny; i++) {
-        if (distances_tmp_buffer[i] < min_dis) {
-            min_dis = distances_tmp_buffer[i];
-            nearest_idx = i;
+    // Max SVE vector length is 2048 bits → 64 fp32 lanes
+    float tmp_buf[64];
+
+    for (size_t k = 0; k < ny; k += lanes) {
+        svbool_t pg = svwhilelt_b32_u64(k, ny);
+        svfloat32_t acc = svdup_n_f32(0.0f);
+
+        for (size_t j = 0; j < d; ++j) {
+            svfloat32_t ychunk = svld1_f32(pg, y + j * d_offset + k);
+            svfloat32_t xj = svdup_n_f32(x[j]);
+            acc = svmla_f32_x(pg, acc, xj, ychunk);
+        }
+
+        svfloat32_t ysq = svld1_f32(pg, y_sqlen + k);
+        svfloat32_t two_acc = svmul_f32_x(pg, acc, svdup_n_f32(2.0f));
+        svfloat32_t sum = svadd_f32_x(pg, svdup_n_f32(x_sq), ysq);
+        svfloat32_t res = svsub_f32_x(pg, sum, two_acc);
+
+        svst1_f32(pg, distances_tmp_buffer + k, res);
+
+        svbool_t less_mask = svcmplt_f32(pg, res, current_min_v);
+
+        if (svptest_any(pg, less_mask)) {
+            float vec_min = svminv_f32(pg, res);
+
+            current_min_v =
+                    svmin_f32_x(pg, current_min_v, svdup_n_f32(vec_min));
+
+            svst1_f32(pg, tmp_buf, res);
+
+            size_t cnt = (size_t)svcntw();
+            for (size_t lane = 0; lane < cnt && (k + lane) < ny; ++lane) {
+                if (tmp_buf[lane] == vec_min) {
+                    current_min_idx = k + lane;
+                    break;
+                }
+            }
         }
     }
 
-    return nearest_idx;
+    return current_min_idx;
 }
 
 template <>

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

@@ -11,7 +11,6 @@
 
 #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

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

@@ -9,6 +9,35 @@
 
 #include <immintrin.h>
 
+#include <memory>
+
+#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::AVX512
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
@@ -954,21 +983,6 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::AVX512>(
             &fvec_L2sqr_ny_nearest_D8<SIMDLevel::AVX512>);
 }
 
-template <>
-size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::AVX512>(
-        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<SIMDLevel::NONE>(
-            distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
-}
-
-// TODO: Following functions are not used in the current codebase. Check AVX2 ,
-// respective implementation has been used
 template <size_t DIM>
 size_t fvec_L2sqr_ny_nearest_y_transposed_D(
         float* /* distances_tmp_buffer */,
@@ -1082,6 +1096,33 @@ size_t fvec_L2sqr_ny_nearest_y_transposed_D(
     return current_min_index;
 }
 
+template <>
+size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::AVX512>(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t ny) {
+    // optimized for a few special cases
+#define DISPATCH(dval)                                     \
+    case dval:                                             \
+        return fvec_L2sqr_ny_nearest_y_transposed_D<dval>( \
+                distances_tmp_buffer, x, y, y_sqlen, d_offset, ny);
+
+    switch (d) {
+        DISPATCH(1)
+        DISPATCH(2)
+        DISPATCH(4)
+        DISPATCH(8)
+        default:
+            return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
+                    distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
+    }
+#undef DISPATCH
+}
+
 template <>
 int fvec_madd_and_argmin<SIMDLevel::AVX512>(
         size_t n,
@@ -1092,4 +1133,211 @@ int fvec_madd_and_argmin<SIMDLevel::AVX512>(
     return fvec_madd_and_argmin_sse(n, a, bf, b, c);
 }
 
+template <>
+void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
+        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;
+    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);
+
+    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++) {
+                float* ip_line = ip_block.get() + (i - i0) * (j1 - j0);
+
+                _mm_prefetch((const char*)ip_line, _MM_HINT_NTA);
+                _mm_prefetch((const char*)(ip_line + 16), _MM_HINT_NTA);
+
+                // constant
+                const __m512 mul_minus2 = _mm512_set1_ps(-2);
+
+                // Track 16 min distances + 16 min indices.
+                // All the distances tracked do not take x_norms[i]
+                // into account in order to get rid of extra
+                // _mm512_add_ps(x_norms[i], ...) instructions
+                // in distance computations.
+                __m512 min_distances =
+                        _mm512_set1_ps(res.dis_tab[i] - x_norms[i]);
+
+                // these indices are local and are relative to j0.
+                // so, value 0 means j0.
+                __m512i min_indices = _mm512_set1_epi32(0);
+
+                __m512i current_indices = _mm512_setr_epi32(
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
+                const __m512i indices_delta = _mm512_set1_epi32(16);
+
+                // current j index
+                size_t idx_j = 0;
+                size_t count = j1 - j0;
+
+                // process 32 elements per loop
+                for (; idx_j < (count / 32) * 32; idx_j += 32, ip_line += 32) {
+                    _mm_prefetch((const char*)(ip_line + 32), _MM_HINT_NTA);
+                    _mm_prefetch((const char*)(ip_line + 48), _MM_HINT_NTA);
+
+                    // load values for norms
+                    const __m512 y_norm_0 =
+                            _mm512_loadu_ps(y_norms + idx_j + j0 + 0);
+                    const __m512 y_norm_1 =
+                            _mm512_loadu_ps(y_norms + idx_j + j0 + 16);
+
+                    // load values for dot products
+                    const __m512 ip_0 = _mm512_loadu_ps(ip_line + 0);
+                    const __m512 ip_1 = _mm512_loadu_ps(ip_line + 16);
+
+                    // 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.
+                    __m512 distances_0 =
+                            _mm512_fmadd_ps(ip_0, mul_minus2, y_norm_0);
+                    __m512 distances_1 =
+                            _mm512_fmadd_ps(ip_1, mul_minus2, y_norm_1);
+
+                    // compare the new distances to the min distances
+                    // for each of the first group of 16 AVX512 components.
+                    const __mmask16 comparison_0 = _mm512_cmp_ps_mask(
+                            min_distances, distances_0, _CMP_LE_OS);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances = _mm512_mask_blend_ps(
+                            comparison_0, distances_0, min_distances);
+                    min_indices = _mm512_mask_blend_epi32(
+                            comparison_0, current_indices, min_indices);
+                    current_indices =
+                            _mm512_add_epi32(current_indices, indices_delta);
+
+                    // compare the new distances to the min distances
+                    // for each of the second group of 16 AVX512 components.
+                    const __mmask16 comparison_1 = _mm512_cmp_ps_mask(
+                            min_distances, distances_1, _CMP_LE_OS);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances = _mm512_mask_blend_ps(
+                            comparison_1, distances_1, min_distances);
+                    min_indices = _mm512_mask_blend_epi32(
+                            comparison_1, current_indices, min_indices);
+                    current_indices =
+                            _mm512_add_epi32(current_indices, indices_delta);
+                }
+
+                // dump values and find the minimum distance / minimum index
+                float min_distances_scalar[16];
+                uint32_t min_indices_scalar[16];
+                _mm512_storeu_ps(min_distances_scalar, min_distances);
+                _mm512_storeu_si512(
+                        (__m512i*)(min_indices_scalar), min_indices);
+
+                float current_min_distance = res.dis_tab[i];
+                uint32_t current_min_index = res.ids_tab[i];
+
+                // if two indices are represented with equal distance values,
+                // then the index with the min value is returned.
+                for (size_t jv = 0; jv < 16; jv++) {
+                    // add missing x_norms[i]
+                    float distance_candidate =
+                            min_distances_scalar[jv] + x_norms[i];
+
+                    // negative values can occur for identical vectors
+                    //    due to roundoff errors.
+                    if (distance_candidate < 0) {
+                        distance_candidate = 0;
+                    }
+
+                    int64_t index_candidate = min_indices_scalar[jv] + j0;
+
+                    if (current_min_distance > distance_candidate) {
+                        current_min_distance = distance_candidate;
+                        current_min_index = index_candidate;
+                    } else if (
+                            current_min_distance == distance_candidate &&
+                            current_min_index > index_candidate) {
+                        current_min_index = index_candidate;
+                    }
+                }
+
+                // process leftovers
+                for (; idx_j < count; idx_j++, ip_line++) {
+                    float ip = *ip_line;
+                    float dis = x_norms[i] + y_norms[idx_j + j0] - 2 * ip;
+                    // negative values can occur for identical vectors
+                    //    due to roundoff errors.
+                    if (dis < 0) {
+                        dis = 0;
+                    }
+
+                    if (current_min_distance > dis) {
+                        current_min_distance = dis;
+                        current_min_index = idx_j + j0;
+                    }
+                }
+
+                //
+                res.add_result(i, current_min_distance, current_min_index);
+            }
+        }
+        // Does nothing for SingleBestResultHandler, but
+        // keeping the call for the consistency.
+        res.end_multiple();
+        InterruptCallback::check();
+    }
+}
+
 } // namespace faiss