faiss 0.6.3 → 0.6.4

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

@@ -16,6 +16,42 @@
 
 namespace faiss {
 
+template <typename Vec, typename Reduce>
+static inline float rvv_reduce(
+        Vec value,
+        size_t vl,
+        float identity,
+        Reduce reduce) {
+    vfloat32m1_t init = __riscv_vfmv_s_f_f32m1(identity, 1);
+    vfloat32m1_t result = reduce(value, init, vl);
+    return __riscv_vfmv_f_s_f32m1_f32(result);
+}
+
+static inline size_t rvv_argmin(const float* values, size_t n) {
+    size_t vlmax = __riscv_vsetvlmax_e32m8();
+    vfloat32m8_t vmin = __riscv_vfmv_v_f_f32m8(__builtin_inff(), vlmax);
+    size_t i = 0;
+    while (i < n) {
+        size_t vl = __riscv_vsetvl_e32m8(n - i);
+        vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
+        vmin = __riscv_vfmin_vv_f32m8_tu(vmin, vmin, vd, vl);
+        i += vl;
+    }
+    float min_val = rvv_reduce(
+            vmin, vlmax, __builtin_inff(), __riscv_vfredmin_vs_f32m8_f32m1);
+    i = 0;
+    while (i < n) {
+        size_t vl = __riscv_vsetvl_e32m8(n - i);
+        vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
+        long j = __riscv_vfirst_m_b4(
+                __riscv_vmfeq_vf_f32m8_b4(vd, min_val, vl), vl);
+        if (j >= 0)
+            return i + static_cast<size_t>(j);
+        i += vl;
+    }
+    return n;
+}
+
 template <>
 float fvec_norm_L2sqr<SIMDLevel::RISCV_RVV>(const float* x, size_t d) {
     size_t vlmax = __riscv_vsetvlmax_e32m8();
@@ -27,9 +63,7 @@ float fvec_norm_L2sqr<SIMDLevel::RISCV_RVV>(const float* x, size_t d) {
         acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
         i += vl;
     }
-    vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
-    return __riscv_vfmv_f_s_f32m1_f32(sum);
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
@@ -48,9 +82,7 @@ float fvec_L2sqr<SIMDLevel::RISCV_RVV>(
         acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
         i += vl;
     }
-    vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
-    return __riscv_vfmv_f_s_f32m1_f32(sum);
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
@@ -68,9 +100,7 @@ float fvec_inner_product<SIMDLevel::RISCV_RVV>(
         acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vy, vl);
         i += vl;
     }
-    vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
-    return __riscv_vfmv_f_s_f32m1_f32(sum);
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
@@ -87,9 +117,7 @@ float fvec_L1<SIMDLevel::RISCV_RVV>(const float* x, const float* y, size_t d) {
         acc = __riscv_vfadd_vv_f32m8_tu(acc, acc, vx, vl);
         i += vl;
     }
-    vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
-    return __riscv_vfmv_f_s_f32m1_f32(sum);
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
@@ -109,9 +137,7 @@ float fvec_Linf<SIMDLevel::RISCV_RVV>(
         vmax = __riscv_vfmax_vv_f32m8_tu(vmax, vmax, vx, vl);
         i += vl;
     }
-    vfloat32m1_t max = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    max = __riscv_vfredmax_vs_f32m8_f32m1(vmax, max, vlmax);
-    return __riscv_vfmv_f_s_f32m1_f32(max);
+    return rvv_reduce(vmax, vlmax, 0.0f, __riscv_vfredmax_vs_f32m8_f32m1);
 }
 
 template <>
