faiss 0.6.2 → 0.6.3

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

@@ -9,6 +9,7 @@
 
 #include <faiss/utils/rabitq_simd.h>
 #include <immintrin.h>
+#include <limits>
 
 namespace faiss::rabitq {
 
@@ -82,8 +83,161 @@ inline uint64_t reduce_add_128(__m128i v) {
     return lanes[0] + lanes[1];
 }
 
+inline float reduce_min_256(__m256 v) {
+    __m128 x =
+            _mm_min_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));
+    x = _mm_min_ps(x, _mm_movehl_ps(x, x));
+    x = _mm_min_ss(x, _mm_shuffle_ps(x, x, 1));
+    return _mm_cvtss_f32(x);
+}
+
+inline float reduce_max_256(__m256 v) {
+    __m128 x =
+            _mm_max_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));
+    x = _mm_max_ps(x, _mm_movehl_ps(x, x));
+    x = _mm_max_ss(x, _mm_shuffle_ps(x, x, 1));
+    return _mm_cvtss_f32(x);
+}
+
+inline __m256i round_nonnegative_ps_to_i32(__m256 x) {
+    return _mm256_cvttps_epi32(_mm256_add_ps(x, _mm256_set1_ps(0.5f)));
+}
+
+inline void store_i32_as_u8_8(__m256i values, uint8_t* out) {
+    const __m128i packed16 = _mm_packus_epi32(
+            _mm256_castsi256_si128(values),
+            _mm256_extracti128_si256(values, 1));
+    const __m128i packed8 = _mm_packus_epi16(packed16, _mm_setzero_si128());
+    _mm_storel_epi64(reinterpret_cast<__m128i*>(out), packed8);
+}
+
+inline void accumulate_i32_as_i64(
+        __m256i values,
+        __m256i& low_acc,
+        __m256i& high_acc) {
+    low_acc = _mm256_add_epi64(
+            low_acc, _mm256_cvtepi32_epi64(_mm256_castsi256_si128(values)));
+    high_acc = _mm256_add_epi64(
+            high_acc,
+            _mm256_cvtepi32_epi64(_mm256_extracti128_si256(values, 1)));
+}
+
 } // namespace
 
