Fuse RaBitQ AND-dot and popcount scan (#5412)

Summary:
Fuse RaBitQ’s AND-dot and doc-side popcount into a single SIMD pass, avoiding a second scan of the same binary code.

Pull Request resolved: #5412

Reviewed By: mnorris11

Differential Revision: D112829641

Pulled By: alibeklfc

fbshipit-source-id: d8015661fe27e656d34a210ddeec7789aed1802e

Author: lyang24

benchs/bench_rabitq_simd.cpp

@@ -82,6 +82,18 @@ void bench_rabitq_and_dot_product_with_sum(benchmark::State& state) {
             });
 }
 
+void bench_rabitq_and_dot_product_with_sum_fused(benchmark::State& state) {
+    bench_rabitq_generic(
+            state,
+            [](const uint8_t* q, const uint8_t* x, size_t size, size_t qb)
+                    -> int64_t {
+                auto result = rabitq::bitwise_and_dot_product_with_popcount(
+                        q, x, size, qb);
+                // Synthetic operation using both inputs for benchmarking.
+                return result.popcount + result.dot_product;
+            });
+}
+
 template <SIMDLevel SL>
 void bench_rabitq_rearrange_impl(benchmark::State& state) {
     size_t qb = state.range(0);
@@ -127,6 +139,9 @@ BENCHMARK(bench_rabitq_xor_dot_product)
 BENCHMARK(bench_rabitq_and_dot_product_with_sum)
         ->ArgsProduct({qbs, dims})
         ->ArgNames({"qb", "d"});
+BENCHMARK(bench_rabitq_and_dot_product_with_sum_fused)
+        ->ArgsProduct({qbs, dims})
+        ->ArgNames({"qb", "d"});
 BENCHMARK(bench_rabitq_rearrange_scalar)
         ->ArgsProduct({qbs, dims})
         ->ArgNames({"qb", "d"});

faiss/impl/RaBitQuantizer.cpp

@@ -468,17 +468,16 @@ struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer {
                                qb);
             final_dot += int_dot * query_fac.int_dot_scale;
         } else {
-            auto dot_qo = rabitq::bitwise_and_dot_product<SL>(
-                    rearranged_rotated_qq.data(), binary_data, size, qb);
-            // It was a willful decision (after the discussion) to not to
-            // pre-cache the sum of all bits, just in order to reduce the
-            // overhead per vector.
-            // process 64-bit popcounts
-            auto sum_q = rabitq::popcount<SL>(binary_data, size);
+            auto bitwise_result =
+                    rabitq::bitwise_and_dot_product_with_popcount<SL>(
+                            rearranged_rotated_qq.data(),
+                            binary_data,
+                            size,
+                            qb);
             // dot-product itself
-            final_dot += query_fac.c1 * dot_qo;
+            final_dot += query_fac.c1 * bitwise_result.dot_product;
             // normalizer coefficients
-            final_dot += query_fac.c2 * sum_q;
+            final_dot += query_fac.c2 * bitwise_result.popcount;
             // normalizer coefficients
             final_dot -= query_fac.c34;
         }

faiss/utils/rabitq_simd.h

@@ -21,6 +21,11 @@ namespace faiss::rabitq {
 constexpr int RABITQ_QUANTIZATION_SIMD_LEVELS = (1 << int(SIMDLevel::NONE)) |
         (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::AVX512));
 
+struct BitwiseAndDotProductResult {
+    uint64_t dot_product;
+    uint64_t popcount;
+};
+
 /**
  * Compute dot product between query and binary data using popcount on AND.
  *
@@ -37,6 +42,17 @@ uint64_t bitwise_and_dot_product(
         size_t size,
         size_t qb);
 
+/**
+ * Compute bitwise_and_dot_product(query, data, size, qb) and popcount(data,
+ * size) in one pass over data.
+ */
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount(
+        const uint8_t* query,
+        const uint8_t* data,
+        size_t size,
+        size_t qb);
+
 /**
  * Compute dot product between query and binary data using popcount on XOR.
  *
@@ -140,6 +156,35 @@ inline uint64_t bitwise_and_dot_product<SIMDLevel::NONE>(
     return sum;
 }
 
+template <>
+inline BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::NONE>(
+        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;
+    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 <>
 inline uint64_t bitwise_xor_dot_product<SIMDLevel::NONE>(
         const uint8_t* query,

faiss/utils/simd_impl/rabitq_avx2.cpp

@@ -291,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,

faiss/utils/simd_impl/rabitq_avx512.cpp

@@ -394,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,

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,

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,