facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

RaBitQ: SIMD helper simplification, use faster popcount for doc-side sum (#4573)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4573

`rabitq_simd.h` had significant duplication across different helper methods, often with the tail of the method being identical.

Reviewed By: luciang

Differential Revision: D81825790

fbshipit-source-id: f3a024889a1be7b4dcc01e84da29e13abcd884e8
This commit is contained in:
Tom Jackson
2025-09-09 19:00:58 -07:00
committed by Facebook GitHub Bot
parent b7c88eaff8
commit bc3e3a1ef1
3 changed files with 189 additions and 436 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchs/bench_rabitq_simd.cpp
+10 -22
View File
@@ -37,7 +37,7 @@ void bench_rabitq_generic(benchmark::State& state, auto distFn) {
for (auto _ : state) {
++r;
for (size_t i = 0; i < n; ++i) {
sum += distFn(q.data(), x.data() + i * size, d, qb);
sum += distFn(q.data(), x.data() + i * size, size, qb);
}
benchmark::DoNotOptimize(sum);
}
@@ -45,41 +45,29 @@ void bench_rabitq_generic(benchmark::State& state, auto distFn) {
state.SetBytesProcessed(r * x.size());
}
uint64_t rabitq_simple_popcnt(const uint8_t* data, size_t size) {
uint64_t sum = 0;
size_t i = 0;
for (; i + 8 <= size; i += 8) {
const auto yv = *(const uint64_t*)(data + i);
sum += __builtin_popcountll(yv);
}
for (; i < size; ++i) {
const auto yv = *(data + i);
sum += __builtin_popcount(yv);
}
return sum;
}
void bench_rabitq_sum(benchmark::State& state) {
bench_rabitq_generic(
state,
[](const uint8_t* /*q*/, const uint8_t* x, size_t d, size_t /* qb*/)
-> int64_t { return rabitq_simple_popcnt(x, d); });
[](const uint8_t*, const uint8_t* x, size_t size, size_t)
-> int64_t { return rabitq::popcount(x, size); });
}
void bench_rabitq_and_dot_product(benchmark::State& state) {
bench_rabitq_generic(
state,
[](const uint8_t* q, const uint8_t* x, size_t d, size_t qb)
-> int64_t { return rabitq_dp_popcnt(q, x, d, qb); });
[](const uint8_t* q, const uint8_t* x, size_t size, size_t qb)
-> int64_t {
return rabitq::bitwise_and_dot_product(q, x, size, qb);
});
}
void bench_rabitq_and_dot_product_with_sum(benchmark::State& state) {
bench_rabitq_generic(
state,
[](const uint8_t* q, const uint8_t* x, size_t d, size_t qb)
[](const uint8_t* q, const uint8_t* x, size_t size, size_t qb)
-> int64_t {
auto sum_q = rabitq_simple_popcnt(x, d);
auto dp = rabitq_dp_popcnt(q, x, d, qb);
auto sum_q = rabitq::popcount(x, size);
auto dp = rabitq::bitwise_and_dot_product(q, x, size, qb);
// Synthetic operation using both inputs for benchmarking.
return sum_q + dp;
});
faiss/impl/RaBitQuantizer.cpp
+7 -39
View File
@@ -329,49 +329,17 @@ float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
metric_type == MetricType::METRIC_INNER_PRODUCT));
// split the code into parts
size_t size = (d + 7) / 8;
const uint8_t* binary_data = code;
const FactorsData* fac =
reinterpret_cast<const FactorsData*>(code + (d + 7) / 8);
const FactorsData* fac = reinterpret_cast<const FactorsData*>(code + size);
// // this is the baseline code
// //
// // compute <q,o> using integers
// size_t dot_qo = 0;
// for (size_t i = 0; i < d; i++) {
// // extract i-th bit
// const uint8_t masker = (1 << (i % 8));
// const uint8_t bit = ((binary_data[i / 8] & masker) == masker) ? 1 :
// 0;
//
// // accumulate dp
// dot_qo += bit * rotated_qq[i];
// }
// this is the scheme for popcount
const size_t di_8b = (d + 7) / 8;
const size_t di_64b = (di_8b / 8) * 8;
// Use the optimized popcount function from rabitq_simd.h
float dot_qo =
rabitq_dp_popcnt(rearranged_rotated_qq.data(), binary_data, d, qb);
// See RaBitDistanceComputerNotQ::distance_to_code() for baseline code.
float dot_qo = rabitq::bitwise_and_dot_product(
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.
uint64_t sum_q = 0;
{
// process 64-bit popcounts
for (size_t i = 0; i < di_64b; i += 8) {
const auto yv = *(const uint64_t*)(binary_data + i);
sum_q += __builtin_popcountll(yv);
}
// process leftovers
for (size_t i = di_64b; i < di_8b; i++) {
const auto yv = *(binary_data + i);
sum_q += __builtin_popcount(yv);
}
}
// the sum of all bits, just in order to reduce the overhead per vector.
auto sum_q = rabitq::popcount(binary_data, size);
float final_dot = 0;
// dot-product itself
final_dot += query_fac.c1 * dot_qo;
faiss/utils/rabitq_simd.h
+172 -375
View File
@@ -14,9 +14,9 @@
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || \
defined(_M_IX86)
#include <immintrin.h>
#endif
#endif // defined(__x86_64__) || defined(_M_X64) || defined(__i386__) ||
namespace faiss {
namespace faiss::rabitq {
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || \
defined(_M_IX86)
@@ -24,6 +24,8 @@ namespace faiss {
* Returns the lookup table for AVX512 popcount operations.
* This table is used for lookup-based popcount implementation.
*
* Source: https://github.com/WojciechMula/sse-popcount.
*
* @return Lookup table as __m512i register
*/
#if defined(__AVX512F__)
@@ -94,7 +96,7 @@ inline __m512i get_lookup_512() {
/* 1 */ 1,
/* 0 */ 0);
}
#endif
#endif // defined(__AVX512F__)
#if defined(__AVX2__)
/**
* Returns the lookup table for AVX2 popcount operations.
@@ -137,409 +139,204 @@ inline __m256i get_lookup_256() {
/* e */ 3,
/* f */ 4);
}
#endif
#endif // defined(__AVX2__)
#if defined(__AVX512F__)
/**
* Performs lookup-based popcount on AVX512 registers.
* Popcount for a 512-bit register, using lookup tables if necessary.
*
* @param v_and Input vector to count bits in
* @return Vector with popcount results
* @param v Input vector to count bits in
* @return Vector int32_t[16] with popcount results.
*/
inline __m512i popcount_lookup_avx512(__m512i v_and) {
inline __m512i popcount_512(__m512i v) {
#if defined(__AVX512VPOPCNTDQ__)
return _mm512_popcnt_epi64(v);
#else
const __m512i lookup = get_lookup_512();
const __m512i low_mask = _mm512_set1_epi8(0x0f);
const __m512i lo = _mm512_and_si512(v_and, low_mask);
const __m512i hi = _mm512_and_si512(_mm512_srli_epi16(v_and, 4), low_mask);
const __m512i popcnt1 = _mm512_shuffle_epi8(lookup, lo);
const __m512i popcnt2 = _mm512_shuffle_epi8(lookup, hi);
return _mm512_add_epi8(popcnt1, popcnt2);
const __m512i lo = _mm512_and_si512(v, low_mask);
const __m512i hi = _mm512_and_si512(_mm512_srli_epi16(v, 4), low_mask);
const __m512i popcnt_lo = _mm512_shuffle_epi8(lookup, lo);
const __m512i popcnt_hi = _mm512_shuffle_epi8(lookup, hi);
const __m512i popcnt = _mm512_add_epi8(popcnt_lo, popcnt_hi);
return _mm512_sad_epu8(_mm512_setzero_si512(), popcnt);
#endif // defined(__AVX512VPOPCNTDQ__)
}
#endif
#endif // defined(__AVX512F__)
#if defined(__AVX2__)
/**
* Performs lookup-based popcount on AVX2 registers.
* Popcount for a 256-bit register, using lookup tables if necessary.
*
* @param v_and Input vector to count bits in
* @return Vector with popcount results
* @param v Input vector to count bits in
* @return uint64_t[4] of popcounts for each portion of the input vector.
*/
inline __m256i popcount_lookup_avx2(__m256i v_and) {
inline __m256i popcount_256(__m256i v) {
const __m256i lookup = get_lookup_256();
const __m256i low_mask = _mm256_set1_epi8(0x0f);
const __m256i lo = _mm256_and_si256(v_and, low_mask);
const __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v_and, 4), low_mask);
const __m256i popcnt1 = _mm256_shuffle_epi8(lookup, lo);
const __m256i popcnt2 = _mm256_shuffle_epi8(lookup, hi);
return _mm256_add_epi8(popcnt1, popcnt2);
const __m256i lo = _mm256_and_si256(v, low_mask);
const __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask);
const __m256i popcnt_lo = _mm256_shuffle_epi8(lookup, lo);
const __m256i popcnt_hi = _mm256_shuffle_epi8(lookup, hi);
const __m256i popcnt = _mm256_add_epi8(popcnt_lo, popcnt_hi);
// Reduce uint8_t[32] into uint64_t[4] by addition.
return _mm256_sad_epu8(_mm256_setzero_si256(), popcnt);
}
#endif
#endif
#if defined(__AVX512F__) && defined(__AVX512VPOPCNTDQ__)
/**
* AVX512-optimized version of dot product computation between query and binary
* data. Requires AVX512F and AVX512VPOPCNTDQ instruction sets.
*
* @param query Pointer to rearranged rotated query data
* @param binary_data Pointer to binary data
* @param d Dimension
* @param qb Number of quantization bits
* @return Dot product result as float
*/
inline float rabitq_dp_popcnt_avx512(
const uint8_t* query,
const uint8_t* binary_data,
size_t d,
size_t qb) {
__m512i sum_512 = _mm512_setzero_si512();
const size_t di_8b = (d + 7) / 8;
const size_t d_512 = (d / 512) * 512;
const size_t d_256 = (d / 256) * 256;
const size_t d_128 = (d / 128) * 128;
for (size_t i = 0; i < d_512; i += 512) {
__m512i v_x = _mm512_loadu_si512((const __m512i*)(binary_data + i / 8));
for (size_t j = 0; j < qb; j++) {
__m512i v_q = _mm512_loadu_si512(
(const __m512i*)(query + j * di_8b + i / 8));
__m512i v_and = _mm512_and_si512(v_q, v_x);
__m512i v_popcnt = _mm512_popcnt_epi32(v_and);
sum_512 = _mm512_add_epi32(sum_512, _mm512_slli_epi32(v_popcnt, j));
}
}
__m256i sum_256 = _mm256_add_epi32(
_mm512_extracti32x8_epi32(sum_512, 0),
_mm512_extracti32x8_epi32(sum_512, 1));
if (d_256 != d_512) {
__m256i v_x =
_mm256_loadu_si256((const __m256i*)(binary_data + d_512 / 8));
for (size_t j = 0; j < qb; j++) {
__m256i v_q = _mm256_loadu_si256(
(const __m256i*)(query + j * di_8b + d_512 / 8));
__m256i v_and = _mm256_and_si256(v_q, v_x);
__m256i v_popcnt = _mm256_popcnt_epi32(v_and);
sum_256 = _mm256_add_epi32(sum_256, _mm256_slli_epi32(v_popcnt, j));
}
}
__m128i sum_128 = _mm_add_epi32(
_mm256_extracti32x4_epi32(sum_256, 0),
_mm256_extracti32x4_epi32(sum_256, 1));
if (d_128 != d_256) {
__m128i v_x =
_mm_loadu_si128((const __m128i*)(binary_data + d_256 / 8));
for (size_t j = 0; j < qb; j++) {
__m128i v_q = _mm_loadu_si128(
(const __m128i*)(query + j * di_8b + d_256 / 8));
__m128i v_and = _mm_and_si128(v_q, v_x);
__m128i v_popcnt = _mm_popcnt_epi32(v_and);
sum_128 = _mm_add_epi32(sum_128, _mm_slli_epi32(v_popcnt, j));
}
}
if (d != d_128) {
const size_t leftovers = d - d_128;
const __mmask16 mask = (1 << ((leftovers + 7) / 8)) - 1;
__m128i v_x = _mm_maskz_loadu_epi8(
mask, (const __m128i*)(binary_data + d_128 / 8));
for (size_t j = 0; j < qb; j++) {
__m128i v_q = _mm_maskz_loadu_epi8(
mask, (const __m128i*)(query + j * di_8b + d_128 / 8));
__m128i v_and = _mm_and_si128(v_q, v_x);
__m128i v_popcnt = _mm_popcnt_epi32(v_and);
sum_128 = _mm_add_epi32(sum_128, _mm_slli_epi32(v_popcnt, j));
}
}
int sum_64le = 0;
sum_64le += _mm_extract_epi32(sum_128, 0);
sum_64le += _mm_extract_epi32(sum_128, 1);
sum_64le += _mm_extract_epi32(sum_128, 2);
sum_64le += _mm_extract_epi32(sum_128, 3);
return static_cast<float>(sum_64le);
}
#endif
#if defined(__AVX512F__) && !defined(__AVX512VPOPCNTDQ__)
/**
* AVX512-optimized version of dot product computation between query and binary
* data. Uses AVX512F instructions but does not require AVX512VPOPCNTDQ.
*
* @param query Pointer to rearranged rotated query data
* @param binary_data Pointer to binary data
* @param d Dimension
* @param qb Number of quantization bits
* @return Dot product result as float
*/
inline float rabitq_dp_popcnt_avx512_fallback(
const uint8_t* query,
const uint8_t* binary_data,
size_t d,
size_t qb) {
const size_t di_8b = (d + 7) / 8;
const size_t d_512 = (d / 512) * 512;
const size_t d_256 = (d / 256) * 256;
const size_t d_128 = (d / 128) * 128;
// Use the lookup-based popcount helper function
__m512i sum_512 = _mm512_setzero_si512();
// Process 512 bits (64 bytes) at a time using lookup-based popcount
for (size_t i = 0; i < d_512; i += 512) {
__m512i v_x = _mm512_loadu_si512((const __m512i*)(binary_data + i / 8));
for (size_t j = 0; j < qb; j++) {
__m512i v_q = _mm512_loadu_si512(
(const __m512i*)(query + j * di_8b + i / 8));
__m512i v_and = _mm512_and_si512(v_q, v_x);
// Use the popcount_lookup_avx512 helper function
__m512i v_popcnt = popcount_lookup_avx512(v_and);
// Sum bytes to 32-bit integers
__m512i v_sad = _mm512_sad_epu8(v_popcnt, _mm512_setzero_si512());
// Shift by j and add to sum
__m512i v_shifted = _mm512_slli_epi64(v_sad, j);
sum_512 = _mm512_add_epi64(sum_512, v_shifted);
}
}
// Handle 256-bit section if needed
__m256i sum_256 = _mm256_setzero_si256();
if (d_256 != d_512) {
__m256i v_x =
_mm256_loadu_si256((const __m256i*)(binary_data + d_512 / 8));
for (size_t j = 0; j < qb; j++) {
__m256i v_q = _mm256_loadu_si256(
(const __m256i*)(query + j * di_8b + d_512 / 8));
__m256i v_and = _mm256_and_si256(v_q, v_x);
// Use the popcount_lookup_avx2 helper function
__m256i v_popcnt = popcount_lookup_avx2(v_and);
// Sum bytes to 64-bit integers
__m256i v_sad = _mm256_sad_epu8(v_popcnt, _mm256_setzero_si256());
// Shift by j and add to sum
__m256i v_shifted = _mm256_slli_epi64(v_sad, j);
sum_256 = _mm256_add_epi64(sum_256, v_shifted);
}
}
// Handle 128-bit section and leftovers
__m128i sum_128 = _mm_setzero_si128();
if (d_128 != d_256) {
__m128i v_x =
_mm_loadu_si128((const __m128i*)(binary_data + d_256 / 8));
for (size_t j = 0; j < qb; j++) {
__m128i v_q = _mm_loadu_si128(
(const __m128i*)(query + j * di_8b + d_256 / 8));
__m128i v_and = _mm_and_si128(v_q, v_x);
// Scalar popcount for each 64-bit lane
uint64_t lane0 = _mm_extract_epi64(v_and, 0);
uint64_t lane1 = _mm_extract_epi64(v_and, 1);
uint64_t pop0 = __builtin_popcountll(lane0) << j;
uint64_t pop1 = __builtin_popcountll(lane1) << j;
sum_128 = _mm_add_epi64(sum_128, _mm_set_epi64x(pop1, pop0));
}
}
// Handle remaining bytes (less than 16)
uint64_t sum_leftover = 0;
size_t d_leftover = d - d_128;
if (d_leftover > 0) {
for (size_t j = 0; j < qb; j++) {
for (size_t k = 0; k < (d_leftover + 7) / 8; ++k) {
uint8_t qv = query[j * di_8b + d_128 / 8 + k];
uint8_t yv = binary_data[d_128 / 8 + k];
sum_leftover += (__builtin_popcount(qv & yv) << j);
}
}
}
// Horizontal sum of all lanes
uint64_t sum = 0;
// Sum from 512-bit registers
alignas(64) uint64_t lanes512[8];
_mm512_store_si512((__m512i*)lanes512, sum_512);
for (int i = 0; i < 8; ++i) {
sum += lanes512[i];
}
// Sum from 256-bit registers
alignas(32) uint64_t lanes256[4];
_mm256_store_si256((__m256i*)lanes256, sum_256);
for (int i = 0; i < 4; ++i) {
sum += lanes256[i];
}
// Sum from 128-bit registers
alignas(16) uint64_t lanes128[2];
_mm_store_si128((__m128i*)lanes128, sum_128);
sum += lanes128[0] + lanes128[1];
// Add leftovers
sum += sum_leftover;
return static_cast<float>(sum);
}
#endif
#ifdef __AVX2__
/**
* AVX2-optimized version of dot product computation between query and binary
* data.
*
* @param query Pointer to rearranged rotated query data
* @param binary_data Pointer to binary data
* @param d Dimension
* @param qb Number of quantization bits
* @return Dot product result as float
*/
inline float rabitq_dp_popcnt_avx2(
const uint8_t* query,
const uint8_t* binary_data,
size_t d,
size_t qb) {
const size_t di_8b = (d + 7) / 8;
const size_t d_256 = (d / 256) * 256;
const size_t d_128 = (d / 128) * 128;
// Use the lookup-based popcount helper function
__m256i sum_256 = _mm256_setzero_si256();
// Process 256 bits (32 bytes) at a time using lookup-based popcount
for (size_t i = 0; i < d_256; i += 256) {
__m256i v_x = _mm256_loadu_si256((const __m256i*)(binary_data + i / 8));
for (size_t j = 0; j < qb; j++) {
__m256i v_q = _mm256_loadu_si256(
(const __m256i*)(query + j * di_8b + i / 8));
__m256i v_and = _mm256_and_si256(v_q, v_x);
// Use the popcount_lookup_avx2 helper function
__m256i v_popcnt = popcount_lookup_avx2(v_and);
// Convert byte counts to 64-bit lanes and shift by j
__m256i v_sad = _mm256_sad_epu8(v_popcnt, _mm256_setzero_si256());
__m256i v_shifted = _mm256_slli_epi64(v_sad, static_cast<int>(j));
sum_256 = _mm256_add_epi64(sum_256, v_shifted);
}
}
// Handle leftovers with 128-bit SIMD
__m128i sum_128 = _mm_setzero_si128();
if (d_128 != d_256) {
__m128i v_x =
_mm_loadu_si128((const __m128i*)(binary_data + d_256 / 8));
for (size_t j = 0; j < qb; j++) {
__m128i v_q = _mm_loadu_si128(
(const __m128i*)(query + j * di_8b + d_256 / 8));
__m128i v_and = _mm_and_si128(v_q, v_x);
// Scalar popcount for each 64-bit lane
uint64_t lane0 = _mm_extract_epi64(v_and, 0);
uint64_t lane1 = _mm_extract_epi64(v_and, 1);
uint64_t pop0 = __builtin_popcountll(lane0) << j;
uint64_t pop1 = __builtin_popcountll(lane1) << j;
sum_128 = _mm_add_epi64(sum_128, _mm_set_epi64x(pop1, pop0));
}
}
// Handle remaining bytes (less than 16)
uint64_t sum_leftover = 0;
size_t d_leftover = d - d_128;
if (d_leftover > 0) {
for (size_t j = 0; j < qb; j++) {
for (size_t k = 0; k < (d_leftover + 7) / 8; ++k) {
uint8_t qv = query[j * di_8b + d_128 / 8 + k];
uint8_t yv = binary_data[d_128 / 8 + k];
sum_leftover += (__builtin_popcount(qv & yv) << j);
}
}
}
// Horizontal sum of all lanes
uint64_t sum = 0;
// sum_256: 4 lanes of 64 bits
inline uint64_t reduce_add_256(__m256i v) {
alignas(32) uint64_t lanes[4];
_mm256_store_si256((__m256i*)lanes, sum_256);
for (int i = 0; i < 4; ++i) {
sum += lanes[i];
}
// sum_128: 2 lanes of 64 bits
alignas(16) uint64_t lanes128[2];
_mm_store_si128((__m128i*)lanes128, sum_128);
sum += lanes128[0] + lanes128[1];
// leftovers
sum += sum_leftover;
return static_cast<float>(sum);
_mm256_store_si256((__m256i*)lanes, v);
return lanes[0] + lanes[1] + lanes[2] + lanes[3];
}
#endif
#endif // defined(__AVX2__)
#if defined(__SSE4_1__)
inline __m128i popcount_128(__m128i v) {
// Scalar popcount for each 64-bit lane
uint64_t lane0 = _mm_extract_epi64(v, 0);
uint64_t lane1 = _mm_extract_epi64(v, 1);
uint64_t pop0 = __builtin_popcountll(lane0);
uint64_t pop1 = __builtin_popcountll(lane1);
return _mm_set_epi64x(pop1, pop0);
}
inline uint64_t reduce_add_128(__m128i v) {
alignas(16) uint64_t lanes[2];
_mm_store_si128((__m128i*)lanes, v);
return lanes[0] + lanes[1];
}
#endif // defined(__SSE4_1__)
#endif // defined(__x86_64__) || defined(_M_X64) || defined(__i386__) ||
/**
* Compute dot product between query and binary data using popcount operations.
*
* @param query Pointer to rearranged rotated query data
* @param binary_data Pointer to binary data
* @param data Pointer to binary data
* @param d Dimension
* @param qb Number of quantization bits
* @return Dot product result as float
* @return Unsigned integer dot product
*/
inline float rabitq_dp_popcnt(
inline uint64_t bitwise_and_dot_product(
const uint8_t* query,
const uint8_t* binary_data,
size_t d,
const uint8_t* data,
size_t size,
size_t qb) {
#if defined(__AVX512F__) && defined(__AVX512VPOPCNTDQ__)
return rabitq_dp_popcnt_avx512(query, binary_data, d, qb);
#elif defined(__AVX512F__)
return rabitq_dp_popcnt_avx512_fallback(query, binary_data, d, qb);
#elif defined(__AVX2__)
return rabitq_dp_popcnt_avx2(query, binary_data, d, qb);
#else
const size_t di_8b = (d + 7) / 8;
const size_t di_64b = (di_8b / 8) * 8;
uint64_t dot_qo = 0;
for (size_t j = 0; j < qb; j++) {
const uint8_t* query_j = query + j * di_8b;
// process 64-bit popcounts
uint64_t count_dot = 0;
for (size_t i = 0; i < di_64b; i += 8) {
const auto qv = *(const uint64_t*)(query_j + i);
const auto yv = *(const uint64_t*)(binary_data + i);
count_dot += __builtin_popcountll(qv & yv);
uint64_t sum = 0;
size_t offset = 0;
#if defined(__AVX512F__)
// Handle 512-bit chunks.
if (size_t step = 512 / 8; offset + step <= size) {
__m512i sum_512 = _mm512_setzero_si512();
for (; offset + step <= size; offset += step) {
__m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset));
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);
sum_512 = _mm512_add_epi64(sum_512, v_shifted);
}
}
// process leftovers
for (size_t i = di_64b; i < di_8b; i++) {
const auto qv = *(query_j + i);
const auto yv = *(binary_data + i);
count_dot += __builtin_popcount(qv & yv);
}
dot_qo += (count_dot << j);
sum += _mm512_reduce_add_epi64(sum_512);
}
return static_cast<float>(dot_qo);
#endif
#endif // defined(__AVX512F__)
#if defined(__AVX2__)
if (size_t step = 256 / 8; offset + step <= size) {
__m256i sum_256 = _mm256_setzero_si256();
for (; offset + step <= size; offset += step) {
__m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset));
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);
sum_256 = _mm256_add_epi64(sum_256, v_shifted);
}
}
sum += reduce_add_256(sum_256);
}
#endif // defined(__AVX2__)
#if defined(__SSE4_1__)
__m128i sum_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));
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);
sum_128 = _mm_add_epi64(sum_128, v_shifted);
}
}
sum += reduce_add_128(sum_128);
#endif // defined(__SSE4_1__)
for (size_t step = 64 / 8; offset + step <= size; offset += step) {
const auto yv = *(const uint64_t*)(data + offset);
for (int j = 0; j < qb; j++) {
const auto qv = *(const uint64_t*)(query + j * size + offset);
sum += __builtin_popcountll(qv & yv) << j;
}
}
for (; offset < size; ++offset) {
const auto yv = *(data + offset);
for (int j = 0; j < qb; j++) {
const auto qv = *(query + j * size + offset);
sum += __builtin_popcount(qv & yv) << j;
}
}
return sum;
}
} // namespace faiss
inline uint64_t popcount(const uint8_t* data, size_t size) {
uint64_t sum = 0;
size_t offset = 0;
#if defined(__AVX512F__)
// Handle 512-bit chunks.
if (offset + 512 / 8 <= size) {
__m512i sum_512 = _mm512_setzero_si512();
for (size_t end; (end = offset + 512 / 8) <= size; offset = end) {
__m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset));
__m512i v_popcnt = popcount_512(v_x);
sum_512 = _mm512_add_epi64(sum_512, v_popcnt);
}
sum += _mm512_reduce_add_epi64(sum_512);
}
#endif // defined(__AVX512F__)
#if defined(__AVX2__)
if (offset + 256 / 8 <= size) {
__m256i sum_256 = _mm256_setzero_si256();
for (size_t end; (end = offset + 256 / 8) <= size; offset = end) {
__m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset));
__m256i v_popcnt = popcount_256(v_x);
sum_256 = _mm256_add_epi64(sum_256, v_popcnt);
}
sum += reduce_add_256(sum_256);
}
#endif // defined(__AVX2__)
#if defined(__SSE4_1__)
__m128i sum_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));
sum_128 = _mm_add_epi64(sum_128, popcount_128(v_x));
}
sum += reduce_add_128(sum_128);
#endif // defined(__SSE4_1__)
for (size_t step = 64 / 8; offset + step <= size; offset += step) {
const auto yv = *(const uint64_t*)(data + offset);
sum += __builtin_popcountll(yv);
}
for (; offset < size; ++offset) {
const auto yv = *(data + offset);
sum += __builtin_popcount(yv);
}
return sum;
}
} // namespace faiss::rabitq