faiss 0.6.2 → 0.6.4

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

@@ -9,6 +9,7 @@
 
 #include <faiss/utils/rabitq_simd.h>
 #include <immintrin.h>
+#include <limits>
 
 namespace faiss::rabitq {
 
@@ -167,8 +168,164 @@ inline uint64_t reduce_add_128(__m128i v) {
     return lanes[0] + lanes[1];
 }
 
+inline __m512i round_nonnegative_ps_to_i32(__m512 x) {
+    return _mm512_cvttps_epi32(_mm512_add_ps(x, _mm512_set1_ps(0.5f)));
+}
+
+inline void store_i32_as_u8_16(__m512i values, uint8_t* out) {
+    const __m128i packed = _mm512_cvtusepi32_epi8(values);
+    _mm_storeu_si128(reinterpret_cast<__m128i*>(out), packed);
+}
+
 } // namespace
 
+template <>
+void lut_minmax_16<SIMDLevel::AVX512>(const float* tab, float& mn, float& mx) {
+    const __m512 v = _mm512_loadu_ps(tab);
+    mn = _mm512_reduce_min_ps(v);
+    mx = _mm512_reduce_max_ps(v);
+}
+
+template <>
+void minmax_values<SIMDLevel::AVX512>(
+        const float* values,
+        size_t n,
+        float& mn,
+        float& mx) {
+    if (n == 0) {
+        return;
+    }
+
+    size_t i = 0;
+    __m512 mn_vec = _mm512_set1_ps(std::numeric_limits<float>::max());
+    __m512 mx_vec = _mm512_set1_ps(std::numeric_limits<float>::lowest());
+    for (; i + 16 <= n; i += 16) {
+        const __m512 v = _mm512_loadu_ps(values + i);
+        mn_vec = _mm512_min_ps(mn_vec, v);
+        mx_vec = _mm512_max_ps(mx_vec, v);
+    }
+
+    if (i < n) {
+        const __mmask16 mask =
+                static_cast<__mmask16>((uint32_t(1) << (n - i)) - 1);
+        const __m512 v = _mm512_maskz_loadu_ps(mask, values + i);
+        mn_vec = _mm512_mask_min_ps(mn_vec, mask, mn_vec, v);
+        mx_vec = _mm512_mask_max_ps(mx_vec, mask, mx_vec, v);
+    }
+
+    mn = _mm512_reduce_min_ps(mn_vec);
+    mx = _mm512_reduce_max_ps(mx_vec);
+}
+
+template <>
+void lut_quantize_16_to_uint8<SIMDLevel::AVX512>(
+        const float* tab,
+        float mn,
+        float a,
+        uint8_t* out) {
+    const __m512 values = _mm512_loadu_ps(tab);
+    const __m512 a_vec = _mm512_set1_ps(a);
+    const __m512 scaled =
+            _mm512_fmsub_ps(values, a_vec, _mm512_set1_ps(mn * a));
+    __m512i rounded = round_nonnegative_ps_to_i32(scaled);
+    rounded = _mm512_max_epi32(rounded, _mm512_setzero_si512());
+    store_i32_as_u8_16(rounded, out);
+}
+
+template <>
+void quantize_query_values<SIMDLevel::AVX512>(
+        const float* rq,
+        size_t d,
+        float v_min,
+        float inv_delta,
+        uint8_t max_code,
+        bool centered,
+        uint8_t* rqq,
+        size_t& sum_qq,
+        int64_t& sum2_signed_odd_int) {
+    const __m512 inv_delta_vec = _mm512_set1_ps(inv_delta);
+    const __m512 v_min_times_inv_delta_vec = _mm512_set1_ps(v_min * inv_delta);
+    const __m512 zero = _mm512_setzero_ps();
+    const __m512 max_code_ps = _mm512_set1_ps(static_cast<float>(max_code));
+    const __m512i max_code_i32 = _mm512_set1_epi32(max_code);
+    const __m512i two = _mm512_set1_epi32(2);
+
+    size_t i = 0;
+    if (centered) {
+        __m512i sum_acc_lo = _mm512_setzero_si512();
+        __m512i sum_acc_hi = _mm512_setzero_si512();
+        __m512i sq_acc_lo = _mm512_setzero_si512();
+        __m512i sq_acc_hi = _mm512_setzero_si512();
+        for (; i + 16 <= d; i += 16) {
+            const __m512 values = _mm512_loadu_ps(rq + i);
+            __m512 scaled = _mm512_fmsub_ps(
+                    values, inv_delta_vec, v_min_times_inv_delta_vec);
+            scaled = _mm512_min_ps(_mm512_max_ps(scaled, zero), max_code_ps);
+            __m512i rounded = round_nonnegative_ps_to_i32(scaled);
+
+            sum_acc_lo = _mm512_add_epi64(
+                    sum_acc_lo,
+                    _mm512_cvtepi32_epi64(_mm512_castsi512_si256(rounded)));
+            sum_acc_hi = _mm512_add_epi64(
+                    sum_acc_hi,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_extracti64x4_epi64(rounded, 1)));
+            const __m512i signed_odd = _mm512_sub_epi32(
+                    _mm512_mullo_epi32(rounded, two), max_code_i32);
+            const __m512i signed_odd_sqr =
+                    _mm512_mullo_epi32(signed_odd, signed_odd);
+            sq_acc_lo = _mm512_add_epi64(
+                    sq_acc_lo,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_castsi512_si256(signed_odd_sqr)));
+            sq_acc_hi = _mm512_add_epi64(
+                    sq_acc_hi,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_extracti64x4_epi64(signed_odd_sqr, 1)));
+            store_i32_as_u8_16(rounded, rqq + i);
+        }
+        sum_qq += static_cast<uint64_t>(
+                _mm512_reduce_add_epi64(sum_acc_lo) +
+                _mm512_reduce_add_epi64(sum_acc_hi));
+        sum2_signed_odd_int += _mm512_reduce_add_epi64(sq_acc_lo);
+        sum2_signed_odd_int += _mm512_reduce_add_epi64(sq_acc_hi);
+    } else {
+        __m512i sum_acc_lo = _mm512_setzero_si512();
+        __m512i sum_acc_hi = _mm512_setzero_si512();
+        for (; i + 16 <= d; i += 16) {
+            const __m512 values = _mm512_loadu_ps(rq + i);
+            __m512 scaled = _mm512_fmsub_ps(
+                    values, inv_delta_vec, v_min_times_inv_delta_vec);
+            scaled = _mm512_min_ps(_mm512_max_ps(scaled, zero), max_code_ps);
+            __m512i rounded = round_nonnegative_ps_to_i32(scaled);
+
+            sum_acc_lo = _mm512_add_epi64(
+                    sum_acc_lo,
+                    _mm512_cvtepi32_epi64(_mm512_castsi512_si256(rounded)));
+            sum_acc_hi = _mm512_add_epi64(
+                    sum_acc_hi,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_extracti64x4_epi64(rounded, 1)));
+            store_i32_as_u8_16(rounded, rqq + i);
+        }
+        sum_qq += static_cast<uint64_t>(
+                _mm512_reduce_add_epi64(sum_acc_lo) +
+                _mm512_reduce_add_epi64(sum_acc_hi));
+    }
+
+    for (; i < d; i++) {
+        const uint8_t v_qq = round_clamped_byte_scalar(
+                (rq[i] - v_min) * inv_delta, max_code);
+        rqq[i] = v_qq;
+        sum_qq += v_qq;
+
+        if (centered) {
+            const int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
+            sum2_signed_odd_int += signed_odd_int * signed_odd_int;
+        }
+    }
+}
+
 template <>
 uint64_t bitwise_and_dot_product<SIMDLevel::AVX512>(
         const uint8_t* query,
@@ -237,6 +394,87 @@ uint64_t bitwise_and_dot_product<SIMDLevel::AVX512>(
     return sum;
 }
 
+template <>
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::AVX512>(
+        const uint8_t* query,
+        const uint8_t* data,
+        size_t size,
+        size_t qb) {
+    uint64_t dot_product = 0;
+    uint64_t popcount_sum = 0;
+    size_t offset = 0;
+    if (size_t step = 512 / 8; offset + step <= size) {
+        __m512i dot_512 = _mm512_setzero_si512();
+        __m512i pop_512 = _mm512_setzero_si512();
+        for (; offset + step <= size; offset += step) {
+            __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset));
+            pop_512 = _mm512_add_epi64(pop_512, popcount_512(v_x));
+            for (int j = 0; j < qb; j++) {
+                __m512i v_q = _mm512_loadu_si512(
+                        (const __m512i*)(query + j * size + offset));
+                __m512i v_and = _mm512_and_si512(v_q, v_x);
+                __m512i v_popcnt = popcount_512(v_and);
+                __m512i v_shifted = _mm512_slli_epi64(v_popcnt, j);
+                dot_512 = _mm512_add_epi64(dot_512, v_shifted);
+            }
+        }
+        dot_product += _mm512_reduce_add_epi64(dot_512);
+        popcount_sum += _mm512_reduce_add_epi64(pop_512);
+    }
+    if (size_t step = 256 / 8; offset + step <= size) {
+        __m256i dot_256 = _mm256_setzero_si256();
+        __m256i pop_256 = _mm256_setzero_si256();
+        for (; offset + step <= size; offset += step) {
+            __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset));
+            pop_256 = _mm256_add_epi64(pop_256, popcount_256(v_x));
+            for (int j = 0; j < qb; j++) {
+                __m256i v_q = _mm256_loadu_si256(
+                        (const __m256i*)(query + j * size + offset));
+                __m256i v_and = _mm256_and_si256(v_q, v_x);
+                __m256i v_popcnt = popcount_256(v_and);
+                __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j);
+                dot_256 = _mm256_add_epi64(dot_256, v_shifted);
+            }
+        }
+        dot_product += reduce_add_256(dot_256);
+        popcount_sum += reduce_add_256(pop_256);
+    }
+    __m128i dot_128 = _mm_setzero_si128();
+    __m128i pop_128 = _mm_setzero_si128();
+    for (size_t step = 128 / 8; offset + step <= size; offset += step) {
+        __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset));
+        pop_128 = _mm_add_epi64(pop_128, popcount_128(v_x));
+        for (int j = 0; j < qb; j++) {
+            __m128i v_q = _mm_loadu_si128(
+                    (const __m128i*)(query + j * size + offset));
+            __m128i v_and = _mm_and_si128(v_q, v_x);
+            __m128i v_popcnt = popcount_128(v_and);
+            __m128i v_shifted = _mm_slli_epi64(v_popcnt, j);
+            dot_128 = _mm_add_epi64(dot_128, v_shifted);
+        }
+    }
+    dot_product += reduce_add_128(dot_128);
+    popcount_sum += reduce_add_128(pop_128);
+    for (size_t step = 64 / 8; offset + step <= size; offset += step) {
+        const auto yv = *(const uint64_t*)(data + offset);
+        popcount_sum += popcount64(yv);
+        for (int j = 0; j < qb; j++) {
+            const auto qv = *(const uint64_t*)(query + j * size + offset);
+            dot_product += popcount64(qv & yv) << j;
+        }
+    }
+    for (; offset < size; ++offset) {
+        const auto yv = *(data + offset);
+        popcount_sum += popcount32(yv);
+        for (int j = 0; j < qb; j++) {
+            const auto qv = *(query + j * size + offset);
+            dot_product += popcount32(qv & yv) << j;
+        }
+    }
+    return {dot_product, popcount_sum};
+}
+
 template <>
 uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512>(
         const uint8_t* query,
@@ -344,22 +582,47 @@ uint64_t popcount<SIMDLevel::AVX512>(const uint8_t* data, size_t size) {
     return sum;
 }
 
+template <>
+void rearrange_bit_planes<SIMDLevel::AVX512>(
+        const uint8_t* rotated_qq,
+        size_t d,
+        size_t qb,
+        uint8_t* out) {
+    const size_t offset = (d + 7) / 8;
+    memset(out, 0, offset * qb);
+    size_t idim = 0;
+    for (; idim + 64 <= d; idim += 64) {
+        __m512i vals = _mm512_loadu_si512((const __m512i*)(rotated_qq + idim));
+        for (size_t iv = 0; iv < qb; iv++) {
+            __m512i mask = _mm512_set1_epi8(static_cast<char>(1 << iv));
+            __mmask64 bits = _mm512_test_epi8_mask(vals, mask);
+            memcpy(&out[iv * offset + idim / 8], &bits, 8);
+        }
+    }
+    for (; idim + 32 <= d; idim += 32) {
+        __m256i vals = _mm256_loadu_si256((const __m256i*)(rotated_qq + idim));
+        for (size_t iv = 0; iv < qb; iv++) {
+            __m256i mask = _mm256_set1_epi8(static_cast<char>(1 << iv));
+            __m256i bits =
+                    _mm256_cmpeq_epi8(_mm256_and_si256(vals, mask), mask);
+            uint32_t packed = static_cast<uint32_t>(_mm256_movemask_epi8(bits));
+            memcpy(&out[iv * offset + idim / 8], &packed, 4);
+        }
+    }
+    for (; idim < d; idim++) {
+        for (size_t iv = 0; iv < qb; iv++) {
+            const bool bit = ((rotated_qq[idim] & (1 << iv)) != 0);
+            out[iv * offset + idim / 8] |= bit ? (1 << (idim % 8)) : 0;
+        }
+    }
+}
+
 } // namespace faiss::rabitq
 
 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,
