faiss 0.6.2 → 0.6.4

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

@@ -12,12 +12,58 @@
 #ifdef COMPILE_SIMD_RISCV_RVV
 
 #include <faiss/utils/extra_distances.h>
+#include <riscv_vector.h>
 
 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) {
-    return fvec_norm_L2sqr<SIMDLevel::NONE>(x, d);
+    size_t vlmax = __riscv_vsetvlmax_e32m8();
+    vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m8(d - i);
+        vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
+        acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
+        i += vl;
+    }
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
@@ -25,7 +71,18 @@ float fvec_L2sqr<SIMDLevel::RISCV_RVV>(
         const float* x,
         const float* y,
         size_t d) {
-    return fvec_L2sqr<SIMDLevel::NONE>(x, y, d);
+    size_t vlmax = __riscv_vsetvlmax_e32m8();
+    vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m8(d - i);
+        vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
+        vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
+        vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
+        acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
+        i += vl;
+    }
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
@@ -33,12 +90,34 @@ float fvec_inner_product<SIMDLevel::RISCV_RVV>(
         const float* x,
         const float* y,
         size_t d) {
-    return fvec_inner_product<SIMDLevel::NONE>(x, y, d);
+    size_t vlmax = __riscv_vsetvlmax_e32m8();
+    vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m8(d - i);
+        vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
+        vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
+        acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vy, vl);
+        i += vl;
+    }
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
 float fvec_L1<SIMDLevel::RISCV_RVV>(const float* x, const float* y, size_t d) {
-    return fvec_L1<SIMDLevel::NONE>(x, y, d);
+    size_t vlmax = __riscv_vsetvlmax_e32m8();
+    vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m8(d - i);
+        vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
+        vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
+        vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
+        vx = __riscv_vfsgnjx_vv_f32m8(vx, vx, vl);
+        acc = __riscv_vfadd_vv_f32m8_tu(acc, acc, vx, vl);
+        i += vl;
+    }
+    return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
 }
 
 template <>
@@ -46,7 +125,19 @@ float fvec_Linf<SIMDLevel::RISCV_RVV>(
         const float* x,
         const float* y,
         size_t d) {
-    return fvec_Linf<SIMDLevel::NONE>(x, y, d);
+    size_t vlmax = __riscv_vsetvlmax_e32m8();
+    vfloat32m8_t vmax = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m8(d - i);
+        vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
+        vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
+        vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
+        vx = __riscv_vfsgnjx_vv_f32m8(vx, vx, vl);
+        vmax = __riscv_vfmax_vv_f32m8_tu(vmax, vmax, vx, vl);
+        i += vl;
+    }
+    return rvv_reduce(vmax, vlmax, 0.0f, __riscv_vfredmax_vs_f32m8_f32m1);
 }
 
 template <>
@@ -61,8 +152,29 @@ void fvec_inner_product_batch_4<SIMDLevel::RISCV_RVV>(
         float& dis1,
         float& dis2,
         float& dis3) {
-    fvec_inner_product_batch_4<SIMDLevel::NONE>(
-            x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    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 <>
@@ -77,8 +189,33 @@ void fvec_L2sqr_batch_4<SIMDLevel::RISCV_RVV>(
         float& dis1,
         float& dis2,
         float& dis3) {
-    fvec_L2sqr_batch_4<SIMDLevel::NONE>(
-            x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    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 <>
@@ -90,8 +227,32 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    fvec_L2sqr_ny_transposed<SIMDLevel::NONE>(
-            dis, x, y, y_sqlen, d, d_offset, ny);
+    size_t vlmax = __riscv_vsetvlmax_e32m8();
+    vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
+    size_t i = 0;
+    while (i < d) {
+        size_t vl = __riscv_vsetvl_e32m8(d - i);
+        vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
+        acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
+        i += vl;
+    }
+    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);
+        acc = __riscv_vfmv_v_f_f32m8(0.0f, vl);
+        for (size_t j = 0; j < d; j++) {
+            vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + j * d_offset + i, vl);
+            acc = __riscv_vfmacc_vf_f32m8(acc, x[j], vy, vl);
+        }
+        vfloat32m8_t vres = __riscv_vle32_v_f32m8(y_sqlen + i, vl);
+        vres = __riscv_vfadd_vf_f32m8(vres, x_sqlen, vl);
+        acc = __riscv_vfmul_vf_f32m8(acc, 2.0f, vl);
+        vres = __riscv_vfsub_vv_f32m8(vres, acc, vl);
+        __riscv_vse32_v_f32m8(dis + i, vres, vl);
+        i += vl;
+    }
 }
 
 template <>
@@ -101,7 +262,10 @@ void fvec_inner_products_ny<SIMDLevel::RISCV_RVV>(
         const float* y,
         size_t d,
         size_t ny) {
-    fvec_inner_products_ny<SIMDLevel::NONE>(ip, x, y, d, ny);
+    for (size_t i = 0; i < ny; i++) {
+        ip[i] = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y, d);
+        y += d;
+    }
 }
 
 template <>
@@ -111,7 +275,10 @@ void fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(
         const float* y,
         size_t d,
         size_t ny) {
-    fvec_L2sqr_ny<SIMDLevel::NONE>(dis, x, y, d, ny);
+    for (size_t i = 0; i < ny; i++) {
+        dis[i] = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y, d);
+        y += d;
+    }
 }
 
 template <>
@@ -121,8 +288,9 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::RISCV_RVV>(
         const float* y,
         size_t d,
         size_t ny) {
-    return fvec_L2sqr_ny_nearest<SIMDLevel::NONE>(
-            distances_tmp_buffer, x, y, d, ny);
+    fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(distances_tmp_buffer, x, y, d, ny);
+    const size_t j = rvv_argmin(distances_tmp_buffer, ny);
+    return j < ny ? j : 0;
 }
 
 template <>
@@ -134,8 +302,10 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::RISCV_RVV>(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
+    fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
             distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
+    const size_t j = rvv_argmin(distances_tmp_buffer, ny);
+    return j < ny ? j : 0;
 }
 
 template <>
@@ -145,7 +315,15 @@ void fvec_madd<SIMDLevel::RISCV_RVV>(
         float bf,
         const float* b,
         float* c) {
-    fvec_madd<SIMDLevel::NONE>(n, a, bf, b, c);
+    size_t i = 0;
+    while (i < n) {
+        size_t vl = __riscv_vsetvl_e32m8(n - i);
+        vfloat32m8_t va = __riscv_vle32_v_f32m8(a + i, vl);
+        vfloat32m8_t vb = __riscv_vle32_v_f32m8(b + i, vl);
+        va = __riscv_vfmacc_vf_f32m8(va, bf, vb, vl);
+        __riscv_vse32_v_f32m8(c + i, va, vl);
+        i += vl;
+    }
 }
 
 template <>
@@ -155,7 +333,9 @@ int fvec_madd_and_argmin<SIMDLevel::RISCV_RVV>(
         float bf,
         const float* b,
         float* c) {
-    return fvec_madd_and_argmin<SIMDLevel::NONE>(n, a, bf, b, c);
+    fvec_madd<SIMDLevel::RISCV_RVV>(n, a, bf, b, c);
+    const size_t j = rvv_argmin(c, n);
+    return j < n ? static_cast<int>(j) : -1;
 }
 
 #define DEFINE_VECTOR_DISTANCE_RVV_FALLBACK(metric)                 \

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 {