faiss 0.6.3 → 0.6.4

data/vendor/faiss/faiss/utils/hamming_distance/{hamming_computer-avx512_spr.h → hamming_computer-avx512_vpopcnt.h}

@@ -5,17 +5,19 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#ifndef HAMMING_COMPUTER_AVX512_SPR_H
-#define HAMMING_COMPUTER_AVX512_SPR_H
+#ifndef HAMMING_COMPUTER_AVX512_VPOPCNT_H
+#define HAMMING_COMPUTER_AVX512_VPOPCNT_H
 
-// AVX512_SPR HammingComputer specializations using VPOPCNTDQ.
-// On Sapphire Rapids+, _mm512_popcnt_epi64 (and _mm256_popcnt_epi64 with VL)
-// are unconditionally available. This gives a faster path than the scalar
-// popcount fallback used in the base AVX512 specializations when compiled
-// without -mavx512vpopcntdq.
+// AVX512_VPOPCNT HammingComputer specializations. The 32/64/Default kernels
+// use VPOPCNTDQ; the batched 20-byte kernel uses AVX512_BITALG. This gives
+// a faster path than the scalar popcount fallback used in the base AVX512
+// specializations when compiled without -mavx512vpopcntdq.
 
 #include <cassert>
 #include <cstdint>
+#include <cstring>
+
+#include <faiss/utils/popcount.h>
 
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/hamming_distance/hamming_computer-avx512.h>
@@ -25,31 +27,90 @@
 namespace faiss {
 
 /***************************************************************************
- * AVX512_SPR inheriting specializations for types without custom SPR code.
+ * AVX512_VPOPCNT inheriting specializations without custom VPOPCNT code.
  ***************************************************************************/
 
-#define FAISS_INHERIT_HAMMING_SPR(Class)                                   \
-    template <>                                                            \
-    struct Class##                                                         \
-            _tpl<SIMDLevel::AVX512_SPR> : Class##_tpl<SIMDLevel::AVX512> { \
-        using Class##_tpl<SIMDLevel::AVX512>::Class##_tpl;                 \
+#define FAISS_INHERIT_HAMMING_VPOPCNT(Class)                                   \
+    template <>                                                                \
+    struct Class##                                                             \
+            _tpl<SIMDLevel::AVX512_VPOPCNT> : Class##_tpl<SIMDLevel::AVX512> { \
+        using Class##_tpl<SIMDLevel::AVX512>::Class##_tpl;                     \
     }
 
-FAISS_INHERIT_HAMMING_SPR(HammingComputer16);
-FAISS_INHERIT_HAMMING_SPR(HammingComputer20);
-FAISS_INHERIT_HAMMING_SPR(GenHammingComputer8);
-FAISS_INHERIT_HAMMING_SPR(GenHammingComputer16);
-FAISS_INHERIT_HAMMING_SPR(GenHammingComputer32);
-FAISS_INHERIT_HAMMING_SPR(GenHammingComputerM8);
+FAISS_INHERIT_HAMMING_VPOPCNT(HammingComputer16);
+FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer8);
+FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer16);
+FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputer32);
+FAISS_INHERIT_HAMMING_VPOPCNT(GenHammingComputerM8);
 
-#undef FAISS_INHERIT_HAMMING_SPR
+#undef FAISS_INHERIT_HAMMING_VPOPCNT
 
 /***************************************************************************
- * Custom AVX512_SPR specializations using VPOPCNTDQ.
+ * Custom AVX512_VPOPCNT specializations using VPOPCNTDQ.
  ***************************************************************************/
 
 template <>