@@ -391,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 ex64 = 0;
-        memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t));
+        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++) {
-            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])) |
+                    (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]);
         }
 
-        __m256 rq = _mm256_loadu_ps(rotated_q + i);
-        acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc);
+        __m512 rq = _mm512_loadu_ps(rotated_q + i);
+        acc = _mm512_fmadd_ps(rq, _mm512_add_ps(recon, v_cb), acc);
     }
 
-    float result = hsum_avx2(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++) {
+            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]);
+        }
+
+        __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 = _mm512_reduce_add_ps(acc);
     result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
     return result;
 }
@@ -466,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,
@@ -186,7 +183,99 @@ uint64_t bitwise_and_dot_product<SIMDLevel::AVX512_SPR>(
 }
 
 template <>
-uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512_SPR>(
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::AVX512_VPOPCNT>(
+        const uint8_t* query,
+        const uint8_t* data,
+        size_t size,
+        size_t qb) {
+    uint64_t dot_product = 0;
+    uint64_t popcount_sum = 0;
+    size_t offset = 0;
+
+    if (size_t step = 512 / 8; offset + step <= size) {
+        __m512i dot_512 = _mm512_setzero_si512();
+        __m512i pop_512 = _mm512_setzero_si512();
+        for (; offset + step <= size; offset += step) {
+            __m512i v_x = _mm512_loadu_si512(
+                    reinterpret_cast<const __m512i*>(data + offset));
+            pop_512 = _mm512_add_epi64(pop_512, popcount_512_vpopcntdq(v_x));
+            for (size_t j = 0; j < qb; j++) {
+                __m512i v_q = _mm512_loadu_si512(
+                        reinterpret_cast<const __m512i*>(
+                                query + j * size + offset));
+                __m512i v_and = _mm512_and_si512(v_q, v_x);
+                __m512i v_popcnt = popcount_512_vpopcntdq(v_and);
+                __m512i v_shifted = _mm512_slli_epi64(v_popcnt, j);
+                dot_512 = _mm512_add_epi64(dot_512, v_shifted);
+            }
+        }
+        dot_product += _mm512_reduce_add_epi64(dot_512);
+        popcount_sum += _mm512_reduce_add_epi64(pop_512);
+    }
+
+    if (size_t step = 256 / 8; offset + step <= size) {
+        __m256i dot_256 = _mm256_setzero_si256();
+        __m256i pop_256 = _mm256_setzero_si256();
+        for (; offset + step <= size; offset += step) {
+            __m256i v_x = _mm256_loadu_si256(
+                    reinterpret_cast<const __m256i*>(data + offset));
+            pop_256 = _mm256_add_epi64(pop_256, popcount_256_vpopcntdq(v_x));
+            for (size_t j = 0; j < qb; j++) {
+                __m256i v_q = _mm256_loadu_si256(
+                        reinterpret_cast<const __m256i*>(
+                                query + j * size + offset));
+                __m256i v_and = _mm256_and_si256(v_q, v_x);
+                __m256i v_popcnt = popcount_256_vpopcntdq(v_and);
+                __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j);
+                dot_256 = _mm256_add_epi64(dot_256, v_shifted);
+            }
+        }
+        dot_product += reduce_add_256(dot_256);
+        popcount_sum += reduce_add_256(pop_256);
+    }
+
+    __m128i dot_128 = _mm_setzero_si128();
+    __m128i pop_128 = _mm_setzero_si128();
+    for (size_t step = 128 / 8; offset + step <= size; offset += step) {
+        __m128i v_x = _mm_loadu_si128(
+                reinterpret_cast<const __m128i*>(data + offset));
+        pop_128 = _mm_add_epi64(pop_128, popcount_128_vpopcntdq(v_x));
+        for (size_t j = 0; j < qb; j++) {
+            __m128i v_q = _mm_loadu_si128(
+                    reinterpret_cast<const __m128i*>(
+                            query + j * size + offset));
+            __m128i v_and = _mm_and_si128(v_q, v_x);
+            __m128i v_popcnt = popcount_128_vpopcntdq(v_and);
+            __m128i v_shifted = _mm_slli_epi64(v_popcnt, j);
+            dot_128 = _mm_add_epi64(dot_128, v_shifted);
+        }
+    }
+    dot_product += reduce_add_128(dot_128);
+    popcount_sum += reduce_add_128(pop_128);
+
+    for (size_t step = 64 / 8; offset + step <= size; offset += step) {
+        const auto yv = *reinterpret_cast<const uint64_t*>(data + offset);
+        popcount_sum += popcount64(yv);
+        for (size_t j = 0; j < qb; j++) {
+            const auto qv = *reinterpret_cast<const uint64_t*>(
+                    query + j * size + offset);
+            dot_product += static_cast<uint64_t>(popcount64(qv & yv)) << j;
+        }
+    }
+    for (; offset < size; ++offset) {
+        const auto yv = *(data + offset);
+        popcount_sum += popcount32(yv);
+        for (size_t j = 0; j < qb; j++) {
+            const auto qv = *(query + j * size + offset);
+            dot_product += static_cast<uint64_t>(popcount32(qv & yv)) << j;
+        }
+    }
+    return {dot_product, popcount_sum};
+}
+
+template <>
+uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512_VPOPCNT>(
         const uint8_t* query,
         const uint8_t* data,
         size_t size,
@@ -265,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;
 
@@ -327,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,
@@ -340,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/rabitq_neon.cpp

@@ -20,6 +20,17 @@ uint64_t bitwise_and_dot_product<SIMDLevel::ARM_NEON>(
     return bitwise_and_dot_product<SIMDLevel::NONE>(query, data, size, qb);
 }
 
+template <>
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::ARM_NEON>(
+        const uint8_t* query,
+        const uint8_t* data,
+        size_t size,
+        size_t qb) {
+    return bitwise_and_dot_product_with_popcount<SIMDLevel::NONE>(
+            query, data, size, qb);
+}
+
 template <>
 uint64_t bitwise_xor_dot_product<SIMDLevel::ARM_NEON>(
         const uint8_t* query,