+template <>
+void lut_minmax_16<SIMDLevel::AVX2>(const float* tab, float& mn, float& mx) {
+    const __m256 lo = _mm256_loadu_ps(tab);
+    const __m256 hi = _mm256_loadu_ps(tab + 8);
+    const __m256 min_vec = _mm256_min_ps(lo, hi);
+    const __m256 max_vec = _mm256_max_ps(lo, hi);
+    mn = reduce_min_256(min_vec);
+    mx = reduce_max_256(max_vec);
+}
+
+template <>
+void minmax_values<SIMDLevel::AVX2>(
+        const float* values,
+        size_t n,
+        float& mn,
+        float& mx) {
+    if (n == 0) {
+        return;
+    }
+
+    size_t i = 0;
+    __m256 min_vec = _mm256_set1_ps(std::numeric_limits<float>::max());
+    __m256 max_vec = _mm256_set1_ps(std::numeric_limits<float>::lowest());
+    for (; i + 8 <= n; i += 8) {
+        const __m256 values_vec = _mm256_loadu_ps(values + i);
+        min_vec = _mm256_min_ps(min_vec, values_vec);
+        max_vec = _mm256_max_ps(max_vec, values_vec);
+    }
+
+    mn = reduce_min_256(min_vec);
+    mx = reduce_max_256(max_vec);
+    for (; i < n; i++) {
+        mn = std::min(mn, values[i]);
+        mx = std::max(mx, values[i]);
+    }
+}
+
+template <>
+void lut_quantize_16_to_uint8<SIMDLevel::AVX2>(
+        const float* tab,
+        float mn,
+        float a,
+        uint8_t* out) {
+    const __m256 a_vec = _mm256_set1_ps(a);
+    const __m256 mn_times_a_vec = _mm256_set1_ps(mn * a);
+    const __m256i zero = _mm256_setzero_si256();
+    for (size_t i = 0; i < 16; i += 8) {
+        const __m256 values = _mm256_loadu_ps(tab + i);
+        const __m256 scaled = _mm256_fmsub_ps(values, a_vec, mn_times_a_vec);
+        const __m256i rounded =
+                _mm256_max_epi32(round_nonnegative_ps_to_i32(scaled), zero);
+        store_i32_as_u8_8(rounded, out + i);
+    }
+}
+
+template <>
+void quantize_query_values<SIMDLevel::AVX2>(
+        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 __m256 inv_delta_vec = _mm256_set1_ps(inv_delta);
+    const __m256 v_min_times_inv_delta_vec = _mm256_set1_ps(v_min * inv_delta);
+    const __m256 zero = _mm256_setzero_ps();
+    const __m256 max_code_ps = _mm256_set1_ps(max_code);
+    const __m256i max_code_i32 = _mm256_set1_epi32(max_code);
+    const __m256i two = _mm256_set1_epi32(2);
+    __m256i sum_acc_lo = _mm256_setzero_si256();
+    __m256i sum_acc_hi = _mm256_setzero_si256();
+    __m256i sq_acc_lo = _mm256_setzero_si256();
+    __m256i sq_acc_hi = _mm256_setzero_si256();
+
+    size_t i = 0;
+    for (; i + 8 <= d; i += 8) {
+        const __m256 values = _mm256_loadu_ps(rq + i);
+        __m256 scaled = _mm256_fmsub_ps(
+                values, inv_delta_vec, v_min_times_inv_delta_vec);
+        scaled = _mm256_min_ps(_mm256_max_ps(scaled, zero), max_code_ps);
+        const __m256i rounded = round_nonnegative_ps_to_i32(scaled);
+        accumulate_i32_as_i64(rounded, sum_acc_lo, sum_acc_hi);
+
+        if (centered) {
+            const __m256i signed_odd = _mm256_sub_epi32(
+                    _mm256_mullo_epi32(rounded, two), max_code_i32);
+            const __m256i signed_odd_sqr =
+                    _mm256_mullo_epi32(signed_odd, signed_odd);
+            accumulate_i32_as_i64(signed_odd_sqr, sq_acc_lo, sq_acc_hi);
+        }
+        store_i32_as_u8_8(rounded, rqq + i);
+    }
+
+    sum_qq += reduce_add_256(sum_acc_lo) + reduce_add_256(sum_acc_hi);
+    if (centered) {
+        sum2_signed_odd_int +=
+                reduce_add_256(sq_acc_lo) + reduce_add_256(sq_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::AVX2>(
         const uint8_t* query,
@@ -137,6 +291,69 @@ uint64_t bitwise_and_dot_product<SIMDLevel::AVX2>(
     return sum;
 }
 
+template <>
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::AVX2>(
+        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 = 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 uint64_t yv = *(const uint64_t*)(data + offset);
+        popcount_sum += popcount64(yv);
+        for (int j = 0; j < qb; j++) {
+            const uint64_t qv = *(const uint64_t*)(query + j * size + offset);
+            dot_product += popcount64(qv & yv) << j;
+        }
+    }
+    for (; offset < size; ++offset) {
+        const uint8_t yv = *(data + offset);
+        popcount_sum += popcount32(yv);
+        for (int j = 0; j < qb; j++) {
+            const uint8_t qv = *(query + j * size + offset);
+            dot_product += popcount32(qv & yv) << j;
+        }
+    }
+    return {dot_product, popcount_sum};
+}
+
 template <>
 uint64_t bitwise_xor_dot_product<SIMDLevel::AVX2>(
         const uint8_t* query,
@@ -220,6 +437,34 @@ uint64_t popcount<SIMDLevel::AVX2>(const uint8_t* data, size_t size) {
     return sum;
 }
 
+template <>
+void rearrange_bit_planes<SIMDLevel::AVX2>(
+        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);
+    const size_t nchunks = d / 32;
+    for (size_t chunk = 0; chunk < nchunks; chunk++) {
+        __m256i vals =
+                _mm256_loadu_si256((const __m256i*)(rotated_qq + chunk * 32));
+        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 + chunk * 4], &packed, 4);
+        }
+    }
+    for (size_t idim = nchunks * 32; 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 {

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,6 +582,41 @@ 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 {

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

@@ -185,6 +185,98 @@ uint64_t bitwise_and_dot_product<SIMDLevel::AVX512_SPR>(
     return sum;
 }
 
+template <>
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::AVX512_SPR>(
+        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_SPR>(
         const uint8_t* query,

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,