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

Convert rabitq_simd.h to runtime SIMD dispatch (#4912)

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

Convert RaBitQ SIMD primitives from compile-time #ifdef cascades to the
DD template specialization pattern.

- Replace inline SIMD functions in rabitq_simd.h with template<SIMDLevel>
  declarations and NONE (scalar) specializations
- Create rabitq_avx2.cpp with AVX2+SSE4.1 specializations for
  bitwise_and_dot_product, bitwise_xor_dot_product, popcount, and
  compute_inner_product (including BMI2 bitplane kernel)
- Create rabitq_avx512.cpp with AVX512 specializations (falls back to
  AVX2 bitplane kernel for compute_inner_product with ex_bits >= 2)
- Add AVX512_SPR forwarding to AVX512 for all four functions
- Create rabitq_neon.cpp forwarding ARM_NEON/ARM_SVE to NONE
- Wrap calls in RaBitQuantizer.cpp and RaBitQUtils.cpp with
  with_simd_level dispatch
- Add -mbmi2 to AVX2 and AVX512 SIMD flags (required by _pext_u64 in
  bitplane kernel; BMI2 is present on all Haswell+ CPUs)
- Register new SIMD files in xplat.bzl and CMakeLists.txt

Reviewed By: mdouze

Differential Revision: D96298743

fbshipit-source-id: 1941bf099664ba2d3d1c1245d5c06d1775614955
This commit is contained in:
Gergely Szilvasy
2026-03-31 14:41:16 -07:00
committed by meta-codesync[bot]
parent 6bca9613b7
commit d1c432fb5f
8 changed files with 1304 additions and 850 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+3
View File
@@ -15,6 +15,7 @@ set(FAISS_SIMD_AVX2_SRC
impl/approx_topk/avx2.cpp
utils/simd_impl/distances_avx2.cpp
utils/distances_fused/simdlib_based.cpp
utils/simd_impl/rabitq_avx2.cpp
)
set(FAISS_SIMD_AVX512_SRC
impl/fast_scan/impl-avx512.cpp
@@ -22,6 +23,7 @@ set(FAISS_SIMD_AVX512_SRC
impl/scalar_quantizer/sq-avx512.cpp
utils/simd_impl/distances_avx512.cpp
utils/distances_fused/avx512.cpp
utils/simd_impl/rabitq_avx512.cpp
)
set(FAISS_SIMD_NEON_SRC
impl/fast_scan/impl-neon.cpp
@@ -29,6 +31,7 @@ set(FAISS_SIMD_NEON_SRC
impl/approx_topk/neon.cpp
utils/simd_impl/distances_aarch64.cpp
utils/distances_fused/simdlib_based_neon.cpp
utils/simd_impl/rabitq_neon.cpp
)
set(FAISS_SIMD_SVE_SRC
impl/pq_code_distance/pq_code_distance-sve.cpp
faiss/impl/RaBitQUtils.cpp
+16 -12
View File
@@ -8,6 +8,7 @@
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/rabitq_simd.h>
#include <algorithm>
@@ -306,6 +307,9 @@ size_t compute_per_vector_storage_size(size_t nb_bits, size_t d) {
}
}
// Non-template wrapper with dynamic dispatch (one dispatch per call).
// The hot path in RaBitQuantizer dispatches once at distance computer
// construction, so per-vector dispatch only affects this utility path.
float compute_full_multibit_distance(
const uint8_t* sign_bits,
const uint8_t* ex_code,
@@ -315,18 +319,18 @@ float compute_full_multibit_distance(
size_t d,
size_t ex_bits,
MetricType metric_type) {
const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
float ex_ip = rabitq::multibit::compute_inner_product(
sign_bits, ex_code, rotated_q, d, ex_bits, cb);
float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
if (metric_type == MetricType::METRIC_L2) {
dist = std::max(0.0f, dist);
}
return dist;
return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
[&]<SIMDLevel SL>() {
return compute_full_multibit_distance<SL>(
sign_bits,
ex_code,
ex_fac,
rotated_q,
qr_base,
d,
ex_bits,
metric_type);
});
}
void populate_block_aux_from_flat_storage(
faiss/impl/RaBitQUtils.h
+30
View File
@@ -10,6 +10,9 @@
#include <faiss/MetricType.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/utils/AlignedTable.h>
#include <faiss/utils/rabitq_simd.h>
#include <faiss/utils/simd_levels.h>
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
@@ -337,6 +340,33 @@ float compute_full_multibit_distance(
size_t ex_bits,
MetricType metric_type);
// SIMDLevel-templatized version — avoids per-call dynamic dispatch.
// Inline so it can be used from templatized distance computers without
// needing explicit instantiations in per-SIMD TUs.
template <SIMDLevel SL>
inline float compute_full_multibit_distance(
const uint8_t* sign_bits,
const uint8_t* ex_code,
const ExtraBitsFactors& ex_fac,
const float* rotated_q,
float qr_base,
size_t d,
size_t ex_bits,
MetricType metric_type) {
const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
float ex_ip = rabitq::multibit::compute_inner_product<SL>(
sign_bits, ex_code, rotated_q, d, ex_bits, cb);
float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
if (metric_type == MetricType::METRIC_L2) {
dist = std::max(0.0f, dist);
}
return dist;
}
/** Compute pointer to a vector's auxiliary data within block layout. */
template <typename T>
inline T* get_block_aux_ptr(
faiss/impl/RaBitQuantizer.cpp
+305 -295
View File
@@ -10,6 +10,7 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/RaBitQuantizerMultiBit.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/rabitq_simd.h>
@@ -221,162 +222,159 @@ void RaBitQuantizer::decode_core(
namespace {
// Distance computers templatized on SIMDLevel to avoid per-call dynamic
// dispatch. The SIMDLevel is baked in at construction time via
// get_distance_computer, so virtual calls through the base class go
// directly to the SIMD-specialized code.
template <SIMDLevel SL>
struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
// the rotated query (qr - c)
std::vector<float> rotated_q;
// some additional numbers for the query
QueryFactorsData query_fac;
RaBitQDistanceComputerNotQ();
RaBitQDistanceComputerNotQ() = default;
// Compute distance using only 1-bit codes (fast)
float distance_to_code_1bit(const uint8_t* code) override;
float distance_to_code_1bit(const uint8_t* code) override {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_ASSERT(rotated_q.size() == d);
// split the code into parts
const uint8_t* binary_data = code;
// Cast to appropriate type based on nb_bits
// For 1-bit: use SignBitFactors (8 bytes)
// For multi-bit: use SignBitFactorsWithError (12 bytes) which includes
// f_error
size_t ex_bits = nb_bits - 1;
const SignBitFactors* base_fac = (ex_bits == 0)
? reinterpret_cast<const SignBitFactors*>(code + (d + 7) / 8)
: reinterpret_cast<const SignBitFactorsWithError*>(
code + (d + 7) / 8);
// this is the baseline code
//
// compute <q,o> using floats
float dot_qo = 0;
// 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;
for (size_t i = 0; i < d; i++) {
// Extract i-th bit
bool bit = rabitq_utils::extract_bit_standard(binary_data, i);
// accumulate dp
dot_qo += bit ? rotated_q[i] : 0;
// accumulate sum-of-bits
sum_q += bit ? 1 : 0;
}
// Apply query factors
float final_dot =
query_fac.c1 * dot_qo + query_fac.c2 * sum_q - query_fac.c34;
// pre_dist = ||or - c||^2 + ||qr - c||^2 -
// 2 * ||or - c|| * ||qr - c|| * <q,o> - (IP ? ||or||^2 : 0)
float pre_dist = base_fac->or_minus_c_l2sqr + query_fac.qr_to_c_L2sqr -
2 * base_fac->dp_multiplier * final_dot;
if (metric_type == MetricType::METRIC_L2) {
// ||or - q||^ 2
return pre_dist;
} else {
// metric == MetricType::METRIC_INNER_PRODUCT
// 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr);
}
}
// Compute full distance using 1-bit + ex-bits (accurate)
float distance_to_code_full(const uint8_t* code) override;
float distance_to_code_full(const uint8_t* code) override {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_ASSERT(rotated_q.size() == d);
void set_query(const float* x) override;
size_t ex_bits = nb_bits - 1;
if (ex_bits == 0) {
// No ex-bits, just return 1-bit distance
return distance_to_code_1bit(code);
}
// Extract pointers to code sections
const uint8_t* binary_data = code;
size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError);
const uint8_t* ex_code = code + offset;
const ExtraBitsFactors* ex_fac =
reinterpret_cast<const ExtraBitsFactors*>(
ex_code + (d * ex_bits + 7) / 8);
float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_fac.q_dot_c
: query_fac.qr_to_c_L2sqr;
return rabitq_utils::compute_full_multibit_distance<SL>(
binary_data,
ex_code,
*ex_fac,
rotated_q.data(),
qr_base,
d,
ex_bits,
metric_type);
}
void set_query(const float* x) override {
q = x;
FAISS_ASSERT(x != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
// compute the distance from the query to the centroid
if (centroid != nullptr) {
query_fac.qr_to_c_L2sqr = fvec_L2sqr(x, centroid, d);
} else {
query_fac.qr_to_c_L2sqr = fvec_norm_L2sqr(x, d);
}
// subtract c, obtain P^(-1)(qr - c)
rotated_q.resize(d);
for (size_t i = 0; i < d; i++) {
rotated_q[i] = x[i] - ((centroid == nullptr) ? 0 : centroid[i]);
}
// Compute g_error = ||qr - c|| (L2 norm of rotated query)
g_error = std::sqrt(query_fac.qr_to_c_L2sqr);
// compute some numbers — do not quantize the query
const float inv_d = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
float sum_q = 0;
for (size_t i = 0; i < d; i++) {
sum_q += rotated_q[i];
}
query_fac.c1 = 2 * inv_d;
query_fac.c2 = 0;
query_fac.c34 = sum_q * inv_d;
if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
query_fac.qr_norm_L2sqr = fvec_norm_L2sqr(x, d);
query_fac.q_dot_c =
centroid ? fvec_inner_product(x, centroid, d) : 0.0f;
}
}
};
RaBitQDistanceComputerNotQ::RaBitQDistanceComputerNotQ() = default;
float RaBitQDistanceComputerNotQ::distance_to_code_1bit(const uint8_t* code) {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_ASSERT(rotated_q.size() == d);
// split the code into parts
const uint8_t* binary_data = code;
// Cast to appropriate type based on nb_bits
// For 1-bit: use SignBitFactors (8 bytes)
// For multi-bit: use SignBitFactorsWithError (12 bytes) which includes
// f_error
size_t ex_bits = nb_bits - 1;
const SignBitFactors* base_fac = (ex_bits == 0)
? reinterpret_cast<const SignBitFactors*>(code + (d + 7) / 8)
: reinterpret_cast<const SignBitFactorsWithError*>(
code + (d + 7) / 8);
// this is the baseline code
//
// compute <q,o> using floats
float dot_qo = 0;
// 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;
for (size_t i = 0; i < d; i++) {
// Extract i-th bit
bool bit = rabitq_utils::extract_bit_standard(binary_data, i);
// accumulate dp
dot_qo += bit ? rotated_q[i] : 0;
// accumulate sum-of-bits
sum_q += bit ? 1 : 0;
}
// Apply query factors
float final_dot =
query_fac.c1 * dot_qo + query_fac.c2 * sum_q - query_fac.c34;
// pre_dist = ||or - c||^2 + ||qr - c||^2 -
// 2 * ||or - c|| * ||qr - c|| * <q,o> - (IP ? ||or||^2 : 0)
float pre_dist = base_fac->or_minus_c_l2sqr + query_fac.qr_to_c_L2sqr -
2 * base_fac->dp_multiplier * final_dot;
if (metric_type == MetricType::METRIC_L2) {
// ||or - q||^ 2
return pre_dist;
} else {
// metric == MetricType::METRIC_INNER_PRODUCT
return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr);
}
}
float RaBitQDistanceComputerNotQ::distance_to_code_full(const uint8_t* code) {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_ASSERT(rotated_q.size() == d);
size_t ex_bits = nb_bits - 1;
if (ex_bits == 0) {
// No ex-bits, just return 1-bit distance
return distance_to_code_1bit(code);
}
// Extract pointers to code sections
const uint8_t* binary_data = code;
size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError);
const uint8_t* ex_code = code + offset;
const ExtraBitsFactors* ex_fac = reinterpret_cast<const ExtraBitsFactors*>(
ex_code + (d * ex_bits + 7) / 8);
// Call shared utility directly with rotated_q pointer
float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_fac.q_dot_c
: query_fac.qr_to_c_L2sqr;
return rabitq_utils::compute_full_multibit_distance(
binary_data,
ex_code,
*ex_fac,
rotated_q.data(),
qr_base,
d,
ex_bits,
metric_type);
}
void RaBitQDistanceComputerNotQ::set_query(const float* x) {
q = x;
FAISS_ASSERT(x != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
// compute the distance from the query to the centroid
if (centroid != nullptr) {
query_fac.qr_to_c_L2sqr = fvec_L2sqr(x, centroid, d);
} else {
query_fac.qr_to_c_L2sqr = fvec_norm_L2sqr(x, d);
}
// subtract c, obtain P^(-1)(qr - c)
rotated_q.resize(d);
for (size_t i = 0; i < d; i++) {
rotated_q[i] = x[i] - ((centroid == nullptr) ? 0 : centroid[i]);
}
// Compute g_error (query norm for lower bound computation)
// g_error = ||qr - c|| (L2 norm of rotated query)
g_error = std::sqrt(query_fac.qr_to_c_L2sqr);
// compute some numbers
const float inv_d = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
// do not quantize the query
float sum_q = 0;
for (size_t i = 0; i < d; i++) {
sum_q += rotated_q[i];
}
query_fac.c1 = 2 * inv_d;
query_fac.c2 = 0;
query_fac.c34 = sum_q * inv_d;
if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
// precompute if needed
query_fac.qr_norm_L2sqr = fvec_norm_L2sqr(x, d);
query_fac.q_dot_c =
centroid ? fvec_inner_product(x, centroid, d) : 0.0f;
}
}
//
template <SIMDLevel SL>
struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
// the rotated and quantized query (qr - c)
std::vector<float> rotated_q;
@@ -394,176 +392,188 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
// the smallest value divisible by 8 that is not smaller than dim
size_t popcount_aligned_dim = 0;
RaBitQDistanceComputerQ();
RaBitQDistanceComputerQ() = default;
// Compute distance using only 1-bit codes (fast)
float distance_to_code_1bit(const uint8_t* code) override;
float distance_to_code_1bit(const uint8_t* code) override {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
// split the code into parts
size_t size = (d + 7) / 8;
const uint8_t* binary_data = code;
// Cast to appropriate type based on nb_bits
// For 1-bit: use SignBitFactors (8 bytes)
// For multi-bit: use SignBitFactorsWithError (12 bytes) which
// includes f_error
size_t ex_bits = nb_bits - 1;
const SignBitFactors* base_fac = (ex_bits == 0)
? reinterpret_cast<const SignBitFactors*>(code + size)
: reinterpret_cast<const SignBitFactorsWithError*>(code + size);
// this is ||or - c||^2 - (IP ? ||or||^2 : 0)
float final_dot = 0;
if (centered) {
int64_t int_dot = ((1 << qb) - 1) * d;
// See RaBitDistanceComputerNotQ::distance_to_code() for
// baseline code.
int_dot -= 2 *
rabitq::bitwise_xor_dot_product<SL>(
rearranged_rotated_qq.data(),
binary_data,
size,
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);
// dot-product itself
final_dot += query_fac.c1 * dot_qo;
// normalizer coefficients
final_dot += query_fac.c2 * sum_q;
// normalizer coefficients
final_dot -= query_fac.c34;
}
const float pre_dist = base_fac->or_minus_c_l2sqr +
query_fac.qr_to_c_L2sqr -
2 * base_fac->dp_multiplier * final_dot;
if (metric_type == MetricType::METRIC_L2) {
// ||or - q||^ 2
return pre_dist;
} else {
// metric == MetricType::METRIC_INNER_PRODUCT
// 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr);
}
}
// Compute full distance using 1-bit + ex-bits (accurate)
float distance_to_code_full(const uint8_t* code) override;
float distance_to_code_full(const uint8_t* code) override {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_ASSERT(rotated_q.size() == d);
void set_query(const float* x) override;
size_t ex_bits = nb_bits - 1;
if (ex_bits == 0) {
// No ex-bits, just return 1-bit distance
return distance_to_code_1bit(code);
}
// Extract pointers to code sections
const uint8_t* binary_data = code;
size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError);
const uint8_t* ex_code = code + offset;
const ExtraBitsFactors* ex_fac =
reinterpret_cast<const ExtraBitsFactors*>(
ex_code + (d * ex_bits + 7) / 8);
float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_fac.q_dot_c
: query_fac.qr_to_c_L2sqr;
return rabitq_utils::compute_full_multibit_distance<SL>(
binary_data,
ex_code,
*ex_fac,
rotated_q.data(),
qr_base,
d,
ex_bits,
metric_type);
}
void set_query(const float* x) override {
q = x;
FAISS_ASSERT(x != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_THROW_IF_NOT(qb <= 8);
FAISS_THROW_IF_NOT(qb > 0);
// Use shared utilities for core query factor computation
// rotated_q is populated directly by compute_query_factors as an
// output parameter
query_fac = rabitq_utils::compute_query_factors(
x,
d,
centroid,
qb,
centered,
metric_type,
rotated_q,
rotated_qq);
// Compute g_error (query norm for lower bound computation)
// g_error = ||qr - c|| (L2 norm of rotated query)
g_error = std::sqrt(query_fac.qr_to_c_L2sqr);
// Rearrange the query vector for SIMD operations
// (RaBitQuantizer-specific)
popcount_aligned_dim = ((d + 7) / 8) * 8;
size_t offset = (d + 7) / 8;
rearranged_rotated_qq.resize(offset * qb);
std::fill(
rearranged_rotated_qq.begin(), rearranged_rotated_qq.end(), 0);
for (size_t idim = 0; idim < d; idim++) {
for (size_t iv = 0; iv < qb; iv++) {
const bool bit = ((rotated_qq[idim] & (1 << iv)) != 0);
rearranged_rotated_qq[iv * offset + idim / 8] |=
bit ? (1 << (idim % 8)) : 0;
}
}
}
};
RaBitQDistanceComputerQ::RaBitQDistanceComputerQ() = default;
float RaBitQDistanceComputerQ::distance_to_code_1bit(const uint8_t* code) {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
// split the code into parts
size_t size = (d + 7) / 8;
const uint8_t* binary_data = code;
// Cast to appropriate type based on nb_bits
// For 1-bit: use SignBitFactors (8 bytes)
// For multi-bit: use SignBitFactorsWithError (12 bytes) which includes
// f_error
size_t ex_bits = nb_bits - 1;
const SignBitFactors* base_fac = (ex_bits == 0)
? reinterpret_cast<const SignBitFactors*>(code + size)
: reinterpret_cast<const SignBitFactorsWithError*>(code + size);
// this is ||or - c||^2 - (IP ? ||or||^2 : 0)
float final_dot = 0;
if (centered) {
int64_t int_dot = ((1 << qb) - 1) * d;
// See RaBitDistanceComputerNotQ::distance_to_code() for baseline code.
int_dot -= 2 *
rabitq::bitwise_xor_dot_product(
rearranged_rotated_qq.data(), binary_data, size, qb);
final_dot += int_dot * query_fac.int_dot_scale;
} else {
auto 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.
// process 64-bit popcounts
auto sum_q = rabitq::popcount(binary_data, size);
// dot-product itself
final_dot += query_fac.c1 * dot_qo;
// normalizer coefficients
final_dot += query_fac.c2 * sum_q;
// normalizer coefficients
final_dot -= query_fac.c34;
}
// pre_dist = ||or - c||^2 + ||qr - c||^2 -
// 2 * ||or - c|| * ||qr - c|| * <q,o> - (IP ? ||or||^2 : 0)
const float pre_dist = base_fac->or_minus_c_l2sqr +
query_fac.qr_to_c_L2sqr - 2 * base_fac->dp_multiplier * final_dot;
if (metric_type == MetricType::METRIC_L2) {
// ||or - q||^ 2
return pre_dist;
} else {
// metric == MetricType::METRIC_INNER_PRODUCT
// 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr);
}
}
float RaBitQDistanceComputerQ::distance_to_code_full(const uint8_t* code) {
FAISS_ASSERT(code != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_ASSERT(rotated_q.size() == d);
size_t ex_bits = nb_bits - 1;
if (ex_bits == 0) {
// No ex-bits, just return 1-bit distance
return distance_to_code_1bit(code);
}
// Extract pointers to code sections
const uint8_t* binary_data = code;
size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError);
const uint8_t* ex_code = code + offset;
const ExtraBitsFactors* ex_fac = reinterpret_cast<const ExtraBitsFactors*>(
ex_code + (d * ex_bits + 7) / 8);
// Call shared utility directly with rotated_q pointer
float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_fac.q_dot_c
: query_fac.qr_to_c_L2sqr;
return rabitq_utils::compute_full_multibit_distance(
binary_data,
ex_code,
*ex_fac,
rotated_q.data(),
qr_base,
d,
ex_bits,
metric_type);
}
// Use shared constant from RaBitQUtils
using rabitq_utils::Z_MAX_BY_QB;
void RaBitQDistanceComputerQ::set_query(const float* x) {
q = x;
FAISS_ASSERT(x != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_THROW_IF_NOT(qb <= 8);
FAISS_THROW_IF_NOT(qb > 0);
// Use shared utilities for core query factor computation
// rotated_q is populated directly by compute_query_factors as an output
// parameter
query_fac = rabitq_utils::compute_query_factors(
x, d, centroid, qb, centered, metric_type, rotated_q, rotated_qq);
// Compute g_error (query norm for lower bound computation)
// g_error = ||qr - c|| (L2 norm of rotated query)
g_error = std::sqrt(query_fac.qr_to_c_L2sqr);
// Rearrange the query vector for SIMD operations (RaBitQuantizer-specific)
popcount_aligned_dim = ((d + 7) / 8) * 8;
size_t offset = (d + 7) / 8;
rearranged_rotated_qq.resize(offset * qb);
std::fill(rearranged_rotated_qq.begin(), rearranged_rotated_qq.end(), 0);
for (size_t idim = 0; idim < d; idim++) {
for (size_t iv = 0; iv < qb; iv++) {
const bool bit = ((rotated_qq[idim] & (1 << iv)) != 0);
rearranged_rotated_qq[iv * offset + idim / 8] |=
bit ? (1 << (idim % 8)) : 0;
}
}
}
} // anonymous namespace
FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
uint8_t qb,
const float* centroid_in,
bool centered) const {
if (qb == 0) {
auto dc = std::make_unique<RaBitQDistanceComputerNotQ>();
dc->metric_type = metric_type;
dc->d = d;
dc->centroid = centroid_in;
dc->nb_bits = nb_bits;
// Dispatch on SIMDLevel once here so the distance computer methods
// call the SIMD-specialized rabitq functions directly (no per-call
// with_simd_level overhead).
return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
[&]<SIMDLevel SL>() -> FlatCodesDistanceComputer* {
if (qb == 0) {
auto dc =
std::make_unique<RaBitQDistanceComputerNotQ<SL>>();
dc->metric_type = metric_type;
dc->d = d;
dc->centroid = centroid_in;
dc->nb_bits = nb_bits;
return dc.release();
} else {
auto dc = std::make_unique<RaBitQDistanceComputerQ>();
dc->metric_type = metric_type;
dc->d = d;
dc->centroid = centroid_in;
dc->qb = qb;
dc->centered = centered;
dc->nb_bits = nb_bits;
return dc.release();
} else {
auto dc = std::make_unique<RaBitQDistanceComputerQ<SL>>();
dc->metric_type = metric_type;
dc->d = d;
dc->centroid = centroid_in;
dc->qb = qb;
dc->centered = centered;
dc->nb_bits = nb_bits;
return dc.release();
}
return dc.release();
}
});
}
} // namespace faiss
faiss/utils/rabitq_simd.h
+63 -543
View File
@@ -11,282 +11,72 @@
#include <cstdint>
#include <cstring>
// Only include x86 SIMD intrinsics on x86/x86_64 architectures
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || \
defined(_M_IX86)
#include <immintrin.h>
#endif // defined(__x86_64__) || defined(_M_X64) || defined(__i386__) ||
#include <faiss/utils/simd_levels.h>
namespace faiss::rabitq {
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || \
defined(_M_IX86)
/**
* Returns the lookup table for AVX512 popcount operations.
* This table is used for lookup-based popcount implementation.
* Compute dot product between query and binary data using popcount on AND.
*
* Source: https://github.com/WojciechMula/sse-popcount.
*
* @return Lookup table as __m512i register
*/
#if defined(__AVX512F__)
inline __m512i get_lookup_512() {
return _mm512_set_epi8(
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0,
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0,
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0,
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0);
}
#endif // defined(__AVX512F__)
#if defined(__AVX2__)
/**
* Returns the lookup table for AVX2 popcount operations.
* This table is used for lookup-based popcount implementation.
*
* @return Lookup table as __m256i register
*/
inline __m256i get_lookup_256() {
return _mm256_setr_epi8(
/* 0 */ 0,
/* 1 */ 1,
/* 2 */ 1,
/* 3 */ 2,
/* 4 */ 1,
/* 5 */ 2,
/* 6 */ 2,
/* 7 */ 3,
/* 8 */ 1,
/* 9 */ 2,
/* a */ 2,
/* b */ 3,
/* c */ 2,
/* d */ 3,
/* e */ 3,
/* f */ 4,
/* 0 */ 0,
/* 1 */ 1,
/* 2 */ 1,
/* 3 */ 2,
/* 4 */ 1,
/* 5 */ 2,
/* 6 */ 2,
/* 7 */ 3,
/* 8 */ 1,
/* 9 */ 2,
/* a */ 2,
/* b */ 3,
/* c */ 2,
/* d */ 3,
/* e */ 3,
/* f */ 4);
}
#endif // defined(__AVX2__)
#if defined(__AVX512F__)
/**
* Popcount for a 512-bit register, using lookup tables if necessary.
*
* @param v Input vector to count bits in
* @return Vector int32_t[16] with popcount results.
*/
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, 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 // defined(__AVX512F__)
#if defined(__AVX2__)
/**
* Popcount for a 256-bit register, using lookup tables if necessary.
*
* @param v Input vector to count bits in
* @return uint64_t[4] of popcounts for each portion of the input vector.
*/
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, 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);
}
inline uint64_t reduce_add_256(__m256i v) {
alignas(32) uint64_t lanes[4];
_mm256_store_si256((__m256i*)lanes, v);
return lanes[0] + lanes[1] + lanes[2] + lanes[3];
}
#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 query Pointer to rearranged rotated query data
* @param data Pointer to binary data
* @param d Dimension
* @param qb Number of quantization bits
* @return Unsigned integer dot product
* @param size Size in bytes
* @param qb Number of quantization bits
* @return Unsigned integer dot product
*/
inline uint64_t bitwise_and_dot_product(
template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
uint64_t bitwise_and_dot_product(
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.
*
* @param query Pointer to rearranged rotated query data
* @param data Pointer to binary data
* @param size Size in bytes
* @param qb Number of quantization bits
* @return Unsigned integer dot product
*/
template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
uint64_t bitwise_xor_dot_product(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb);
/**
* Count total set bits in data.
*
* @param data Pointer to binary data
* @param size Size in bytes
* @return Total popcount
*/
template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
uint64_t popcount(const uint8_t* data, size_t size);
// NONE specializations — scalar fallbacks
template <>
inline uint64_t bitwise_and_dot_product<SIMDLevel::NONE>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
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 (size_t 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);
}
}
sum += _mm512_reduce_add_epi64(sum_512);
}
#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 (size_t 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 (size_t 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 (size_t j = 0; j < qb; j++) {
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 (size_t j = 0; j < qb; j++) {
for (int j = 0; j < qb; j++) {
const auto qv = *(query + j * size + offset);
sum += __builtin_popcount(qv & yv) << j;
}
@@ -294,82 +84,24 @@ inline uint64_t bitwise_and_dot_product(
return sum;
}
/**
* Compute dot product between query and binary data using popcount operations.
*
* @param query Pointer to rearranged rotated query data
* @param data Pointer to binary data
* @param d Dimension
* @param qb Number of quantization bits
* @return Unsigned integer dot product
*/
inline uint64_t bitwise_xor_dot_product(
template <>
inline uint64_t bitwise_xor_dot_product<SIMDLevel::NONE>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
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 (size_t j = 0; j < qb; j++) {
__m512i v_q = _mm512_loadu_si512(
(const __m512i*)(query + j * size + offset));
__m512i v_xor = _mm512_xor_si512(v_q, v_x);
__m512i v_popcnt = popcount_512(v_xor);
__m512i v_shifted = _mm512_slli_epi64(v_popcnt, j);
sum_512 = _mm512_add_epi64(sum_512, v_shifted);
}
}
sum += _mm512_reduce_add_epi64(sum_512);
}
#endif
#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 (size_t j = 0; j < qb; j++) {
__m256i v_q = _mm256_loadu_si256(
(const __m256i*)(query + j * size + offset));
__m256i v_xor = _mm256_xor_si256(v_q, v_x);
__m256i v_popcnt = popcount_256(v_xor);
__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
#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 (size_t j = 0; j < qb; j++) {
__m128i v_q = _mm_loadu_si128(
(const __m128i*)(query + j * size + offset));
__m128i v_xor = _mm_xor_si128(v_q, v_x);
__m128i v_popcnt = popcount_128(v_xor);
__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
for (size_t step = 64 / 8; offset + step <= size; offset += step) {
const auto yv = *(const uint64_t*)(data + offset);
for (size_t j = 0; j < qb; j++) {
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 (size_t j = 0; j < qb; j++) {
for (int j = 0; j < qb; j++) {
const auto qv = *(query + j * size + offset);
sum += __builtin_popcount(qv ^ yv) << j;
}
@@ -377,41 +109,10 @@ inline uint64_t bitwise_xor_dot_product(
return sum;
}
inline uint64_t popcount(const uint8_t* data, size_t size) {
template <>
inline uint64_t popcount<SIMDLevel::NONE>(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);
@@ -469,186 +170,9 @@ inline float ip_scalar(
return result;
}
#if defined(__x86_64__) || defined(_M_X64)
#if defined(__AVX2__)
/// Horizontal sum of 8 floats in a __m256 register.
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));
}
#endif // __AVX2__
/*********************************************************
* Specialized 1-bit kernels (ex_bits == 1).
*
* For 1 extra bit, both sign_bits and ex_code are 1-bit-per-dim packed,
* so we convert bits to floats directly — no extraction loops needed.
*********************************************************/
#if defined(__AVX512F__)
/// AVX-512: 16 dims/iter, ex_bits == 1.
inline float ip_1exbit_avx512(
const uint8_t* __restrict sign_bits,
const uint8_t* __restrict ex_code,
const float* __restrict rotated_q,
size_t d,
float cb) {
__m512 acc = _mm512_setzero_ps();
const __m512 v_cb = _mm512_set1_ps(cb);
const __m512 v_two = _mm512_set1_ps(2.0f);
const __m512 v_one = _mm512_set1_ps(1.0f);
size_t i = 0;
for (; i + 16 <= d; i += 16) {
uint16_t sb16;
memcpy(&sb16, sign_bits + i / 8, sizeof(uint16_t));
uint16_t eb16;
memcpy(&eb16, ex_code + i / 8, sizeof(uint16_t));
__m512 sb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(sb16), v_one);
__m512 eb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(eb16), v_one);
__m512 recon = _mm512_add_ps(_mm512_fmadd_ps(sb_f, v_two, eb_f), v_cb);
__m512 rq = _mm512_loadu_ps(rotated_q + i);
acc = _mm512_fmadd_ps(rq, recon, acc);
}
float result = _mm512_reduce_add_ps(acc);
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb);
return result;
}
#endif // __AVX512F__
#if defined(__AVX2__)
/// AVX2: 8 dims/iter, ex_bits == 1.
inline float ip_1exbit_avx2(
const uint8_t* __restrict sign_bits,
const uint8_t* __restrict ex_code,
const float* __restrict rotated_q,
size_t d,
float cb) {
__m256 acc = _mm256_setzero_ps();
const __m256 v_cb = _mm256_set1_ps(cb);
const __m256 v_two = _mm256_set1_ps(2.0f);
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();
size_t i = 0;
for (; i + 8 <= d; i += 8) {
uint8_t sb = sign_bits[i / 8];
uint8_t eb = ex_code[i / 8];
__m256i sb_cmp = _mm256_cmpgt_epi32(
_mm256_and_si256(_mm256_set1_epi32(sb), bit_pos), zero);
__m256 sb_f = _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one);
__m256i eb_cmp = _mm256_cmpgt_epi32(
_mm256_and_si256(_mm256_set1_epi32(eb), bit_pos), zero);
__m256 eb_f = _mm256_and_ps(_mm256_castsi256_ps(eb_cmp), v_one);
__m256 recon = _mm256_add_ps(_mm256_fmadd_ps(sb_f, v_two, eb_f), v_cb);
__m256 rq = _mm256_loadu_ps(rotated_q + i);
acc = _mm256_fmadd_ps(rq, recon, acc);
}
float result = hsum_avx2(acc);
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb);
return result;
}
#endif // __AVX2__
/*********************************************************
* Bit-plane decomposition kernels (ex_bits >= 2, BMI2 required).
*
* Decomposes the inner product as:
* ex_ip = (1 << ex_bits) * sign_dot
* + Σ_{b=0}^{ex_bits-1} (1 << b) * plane_dot_b
* + cb * total_q
*
* Each plane_dot_b is a float × bit-vector dot product, computed using
* the same bit→mask→float conversion as the 1-bit kernel. PEXT
* extracts each bit plane from the packed ex_code in one instruction
* per 8 dimensions.
*********************************************************/
#if defined(__AVX2__) && defined(__BMI2__)
/// AVX2 + BMI2 bit-plane decomposition: 8 dims/iter, ex_bits in [2, 7].
/// Caller must ensure ex_bits <= 7 (pext_masks[7] / v_weights[8]).
inline float ip_bitplane_avx2(
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);
// Precompute PEXT masks and plane weights
uint64_t pext_masks[7];
__m256 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[ex_bits] = _mm256_set1_ps(static_cast<float>(1u << ex_bits));
size_t i = 0;
for (; i + 8 <= d; i += 8) {
// Sign bit → float via bit mask comparison
__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]);
// Load packed ex_code for 8 dims (8 × ex_bits bits = ex_bits bytes)
uint64_t ex64 = 0;
memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t));
// Extract each bit plane via PEXT → bit mask → float
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);
}
__m256 rq = _mm256_loadu_ps(rotated_q + i);
acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc);
}
float result = hsum_avx2(acc);
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
return result;
}
#endif // __AVX2__ && __BMI2__
#endif // x86_64
/**
* Dispatch to the best available kernel for the given ex_bits.
*
* Routing (compile-time):
* ex_bits == 1: specialized 1-bit kernel (AVX-512 > AVX2 > scalar)
* ex_bits >= 2: bit-plane decomposition (AVX2+BMI2 > scalar)
*
* @param sign_bits packed sign bits (1 bit/dim, standard byte packing)
* @param ex_code packed extra-bit codes (ex_bits bits/dim)
* @param rotated_q rotated query vector (float[d])
@@ -657,28 +181,24 @@ inline float ip_bitplane_avx2(
* @param cb constant bias: -(2^ex_bits - 0.5)
* @return inner product value
*/
inline float compute_inner_product(
template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
float compute_inner_product(
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);
// NONE specialization — pure scalar
template <>
inline float compute_inner_product<SIMDLevel::NONE>(
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) {
if (ex_bits == 1) {
#if defined(__AVX512F__)
return ip_1exbit_avx512(sign_bits, ex_code, rotated_q, d, cb);
#elif defined(__AVX2__)
return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb);
#else
return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, 1, cb);
#endif
}
#if defined(__AVX2__) && defined(__BMI2__)
if (ex_bits <= 7) {
return ip_bitplane_avx2(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);
}
faiss/utils/simd_impl/rabitq_avx2.cpp
+355
View File
@@ -0,0 +1,355 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#ifdef COMPILE_SIMD_AVX2
#include <faiss/utils/rabitq_simd.h>
#include <immintrin.h>
namespace faiss::rabitq {
namespace {
inline __m256i get_lookup_256() {
return _mm256_setr_epi8(
/* 0 */ 0,
/* 1 */ 1,
/* 2 */ 1,
/* 3 */ 2,
/* 4 */ 1,
/* 5 */ 2,
/* 6 */ 2,
/* 7 */ 3,
/* 8 */ 1,
/* 9 */ 2,
/* a */ 2,
/* b */ 3,
/* c */ 2,
/* d */ 3,
/* e */ 3,
/* f */ 4,
/* 0 */ 0,
/* 1 */ 1,
/* 2 */ 1,
/* 3 */ 2,
/* 4 */ 1,
/* 5 */ 2,
/* 6 */ 2,
/* 7 */ 3,
/* 8 */ 1,
/* 9 */ 2,
/* a */ 2,
/* b */ 3,
/* c */ 2,
/* d */ 3,
/* e */ 3,
/* f */ 4);
}
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, 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);
return _mm256_sad_epu8(_mm256_setzero_si256(), popcnt);
}
inline uint64_t reduce_add_256(__m256i v) {
alignas(32) uint64_t lanes[4];
_mm256_store_si256((__m256i*)lanes, v);
return lanes[0] + lanes[1] + lanes[2] + lanes[3];
}
inline __m128i popcount_128(__m128i v) {
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];
}
} // namespace
template <>
uint64_t bitwise_and_dot_product<SIMDLevel::AVX2>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
uint64_t sum = 0;
size_t offset = 0;
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);
}
__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);
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;
}
template <>
uint64_t bitwise_xor_dot_product<SIMDLevel::AVX2>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
uint64_t sum = 0;
size_t offset = 0;
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_xor = _mm256_xor_si256(v_q, v_x);
__m256i v_popcnt = popcount_256(v_xor);
__m256i v_shifted = _mm256_slli_epi64(v_popcnt, j);
sum_256 = _mm256_add_epi64(sum_256, v_shifted);
}
}
sum += reduce_add_256(sum_256);
}
__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_xor = _mm_xor_si128(v_q, v_x);
__m128i v_popcnt = popcount_128(v_xor);
__m128i v_shifted = _mm_slli_epi64(v_popcnt, j);
sum_128 = _mm_add_epi64(sum_128, v_shifted);
}
}
sum += reduce_add_128(sum_128);
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;
}
template <>
uint64_t popcount<SIMDLevel::AVX2>(const uint8_t* data, size_t size) {
uint64_t sum = 0;
size_t offset = 0;
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);
}
__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);
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
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_avx2(
const uint8_t* __restrict sign_bits,
const uint8_t* __restrict ex_code,
const float* __restrict rotated_q,
size_t d,
float cb) {
__m256 acc = _mm256_setzero_ps();
const __m256 v_cb = _mm256_set1_ps(cb);
const __m256 v_two = _mm256_set1_ps(2.0f);
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();
size_t i = 0;
for (; i + 8 <= d; i += 8) {
uint8_t sb = sign_bits[i / 8];
uint8_t eb = ex_code[i / 8];
__m256i sb_cmp = _mm256_cmpgt_epi32(
_mm256_and_si256(_mm256_set1_epi32(sb), bit_pos), zero);
__m256 sb_f = _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one);
__m256i eb_cmp = _mm256_cmpgt_epi32(
_mm256_and_si256(_mm256_set1_epi32(eb), bit_pos), zero);
__m256 eb_f = _mm256_and_ps(_mm256_castsi256_ps(eb_cmp), v_one);
__m256 recon = _mm256_add_ps(_mm256_fmadd_ps(sb_f, v_two, eb_f), v_cb);
__m256 rq = _mm256_loadu_ps(rotated_q + i);
acc = _mm256_fmadd_ps(rq, recon, acc);
}
float result = hsum_avx2(acc);
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb);
return result;
}
#ifdef __BMI2__
inline float ip_bitplane_avx2(
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);
uint64_t pext_masks[7];
__m256 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[ex_bits] = _mm256_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]);
uint64_t ex64 = 0;
memcpy(&ex64, ex_code + (i / 8) * 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);
}
__m256 rq = _mm256_loadu_ps(rotated_q + i);
acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc);
}
float result = hsum_avx2(acc);
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
return result;
}
#endif // __BMI2__
} // namespace
template <>
float compute_inner_product<SIMDLevel::AVX2>(
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) {
if (ex_bits == 1) {
return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb);
}
#ifdef __BMI2__
if (ex_bits <= 7) {
return ip_bitplane_avx2(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);
}
} // namespace faiss::rabitq::multibit
#endif // COMPILE_SIMD_AVX2
faiss/utils/simd_impl/rabitq_avx512.cpp
+477
View File
@@ -0,0 +1,477 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#ifdef COMPILE_SIMD_AVX512
#include <faiss/utils/rabitq_simd.h>
#include <immintrin.h>
namespace faiss::rabitq {
namespace {
inline __m512i get_lookup_512() {
return _mm512_set_epi8(
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0,
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0,
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0,
/* f */ 4,
/* e */ 3,
/* d */ 3,
/* c */ 2,
/* b */ 3,
/* a */ 2,
/* 9 */ 2,
/* 8 */ 1,
/* 7 */ 3,
/* 6 */ 2,
/* 5 */ 2,
/* 4 */ 1,
/* 3 */ 2,
/* 2 */ 1,
/* 1 */ 1,
/* 0 */ 0);
}
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, 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
}
// AVX2 helpers needed for AVX512 fallback paths (compute_inner_product)
inline __m256i get_lookup_256() {
return _mm256_setr_epi8(
/* 0 */ 0,
/* 1 */ 1,
/* 2 */ 1,
/* 3 */ 2,
/* 4 */ 1,
/* 5 */ 2,
/* 6 */ 2,
/* 7 */ 3,
/* 8 */ 1,
/* 9 */ 2,
/* a */ 2,
/* b */ 3,
/* c */ 2,
/* d */ 3,
/* e */ 3,
/* f */ 4,
/* 0 */ 0,
/* 1 */ 1,
/* 2 */ 1,
/* 3 */ 2,
/* 4 */ 1,
/* 5 */ 2,
/* 6 */ 2,
/* 7 */ 3,
/* 8 */ 1,
/* 9 */ 2,
/* a */ 2,
/* b */ 3,
/* c */ 2,
/* d */ 3,
/* e */ 3,
/* f */ 4);
}
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, 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);
return _mm256_sad_epu8(_mm256_setzero_si256(), popcnt);
}
inline uint64_t reduce_add_256(__m256i v) {
alignas(32) uint64_t lanes[4];
_mm256_store_si256((__m256i*)lanes, v);
return lanes[0] + lanes[1] + lanes[2] + lanes[3];
}
inline __m128i popcount_128(__m128i v) {
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];
}
} // namespace
template <>
uint64_t bitwise_and_dot_product<SIMDLevel::AVX512>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
uint64_t sum = 0;
size_t offset = 0;
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);
}
}
sum += _mm512_reduce_add_epi64(sum_512);
}
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);
}
__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);
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;
}
template <>
uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
uint64_t sum = 0;
size_t offset = 0;
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_xor = _mm512_xor_si512(v_q, v_x);
__m512i v_popcnt = popcount_512(v_xor);
__m512i v_shifted = _mm512_slli_epi64(v_popcnt, j);
sum_512 = _mm512_add_epi64(sum_512, v_shifted);
}
}
sum += _mm512_reduce_add_epi64(sum_512);
}
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_xor = _mm256_xor_si256(v_q, v_x);
__m256i v_popcnt = popcount_256(v_xor);
__m256i v_shifted = _mm256_slli_epi64(v_popcnt, j);
sum_256 = _mm256_add_epi64(sum_256, v_shifted);
}
}
sum += reduce_add_256(sum_256);
}
__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_xor = _mm_xor_si128(v_q, v_x);
__m128i v_popcnt = popcount_128(v_xor);
__m128i v_shifted = _mm_slli_epi64(v_popcnt, j);
sum_128 = _mm_add_epi64(sum_128, v_shifted);
}
}
sum += reduce_add_128(sum_128);
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;
}
template <>
uint64_t popcount<SIMDLevel::AVX512>(const uint8_t* data, size_t size) {
uint64_t sum = 0;
size_t offset = 0;
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);
}
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);
}
__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);
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
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,
const float* __restrict rotated_q,
size_t d,
float cb) {
__m512 acc = _mm512_setzero_ps();
const __m512 v_cb = _mm512_set1_ps(cb);
const __m512 v_two = _mm512_set1_ps(2.0f);
const __m512 v_one = _mm512_set1_ps(1.0f);
size_t i = 0;
for (; i + 16 <= d; i += 16) {
uint16_t sb16;
memcpy(&sb16, sign_bits + i / 8, sizeof(uint16_t));
uint16_t eb16;
memcpy(&eb16, ex_code + i / 8, sizeof(uint16_t));
__m512 sb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(sb16), v_one);
__m512 eb_f = _mm512_maskz_mov_ps(_cvtu32_mask16(eb16), v_one);
__m512 recon = _mm512_add_ps(_mm512_fmadd_ps(sb_f, v_two, eb_f), v_cb);
__m512 rq = _mm512_loadu_ps(rotated_q + i);
acc = _mm512_fmadd_ps(rq, recon, acc);
}
float result = _mm512_reduce_add_ps(acc);
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, 1, cb);
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.
#ifdef __BMI2__
inline float ip_bitplane_avx2(
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);
uint64_t pext_masks[7];
__m256 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[ex_bits] = _mm256_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]);
uint64_t ex64 = 0;
memcpy(&ex64, ex_code + (i / 8) * 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);
}
__m256 rq = _mm256_loadu_ps(rotated_q + i);
acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc);
}
float result = hsum_avx2(acc);
result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
return result;
}
#endif // __BMI2__
} // namespace
template <>
float compute_inner_product<SIMDLevel::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) {
if (ex_bits == 1) {
return ip_1exbit_avx512(sign_bits, ex_code, rotated_q, d, cb);
}
#ifdef __BMI2__
if (ex_bits <= 7) {
return ip_bitplane_avx2(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);
}
} // namespace faiss::rabitq::multibit
#endif // COMPILE_SIMD_AVX512
faiss/utils/simd_impl/rabitq_neon.cpp
+55
View File
@@ -0,0 +1,55 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <faiss/utils/rabitq_simd.h>
#ifdef COMPILE_SIMD_ARM_NEON
namespace faiss::rabitq {
template <>
uint64_t bitwise_and_dot_product<SIMDLevel::ARM_NEON>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
return bitwise_and_dot_product<SIMDLevel::NONE>(query, data, size, qb);
}
template <>
uint64_t bitwise_xor_dot_product<SIMDLevel::ARM_NEON>(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
return bitwise_xor_dot_product<SIMDLevel::NONE>(query, data, size, qb);
}
template <>
uint64_t popcount<SIMDLevel::ARM_NEON>(const uint8_t* data, size_t size) {
return popcount<SIMDLevel::NONE>(data, size);
}
} // namespace faiss::rabitq
namespace faiss::rabitq::multibit {
template <>
float compute_inner_product<SIMDLevel::ARM_NEON>(
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) {
return compute_inner_product<SIMDLevel::NONE>(
sign_bits, ex_code, rotated_q, d, ex_bits, cb);
}
} // namespace faiss::rabitq::multibit
#endif // COMPILE_SIMD_ARM_NEON