mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Simplify RaBitQ slightly, improving speed and recall (#4550)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4550 ## Introduction The [RaBitQ paper](https://arxiv.org/abs/2405.12497) describes a method for implementing vector inner products with bitwise operations, quantizing the query with a uniform scalar quantizer, then performing a **bitwise and** of a single bit position of the query vector with the document vector, then using a `popcount()` to accumulate the inner product. This is done for each query bit position in `[0, B_q)`, and an additional `popcount()` on the document vector alone. This additional `popcount()` can be avoided by rearranging the arithmetic. The quantization of the query, `q_u`, is specified in equation 18 of the paper, reproduced here: ```lang=txt q' = q @ P.T v_l, v_h = min(q'), max(q') delta = v_h - v_l u_i[i] = sample(uniform(0,1)) q_u[i] = floor((q'[i] - v_l) / (v_h - v_l) + u[i]) ``` The *unsigned dot product* is then estimated as (eq. 22): ```lang=txt x_b[i] = (o @ P.T)[i] > 0 ? 1 : 0 q_b[i, j] = (q_u[i] >> j) & 1 unsigned_inner_product = sum_j(2**j * sum_i(x_b[i] * q[i, j]) ``` Note that `x_b[i] * q[i, j]` is implemented as a bitwise AND. This is then used to estimate the real inner product (eq. 20): ```lang=txt inner_product_estimate = (2 * delta * unsigned_inner_product + 2 * v_l * sum_i(x_b[i]) - delta * sum(q_u[i]) - 1) / sqrt(D) ``` ## Adjustment The presence of `sum_i(x_b[i])` can be eliminated by redefining the quantization scheme used to construct `q_b[i,j]`. Instead of mapping the query to an unsigned integer, the original value is scaled, then rounded to a **signed, odd integer** in a manner very similar to a paper which builds upon RaBitQ (https://arxiv.org/abs/2409.09913). ```lang=txt |-------|-------|-------|-------|-------|-------|-------|-------| q'[i]*m: -8 v -6 v -4 v -2 v 0 v +2 v +4 v +6 v +8 | | | | | | | | odd_int: -7 -5 -3 -1 +1 +3 +5 +7 | | | | | | | | v v v v v v v v q_s[i]: 0 1 2 3 4 5 6 7 ``` This has the convenient property that `odd_int = sum(2**j * (2 * q_s[i, j] - 1))`. That is, each component is encoded as `(±4) + (±2) + (±1)`. With both queries and documents now mapped to the signed, odd integer domain, the signed integer dot product can be computed more directly. All query bits represent `±2^j`, and all document bits represent `±1`, so the maximum possible dot product is `(2^0 + 2^1 + ... + 2^j) * D = (2^(B_q-1) - 1) * D`. Every *mismatched bit* represents a *mismatched sign*, so each such bit mismatch represents a replacement of an inner product term `+2^j` with `-2^j`, a reduction of `2 * 2^j`. So, the signed dot product can be estimated as: ```lang=txt signed_inner_product = (2 ** B_q - 1) * D - 2 * sum_j(2**j * sum_i(xor(q_s[i,j], x_b[i]))) ``` Which can then be simply rescaled to to estimate the real inner product: ``` scalar = norm(q) / (norm(q_s) * sqrt(D)) inner_product_estimate = signed_inner_product * scalar ``` ## Results But because queries are expected to be centered around 0 due to random rotation, there is minimal loss in accuracy by defining the scalar quantizer such that the bias is fixed at zero. To the contrary, the centering of the scalar quantizer appears to **correct a systematic bias** in the RaBitQ quantizer which may lead to query components being disproportionately rounded up or down depending on the value of the `bias`. Particularly for smaller values of `B_q`, this imbalance can lead to significant systematic bias, adversely affecting recall. For `B_q = 1` in particular, experiments show that this centering **halves recall loss**. Costs and recall losses are consistently reduced across all `B_q` values tested, for both a 1024-d and 128-d test corpus. Results for a 1024-d test: * `B_q = 1`: 60.3% -> 86.16% recall (2.9x loss reduction) * `B_q = 2`: 87.3% -> 87.5% recall (+0.2 pp) * `B_q = 3`: 87.9% -> 88.1% recall (+0.2 pp) ## Interfaces / Compatibility This new mode is gated by a parameter on the distance computer, exposed through `IVFRaBitQSearchParameters::centered` among other APIs. For fairer performance comparisons, the **existing implementation** is also sped up with a specialized SIMD implementations for the doc-side `popcount()`. Reviewed By: luciang Differential Revision: D80904214 fbshipit-source-id: 6075deeda13420ff08adac80d7ad5bf49552aff4
This commit is contained in:
committed by
Facebook GitHub Bot
parent
bc3e3a1ef1
commit
6470b8d9d0
@@ -61,6 +61,15 @@ void bench_rabitq_and_dot_product(benchmark::State& state) {
|
||||
});
|
||||
}
|
||||
|
||||
void bench_rabitq_xor_dot_product(benchmark::State& state) {
|
||||
bench_rabitq_generic(
|
||||
state,
|
||||
[](const uint8_t* q, const uint8_t* x, size_t size, size_t qb)
|
||||
-> int64_t {
|
||||
return rabitq::bitwise_xor_dot_product(q, x, size, qb);
|
||||
});
|
||||
}
|
||||
|
||||
void bench_rabitq_and_dot_product_with_sum(benchmark::State& state) {
|
||||
bench_rabitq_generic(
|
||||
state,
|
||||
@@ -80,6 +89,9 @@ BENCHMARK(bench_rabitq_sum)->ArgsProduct({{0}, dims})->ArgNames({"qb", "d"});
|
||||
BENCHMARK(bench_rabitq_and_dot_product)
|
||||
->ArgsProduct({qbs, dims})
|
||||
->ArgNames({"qb", "d"});
|
||||
BENCHMARK(bench_rabitq_xor_dot_product)
|
||||
->ArgsProduct({qbs, dims})
|
||||
->ArgNames({"qb", "d"});
|
||||
BENCHMARK(bench_rabitq_and_dot_product_with_sum)
|
||||
->ArgsProduct({qbs, dims})
|
||||
->ArgNames({"qb", "d"});
|
||||
|
||||
@@ -155,15 +155,18 @@ struct RaBitInvertedListScanner : InvertedListScanner {
|
||||
std::unique_ptr<FlatCodesDistanceComputer> dc;
|
||||
|
||||
uint8_t qb = 0;
|
||||
bool centered = false;
|
||||
|
||||
explicit RaBitInvertedListScanner(
|
||||
const IndexIVFRaBitQ& ivf_rabitq_in,
|
||||
bool store_pairs = false,
|
||||
const IDSelector* sel = nullptr,
|
||||
uint8_t qb_in = 0)
|
||||
uint8_t qb_in = 0,
|
||||
bool centered = false)
|
||||
: InvertedListScanner(store_pairs, sel),
|
||||
ivf_rabitq{ivf_rabitq_in},
|
||||
qb{qb_in} {
|
||||
qb{qb_in},
|
||||
centered(centered) {
|
||||
keep_max = is_similarity_metric(ivf_rabitq.metric_type);
|
||||
code_size = ivf_rabitq.code_size;
|
||||
}
|
||||
@@ -196,7 +199,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
|
||||
// both query_vector and centroid are available!
|
||||
// set up DistanceComputer
|
||||
dc.reset(ivf_rabitq.rabitq.get_distance_computer(
|
||||
qb, reconstructed_centroid.data()));
|
||||
qb, reconstructed_centroid.data(), centered));
|
||||
|
||||
dc->set_query(query_vector.data());
|
||||
}
|
||||
@@ -208,12 +211,15 @@ InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
|
||||
const IDSelector* sel,
|
||||
const IVFSearchParameters* search_params_in) const {
|
||||
uint8_t used_qb = qb;
|
||||
bool centered = false;
|
||||
if (auto params = dynamic_cast<const IVFRaBitQSearchParameters*>(
|
||||
search_params_in)) {
|
||||
used_qb = params->qb;
|
||||
centered = params->centered;
|
||||
}
|
||||
|
||||
return new RaBitInvertedListScanner(*this, store_pairs, sel, used_qb);
|
||||
return new RaBitInvertedListScanner(
|
||||
*this, store_pairs, sel, used_qb, centered);
|
||||
}
|
||||
|
||||
void IndexIVFRaBitQ::reconstruct_from_offset(
|
||||
@@ -278,7 +284,8 @@ float IVFRaBitDistanceComputer::operator()(idx_t i) {
|
||||
float distance = 0;
|
||||
|
||||
std::unique_ptr<FlatCodesDistanceComputer> dc(
|
||||
parent->rabitq.get_distance_computer(parent->qb, centroid.data()));
|
||||
parent->rabitq.get_distance_computer(
|
||||
parent->qb, centroid.data(), /*centered=*/false));
|
||||
dc->set_query(q);
|
||||
distance = dc->distance_to_code(code);
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ namespace faiss {
|
||||
|
||||
struct IVFRaBitQSearchParameters : IVFSearchParameters {
|
||||
uint8_t qb = 0;
|
||||
bool centered = false;
|
||||
};
|
||||
|
||||
// * by_residual is true, just by design
|
||||
|
||||
+13
-10
@@ -55,16 +55,17 @@ void IndexRaBitQ::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
|
||||
|
||||
FlatCodesDistanceComputer* IndexRaBitQ::get_FlatCodesDistanceComputer() const {
|
||||
FlatCodesDistanceComputer* dc =
|
||||
rabitq.get_distance_computer(qb, center.data());
|
||||
rabitq.get_distance_computer(qb, center.data(), centered);
|
||||
dc->code_size = rabitq.code_size;
|
||||
dc->codes = codes.data();
|
||||
return dc;
|
||||
}
|
||||
|
||||
FlatCodesDistanceComputer* IndexRaBitQ::get_quantized_distance_computer(
|
||||
const uint8_t qb) const {
|
||||
const uint8_t qb,
|
||||
bool centered) const {
|
||||
FlatCodesDistanceComputer* dc =
|
||||
rabitq.get_distance_computer(qb, center.data());
|
||||
rabitq.get_distance_computer(qb, center.data(), centered);
|
||||
dc->code_size = rabitq.code_size;
|
||||
dc->codes = codes.data();
|
||||
return dc;
|
||||
@@ -76,6 +77,7 @@ struct Run_search_with_dc_res {
|
||||
using T = void;
|
||||
|
||||
uint8_t qb = 0;
|
||||
bool centered = false;
|
||||
|
||||
template <class BlockResultHandler>
|
||||
void f(BlockResultHandler& res, const IndexRaBitQ* index, const float* xq) {
|
||||
@@ -87,7 +89,7 @@ struct Run_search_with_dc_res {
|
||||
#pragma omp parallel // if (res.nq > 100)
|
||||
{
|
||||
std::unique_ptr<FlatCodesDistanceComputer> dc(
|
||||
index->get_quantized_distance_computer(qb));
|
||||
index->get_quantized_distance_computer(qb, centered));
|
||||
SingleResultHandler resi(res);
|
||||
#pragma omp for
|
||||
for (int64_t q = 0; q < res.nq; q++) {
|
||||
@@ -114,14 +116,15 @@ void IndexRaBitQ::search(
|
||||
float* distances,
|
||||
idx_t* labels,
|
||||
const SearchParameters* params_in) const {
|
||||
uint8_t used_qb = qb;
|
||||
if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
|
||||
used_qb = params->qb;
|
||||
}
|
||||
|
||||
const IDSelector* sel = (params_in != nullptr) ? params_in->sel : nullptr;
|
||||
Run_search_with_dc_res r;
|
||||
r.qb = used_qb;
|
||||
if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
|
||||
r.qb = params->qb;
|
||||
r.centered = params->centered;
|
||||
} else {
|
||||
r.qb = this->qb;
|
||||
r.centered = this->centered;
|
||||
}
|
||||
|
||||
dispatch_knn_ResultHandler(
|
||||
n, distances, labels, k, metric_type, sel, r, this, x);
|
||||
|
||||
+6
-1
@@ -14,6 +14,7 @@ namespace faiss {
|
||||
|
||||
struct RaBitQSearchParameters : SearchParameters {
|
||||
uint8_t qb = 0;
|
||||
bool centered = false;
|
||||
};
|
||||
|
||||
struct IndexRaBitQ : IndexFlatCodes {
|
||||
@@ -26,6 +27,9 @@ struct IndexRaBitQ : IndexFlatCodes {
|
||||
// use '0' to disable quantization and use raw fp32 values.
|
||||
uint8_t qb = 0;
|
||||
|
||||
// quantize the query with a zero-centered scalar quantizer.
|
||||
bool centered = false;
|
||||
|
||||
IndexRaBitQ();
|
||||
|
||||
explicit IndexRaBitQ(idx_t d, MetricType metric = METRIC_L2);
|
||||
@@ -42,7 +46,8 @@ struct IndexRaBitQ : IndexFlatCodes {
|
||||
// returns a quantized-to-qb bits DC if qb_in > 0
|
||||
// returns a default fp32-based DC if qb_in == 0
|
||||
FlatCodesDistanceComputer* get_quantized_distance_computer(
|
||||
const uint8_t qb_in) const;
|
||||
const uint8_t qb_in,
|
||||
bool centered) const;
|
||||
|
||||
// Don't rely on sa_decode(), bcz it is good for IP, but not for L2.
|
||||
// As a result, use get_FlatCodesDistanceComputer() for the search.
|
||||
|
||||
@@ -32,6 +32,8 @@ struct QueryFactorsData {
|
||||
|
||||
float qr_to_c_L2sqr = 0;
|
||||
float qr_norm_L2sqr = 0;
|
||||
|
||||
float int_dot_scale = 1;
|
||||
};
|
||||
|
||||
static size_t get_code_size(const size_t d) {
|
||||
@@ -310,6 +312,7 @@ struct RaBitDistanceComputerQ : RaBitDistanceComputer {
|
||||
|
||||
// the number of bits for SQ quantization of the query (qb > 0)
|
||||
uint8_t qb = 8;
|
||||
bool centered = false;
|
||||
// the smallest value divisible by 8 that is not smaller than dim
|
||||
size_t popcount_aligned_dim = 0;
|
||||
|
||||
@@ -333,20 +336,30 @@ float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
|
||||
const uint8_t* binary_data = code;
|
||||
const FactorsData* fac = reinterpret_cast<const FactorsData*>(code + size);
|
||||
|
||||
// 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.
|
||||
auto sum_q = rabitq::popcount(binary_data, size);
|
||||
// this is ||or - c||^2 - (IP ? ||or||^2 : 0)
|
||||
float final_dot = 0;
|
||||
// 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;
|
||||
if (centered) {
|
||||
int64_t int_dot = ((1 << qb) - 1) * d;
|
||||
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 {
|
||||
// See RaBitDistanceComputerNotQ::distance_to_code() for baseline code.
|
||||
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;
|
||||
}
|
||||
|
||||
// this is ||or - c||^2 - (IP ? ||or||^2 : 0)
|
||||
const float or_c_l2sqr = fac->or_minus_c_l2sqr;
|
||||
@@ -370,11 +383,30 @@ float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
|
||||
}
|
||||
}
|
||||
|
||||
namespace {
|
||||
|
||||
// Ideal quantizer radii for quantizers of 1..8 bits, optimized to minimize L2
|
||||
// reconstruction error.
|
||||
const float z_max_by_qb[8] = {
|
||||
0.79688, // qb = 1.
|
||||
1.49375,
|
||||
2.05078,
|
||||
2.50938,
|
||||
2.91250,
|
||||
3.26406,
|
||||
3.59844,
|
||||
3.91016, // qb = 8.
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
void RaBitDistanceComputerQ::set_query(const float* 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);
|
||||
|
||||
// compute the distance from the query to the centroid
|
||||
if (centroid != nullptr) {
|
||||
@@ -398,26 +430,36 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
|
||||
// quantize the query. compute min and max
|
||||
float v_min = std::numeric_limits<float>::max();
|
||||
float v_max = std::numeric_limits<float>::lowest();
|
||||
for (size_t i = 0; i < d; i++) {
|
||||
const float v_q = rotated_q[i];
|
||||
v_min = std::min(v_min, v_q);
|
||||
v_max = std::max(v_max, v_q);
|
||||
if (centered) {
|
||||
float z_max = z_max_by_qb[qb - 1];
|
||||
float v_radius = z_max * std::sqrt(query_fac.qr_to_c_L2sqr / d);
|
||||
v_min = -v_radius;
|
||||
v_max = v_radius;
|
||||
} else {
|
||||
for (size_t i = 0; i < d; i++) {
|
||||
const float v_q = rotated_q[i];
|
||||
v_min = std::min(v_min, v_q);
|
||||
v_max = std::max(v_max, v_q);
|
||||
}
|
||||
}
|
||||
|
||||
const float pow_2_qb = 1 << qb;
|
||||
|
||||
const float delta = (v_max - v_min) / (pow_2_qb - 1);
|
||||
const uint8_t max_code = (1 << qb) - 1;
|
||||
const float delta = (v_max - v_min) / max_code;
|
||||
const float inv_delta = 1.0f / delta;
|
||||
|
||||
size_t sum_qq = 0;
|
||||
int64_t sum2_signed_odd_int = 0;
|
||||
for (int32_t i = 0; i < d; i++) {
|
||||
const float v_q = rotated_q[i];
|
||||
|
||||
// a default non-randomized SQ
|
||||
const int v_qq = std::round((v_q - v_min) * inv_delta);
|
||||
|
||||
rotated_qq[i] = std::min(255, std::max(0, v_qq));
|
||||
const uint8_t v_qq = std::clamp<float>(
|
||||
std::round((v_q - v_min) * inv_delta), 0, max_code);
|
||||
rotated_qq[i] = v_qq;
|
||||
sum_qq += v_qq;
|
||||
if (centered) {
|
||||
int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
|
||||
sum2_signed_odd_int += signed_odd_int * signed_odd_int;
|
||||
}
|
||||
}
|
||||
|
||||
// rearrange the query vector
|
||||
@@ -438,6 +480,8 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
|
||||
query_fac.c1 = 2 * delta * inv_d;
|
||||
query_fac.c2 = 2 * v_min * inv_d;
|
||||
query_fac.c34 = inv_d * (delta * sum_qq + d * v_min);
|
||||
query_fac.int_dot_scale =
|
||||
std::sqrt(query_fac.qr_to_c_L2sqr / (sum2_signed_odd_int * d));
|
||||
|
||||
if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
|
||||
// precompute if needed
|
||||
@@ -447,7 +491,8 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
|
||||
|
||||
FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
|
||||
uint8_t qb,
|
||||
const float* centroid_in) const {
|
||||
const float* centroid_in,
|
||||
bool centered) const {
|
||||
if (qb == 0) {
|
||||
auto dc = std::make_unique<RaBitDistanceComputerNotQ>();
|
||||
dc->metric_type = metric_type;
|
||||
@@ -461,6 +506,7 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
|
||||
dc->d = d;
|
||||
dc->centroid = centroid_in;
|
||||
dc->qb = qb;
|
||||
dc->centered = centered;
|
||||
|
||||
return dc.release();
|
||||
}
|
||||
|
||||
@@ -72,7 +72,8 @@ struct RaBitQuantizer : Quantizer {
|
||||
// specify qb > 0 to have SQ qb-bits query
|
||||
FlatCodesDistanceComputer* get_distance_computer(
|
||||
uint8_t qb,
|
||||
const float* centroid_in = nullptr) const;
|
||||
const float* centroid_in = nullptr,
|
||||
bool centered = false) const;
|
||||
};
|
||||
|
||||
} // namespace faiss
|
||||
|
||||
@@ -293,6 +293,89 @@ 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(
|
||||
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 (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);
|
||||
}
|
||||
#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 (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);
|
||||
}
|
||||
#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 (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);
|
||||
#endif
|
||||
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;
|
||||
}
|
||||
|
||||
inline uint64_t popcount(const uint8_t* data, size_t size) {
|
||||
uint64_t sum = 0;
|
||||
size_t offset = 0;
|
||||
|
||||
+18
-6
@@ -223,21 +223,25 @@ class TestRaBitQ(unittest.TestCase):
|
||||
index_rbq.train(ds.get_train())
|
||||
index_rbq.add(ds.get_database())
|
||||
|
||||
for qb in [1, 2, 3, 4, 8]:
|
||||
params = faiss.RaBitQSearchParameters(qb=qb)
|
||||
def test(params):
|
||||
_, I_rbq = index_rbq.search(ds.get_queries(), k, params=params)
|
||||
|
||||
# ensure that RaBitQ and PQ are relatively close
|
||||
loss_rbq = 1 - eval_I(I_rbq)
|
||||
ratio_threshold = 2 ** (1 / qb)
|
||||
print(
|
||||
f"{random_rotate=:1}, {params.qb=}: "
|
||||
f"{random_rotate=:1}, {params.qb=}, {params.centered=:1}: "
|
||||
f"{loss_rbq=:5.3f} = loss_pq * {loss_rbq / loss_pq:5.3f}"
|
||||
f" < {ratio_threshold=:.2f}"
|
||||
)
|
||||
|
||||
np.testing.assert_(loss_rbq < loss_pq * ratio_threshold)
|
||||
|
||||
for qb in [1, 2, 3, 4, 8]:
|
||||
print()
|
||||
for centered in [False, True]:
|
||||
test(faiss.RaBitQSearchParameters(qb=qb, centered=centered))
|
||||
|
||||
def test_comparison_vs_pq_L2(self):
|
||||
self.do_comparison_vs_pq_test(faiss.METRIC_L2)
|
||||
|
||||
@@ -387,21 +391,29 @@ class TestIVFRaBitQ(unittest.TestCase):
|
||||
print(f"{random_rotate=:1}, {loss_pq=:5.3f}")
|
||||
np.testing.assert_(loss_pq < 0.25, f"{loss_pq}")
|
||||
|
||||
for qb in [1, 2, 3, 4, 8]:
|
||||
params = faiss.IVFRaBitQSearchParameters(nprobe=nprobe, qb=qb)
|
||||
def test(params):
|
||||
D_rbq, I_rbq = index_rbq.search(xq, k, params=params)
|
||||
|
||||
# ensure that RaBitQ and PQ are relatively close
|
||||
loss_rbq = 1 - eval_I(I_rbq)
|
||||
ratio_threshold = 2 ** (1 / qb)
|
||||
print(
|
||||
f"{random_rotate=:1}, {params.qb=}: "
|
||||
f"{random_rotate=:1}, {params.qb=}, {params.centered=:1}: "
|
||||
f"{loss_rbq=:5.3f} = loss_pq * {loss_rbq / loss_pq:5.3f}"
|
||||
f" < {ratio_threshold=:.2f}"
|
||||
)
|
||||
|
||||
np.testing.assert_(loss_rbq < loss_pq * ratio_threshold)
|
||||
|
||||
for qb in [1, 2, 3, 4, 8]:
|
||||
print()
|
||||
for centered in [False, True]:
|
||||
test(
|
||||
faiss.IVFRaBitQSearchParameters(
|
||||
nprobe=nprobe, qb=qb, centered=centered
|
||||
)
|
||||
)
|
||||
|
||||
def test_comparison_vs_pq_L2(self):
|
||||
self.do_comparison_vs_pq_test(faiss.METRIC_L2)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user