@@ -126,10 +152,29 @@ void fvec_inner_product_batch_4<SIMDLevel::RISCV_RVV>(
         float& dis1,
         float& dis2,
         float& dis3) {
-    dis0 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y0, d);
-    dis1 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y1, d);
-    dis2 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y2, d);
-    dis3 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y3, d);
+    size_t vlmax = __riscv_vsetvlmax_e32m4();
+    vfloat32m4_t vacc0 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    vfloat32m4_t vacc1 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    vfloat32m4_t vacc2 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    vfloat32m4_t vacc3 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m4(d - i);
+        vfloat32m4_t vx = __riscv_vle32_v_f32m4(x + i, vl);
+        vfloat32m4_t vy = __riscv_vle32_v_f32m4(y0 + i, vl);
+        vacc0 = __riscv_vfmacc_vv_f32m4_tu(vacc0, vx, vy, vl);
+        vy = __riscv_vle32_v_f32m4(y1 + i, vl);
+        vacc1 = __riscv_vfmacc_vv_f32m4_tu(vacc1, vx, vy, vl);
+        vy = __riscv_vle32_v_f32m4(y2 + i, vl);
+        vacc2 = __riscv_vfmacc_vv_f32m4_tu(vacc2, vx, vy, vl);
+        vy = __riscv_vle32_v_f32m4(y3 + i, vl);
+        vacc3 = __riscv_vfmacc_vv_f32m4_tu(vacc3, vx, vy, vl);
+        i += vl;
+    }
+    dis0 = rvv_reduce(vacc0, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
+    dis1 = rvv_reduce(vacc1, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
+    dis2 = rvv_reduce(vacc2, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
+    dis3 = rvv_reduce(vacc3, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
 }
 
 template <>
@@ -144,10 +189,33 @@ void fvec_L2sqr_batch_4<SIMDLevel::RISCV_RVV>(
         float& dis1,
         float& dis2,
         float& dis3) {
-    dis0 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y0, d);
-    dis1 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y1, d);
-    dis2 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y2, d);
-    dis3 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y3, d);
+    size_t vlmax = __riscv_vsetvlmax_e32m4();
+    vfloat32m4_t vacc0 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    vfloat32m4_t vacc1 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    vfloat32m4_t vacc2 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    vfloat32m4_t vacc3 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m4(d - i);
+        vfloat32m4_t vx = __riscv_vle32_v_f32m4(x + i, vl);
+        vfloat32m4_t vy = __riscv_vle32_v_f32m4(y0 + i, vl);
+        vy = __riscv_vfsub_vv_f32m4(vx, vy, vl);
+        vacc0 = __riscv_vfmacc_vv_f32m4_tu(vacc0, vy, vy, vl);
+        vy = __riscv_vle32_v_f32m4(y1 + i, vl);
+        vy = __riscv_vfsub_vv_f32m4(vx, vy, vl);
+        vacc1 = __riscv_vfmacc_vv_f32m4_tu(vacc1, vy, vy, vl);
+        vy = __riscv_vle32_v_f32m4(y2 + i, vl);
+        vy = __riscv_vfsub_vv_f32m4(vx, vy, vl);
+        vacc2 = __riscv_vfmacc_vv_f32m4_tu(vacc2, vy, vy, vl);
+        vy = __riscv_vle32_v_f32m4(y3 + i, vl);
+        vy = __riscv_vfsub_vv_f32m4(vx, vy, vl);
+        vacc3 = __riscv_vfmacc_vv_f32m4_tu(vacc3, vy, vy, vl);
+        i += vl;
+    }
+    dis0 = rvv_reduce(vacc0, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
+    dis1 = rvv_reduce(vacc1, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
+    dis2 = rvv_reduce(vacc2, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
+    dis3 = rvv_reduce(vacc3, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
 }
 
 template <>
@@ -168,9 +236,8 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
         acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
         i += vl;
     }
-    vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
-    sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
-    float x_sqlen = __riscv_vfmv_f_s_f32m1_f32(sum);
+    float x_sqlen =
+            rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
     i = 0;
     while (i < ny) {
         size_t vl = __riscv_vsetvl_e32m8(ny - i);
@@ -214,35 +281,6 @@ void fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(
     }
 }
 
-// Index of the first element equal to the minimum of values[0..n), or n when
-// there is none (e.g. n == 0). Shared by the *_nearest and madd_and_argmin
-// kernels so the vfmin/vfredmin/vfirst sequence lives in one place.
-static size_t rvv_argmin(const float* values, size_t n) {
-    size_t vlmax = __riscv_vsetvlmax_e32m8();
-    vfloat32m8_t vmin = __riscv_vfmv_v_f_f32m8(__builtin_inff(), vlmax);
-    size_t i = 0;
-    while (i < n) {
-        size_t vl = __riscv_vsetvl_e32m8(n - i);
-        vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
-        vmin = __riscv_vfmin_vv_f32m8_tu(vmin, vmin, vd, vl);
-        i += vl;
-    }
-    vfloat32m1_t rmin = __riscv_vfmv_s_f_f32m1(__builtin_inff(), 1);
-    rmin = __riscv_vfredmin_vs_f32m8_f32m1(vmin, rmin, vlmax);
-    float min_val = __riscv_vfmv_f_s_f32m1_f32(rmin);
-    i = 0;
-    while (i < n) {
-        size_t vl = __riscv_vsetvl_e32m8(n - i);
-        vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
-        long j = __riscv_vfirst_m_b4(
-                __riscv_vmfeq_vf_f32m8_b4(vd, min_val, vl), vl);
-        if (j >= 0)
-            return i + static_cast<size_t>(j);
-        i += vl;
-    }
-    return n;
-}
-
 template <>
 size_t fvec_L2sqr_ny_nearest<SIMDLevel::RISCV_RVV>(
         float* distances_tmp_buffer,

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

@@ -623,16 +623,6 @@ namespace faiss::rabitq::multibit {
 
 namespace {
 
-inline float hsum_avx2(__m256 v) {
-    __m128 hi = _mm256_extractf128_ps(v, 1);
-    __m128 lo = _mm256_castps256_ps128(v);
-    lo = _mm_add_ps(lo, hi);
-    __m128 shuf = _mm_movehdup_ps(lo);
-    lo = _mm_add_ps(lo, shuf);
-    shuf = _mm_movehl_ps(shuf, lo);
-    return _mm_cvtss_f32(_mm_add_ss(lo, shuf));
-}
-
 inline float ip_1exbit_avx512(
         const uint8_t* __restrict sign_bits,
         const uint8_t* __restrict ex_code,
@@ -664,60 +654,86 @@ inline float ip_1exbit_avx512(
     return result;
 }
 
-// AVX2+BMI2 bitplane kernel used as fallback for ex_bits >= 2.
-// AVX512 TU has AVX2 available. BMI2 guarded separately since
-// VIA Eden X4 has AVX2 without BMI2.
+// Needs BMI2 for _pext_u64. Some AVX2 CPUs lack it, and FAISS_BMI2_FLAGS can
+// be empty, so the dispatcher falls back to the scalar path without it.
 #ifdef __BMI2__
-inline float ip_bitplane_avx2(
+// Bitplane kernel for ex_bits >= 2, 16 dims per iteration. A bitplane is
+// already a bitmask, so it goes into a mask register and one masked add
+// applies its weight. Reads of ex_code run a few bytes past the ex-code
+// section into the record's own trailing factors, so they stay in bounds.
+inline float ip_bitplane_avx512(
         const uint8_t* __restrict sign_bits,
         const uint8_t* __restrict ex_code,
         const float* __restrict rotated_q,
         size_t d,
         size_t ex_bits,
         float cb) {
-    __m256 acc = _mm256_setzero_ps();
-    const __m256 v_one = _mm256_set1_ps(1.0f);
-    const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
-    const __m256i zero = _mm256_setzero_si256();
-    const __m256 v_cb = _mm256_set1_ps(cb);
+    __m512 acc = _mm512_setzero_ps();
+    const __m512 v_cb = _mm512_set1_ps(cb);
 
     uint64_t pext_masks[7];
-    __m256 v_weights[8];
+    __m512 v_weights[8];
     for (size_t b = 0; b < ex_bits; b++) {
         uint64_t m = 0;
         for (int j = 0; j < 8; j++) {
             m |= (1ULL << (b + j * ex_bits));
         }
         pext_masks[b] = m;
-        v_weights[b] = _mm256_set1_ps(static_cast<float>(1u << b));
+        v_weights[b] = _mm512_set1_ps(static_cast<float>(1u << b));
     }
-    v_weights[ex_bits] = _mm256_set1_ps(static_cast<float>(1u << ex_bits));
+    v_weights[ex_bits] = _mm512_set1_ps(static_cast<float>(1u << ex_bits));
 
     size_t i = 0;
-    for (; i + 8 <= d; i += 8) {
-        __m256i sb_cmp = _mm256_cmpgt_epi32(
-                _mm256_and_si256(_mm256_set1_epi32(sign_bits[i / 8]), bit_pos),
-                zero);
-        __m256 recon = _mm256_mul_ps(
-                _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one),
-                v_weights[ex_bits]);
+    for (; i + 16 <= d; i += 16) {
+        uint16_t sb = 0;
+        memcpy(&sb, sign_bits + (i / 8), sizeof(uint16_t));
+        __m512 recon = _mm512_maskz_mov_ps(
+                static_cast<__mmask16>(sb), v_weights[ex_bits]);
+
+        uint64_t lo64 = 0;
+        uint64_t hi64 = 0;
+        memcpy(&lo64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t));
+        memcpy(&hi64, ex_code + ((i / 8) + 1) * ex_bits, sizeof(uint64_t));
+
+        for (size_t b = 0; b < ex_bits; b++) {
+            const uint32_t plane =
+                    static_cast<uint32_t>(_pext_u64(lo64, pext_masks[b])) |
+                    (static_cast<uint32_t>(_pext_u64(hi64, pext_masks[b]))
+                     << 8);
+            recon = _mm512_mask_add_ps(
+                    recon, static_cast<__mmask16>(plane), recon, v_weights[b]);
+        }
 
-        uint64_t ex64 = 0;
-        memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t));
+        __m512 rq = _mm512_loadu_ps(rotated_q + i);
+        acc = _mm512_fmadd_ps(rq, _mm512_add_ps(recon, v_cb), acc);
+    }
+
+    // Half-width step: keeps the scalar tail under 8 dims when d is a multiple
+    // of 8 but not of 16 (e.g. 200, 1000). The upper 8 lanes are masked off
+    // throughout, and rotated_q is loaded masked so nothing is read past the
+    // end.
+    if (i + 8 <= d) {
+        const __mmask16 low8 = static_cast<__mmask16>(0x00ff);
+        __m512 recon = _mm512_maskz_mov_ps(
+                static_cast<__mmask16>(sign_bits[i / 8]), v_weights[ex_bits]);
+
+        uint64_t lo64 = 0;
+        memcpy(&lo64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t));
 
         for (size_t b = 0; b < ex_bits; b++) {
-            auto plane = static_cast<uint8_t>(_pext_u64(ex64, pext_masks[b]));
-            __m256i p_cmp = _mm256_cmpgt_epi32(
-                    _mm256_and_si256(_mm256_set1_epi32(plane), bit_pos), zero);
-            __m256 p_f = _mm256_and_ps(_mm256_castsi256_ps(p_cmp), v_one);
-            recon = _mm256_fmadd_ps(p_f, v_weights[b], recon);
+            const uint32_t plane =
+                    static_cast<uint32_t>(_pext_u64(lo64, pext_masks[b]));
+            recon = _mm512_mask_add_ps(
+                    recon, static_cast<__mmask16>(plane), recon, v_weights[b]);
         }
 
-        __m256 rq = _mm256_loadu_ps(rotated_q + i);
-        acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc);
+        __m512 rq = _mm512_maskz_loadu_ps(low8, rotated_q + i);
+        acc = _mm512_fmadd_ps(
+                rq, _mm512_mask_add_ps(recon, low8, recon, v_cb), acc);
+        i += 8;
     }
 
-    float result = hsum_avx2(acc);
+    float result = _mm512_reduce_add_ps(acc);
     result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
     return result;
 }
@@ -739,7 +755,8 @@ float compute_inner_product<SIMDLevel::AVX512>(
 
 #ifdef __BMI2__
     if (ex_bits <= 7) {
-        return ip_bitplane_avx2(sign_bits, ex_code, rotated_q, d, ex_bits, cb);
+        return ip_bitplane_avx512(
+                sign_bits, ex_code, rotated_q, d, ex_bits, cb);
     }
 #endif
     return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, ex_bits, cb);

data/vendor/faiss/faiss/utils/simd_impl/{rabitq_avx512_spr.cpp → rabitq_avx512_vpopcnt.cpp}

@@ -6,34 +6,31 @@
  */
 
 /**
- * @file rabitq_avx512_spr.cpp
+ * @file rabitq_avx512_vpopcnt.cpp
  *
- * RaBitQ SIMD kernels specialized for SIMDLevel::AVX512_SPR.
+ * RaBitQ SIMD kernels specialized for SIMDLevel::AVX512_VPOPCNT.
  *
- * Sapphire Rapids (SPR) and later Intel microarchitectures expose
- * AVX-512 VPOPCNTDQ (vpopcntq), which performs a per-lane 64-bit
- * popcount in a single instruction. This is used here to replace the
- * multi-step shuffle/pshufb-based popcount used by the generic AVX-512
- * specialization in rabitq_avx512.cpp. The popcount-heavy kernels
- * (bitwise_and_dot_product, bitwise_xor_dot_product, popcount) become
- * substantially shorter and faster on SPR+ as a result.
+ * AVX-512 VPOPCNTDQ performs a per-lane 64-bit popcount in a single
+ * instruction. It is available on CPUs including Ice Lake, Zen 4, and
+ * Sapphire Rapids, independently of the other SPR-only extensions. This
+ * replaces the multi-step shuffle-based popcount used by the generic
+ * AVX-512 specialization in rabitq_avx512.cpp.
  *
  * Build / dispatch behavior:
  *   - faiss_avx512 (AVX-512 only, no SPR features): NOT compiled.
  *     The existing AVX512 specialization in rabitq_avx512.cpp is used.
- *   - faiss_avx512_spr (statically built for SPR+): compiled. The
- *     SINGLE_SIMD_LEVEL is AVX512_SPR, so this specialization is
- *     selected by static dispatch.
+ *   - faiss_avx512_spr: compiled alongside the full SPR specialization and
+ *     selected through the SPR -> VPOPCNT fallback.
  *   - faiss with FAISS_OPT_LEVEL=dd (dynamic dispatch): compiled with
  *     -mavx512vpopcntdq as a per-file flag. Selected at runtime when
- *     SIMDConfig::level == SIMDLevel::AVX512_SPR.
+ *     the CPU exposes AVX512_VPOPCNTDQ.
  *
  * The floating-point multi-bit inner-product kernel does not benefit
- * from VPOPCNTDQ, so this TU forwards compute_inner_product<SPR> to
+ * from VPOPCNTDQ, so this TU forwards compute_inner_product<VPOPCNT> to
  * the AVX512 implementation to avoid duplicating that code path.
  */
 
-#ifdef COMPILE_SIMD_AVX512_SPR
+#ifdef COMPILE_SIMD_AVX512_VPOPCNT
 
 #include <faiss/utils/popcount.h>
 #include <faiss/utils/rabitq_simd.h>
@@ -47,7 +44,7 @@
 namespace faiss::rabitq {
 
 // Forward declarations for the AVX512 specializations defined in
-// rabitq_avx512.cpp. They live in the same TU group on SPR builds, so
+// rabitq_avx512.cpp. They live in the same TU group in supported builds, so
 // we can reuse them as a tail handler / fallback. Declaring rather
 // than redefining avoids ODR risk and keeps a single source of truth
 // for the floating-point kernel.
@@ -75,8 +72,8 @@ inline __m512i popcount_512_vpopcntdq(__m512i v) {
 }
 
 // 256-bit popcount using AVX-512VL VPOPCNTDQ.
-// AVX512VL is part of the SPR feature set, so vpopcntq is available
-// on 256-bit registers via _mm256_popcnt_epi64.
+// Baseline AVX-512 includes AVX512VL, so VPOPCNTDQ is also available on
+// 256-bit registers via _mm256_popcnt_epi64.
 inline __m256i popcount_256_vpopcntdq(__m256i v) {
     return _mm256_popcnt_epi64(v);
 }
@@ -101,7 +98,7 @@ inline uint64_t reduce_add_128(__m128i v) {
 } // namespace
 
 template <>
-uint64_t bitwise_and_dot_product<SIMDLevel::AVX512_SPR>(
+uint64_t bitwise_and_dot_product<SIMDLevel::AVX512_VPOPCNT>(
         const uint8_t* query,
         const uint8_t* data,
         size_t size,
@@ -187,7 +184,7 @@ uint64_t bitwise_and_dot_product<SIMDLevel::AVX512_SPR>(
 
 template <>
 BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
-        SIMDLevel::AVX512_SPR>(
+        SIMDLevel::AVX512_VPOPCNT>(
         const uint8_t* query,
         const uint8_t* data,
         size_t size,
@@ -278,7 +275,7 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
 }
 
 template <>
-uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512_SPR>(
+uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512_VPOPCNT>(
         const uint8_t* query,
         const uint8_t* data,
         size_t size,
@@ -357,7 +354,7 @@ uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512_SPR>(
 }
 
 template <>
-uint64_t popcount<SIMDLevel::AVX512_SPR>(const uint8_t* data, size_t size) {
+uint64_t popcount<SIMDLevel::AVX512_VPOPCNT>(const uint8_t* data, size_t size) {
     uint64_t sum = 0;
     size_t offset = 0;
 
@@ -419,7 +416,7 @@ float compute_inner_product<SIMDLevel::AVX512>(
         float cb);
 
 template <>
-float compute_inner_product<SIMDLevel::AVX512_SPR>(
+float compute_inner_product<SIMDLevel::AVX512_VPOPCNT>(
         const uint8_t* __restrict sign_bits,
         const uint8_t* __restrict ex_code,
         const float* __restrict rotated_q,
@@ -432,4 +429,4 @@ float compute_inner_product<SIMDLevel::AVX512_SPR>(
 
 } // namespace faiss::rabitq::multibit
 
-#endif // COMPILE_SIMD_AVX512_SPR
+#endif // COMPILE_SIMD_AVX512_VPOPCNT

data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h

@@ -9,13 +9,8 @@
 
 // Private dispatch wrapper for SuperKMeans's block_l2. Routes to the
 // highest available SIMD specialization at runtime (DD mode) or the
-// compiled-in level (static mode). aarch64 currently falls through to the
-// scalar primary template; adding NEON/SVE means just adding a new
-// specialization file alongside the AVX ones.
+// compiled-in level (static mode).
 //
-// Known perf gap: aarch64 (NEON/SVE) specializations are not implemented yet.
-// aarch64 falls through to the scalar primary template. Validating SVE requires
-// a Graviton-class host; deferred to a focused follow-up.
 
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/simd_impl/super_kmeans_kernels.h>
@@ -24,7 +19,7 @@ namespace faiss {
 namespace detail {
 
 inline float block_l2_dispatch(const float* x, const float* y, int n) {
-    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+    return with_simd_level_with_sve(
             [&]<SIMDLevel SL>() { return block_l2<SL>(x, y, n); });
 }
 

data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h

@@ -16,7 +16,7 @@ namespace detail {
 
 // Squared L2 over `n` dimensions; n in [1, pdx_block_size].
 // Primary template is the scalar fallback; SIMDLevels without a dedicated
-// specialization (ARM_NEON, ARM_SVE, NONE, ...) use it directly.
+// specialization (ARM_NEON, NONE, ...) use it directly.
 template <SIMDLevel Level>
 inline float block_l2(const float* x, const float* y, int n) {
     float s = 0.0f;
@@ -39,5 +39,10 @@ template <>
 float block_l2<SIMDLevel::AVX512>(const float* x, const float* y, int n);
 #endif
 
+#ifdef COMPILE_SIMD_ARM_SVE
+template <>
+float block_l2<SIMDLevel::ARM_SVE>(const float* x, const float* y, int n);
+#endif
+
 } // namespace detail
 } // namespace faiss

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

@@ -0,0 +1,34 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_ARM_SVE
+
+#include <faiss/utils/simd_impl/super_kmeans_kernels.h>
+
+#include <arm_sve.h>
+
+namespace faiss {
+namespace detail {
+
+template <>
+float block_l2<SIMDLevel::ARM_SVE>(const float* x, const float* y, int n) {
+    svfloat32_t acc = svdup_n_f32(0.0f);
+    const int lanes = static_cast<int>(svcntw());
+    for (int m = 0; m < n; m += lanes) {
+        const svbool_t pg = svwhilelt_b32(m, n);
+        const svfloat32_t xv = svld1_f32(pg, x + m);
+        const svfloat32_t yv = svld1_f32(pg, y + m);
+        const svfloat32_t diff = svsub_f32_x(pg, xv, yv);
+        acc = svmla_f32_m(pg, acc, diff, diff);
+    }
+    return svaddv_f32(svptrue_b32(), acc);
+}
+
+} // namespace detail
+} // namespace faiss
+
+#endif // COMPILE_SIMD_ARM_SVE