-struct HammingComputer32_tpl<SIMDLevel::AVX512_SPR> {
+struct HammingComputer20_tpl<SIMDLevel::AVX512_VPOPCNT>
+        : HammingComputer20_tpl<SIMDLevel::AVX512> {
+    using HammingComputer20_tpl<SIMDLevel::AVX512>::HammingComputer20_tpl;
+
+    static constexpr size_t batch_size = 8;
+    static constexpr size_t kStride = get_code_size();
+    // 160 bytes is what the three loads in hamming_batch() cover, and the
+    // 16+4 or 4+16 split it applies per lane is written out for 20 bytes:
+    // the literal offsets and the two group indices are not derived from
+    // kStride, so another width needs the body reworked, not just retuned.
+    static_assert(batch_size * kStride == 160);
+    static_assert(kStride == 20, "hamming_batch() hardcodes the 16+4 split");
+    static constexpr __mmask64 kTailMask = 0xFFFFFFFFull;
+
+    /// Writes the query repeated batch_size times. The caller owns the buffer,
+    /// so a computer used only through hamming() carries no batch state.
+    static void build_batch_query(const uint8_t* a8, uint8_t* tile) {
+        for (size_t k = 0; k < batch_size; k++) {
+            memcpy(tile + k * kStride, a8, kStride);
+        }
+    }
+
+    static void hamming_batch(
+            const uint8_t* tile,
+            const uint8_t* codes,
+            int32_t* dis) {
+        const __m512i zero = _mm512_setzero_si512();
+        const __m512i p0 = _mm512_popcnt_epi8(_mm512_xor_si512(
+                _mm512_loadu_si512(codes), _mm512_loadu_si512(tile)));
+        const __m512i p1 = _mm512_popcnt_epi8(_mm512_xor_si512(
+                _mm512_loadu_si512(codes + 64), _mm512_loadu_si512(tile + 64)));
+        const __m512i p2 = _mm512_popcnt_epi8(_mm512_xor_si512(
+                _mm512_maskz_loadu_epi8(kTailMask, codes + 128),
+                _mm512_maskz_loadu_epi8(kTailMask, tile + 128)));
+
+        alignas(64) uint64_t grp[24];
+        _mm512_store_si512(grp, _mm512_sad_epu8(p0, zero));
+        _mm512_store_si512(grp + 8, _mm512_sad_epu8(p1, zero));
+        _mm512_store_si512(grp + 16, _mm512_sad_epu8(p2, zero));
+
+        for (size_t k = 0; k < batch_size; k++) {
+            const size_t s = k * kStride;
+            const size_t g = s / 8;
+            uint32_t xh, qh;
+            if (s % 8 == 0) {
+                memcpy(&xh, codes + s + 16, 4);
+                memcpy(&qh, tile + s + 16, 4);
+                dis[k] = static_cast<int32_t>(
+                        grp[g] + grp[g + 1] + popcount32(xh ^ qh));
+            } else {
+                memcpy(&xh, codes + s, 4);
+                memcpy(&qh, tile + s, 4);
+                dis[k] = static_cast<int32_t>(
+                        popcount32(xh ^ qh) + grp[g + 1] + grp[g + 2]);
+            }
+        }
+    }
+};
+
+template <>
+struct HammingComputer32_tpl<SIMDLevel::AVX512_VPOPCNT> {
     const uint8_t* a8;
 
     HammingComputer32_tpl() {}
@@ -81,7 +142,7 @@ struct HammingComputer32_tpl<SIMDLevel::AVX512_SPR> {
 };
 
 template <>
-struct HammingComputer64_tpl<SIMDLevel::AVX512_SPR> {
+struct HammingComputer64_tpl<SIMDLevel::AVX512_VPOPCNT> {
     const uint8_t* a8;
 
     HammingComputer64_tpl() {}
@@ -108,7 +169,7 @@ struct HammingComputer64_tpl<SIMDLevel::AVX512_SPR> {
 };
 
 template <>
-struct HammingComputerDefault_tpl<SIMDLevel::AVX512_SPR> {
+struct HammingComputerDefault_tpl<SIMDLevel::AVX512_VPOPCNT> {
     const uint8_t* a8;
     int quotient8;
     int remainder8;

data/vendor/faiss/faiss/utils/hamming_distance/hamming_impl.h

@@ -36,6 +36,7 @@
 
 #include <algorithm>
 #include <cstdio>
+#include <cstring>
 #include <limits>
 #include <memory>
 #include <vector>
@@ -148,6 +149,7 @@ void hammings_knn_mc_impl(
     std::unique_ptr<int64_t[]> all_ids_per_dis(new int64_t[na * nBuckets * k]);
 
     std::vector<HCounterState<HammingComputer>> cs;
+    cs.reserve(na);
     for (size_t i = 0; i < na; ++i) {
         cs.push_back(
                 HCounterState<HammingComputer>(
@@ -298,6 +300,132 @@ void generalized_hammings_knn_hc_impl(
     }
 }
 
+/* Both number of words and remainder are constants, so the compiler flattens
+ * the per-chunk loop. The query is copied into qw/qt once, not re-read for
+ * every candidate. */
+template <int NW, int REM>
+void hammings_ragged_fixed(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    for (size_t i = 0; i < na; i++) {
+        const uint8_t* __restrict ai = a + i * ncodes;
+        hamdis_t* __restrict dis_ = dis + i * nb;
+        // Read the query once here, not once per candidate.
+        uint64_t qw[NW == 0 ? 1 : NW];
+        for (int k = 0; k < NW; k++) {
+            memcpy(&qw[k], ai + k * 8, 8);
+        }
+        uint64_t qt = 0;
+        if constexpr (REM > 0) {
+            memcpy(&qt, ai + NW * 8, REM);
+        }
+        for (size_t j = 0; j < nb; j++) {
+            const uint8_t* __restrict bj = b + j * ncodes;
+            hamdis_t h = 0;
+            uint64_t y;
+            for (int k = 0; k < NW; k++) {
+                memcpy(&y, bj + k * 8, 8);
+                h += popcount64(qw[k] ^ y);
+            }
+            if constexpr (REM > 0) {
+                y = 0;
+                memcpy(&y, bj + NW * 8, REM);
+                h += popcount64(qt ^ y);
+            }
+            dis_[j] = h;
+        }
+    }
+}
+
+/* Same loop with the word count read at run time, for codes long enough that
+ * a constant count gives no performance improvement. */
+template <int REM>
+void hammings_ragged_var(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    const size_t nwords = ncodes / 8;
+    for (size_t i = 0; i < na; i++) {
+        const uint8_t* __restrict ai = a + i * ncodes;
+        hamdis_t* __restrict dis_ = dis + i * nb;
+        for (size_t j = 0; j < nb; j++) {
+            const uint8_t* __restrict bj = b + j * ncodes;
+            hamdis_t h = 0;
+            uint64_t x, y;
+            for (size_t k = 0; k < nwords; k++) {
+                memcpy(&x, ai + k * 8, 8);
+                memcpy(&y, bj + k * 8, 8);
+                h += popcount64(x ^ y);
+            }
+            if constexpr (REM > 0) {
+                x = 0;
+                y = 0;
+                memcpy(&x, ai + nwords * 8, REM);
+                memcpy(&y, bj + nwords * 8, REM);
+                h += popcount64(x ^ y);
+            }
+            dis_[j] = h;
+        }
+    }
+}
+
+/* Makes the word count a constant for short codes when calling
+ * hammings_ragged_fixed, worth roughly 2x perf improvement. Longer codes
+ * gain nothing from it, so they share hammings_ragged_var. */
+template <int REM>
+void hammings_ragged_by_nwords(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    switch (ncodes / 8) {
+        case 0:
+            return hammings_ragged_fixed<0, REM>(a, b, na, nb, ncodes, dis);
+        case 1:
+            return hammings_ragged_fixed<1, REM>(a, b, na, nb, ncodes, dis);
+        case 2:
+            return hammings_ragged_fixed<2, REM>(a, b, na, nb, ncodes, dis);
+        case 3:
+            return hammings_ragged_fixed<3, REM>(a, b, na, nb, ncodes, dis);
+        default:
+            return hammings_ragged_var<REM>(a, b, na, nb, ncodes, dis);
+    }
+}
+
+void hammings_ragged_dispatch(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    switch (ncodes % 8) {
+        case 1:
+            return hammings_ragged_by_nwords<1>(a, b, na, nb, ncodes, dis);
+        case 2:
+            return hammings_ragged_by_nwords<2>(a, b, na, nb, ncodes, dis);
+        case 3:
+            return hammings_ragged_by_nwords<3>(a, b, na, nb, ncodes, dis);
+        case 4:
+            return hammings_ragged_by_nwords<4>(a, b, na, nb, ncodes, dis);
+        case 5:
+            return hammings_ragged_by_nwords<5>(a, b, na, nb, ncodes, dis);
+        case 6:
+            return hammings_ragged_by_nwords<6>(a, b, na, nb, ncodes, dis);
+        default:
+            return hammings_ragged_by_nwords<7>(a, b, na, nb, ncodes, dis);
+    }
+}
+
 } // anonymous namespace
 
 /******************************************************************
@@ -356,6 +484,19 @@ void hamming_range_search_fixSL<THE_SIMD_LEVEL>(
             });
 }
 
+/* Its own entry point, so the word-multiple kernels keep the code the
+ * compiler already generates for them. */
+template <>
+void hammings_ragged_fixSL<THE_SIMD_LEVEL>(
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* dis) {
+    hammings_ragged_dispatch(a, b, na, nb, ncodes, dis);
+}
+
 template <>
 void hammings_fixSL<THE_SIMD_LEVEL>(
         const uint8_t* a,

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 <>