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 Fast Scan (#4595)

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

**Introduction**

This diff adds a new index called the IndexRaBitQFastScan algorithm. The algorithm is based on the existing IndexRaBitQ but achieves higher speed as it processes batches of 32 data vectors concurrently. It leverages the established IndexFastScan architecture to enable efficient batch processing and parallelism.

**Implementation**

* **New Source and Header Files**: Added implementations for IndexRaBitQFastScan, following a similar interface to IndexRaBitQ.

* **Batched Processing**: The search operation processes multiple (32) data vectors in a single batch, taking advantage of low-level parallelism to improve throughput.

* **Specialized Post-processing Handler**:  A dedicated handler was added for IndexRaBitQFastScan to perform necessary post-processing during search because the LUT accumulates only partial distances. Unlike AQ Fast Scan's simple scalar post-processing, RaBitQ requires complex distance adjustments depending on both query and database vector factors.

* **LUT**: IndexRaBitQFastScan produces slightly different results than IndexRaBitQ due to an extra quantization step in the IndexFastScan architecture. Specifically:
  * The LUT computes a float value as c1 * inner_product + c2 * popcount, which is then quantized. This quantization can cause the results to differ slightly from those of IndexRaBitQ.
  * It is possible to avoid this by storing only the inner_product in the LUT, but doing so would require calculating all data vector popcounts during search, introducing a tradeoff between speed and accuracy.
  * With the idea proposed in diff D80904214, the algorithm can be modified in the future to eliminate the popcount calculation step, potentially improving both efficiency and accuracy.
* **Query Offset Parameter**: RaBitQ uses query factors in distance calculations that should be computed in `compute_float_LUT` method (the most efficient place since we are calculating `rotated_qq` anyways) and used for final distance calculations in handlers. However, the previous version of `compute_quantized_LUT` that calls `compute_float_LUT` did not know the query_offset, preventing proper storage of query factors at their global indices. To solve this, I added the extra parameter `query_offset` to both `compute_quantized_LUT` and `compute_float_LUT` methods. After this change, computed query factors can be accessed by the correct global query index during distance calculations, avoiding expensive recalculation.

**Testing**

* Conducted comprehensive tests in the test_rabitq suite covering accuracy comparisons with IndexRaBitQ for L2 and Inner Product metrics, encoding/decoding consistency, query quantization bit settings, small dataset functionality, performance against PQFastScan, serialization, memory management, error handling, and thread safety.
* All tests passed successfully, validating the correctness and robustness of IndexRaBitQFastScan.

**Results**
results_rabitq
* **Performance Dependency**: Performance measurements confirm that IndexRaBitQFastScan is notably faster than IndexRaBitQ when the qb value is high. While the original IndexRaBitQ experiences increased runtime with higher qb values, the fast scan variant maintains consistent runtime regardless of qb.
* **Parallelized Training Loop**: The training loop is parallelized, greatly reducing training time. This parallelism should also be added to the original IndexRaBitQ.
* **Consistency Across Metrics**: The performance advantages of IndexRaBitQFastScan hold true for both L2 and Inner Product metrics, demonstrating robustness across different distance measures.
* One of the next steps is to benchmark IndexRaBitQFastScan against other algorithms to evaluate its performance in a broader context.

Reviewed By: mdouze

Differential Revision: D81787307

fbshipit-source-id: d30827b990116254a88e80df0c08a2f5e086b64c
This commit is contained in:
Alibek Zhakubayev
2025-10-14 09:55:44 -07:00
committed by meta-codesync[bot]
parent 041ac84318
commit 01d394e583
20 changed files with 1899 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+5
View File
@@ -41,6 +41,7 @@ set(FAISS_SRC
IndexPQFastScan.cpp
IndexPreTransform.cpp
IndexRaBitQ.cpp
IndexRaBitQFastScan.cpp
IndexRefine.cpp
IndexReplicas.cpp
IndexRowwiseMinMax.cpp
@@ -63,6 +64,7 @@ set(FAISS_SRC
impl/ProductQuantizer.cpp
impl/AdditiveQuantizer.cpp
impl/RaBitQuantizer.cpp
impl/RaBitQUtils.cpp
impl/ResidualQuantizer.cpp
impl/LocalSearchQuantizer.cpp
impl/ProductAdditiveQuantizer.cpp
@@ -141,6 +143,7 @@ set(FAISS_HEADERS
IndexRefine.h
IndexReplicas.h
IndexRaBitQ.h
IndexRaBitQFastScan.h
IndexRowwiseMinMax.h
IndexScalarQuantizer.h
IndexShards.h
@@ -163,6 +166,7 @@ set(FAISS_HEADERS
impl/LocalSearchQuantizer.h
impl/ProductAdditiveQuantizer.h
impl/LookupTableScaler.h
impl/FastScanDistancePostProcessing.h
impl/maybe_owned_vector.h
impl/NNDescent.h
impl/NSG.h
@@ -171,6 +175,7 @@ set(FAISS_HEADERS
impl/ProductQuantizer.h
impl/Quantizer.h
impl/RaBitQuantizer.h
impl/RaBitQUtils.h
impl/ResidualQuantizer.h
impl/ResultHandler.h
impl/ScalarQuantizer.h
faiss/IndexAdditiveQuantizerFastScan.cpp
+9 -4
View File
@@ -11,6 +11,7 @@
#include <memory>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/FastScanDistancePostProcessing.h>
#include <faiss/impl/LocalSearchQuantizer.h>
#include <faiss/impl/LookupTableScaler.h>
#include <faiss/impl/ResidualQuantizer.h>
@@ -123,7 +124,8 @@ void IndexAdditiveQuantizerFastScan::estimate_norm_scale(
}
std::vector<float> dis_tables(n * M * ksub);
compute_float_LUT(dis_tables.data(), n, x);
FastScanDistancePostProcessing empty_context;
compute_float_LUT(dis_tables.data(), n, x, empty_context);
// here we compute the mean of scales for each query
// TODO: try max of scales
@@ -153,7 +155,8 @@ void IndexAdditiveQuantizerFastScan::compute_codes(
void IndexAdditiveQuantizerFastScan::compute_float_LUT(
float* lut,
idx_t n,
const float* x) const {
const float* x,
const FastScanDistancePostProcessing&) const {
if (metric_type == METRIC_INNER_PRODUCT) {
aq->compute_LUT(n, x, lut, 1.0f);
} else {
@@ -200,10 +203,12 @@ void IndexAdditiveQuantizerFastScan::search(
}
NormTableScaler scaler(norm_scale);
FastScanDistancePostProcessing context;
context.norm_scaler = &scaler;
if (metric_type == METRIC_L2) {
search_dispatch_implem<true>(n, x, k, distances, labels, &scaler);
search_dispatch_implem<true>(n, x, k, distances, labels, context);
} else {
search_dispatch_implem<false>(n, x, k, distances, labels, &scaler);
search_dispatch_implem<false>(n, x, k, distances, labels, context);
}
}
faiss/IndexAdditiveQuantizerFastScan.h
+5 -1
View File
@@ -62,7 +62,11 @@ struct IndexAdditiveQuantizerFastScan : IndexFastScan {
void compute_codes(uint8_t* codes, idx_t n, const float* x) const override;
void compute_float_LUT(float* lut, idx_t n, const float* x) const override;
void compute_float_LUT(
float* lut,
idx_t n,
const float* x,
const FastScanDistancePostProcessing& context) const override;
void search(
idx_t n,
faiss/IndexFastScan.cpp
+124 -50
View File
@@ -7,17 +7,20 @@
#include <faiss/IndexFastScan.h>
#include <cassert>
#include <climits>
#include <omp.h>
#include <cstring>
#include <memory>
#include <omp.h>
#include <faiss/impl/CodePacker.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/FastScanDistancePostProcessing.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/LookupTableScaler.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/pq4_fast_scan.h>
#include <faiss/impl/simd_result_handlers.h>
#include <faiss/utils/hamming.h>
#include <faiss/utils/utils.h>
#include <faiss/impl/pq4_fast_scan.h>
#include <faiss/impl/simd_result_handlers.h>
@@ -163,14 +166,14 @@ void estimators_from_tables_generic(
size_t k,
typename C::T* heap_dis,
int64_t* heap_ids,
const NormTableScaler* scaler) {
const FastScanDistancePostProcessing& context) {
using accu_t = typename C::T;
for (size_t j = 0; j < ncodes; ++j) {
BitstringReader bsr(codes + j * index.code_size, index.code_size);
accu_t dis = 0;
const dis_t* dt = dis_table;
int nscale = scaler ? scaler->nscale : 0;
int nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
for (size_t m = 0; m < index.M - nscale; m++) {
uint64_t c = bsr.read(index.nbits);
@@ -178,10 +181,10 @@ void estimators_from_tables_generic(
dt += index.ksub;
}
if (nscale) {
if (nscale && context.norm_scaler) {
for (size_t m = 0; m < nscale; m++) {
uint64_t c = bsr.read(index.nbits);
dis += scaler->scale_one(dt[c]);
dis += context.norm_scaler->scale_one(dt[c]);
dt += index.ksub;
}
}
@@ -193,40 +196,58 @@ void estimators_from_tables_generic(
}
}
template <class C>
ResultHandlerCompare<C, false>* make_knn_handler(
} // anonymous namespace
// Default implementation of make_knn_handler with centralized fallback logic
void* IndexFastScan::make_knn_handler(
bool is_max,
int impl,
idx_t n,
idx_t k,
size_t ntotal,
float* distances,
idx_t* labels,
const IDSelector* sel = nullptr) {
using HeapHC = HeapHandler<C, false>;
using ReservoirHC = ReservoirHandler<C, false>;
using SingleResultHC = SingleResultHandler<C, false>;
const IDSelector* sel,
const FastScanDistancePostProcessing&) const {
// Create default handlers based on k and impl
if (is_max) {
using HeapHC = HeapHandler<CMax<uint16_t, int>, false>;
using ReservoirHC = ReservoirHandler<CMax<uint16_t, int>, false>;
using SingleResultHC = SingleResultHandler<CMax<uint16_t, int>, false>;
if (k == 1) {
return new SingleResultHC(n, ntotal, distances, labels, sel);
} else if (impl % 2 == 0) {
return new HeapHC(n, ntotal, k, distances, labels, sel);
} else /* if (impl % 2 == 1) */ {
return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels, sel);
if (k == 1) {
return new SingleResultHC(n, ntotal, distances, labels, sel);
} else if (impl % 2 == 0) {
return new HeapHC(n, ntotal, k, distances, labels, sel);
} else {
return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels, sel);
}
} else {
using HeapHC = HeapHandler<CMin<uint16_t, int>, false>;
using ReservoirHC = ReservoirHandler<CMin<uint16_t, int>, false>;
using SingleResultHC = SingleResultHandler<CMin<uint16_t, int>, false>;
if (k == 1) {
return new SingleResultHC(n, ntotal, distances, labels, sel);
} else if (impl % 2 == 0) {
return new HeapHC(n, ntotal, k, distances, labels, sel);
} else {
return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels, sel);
}
}
}
} // anonymous namespace
using namespace quantize_lut;
void IndexFastScan::compute_quantized_LUT(
idx_t n,
const float* x,
uint8_t* lut,
float* normalizers) const {
float* normalizers,
const FastScanDistancePostProcessing& context) const {
size_t dim12 = ksub * M;
std::unique_ptr<float[]> dis_tables(new float[n * dim12]);
compute_float_LUT(dis_tables.get(), n, x);
compute_float_LUT(dis_tables.get(), n, x, context);
for (uint64_t i = 0; i < n; i++) {
round_uint8_per_column(
@@ -263,10 +284,12 @@ void IndexFastScan::search(
!params, "search params not supported for this index");
FAISS_THROW_IF_NOT(k > 0);
FastScanDistancePostProcessing empty_context{};
if (metric_type == METRIC_L2) {
search_dispatch_implem<true>(n, x, k, distances, labels, nullptr);
search_dispatch_implem<true>(n, x, k, distances, labels, empty_context);
} else {
search_dispatch_implem<false>(n, x, k, distances, labels, nullptr);
search_dispatch_implem<false>(
n, x, k, distances, labels, empty_context);
}
}
@@ -277,7 +300,7 @@ void IndexFastScan::search_dispatch_implem(
idx_t k,
float* distances,
idx_t* labels,
const NormTableScaler* scaler) const {
const FastScanDistancePostProcessing& context) const {
using Cfloat = typename std::conditional<
is_max,
CMax<float, int64_t>,
@@ -308,15 +331,20 @@ void IndexFastScan::search_dispatch_implem(
FAISS_THROW_MSG("not implemented");
} else if (implem == 2 || implem == 3 || implem == 4) {
FAISS_THROW_IF_NOT(orig_codes != nullptr);
search_implem_234<Cfloat>(n, x, k, distances, labels, scaler);
search_implem_234<Cfloat>(n, x, k, distances, labels, context);
} else if (impl >= 12 && impl <= 15) {
FAISS_THROW_IF_NOT(ntotal < INT_MAX);
int nt = std::min(omp_get_max_threads(), int(n));
// Fall back to single-threaded implementations when parallelization not
// beneficial:
// - Single-core system (omp_get_max_threads() = 1)
// - Single query (n = 1)
// - OpenMP disabled (omp_get_max_threads() = 1)
if (nt < 2) {
if (impl == 12 || impl == 13) {
search_implem_12<C>(n, x, k, distances, labels, impl, scaler);
search_implem_12<C>(n, x, k, distances, labels, impl, context);
} else {
search_implem_14<C>(n, x, k, distances, labels, impl, scaler);
search_implem_14<C>(n, x, k, distances, labels, impl, context);
}
} else {
// explicitly slice over threads
@@ -324,14 +352,33 @@ void IndexFastScan::search_dispatch_implem(
for (int slice = 0; slice < nt; slice++) {
idx_t i0 = n * slice / nt;
idx_t i1 = n * (slice + 1) / nt;
// Create per-thread context with adjusted query_factors pointer
FastScanDistancePostProcessing thread_context = context;
if (thread_context.query_factors != nullptr) {
thread_context.query_factors += i0;
}
float* dis_i = distances + i0 * k;
idx_t* lab_i = labels + i0 * k;
if (impl == 12 || impl == 13) {
search_implem_12<C>(
i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler);
i1 - i0,
x + i0 * d,
k,
dis_i,
lab_i,
impl,
thread_context);
} else {
search_implem_14<C>(
i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler);
i1 - i0,
x + i0 * d,
k,
dis_i,
lab_i,
impl,
thread_context);
}
}
}
@@ -347,12 +394,12 @@ void IndexFastScan::search_implem_234(
idx_t k,
float* distances,
idx_t* labels,
const NormTableScaler* scaler) const {
const FastScanDistancePostProcessing& context) const {
FAISS_THROW_IF_NOT(implem == 2 || implem == 3 || implem == 4);
const size_t dim12 = ksub * M;
std::unique_ptr<float[]> dis_tables(new float[n * dim12]);
compute_float_LUT(dis_tables.get(), n, x);
compute_float_LUT(dis_tables.get(), n, x, context);
std::vector<float> normalizers(n * 2);
@@ -384,7 +431,7 @@ void IndexFastScan::search_implem_234(
k,
heap_dis,
heap_ids,
scaler);
context);
heap_reorder<Cfloat>(k, heap_dis, heap_ids);
@@ -407,7 +454,7 @@ void IndexFastScan::search_implem_12(
float* distances,
idx_t* labels,
int impl,
const NormTableScaler* scaler) const {
const FastScanDistancePostProcessing& context) const {
using RH = ResultHandlerCompare<C, false>;
FAISS_THROW_IF_NOT(bbs == 32);
@@ -416,6 +463,11 @@ void IndexFastScan::search_implem_12(
if (n > qbs2) {
for (int64_t i0 = 0; i0 < n; i0 += qbs2) {
int64_t i1 = std::min(i0 + qbs2, n);
// Create sub-context with adjusted query_factors pointer
FastScanDistancePostProcessing sub_context = context;
if (sub_context.query_factors != nullptr) {
sub_context.query_factors += i0;
}
search_implem_12<C>(
i1 - i0,
x + d * i0,
@@ -423,7 +475,7 @@ void IndexFastScan::search_implem_12(
distances + i0 * k,
labels + i0 * k,
impl,
scaler);
sub_context);
}
return;
}
@@ -436,7 +488,7 @@ void IndexFastScan::search_implem_12(
quantized_dis_tables.clear();
} else {
compute_quantized_LUT(
n, x, quantized_dis_tables.get(), normalizers.get());
n, x, quantized_dis_tables.get(), normalizers.get(), context);
}
AlignedTable<uint8_t> LUT(n * dim12);
@@ -454,8 +506,17 @@ void IndexFastScan::search_implem_12(
pq4_pack_LUT_qbs(qbs, M2, quantized_dis_tables.get(), LUT.get());
FAISS_THROW_IF_NOT(LUT_nq == n);
std::unique_ptr<RH> handler(
make_knn_handler<C>(impl, n, k, ntotal, distances, labels));
std::unique_ptr<RH> handler(static_cast<RH*>(make_knn_handler(
C::is_max,
impl,
n,
k,
ntotal,
distances,
labels,
nullptr,
context)));
handler->disable = bool(skip & 2);
handler->normalizers = normalizers.get();
@@ -469,7 +530,7 @@ void IndexFastScan::search_implem_12(
codes.get(),
LUT.get(),
*handler.get(),
scaler);
context.norm_scaler);
}
if (!(skip & 8)) {
handler->end();
@@ -486,7 +547,7 @@ void IndexFastScan::search_implem_14(
float* distances,
idx_t* labels,
int impl,
const NormTableScaler* scaler) const {
const FastScanDistancePostProcessing& context) const {
using RH = ResultHandlerCompare<C, false>;
FAISS_THROW_IF_NOT(bbs % 32 == 0);
@@ -496,6 +557,11 @@ void IndexFastScan::search_implem_14(
if (n > qbs2) {
for (int64_t i0 = 0; i0 < n; i0 += qbs2) {
int64_t i1 = std::min(i0 + qbs2, n);
// Create sub-context with adjusted query_factors pointer
FastScanDistancePostProcessing sub_context = context;
if (sub_context.query_factors != nullptr) {
sub_context.query_factors += i0;
}
search_implem_14<C>(
i1 - i0,
x + d * i0,
@@ -503,7 +569,7 @@ void IndexFastScan::search_implem_14(
distances + i0 * k,
labels + i0 * k,
impl,
scaler);
sub_context);
}
return;
}
@@ -516,14 +582,22 @@ void IndexFastScan::search_implem_14(
quantized_dis_tables.clear();
} else {
compute_quantized_LUT(
n, x, quantized_dis_tables.get(), normalizers.get());
n, x, quantized_dis_tables.get(), normalizers.get(), context);
}
AlignedTable<uint8_t> LUT(n * dim12);
pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get());
std::unique_ptr<RH> handler(
make_knn_handler<C>(impl, n, k, ntotal, distances, labels));
std::unique_ptr<RH> handler(static_cast<RH*>(make_knn_handler(
C::is_max,
impl,
n,
k,
ntotal,
distances,
labels,
nullptr,
context)));
handler->disable = bool(skip & 2);
handler->normalizers = normalizers.get();
@@ -538,7 +612,7 @@ void IndexFastScan::search_implem_14(
codes.get(),
LUT.get(),
*handler.get(),
scaler);
context.norm_scaler);
}
if (!(skip & 8)) {
handler->end();
@@ -551,7 +625,7 @@ template void IndexFastScan::search_dispatch_implem<true>(
idx_t k,
float* distances,
idx_t* labels,
const NormTableScaler* scaler) const;
const FastScanDistancePostProcessing& context) const;
template void IndexFastScan::search_dispatch_implem<false>(
idx_t n,
@@ -559,7 +633,7 @@ template void IndexFastScan::search_dispatch_implem<false>(
idx_t k,
float* distances,
idx_t* labels,
const NormTableScaler* scaler) const;
const FastScanDistancePostProcessing& context) const;
void IndexFastScan::reconstruct(idx_t key, float* recons) const {
std::vector<uint8_t> code(code_size, 0);
faiss/IndexFastScan.h
+107 -7
View File
@@ -8,6 +8,7 @@
#pragma once
#include <faiss/Index.h>
#include <faiss/impl/FastScanDistancePostProcessing.h>
#include <faiss/utils/AlignedTable.h>
namespace faiss {
@@ -15,6 +16,13 @@ namespace faiss {
struct CodePacker;
struct NormTableScaler;
// Forward declarations for result handlers
namespace simd_result_handlers {
template <class C, bool with_id_map>
struct ResultHandlerCompare;
}
struct IDSelector;
/** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now.
*
* The codes are not stored sequentially but grouped in blocks of size bbs.
@@ -54,6 +62,14 @@ struct IndexFastScan : Index {
// (set when initialized by IndexPQ or IndexAQ)
const uint8_t* orig_codes = nullptr;
/** Initialize the fast scan index
*
* @param d dimensionality of vectors
* @param M number of subquantizers
* @param nbits number of bits per subquantizer
* @param metric distance metric to use
* @param bbs block size for SIMD processing
*/
void init_fastscan(
int d,
size_t M,
@@ -65,6 +81,15 @@ struct IndexFastScan : Index {
void reset() override;
/** Search for k nearest neighbors
*
* @param n number of query vectors
* @param x query vectors (n * d)
* @param k number of nearest neighbors to find
* @param distances output distances (n * k)
* @param labels output labels/indices (n * k)
* @param params optional search parameters
*/
void search(
idx_t n,
const float* x,
@@ -73,20 +98,70 @@ struct IndexFastScan : Index {
idx_t* labels,
const SearchParameters* params = nullptr) const override;
/** Add vectors to the index
*
* @param n number of vectors to add
* @param x vectors to add (n * d)
*/
void add(idx_t n, const float* x) override;
/** Compute codes for vectors
*
* @param codes output codes
* @param n number of vectors to encode
* @param x vectors to encode (n * d)
*/
virtual void compute_codes(uint8_t* codes, idx_t n, const float* x)
const = 0;
virtual void compute_float_LUT(float* lut, idx_t n, const float* x)
const = 0;
/** Compute floating-point lookup table for distance computation
*
* @param lut output lookup table
* @param n number of query vectors
* @param x query vectors (n * d)
* @param context processing context containing all processors
*/
virtual void compute_float_LUT(
float* lut,
idx_t n,
const float* x,
const FastScanDistancePostProcessing& context) const = 0;
/** Create a KNN handler for this index type
*
* This method can be overridden by derived classes to provide
* specialized handlers (e.g., RaBitQHeapHandler for RaBitQ indexes).
* Base implementation creates standard handlers based on k and impl.
*
* @param is_max whether to use CMax comparator (true) or CMin (false)
* @param impl implementation number
* @param n number of queries
* @param k number of neighbors to find
* @param ntotal total number of vectors in database
* @param distances output distances array
* @param labels output labels array
* @param sel optional ID selector
* @param query_offset query offset for batch processing
* @return pointer to created handler (never returns nullptr)
*/
virtual void* make_knn_handler(
bool is_max,
int impl,
idx_t n,
idx_t k,
size_t ntotal,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context) const;
// called by search function
void compute_quantized_LUT(
idx_t n,
const float* x,
uint8_t* lut,
float* normalizers) const;
float* normalizers,
const FastScanDistancePostProcessing& context) const;
template <bool is_max>
void search_dispatch_implem(
@@ -95,7 +170,7 @@ struct IndexFastScan : Index {
idx_t k,
float* distances,
idx_t* labels,
const NormTableScaler* scaler) const;
const FastScanDistancePostProcessing& context) const;
template <class Cfloat>
void search_implem_234(
@@ -104,7 +179,7 @@ struct IndexFastScan : Index {
idx_t k,
float* distances,
idx_t* labels,
const NormTableScaler* scaler) const;
const FastScanDistancePostProcessing& context) const;
template <class C>
void search_implem_12(
@@ -114,7 +189,7 @@ struct IndexFastScan : Index {
float* distances,
idx_t* labels,
int impl,
const NormTableScaler* scaler) const;
const FastScanDistancePostProcessing& context) const;
template <class C>
void search_implem_14(
@@ -124,14 +199,39 @@ struct IndexFastScan : Index {
float* distances,
idx_t* labels,
int impl,
const NormTableScaler* scaler) const;
const FastScanDistancePostProcessing& context) const;
/** Reconstruct a vector from its code
*
* @param key index of vector to reconstruct
* @param recons output reconstructed vector
*/
void reconstruct(idx_t key, float* recons) const override;
/** Remove vectors by ID selector
*
* @param sel selector defining which vectors to remove
* @return number of vectors removed
*/
size_t remove_ids(const IDSelector& sel) override;
/** Get the code packer for this index
*
* @return pointer to the code packer
*/
CodePacker* get_CodePacker() const;
/** Merge another index into this one
*
* @param otherIndex index to merge from
* @param add_id ID offset to add to merged vectors
*/
void merge_from(Index& otherIndex, idx_t add_id = 0) override;
/** Check if another index is compatible for merging
*
* @param otherIndex index to check compatibility with
*/
void check_compatible_for_merge(const Index& otherIndex) const override;
/// standalone codes interface (but the codes are flattened)
faiss/IndexIVFAdditiveQuantizerFastScan.h
+1
View File
@@ -12,6 +12,7 @@
#include <faiss/IndexIVFAdditiveQuantizer.h>
#include <faiss/IndexIVFFastScan.h>
#include <faiss/impl/AdditiveQuantizer.h>
#include <faiss/impl/FastScanDistancePostProcessing.h>
#include <faiss/impl/ProductAdditiveQuantizer.h>
#include <faiss/utils/AlignedTable.h>
faiss/IndexPQFastScan.cpp
+6 -2
View File
@@ -9,6 +9,7 @@
#include <memory>
#include <faiss/impl/FastScanDistancePostProcessing.h>
#include <faiss/impl/pq4_fast_scan.h>
#include <faiss/utils/utils.h>
@@ -53,8 +54,11 @@ void IndexPQFastScan::compute_codes(uint8_t* codes, idx_t n, const float* x)
pq.compute_codes(x, codes, n);
}
void IndexPQFastScan::compute_float_LUT(float* lut, idx_t n, const float* x)
const {
void IndexPQFastScan::compute_float_LUT(
float* lut,
idx_t n,
const float* x,
const FastScanDistancePostProcessing&) const {
if (metric_type == METRIC_L2) {
pq.compute_distance_tables(n, x, lut);
} else {
faiss/IndexPQFastScan.h
+5 -1
View File
@@ -45,7 +45,11 @@ struct IndexPQFastScan : IndexFastScan {
void compute_codes(uint8_t* codes, idx_t n, const float* x) const override;
void compute_float_LUT(float* lut, idx_t n, const float* x) const override;
void compute_float_LUT(
float* lut,
idx_t n,
const float* x,
const FastScanDistancePostProcessing& context) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
};
faiss/IndexRaBitQFastScan.cpp
+586
View File
@@ -0,0 +1,586 @@
/*
* 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/IndexRaBitQFastScan.h>
#include <faiss/impl/FastScanDistancePostProcessing.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/pq4_fast_scan.h>
#include <faiss/utils/utils.h>
#include <algorithm>
#include <cmath>
namespace faiss {
static inline size_t roundup(size_t a, size_t b) {
return (a + b - 1) / b * b;
}
IndexRaBitQFastScan::IndexRaBitQFastScan() = default;
IndexRaBitQFastScan::IndexRaBitQFastScan(idx_t d, MetricType metric, int bbs)
: rabitq(d, metric) {
// RaBitQ-specific validation
FAISS_THROW_IF_NOT_MSG(d > 0, "Dimension must be positive");
FAISS_THROW_IF_NOT_MSG(
metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT,
"RaBitQ FastScan only supports L2 and Inner Product metrics");
// RaBitQ uses 1 bit per dimension packed into 4-bit FastScan sub-quantizers
// Each FastScan sub-quantizer handles 4 RaBitQ dimensions
const size_t M_fastscan = (d + 3) / 4;
constexpr size_t nbits_fastscan = 4;
// init_fastscan will validate bbs % 32 == 0 and nbits_fastscan == 4
init_fastscan(static_cast<int>(d), M_fastscan, nbits_fastscan, metric, bbs);
// Override code_size to include space for factors after bit patterns
// RaBitQ stores 1 bit per dimension, requiring (d + 7) / 8 bytes
const size_t bit_pattern_size = (d + 7) / 8;
code_size = bit_pattern_size + sizeof(FactorsData);
// Set RaBitQ-specific parameters
qb = 8;
center.resize(d, 0.0f);
// Pre-allocate storage vectors for efficiency
factors_storage.clear();
}
IndexRaBitQFastScan::IndexRaBitQFastScan(const IndexRaBitQ& orig, int bbs)
: rabitq(orig.rabitq) {
// RaBitQ-specific validation
FAISS_THROW_IF_NOT_MSG(orig.d > 0, "Dimension must be positive");
FAISS_THROW_IF_NOT_MSG(
orig.metric_type == METRIC_L2 ||
orig.metric_type == METRIC_INNER_PRODUCT,
"RaBitQ FastScan only supports L2 and Inner Product metrics");
// RaBitQ uses 1 bit per dimension packed into 4-bit FastScan sub-quantizers
// Each FastScan sub-quantizer handles 4 RaBitQ dimensions
const size_t M_fastscan = (orig.d + 3) / 4;
constexpr size_t nbits_fastscan = 4;
// Initialize FastScan base with the original index's parameters
init_fastscan(
static_cast<int>(orig.d),
M_fastscan,
nbits_fastscan,
orig.metric_type,
bbs);
// Override code_size to include space for factors after bit patterns
// RaBitQ stores 1 bit per dimension, requiring (d + 7) / 8 bytes
const size_t bit_pattern_size = (orig.d + 7) / 8;
code_size = bit_pattern_size + sizeof(FactorsData);
// Copy properties from original index
ntotal = orig.ntotal;
ntotal2 = roundup(ntotal, bbs);
is_trained = orig.is_trained;
orig_codes = orig.codes.data();
qb = orig.qb;
centered = orig.centered;
center = orig.center;
// If the original index has data, extract factors and pack codes
if (ntotal > 0) {
// Allocate space for factors
factors_storage.resize(ntotal);
// Extract factors from original codes for each vector
const float* centroid_data = center.data();
// Use the original RaBitQ quantizer to decode and compute factors
std::vector<float> decoded_vectors(ntotal * orig.d);
orig.sa_decode(ntotal, orig.codes.data(), decoded_vectors.data());
for (idx_t i = 0; i < ntotal; i++) {
FactorsData& fac = factors_storage[i];
const float* x_row = decoded_vectors.data() + i * orig.d;
// Use shared utilities for computing factors
fac = rabitq_utils::compute_vector_factors(
x_row, orig.d, centroid_data, orig.metric_type);
}
// Convert RaBitQ bit format to FastScan 4-bit sub-quantizer format
// This follows the same pattern as IndexPQFastScan constructor
AlignedTable<uint8_t> fastscan_codes(ntotal * code_size);
memset(fastscan_codes.get(), 0, ntotal * code_size);
// Convert from RaBitQ 1-bit-per-dimension to FastScan
// 4-bit-per-sub-quantizer
for (idx_t i = 0; i < ntotal; i++) {
const uint8_t* orig_code = orig.codes.data() + i * orig.code_size;
uint8_t* fs_code = fastscan_codes.get() + i * code_size;
// Convert each dimension's bit (same logic as compute_codes)
for (size_t j = 0; j < orig.d; j++) {
// Extract bit from original RaBitQ format
const size_t orig_byte_idx = j / 8;
const size_t orig_bit_offset = j % 8;
const bool bit_value =
(orig_code[orig_byte_idx] >> orig_bit_offset) & 1;
// Use RaBitQUtils for consistent bit setting
if (bit_value) {
rabitq_utils::set_bit_fastscan(fs_code, j);
}
}
}
// Pack the converted codes using pq4_pack_codes with custom stride
codes.resize(ntotal2 * M2 / 2);
pq4_pack_codes(
fastscan_codes.get(),
ntotal,
M,
ntotal2,
bbs,
M2,
codes.get(),
code_size);
}
}
void IndexRaBitQFastScan::train(idx_t n, const float* x) {
// compute a centroid
std::vector<float> centroid(d, 0);
for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
for (size_t j = 0; j < d; j++) {
centroid[j] += x[i * d + j];
}
}
if (n != 0) {
for (size_t j = 0; j < d; j++) {
centroid[j] /= (float)n;
}
}
center = std::move(centroid);
rabitq.train(n, x);
is_trained = true;
}
void IndexRaBitQFastScan::add(idx_t n, const float* x) {
FAISS_THROW_IF_NOT(is_trained);
// Handle blocking to avoid excessive allocations
constexpr idx_t bs = 65536;
if (n > bs) {
for (idx_t i0 = 0; i0 < n; i0 += bs) {
idx_t i1 = std::min(n, i0 + bs);
if (verbose) {
printf("IndexRaBitQFastScan::add %zd/%zd\n",
size_t(i1),
size_t(n));
}
add(i1 - i0, x + i0 * d);
}
return;
}
InterruptCallback::check();
// Create codes with embedded factors using our compute_codes
AlignedTable<uint8_t> tmp_codes(n * code_size);
compute_codes(tmp_codes.get(), n, x);
// Extract and store factors from embedded codes for handler access
const size_t bit_pattern_size = (d + 7) / 8;
factors_storage.resize(ntotal + n);
for (idx_t i = 0; i < n; i++) {
const uint8_t* code = tmp_codes.get() + i * code_size;
const uint8_t* factors_ptr = code + bit_pattern_size;
const FactorsData& embedded_factors =
*reinterpret_cast<const FactorsData*>(factors_ptr);
factors_storage[ntotal + i] = embedded_factors;
}
// Resize main storage (same logic as parent)
ntotal2 = roundup(ntotal + n, bbs);
size_t new_size = ntotal2 * M2 / 2; // assume nbits = 4
size_t old_size = codes.size();
if (new_size > old_size) {
codes.resize(new_size);
memset(codes.get() + old_size, 0, new_size - old_size);
}
// Use our custom packing function with correct stride
pq4_pack_codes_range(
tmp_codes.get(),
M, // Number of sub-quantizers (bit patterns only)
ntotal,
ntotal + n, // Range to pack
bbs,
M2, // Block parameters
codes.get(), // Output
code_size); // CUSTOM STRIDE: includes factor space
ntotal += n;
}
void IndexRaBitQFastScan::compute_codes(uint8_t* codes, idx_t n, const float* x)
const {
FAISS_ASSERT(codes != nullptr);
FAISS_ASSERT(x != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
if (n == 0) {
return;
}
// Hoist loop-invariant computations
const float* centroid_data = center.data();
const size_t bit_pattern_size = (d + 7) / 8;
memset(codes, 0, n * code_size);
#pragma omp parallel for if (n > 1000)
for (int64_t i = 0; i < n; i++) {
uint8_t* const code = codes + i * code_size;
const float* const x_row = x + i * d;
// Pack bits directly into FastScan format
for (size_t j = 0; j < d; j++) {
const float x_val = x_row[j];
const float centroid_val = centroid_data ? centroid_data[j] : 0.0f;
const float or_minus_c = x_val - centroid_val;
const bool xb = (or_minus_c > 0.0f);
if (xb) {
rabitq_utils::set_bit_fastscan(code, j);
}
}
// Calculate and append factors after the bit data
FactorsData factors = rabitq_utils::compute_vector_factors(
x_row, d, centroid_data, metric_type);
// Append factors at the end of the code
uint8_t* factors_ptr = code + bit_pattern_size;
*reinterpret_cast<FactorsData*>(factors_ptr) = factors;
}
}
void IndexRaBitQFastScan::compute_float_LUT(
float* lut,
idx_t n,
const float* x,
const FastScanDistancePostProcessing& context) const {
FAISS_THROW_IF_NOT(is_trained);
// Pre-allocate working buffers to avoid repeated allocations
std::vector<float> rotated_q(d);
std::vector<uint8_t> rotated_qq(d);
// Compute lookup tables for FastScan SIMD operations
// For each query vector, computes distance contributions for all
// possible 4-bit codes per sub-quantizer. Also computes and stores
// query factors for distance reconstruction.
for (idx_t i = 0; i < n; i++) {
const float* query = x + i * d;
// Compute query factors and store in array if available
rabitq_utils::QueryFactorsData query_factors_data =
rabitq_utils::compute_query_factors(
query,
d,
center.data(),
qb,
centered,
metric_type,
rotated_q,
rotated_qq);
// Store query factors in context array if provided
if (context.query_factors) {
context.query_factors[i] = query_factors_data;
}
// Create lookup table storing distance contributions for all possible
// 4-bit codes per sub-quantizer for FastScan SIMD operations
float* query_lut = lut + i * M * 16;
if (centered) {
// For centered mode, we use the signed odd integer quantization
// scheme.
// Formula:
// int_dot = ((1 << qb) - 1) * d - 2 * xor_dot_product
// We precompute the XOR contribution for each
// sub-quantizer
const float max_code_value = (1 << qb) - 1;
for (size_t m = 0; m < M; m++) {
const size_t dim_start = m * 4;
for (int code_val = 0; code_val < 16; code_val++) {
float xor_contribution = 0.0f;
// Process 4 bits per sub-quantizer
for (size_t dim_offset = 0; dim_offset < 4; dim_offset++) {
const size_t dim_idx = dim_start + dim_offset;
if (dim_idx < d) {
const bool db_bit = (code_val >> dim_offset) & 1;
const float query_value = rotated_qq[dim_idx];
// XOR contribution:
// If db_bit == 0: XOR result = query_value
// If db_bit == 1: XOR result = (2^qb - 1) -
// query_value
xor_contribution += db_bit
? (max_code_value - query_value)
: query_value;
}
}
// Store the XOR contribution (will be scaled by -2 *
// int_dot_scale during distance computation)
query_lut[m * 16 + code_val] = xor_contribution;
}
}
} else {
// For non-centered quantization, use traditional AND dot
// product Compute lookup table entries by processing popcount
// and inner product together
for (size_t m = 0; m < M; m++) {
const size_t dim_start = m * 4;
for (int code_val = 0; code_val < 16; code_val++) {
float inner_product = 0.0f;
int popcount = 0;
// Process 4 bits per sub-quantizer
for (size_t dim_offset = 0; dim_offset < 4; dim_offset++) {
const size_t dim_idx = dim_start + dim_offset;
if (dim_idx < d && ((code_val >> dim_offset) & 1)) {
inner_product += rotated_qq[dim_idx];
popcount++;
}
}
// Store pre-computed distance contribution
query_lut[m * 16 + code_val] =
query_factors_data.c1 * inner_product +
query_factors_data.c2 * popcount;
}
}
}
}
}
void IndexRaBitQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
const {
const float* centroid_in =
(center.data() == nullptr) ? nullptr : center.data();
const uint8_t* codes = bytes;
FAISS_ASSERT(codes != nullptr);
FAISS_ASSERT(x != nullptr);
const float inv_d_sqrt = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
const size_t bit_pattern_size = (d + 7) / 8;
#pragma omp parallel for if (n > 1000)
for (int64_t i = 0; i < n; i++) {
// Access code using correct FastScan format
const uint8_t* code = codes + i * code_size;
// Extract factors directly from embedded codes
const uint8_t* factors_ptr = code + bit_pattern_size;
const FactorsData& fac =
*reinterpret_cast<const FactorsData*>(factors_ptr);
for (size_t j = 0; j < d; j++) {
// Use RaBitQUtils for consistent bit extraction
bool bit_value = rabitq_utils::extract_bit_fastscan(code, j);
float bit = bit_value ? 1.0f : 0.0f;
// Compute the output using RaBitQ reconstruction formula
x[i * d + j] = (bit - 0.5f) * fac.dp_multiplier * 2 * inv_d_sqrt +
((centroid_in == nullptr) ? 0 : centroid_in[j]);
}
}
}
void IndexRaBitQFastScan::search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels,
const SearchParameters* params) const {
FAISS_THROW_IF_NOT_MSG(
!params, "search params not supported for this index");
// Create query factors array on stack - memory managed by caller
std::vector<rabitq_utils::QueryFactorsData> query_factors_storage(n);
// Use the faster search_dispatch_implem flow from IndexFastScan
// Pass the query factors array - factors will be computed during LUT
// computation
FastScanDistancePostProcessing context;
context.query_factors = query_factors_storage.data();
if (metric_type == METRIC_L2) {
search_dispatch_implem<true>(n, x, k, distances, labels, context);
} else {
search_dispatch_implem<false>(n, x, k, distances, labels, context);
}
}
// Template implementations for RaBitQHeapHandler
template <class C, bool with_id_map>
RaBitQHeapHandler<C, with_id_map>::RaBitQHeapHandler(
const IndexRaBitQFastScan* index,
size_t nq_val,
size_t k_val,
float* distances,
int64_t* labels,
const IDSelector* sel_in,
const FastScanDistancePostProcessing& ctx)
: RHC(nq_val, index->ntotal, sel_in),
rabitq_index(index),
heap_distances(distances),
heap_labels(labels),
nq(nq_val),
k(k_val),
context(ctx) {
// Initialize heaps for all queries in constructor
// This allows us to support direct normalizer assignment
#pragma omp parallel for if (nq > 100)
for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
float* heap_dis = heap_distances + q * k;
int64_t* heap_ids = heap_labels + q * k;
heap_heapify<Cfloat>(k, heap_dis, heap_ids);
}
}
template <class C, bool with_id_map>
void RaBitQHeapHandler<C, with_id_map>::handle(
size_t q,
size_t b,
simd16uint16 d0,
simd16uint16 d1) {
ALIGNED(32) uint16_t d32tab[32];
d0.store(d32tab);
d1.store(d32tab + 16);
// Get heap pointers and query factors (computed once per batch)
float* const heap_dis = heap_distances + q * k;
int64_t* const heap_ids = heap_labels + q * k;
// Access query factors from query_factors pointer
rabitq_utils::QueryFactorsData query_factors_data = {};
if (context.query_factors) {
query_factors_data = context.query_factors[q];
}
// Compute normalizers once per batch
const float one_a = normalizers ? (1.0f / normalizers[2 * q]) : 1.0f;
const float bias = normalizers ? normalizers[2 * q + 1] : 0.0f;
// Compute loop bounds to avoid redundant bounds checking
const size_t base_db_idx = this->j0 + b * 32;
const size_t max_vectors = (base_db_idx < rabitq_index->ntotal)
? std::min(size_t(32), rabitq_index->ntotal - base_db_idx)
: 0;
// Process distances in batch
for (size_t i = 0; i < max_vectors; i++) {
const size_t db_idx = base_db_idx + i;
// Normalize distance from LUT lookup
const float normalized_distance = d32tab[i] * one_a + bias;
// Access factors from storage (populated from embedded codes during
// add())
const auto& db_factors = rabitq_index->factors_storage[db_idx];
float adjusted_distance;
if (rabitq_index->centered) {
// For centered mode: normalized_distance contains the raw XOR
// contribution. Apply the signed odd integer quantization formula:
// int_dot = ((1 << qb) - 1) * d - 2 * xor_dot_product
int64_t int_dot = ((1 << rabitq_index->qb) - 1) * rabitq_index->d;
int_dot -= 2 * static_cast<int64_t>(normalized_distance);
adjusted_distance = query_factors_data.qr_to_c_L2sqr +
db_factors.or_minus_c_l2sqr -
2 * db_factors.dp_multiplier * int_dot *
query_factors_data.int_dot_scale;
} else {
// For non-centered quantization: use traditional formula
float final_dot = normalized_distance - query_factors_data.c34;
adjusted_distance = db_factors.or_minus_c_l2sqr +
query_factors_data.qr_to_c_L2sqr -
2 * db_factors.dp_multiplier * final_dot;
}
// Apply inner product correction if needed
if (query_factors_data.qr_norm_L2sqr != 0.0f) {
adjusted_distance = -0.5f *
(adjusted_distance - query_factors_data.qr_norm_L2sqr);
}
// Add to heap if better than current worst
if (Cfloat::cmp(heap_dis[0], adjusted_distance)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, adjusted_distance, db_idx);
}
}
}
template <class C, bool with_id_map>
void RaBitQHeapHandler<C, with_id_map>::begin(const float* norms) {
normalizers = norms;
// Heap initialization is now done in constructor
}
template <class C, bool with_id_map>
void RaBitQHeapHandler<C, with_id_map>::end() {
// Reorder final results
#pragma omp parallel for if (nq > 100)
for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
float* heap_dis = heap_distances + q * k;
int64_t* heap_ids = heap_labels + q * k;
heap_reorder<Cfloat>(k, heap_dis, heap_ids);
}
}
// Implementation of virtual make_knn_handler method
void* IndexRaBitQFastScan::make_knn_handler(
bool is_max,
int /*impl*/,
idx_t n,
idx_t k,
size_t /*ntotal*/,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context) const {
if (is_max) {
return new RaBitQHeapHandler<CMax<uint16_t, int>, false>(
this, n, k, distances, labels, sel, context);
} else {
return new RaBitQHeapHandler<CMin<uint16_t, int>, false>(
this, n, k, distances, labels, sel, context);
}
}
// Explicit template instantiations for the required comparator types
template struct RaBitQHeapHandler<CMin<uint16_t, int>, false>;
template struct RaBitQHeapHandler<CMax<uint16_t, int>, false>;
template struct RaBitQHeapHandler<CMin<uint16_t, int>, true>;
template struct RaBitQHeapHandler<CMax<uint16_t, int>, true>;
} // namespace faiss
faiss/IndexRaBitQFastScan.h
+149
View File
@@ -0,0 +1,149 @@
/*
* 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.
*/
#pragma once
#include <vector>
#include <faiss/IndexFastScan.h>
#include <faiss/IndexRaBitQ.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/RaBitQuantizer.h>
#include <faiss/impl/simd_result_handlers.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/simdlib.h>
namespace faiss {
// Import shared utilities from RaBitQUtils
using rabitq_utils::FactorsData;
using rabitq_utils::QueryFactorsData;
/** Fast-scan version of RaBitQ index that processes 32 database vectors at a
* time using SIMD operations. Similar to IndexPQFastScan but adapted for
* RaBitQ's bit-level quantization with factors.
*
* The key differences from IndexRaBitQ:
* - Processes vectors in batches of 32
* - Uses 4-bit groupings for SIMD optimization (4 dimensions per 4-bit unit)
* - Separates factors from quantized bits for efficient processing
* - Leverages existing PQ4 FastScan infrastructure where possible
*/
struct IndexRaBitQFastScan : IndexFastScan {
/// RaBitQ quantizer for encoding/decoding
RaBitQuantizer rabitq;
/// Center of all points (same as IndexRaBitQ)
std::vector<float> center;
/// Extracted factors storage for batch processing
/// Size: ntotal, stores factors separately from packed codes
std::vector<FactorsData> factors_storage;
/// Default number of bits to quantize a query with
uint8_t qb = 8;
// quantize the query with a zero-centered scalar quantizer.
bool centered = false;
IndexRaBitQFastScan();
explicit IndexRaBitQFastScan(
idx_t d,
MetricType metric = METRIC_L2,
int bbs = 32);
/// build from an existing IndexRaBitQ
explicit IndexRaBitQFastScan(const IndexRaBitQ& orig, int bbs = 32);
void train(idx_t n, const float* x) override;
void add(idx_t n, const float* x) override;
void compute_codes(uint8_t* codes, idx_t n, const float* x) const override;
void compute_float_LUT(
float* lut,
idx_t n,
const float* x,
const FastScanDistancePostProcessing& context) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
void search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels,
const SearchParameters* params = nullptr) const override;
/// Override to create RaBitQ-specific handlers
void* make_knn_handler(
bool is_max,
int /*impl*/,
idx_t n,
idx_t k,
size_t /*ntotal*/,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context) const override;
};
/** SIMD result handler for RaBitQ FastScan that applies distance corrections
* and maintains heaps directly during SIMD operations.
*
* This handler processes batches of 32 distance computations from SIMD kernels,
* applies RaBitQ-specific adjustments (factors and normalizers), and
* immediately updates result heaps without intermediate storage. This
* eliminates the need for post-processing and provides significant memory and
* performance benefits.
*
* Key optimizations:
* - Direct heap integration (no intermediate result storage)
* - Batch-level computation of normalizers and query factors
* - Preserves exact mathematical equivalence to original RaBitQ distances
* @tparam C Comparator type (CMin/CMax) for heap operations
* @tparam with_id_map Whether to use id mapping (similar to HeapHandler)
*/
template <class C, bool with_id_map = false>
struct RaBitQHeapHandler
: simd_result_handlers::ResultHandlerCompare<C, with_id_map> {
using RHC = simd_result_handlers::ResultHandlerCompare<C, with_id_map>;
using RHC::normalizers;
const IndexRaBitQFastScan* rabitq_index;
float* heap_distances; // [nq * k]
int64_t* heap_labels; // [nq * k]
const size_t nq, k;
const FastScanDistancePostProcessing&
context; // Processing context with query offset
// Use float-based comparator for heap operations
using Cfloat = typename std::conditional<
C::is_max,
CMax<float, int64_t>,
CMin<float, int64_t>>::type;
RaBitQHeapHandler(
const IndexRaBitQFastScan* index,
size_t nq_val,
size_t k_val,
float* distances,
int64_t* labels,
const IDSelector* sel_in,
const FastScanDistancePostProcessing& context);
void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final;
void begin(const float* norms);
void end();
};
} // namespace faiss
faiss/impl/FastScanDistancePostProcessing.h
+45
View File
@@ -0,0 +1,45 @@
/*
* 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.
*/
#pragma once
namespace faiss {
// Forward declarations
struct NormTableScaler;
namespace rabitq_utils {
struct QueryFactorsData;
}
/**
* Simple context object that holds processors for FastScan operations.
* */
struct FastScanDistancePostProcessing {
/// Norm scaling processor for Additive Quantizers (nullptr if not needed)
const NormTableScaler* norm_scaler = nullptr;
/// Query factors data pointer for RaBitQ (nullptr if not needed)
/// This pointer should point to the beginning of the relevant
/// QueryFactorsData subset for this context.
rabitq_utils::QueryFactorsData* query_factors = nullptr;
/// Default constructor - no processing
FastScanDistancePostProcessing() = default;
/// Check if norm scaling is enabled
bool has_norm_scaling() const {
return norm_scaler != nullptr;
}
/// Check if query factors processing is enabled
bool has_query_processing() const {
return query_factors != nullptr;
}
};
} // namespace faiss
faiss/impl/RaBitQUtils.cpp
+246
View File
@@ -0,0 +1,246 @@
/*
* 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/impl/RaBitQUtils.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/distances.h>
#include <algorithm>
#include <cmath>
#include <limits>
namespace faiss {
namespace rabitq_utils {
// 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.
};
void compute_vector_intermediate_values(
const float* x,
size_t d,
const float* centroid,
float& norm_L2sqr,
float& or_L2sqr,
float& dp_oO) {
norm_L2sqr = 0.0f;
or_L2sqr = 0.0f;
dp_oO = 0.0f;
for (size_t j = 0; j < d; j++) {
const float x_val = x[j];
const float centroid_val = (centroid != nullptr) ? centroid[j] : 0.0f;
const float or_minus_c = x_val - centroid_val;
const float or_minus_c_sq = or_minus_c * or_minus_c;
norm_L2sqr += or_minus_c_sq;
or_L2sqr += x_val * x_val;
const bool xb = (or_minus_c > 0.0f);
dp_oO += xb ? or_minus_c : -or_minus_c;
}
}
FactorsData compute_factors_from_intermediates(
float norm_L2sqr,
float or_L2sqr,
float dp_oO,
size_t d,
MetricType metric_type) {
constexpr float epsilon = std::numeric_limits<float>::epsilon();
const float inv_d_sqrt =
(d == 0) ? 1.0f : (1.0f / std::sqrt(static_cast<float>(d)));
const float sqrt_norm_L2 = std::sqrt(norm_L2sqr);
const float inv_norm_L2 =
(norm_L2sqr < epsilon) ? 1.0f : (1.0f / sqrt_norm_L2);
const float normalized_dp = dp_oO * inv_norm_L2 * inv_d_sqrt;
const float inv_dp_oO =
(std::abs(normalized_dp) < epsilon) ? 1.0f : (1.0f / normalized_dp);
FactorsData factors;
factors.or_minus_c_l2sqr = (metric_type == MetricType::METRIC_INNER_PRODUCT)
? (norm_L2sqr - or_L2sqr)
: norm_L2sqr;
factors.dp_multiplier = inv_dp_oO * sqrt_norm_L2;
return factors;
}
FactorsData compute_vector_factors(
const float* x,
size_t d,
const float* centroid,
MetricType metric_type) {
float norm_L2sqr, or_L2sqr, dp_oO;
compute_vector_intermediate_values(
x, d, centroid, norm_L2sqr, or_L2sqr, dp_oO);
return compute_factors_from_intermediates(
norm_L2sqr, or_L2sqr, dp_oO, d, metric_type);
}
QueryFactorsData compute_query_factors(
const float* query,
size_t d,
const float* centroid,
uint8_t qb,
bool centered,
MetricType metric_type,
std::vector<float>& rotated_q,
std::vector<uint8_t>& rotated_qq) {
FAISS_THROW_IF_NOT(qb <= 8);
FAISS_THROW_IF_NOT(qb > 0);
QueryFactorsData query_factors;
// Compute distance from query to centroid
if (centroid != nullptr) {
query_factors.qr_to_c_L2sqr = fvec_L2sqr(query, centroid, d);
} else {
query_factors.qr_to_c_L2sqr = fvec_norm_L2sqr(query, d);
}
// Rotate the query (subtract centroid)
rotated_q.resize(d);
for (size_t i = 0; i < d; i++) {
if (i < rotated_q.size()) {
rotated_q[i] =
query[i] - ((centroid == nullptr) ? 0.0f : centroid[i]);
}
}
const float inv_d_sqrt =
(d == 0) ? 1.0f : (1.0f / std::sqrt(static_cast<float>(d)));
// Compute quantization range
float v_min = std::numeric_limits<float>::max();
float v_max = std::numeric_limits<float>::lowest();
if (centered) {
float z_max = Z_MAX_BY_QB[qb - 1];
float v_radius = z_max * std::sqrt(query_factors.qr_to_c_L2sqr / d);
v_min = -v_radius;
v_max = v_radius;
} else {
// Only compute min/max if we have dimensions to process
if (d > 0 && !rotated_q.empty()) {
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);
}
} else {
// For empty dimensions, use default range
v_min = 0.0f;
v_max = 1.0f;
}
}
// Quantize the query
const uint8_t max_code = (1 << qb) - 1;
const float delta = (v_max - v_min) / max_code;
const float inv_delta = 1.0f / delta;
rotated_qq.resize(d);
size_t sum_qq = 0;
int64_t sum2_signed_odd_int = 0;
// Process arrays - throw error if they are unexpectedly empty
if (d > 0 && !rotated_q.empty() && !rotated_qq.empty()) {
for (size_t i = 0; i < d; i++) {
const float v_q = rotated_q[i];
// Non-randomized scalar quantization
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;
}
}
} else {
FAISS_THROW_MSG(
"Arrays unexpectedly empty when d=" + std::to_string(d) +
"or d is incorrectly set");
}
// Compute query factors
query_factors.c1 = 2 * delta * inv_d_sqrt;
query_factors.c2 = 2 * v_min * inv_d_sqrt;
query_factors.c34 = inv_d_sqrt * (delta * sum_qq + d * v_min);
if (centered) {
query_factors.int_dot_scale = std::sqrt(
query_factors.qr_to_c_L2sqr / (sum2_signed_odd_int * d));
} else {
query_factors.int_dot_scale = 1.0f;
}
// Compute query norm for inner product metric
query_factors.qr_norm_L2sqr = 0.0f;
if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
query_factors.qr_norm_L2sqr = fvec_norm_L2sqr(query, d);
}
return query_factors;
}
bool extract_bit_standard(const uint8_t* code, size_t bit_index) {
const size_t byte_idx = bit_index / 8;
const size_t bit_offset = bit_index % 8;
return (code[byte_idx] >> bit_offset) & 1;
}
bool extract_bit_fastscan(const uint8_t* code, size_t bit_index) {
const size_t m = bit_index / 4; // Sub-quantizer index
const size_t dim_offset =
bit_index % 4; // Bit position within sub-quantizer
const size_t byte_idx = m / 2; // Byte index (2 sub-quantizers per byte)
const uint8_t bit_mask = static_cast<uint8_t>(1 << dim_offset);
if (m % 2 == 0) {
// Lower 4 bits of byte
return (code[byte_idx] & bit_mask) != 0;
} else {
// Upper 4 bits of byte (shifted)
return (code[byte_idx] & (bit_mask << 4)) != 0;
}
}
void set_bit_standard(uint8_t* code, size_t bit_index) {
const size_t byte_idx = bit_index / 8;
const size_t bit_offset = bit_index % 8;
code[byte_idx] |= (1 << bit_offset);
}
void set_bit_fastscan(uint8_t* code, size_t bit_index) {
const size_t m = bit_index / 4;
const size_t dim_offset = bit_index % 4;
const uint8_t bit_mask = static_cast<uint8_t>(1 << dim_offset);
const size_t byte_idx = m / 2;
if (m % 2 == 0) {
code[byte_idx] |= bit_mask;
} else {
code[byte_idx] |= (bit_mask << 4);
}
}
} // namespace rabitq_utils
} // namespace faiss
faiss/impl/RaBitQUtils.h
+153
View File
@@ -0,0 +1,153 @@
/*
* 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.
*/
#pragma once
#include <faiss/MetricType.h>
#include <faiss/impl/platform_macros.h>
#include <cstddef>
#include <cstdint>
#include <vector>
namespace faiss {
namespace rabitq_utils {
/** Factors computed per database vector for RaBitQ distance computation.
* These can be stored either embedded in codes (IndexRaBitQ) or separately
* (IndexRaBitQFastScan).
*/
struct FactorsData {
// ||or - c||^2 - ((metric==IP) ? ||or||^2 : 0)
float or_minus_c_l2sqr = 0;
float dp_multiplier = 0;
};
/** Query-specific factors computed during search for RaBitQ distance
* computation. Used by both IndexRaBitQ and IndexRaBitQFastScan
* implementations.
*/
struct QueryFactorsData {
float c1 = 0;
float c2 = 0;
float c34 = 0;
float qr_to_c_L2sqr = 0;
float qr_norm_L2sqr = 0;
float int_dot_scale = 1;
};
/** Ideal quantizer radii for quantizers of 1..8 bits, optimized to minimize
* L2 reconstruction error. Shared between all RaBitQ implementations.
*/
FAISS_API extern const float Z_MAX_BY_QB[8];
/** Compute factors for a single database vector using RaBitQ algorithm.
* This function consolidates the mathematical logic that was duplicated
* between IndexRaBitQ and IndexRaBitQFastScan.
*
* @param x input vector (d dimensions)
* @param d dimensionality
* @param centroid database centroid (nullptr if not used)
* @param metric_type distance metric (L2 or Inner Product)
* @return computed factors for distance computation
*/
FactorsData compute_vector_factors(
const float* x,
size_t d,
const float* centroid,
MetricType metric_type);
/** Compute intermediate values needed for vector factor computation.
* Separated out to allow different bit packing strategies while sharing
* the core mathematical computation.
*
* @param x input vector (d dimensions)
* @param d dimensionality
* @param centroid database centroid (nullptr if not used)
* @param norm_L2sqr output: ||or - c||^2
* @param or_L2sqr output: ||or||^2
* @param dp_oO output: sum of |or_i - c_i| (absolute deviations)
*/
void compute_vector_intermediate_values(
const float* x,
size_t d,
const float* centroid,
float& norm_L2sqr,
float& or_L2sqr,
float& dp_oO);
/** Compute final factors from intermediate values.
* @param norm_L2sqr ||or - c||^2
* @param or_L2sqr ||or||^2
* @param dp_oO sum of |or_i - c_i|
* @param d dimensionality
* @param metric_type distance metric
* @return computed factors
*/
FactorsData compute_factors_from_intermediates(
float norm_L2sqr,
float or_L2sqr,
float dp_oO,
size_t d,
MetricType metric_type);
/** Compute query factors for RaBitQ distance computation.
* This consolidates the query processing logic shared between implementations.
*
* @param query query vector (d dimensions)
* @param d dimensionality
* @param centroid database centroid (nullptr if not used)
* @param qb number of quantization bits (1-8)
* @param centered whether to use centered quantization
* @param metric_type distance metric
* @param rotated_q output: query - centroid
* @param rotated_qq output: quantized query values
* @return computed query factors
*/
QueryFactorsData compute_query_factors(
const float* query,
size_t d,
const float* centroid,
uint8_t qb,
bool centered,
MetricType metric_type,
std::vector<float>& rotated_q,
std::vector<uint8_t>& rotated_qq);
/** Extract bit value from RaBitQ code in standard format.
* Used by IndexRaBitQ which stores bits sequentially.
*
* @param code RaBitQ code data
* @param bit_index which bit to extract (0 to d-1)
* @return bit value (true/false)
*/
bool extract_bit_standard(const uint8_t* code, size_t bit_index);
/** Extract bit value from FastScan code format.
* Used by IndexRaBitQFastScan which packs bits into 4-bit sub-quantizers.
*
* @param code FastScan code data
* @param bit_index which bit to extract (0 to d-1)
* @return bit value (true/false)
*/
bool extract_bit_fastscan(const uint8_t* code, size_t bit_index);
/** Set bit value in standard RaBitQ code format.
* @param code RaBitQ code data to modify
* @param bit_index which bit to set (0 to d-1)
*/
void set_bit_standard(uint8_t* code, size_t bit_index);
/** Set bit value in FastScan code format.
* @param code FastScan code data to modify
* @param bit_index which bit to set (0 to d-1)
*/
void set_bit_fastscan(uint8_t* code, size_t bit_index);
} // namespace rabitq_utils
} // namespace faiss
faiss/impl/RaBitQuantizer.cpp
+25 -143
View File
@@ -8,33 +8,20 @@
#include <faiss/impl/RaBitQuantizer.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/rabitq_simd.h>
#include <algorithm>
#include <cmath>
#include <cstring>
#include <limits>
#include <memory>
#include <vector>
namespace faiss {
struct FactorsData {
// ||or - c||^2 - ((metric==IP) ? ||or||^2 : 0)
float or_minus_c_l2sqr = 0;
float dp_multiplier = 0;
};
struct QueryFactorsData {
float c1 = 0;
float c2 = 0;
float c34 = 0;
float qr_to_c_L2sqr = 0;
float qr_norm_L2sqr = 0;
float int_dot_scale = 1;
};
// Import shared utilities from RaBitQUtils
using rabitq_utils::FactorsData;
using rabitq_utils::QueryFactorsData;
static size_t get_code_size(const size_t d) {
return (d + 7) / 8 + sizeof(FactorsData);
@@ -67,19 +54,9 @@ void RaBitQuantizer::compute_codes_core(
return;
}
// compute some helper constants
const float inv_d_sqrt = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
// compute codes
#pragma omp parallel for if (n > 1000)
for (int64_t i = 0; i < n; i++) {
// ||or - c||^2
float norm_L2sqr = 0;
// ||or||^2, which is equal to ||P(or)||^2 and ||P^(-1)(or)||^2
float or_L2sqr = 0;
// dot product
float dp_oO = 0;
// the code
uint8_t* code = codes + i * code_size;
FactorsData* fac = reinterpret_cast<FactorsData*>(code + (d + 7) / 8);
@@ -89,46 +66,25 @@ void RaBitQuantizer::compute_codes_core(
memset(code, 0, code_size);
}
const float* x_row = x + i * d;
// Use shared utilities for computing factors
*fac = rabitq_utils::compute_vector_factors(
x_row, d, centroid_in, metric_type);
// Pack bits into standard RaBitQ format
for (size_t j = 0; j < d; j++) {
const float or_minus_c = x[i * d + j] -
((centroid_in == nullptr) ? 0 : centroid_in[j]);
norm_L2sqr += or_minus_c * or_minus_c;
or_L2sqr += x[i * d + j] * x[i * d + j];
const bool xb = (or_minus_c > 0);
dp_oO += xb ? or_minus_c : (-or_minus_c);
const float x_val = x_row[j];
const float centroid_val =
(centroid_in == nullptr) ? 0.0f : centroid_in[j];
const float or_minus_c = x_val - centroid_val;
const bool xb = (or_minus_c > 0.0f);
// store the output data
if (code != nullptr) {
if (xb) {
// enable a particular bit
code[j / 8] |= (1 << (j % 8));
}
if (code != nullptr && xb) {
rabitq_utils::set_bit_standard(code, j);
}
}
// compute factors
// compute the inverse norm
const float inv_norm_L2 =
(std::abs(norm_L2sqr) < std::numeric_limits<float>::epsilon())
? 1.0f
: (1.0f / std::sqrt(norm_L2sqr));
dp_oO *= inv_norm_L2;
dp_oO *= inv_d_sqrt;
const float inv_dp_oO =
(std::abs(dp_oO) < std::numeric_limits<float>::epsilon())
? 1.0f
: (1.0f / dp_oO);
fac->or_minus_c_l2sqr = norm_L2sqr;
if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
fac->or_minus_c_l2sqr -= or_L2sqr;
}
fac->dp_multiplier = inv_dp_oO * std::sqrt(norm_L2sqr);
}
}
@@ -383,22 +339,8 @@ 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
// Use shared constant from RaBitQUtils
using rabitq_utils::Z_MAX_BY_QB;
void RaBitDistanceComputerQ::set_query(const float* x) {
FAISS_ASSERT(x != nullptr);
@@ -408,61 +350,12 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
FAISS_THROW_IF_NOT(qb <= 8);
FAISS_THROW_IF_NOT(qb > 0);
// 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);
}
// Use shared utilities for core query factor computation
std::vector<float> rotated_q;
query_fac = rabitq_utils::compute_query_factors(
x, d, centroid, qb, centered, metric_type, rotated_q, rotated_qq);
// allocate space
rotated_qq.resize(d);
// rotate the query
std::vector<float> rotated_q(d);
for (size_t i = 0; i < d; i++) {
rotated_q[i] = x[i] - ((centroid == nullptr) ? 0 : centroid[i]);
}
// compute some numbers
const float inv_d = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
// quantize the query. compute min and max
float v_min = std::numeric_limits<float>::max();
float v_max = std::numeric_limits<float>::lowest();
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 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 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
// Rearrange the query vector for SIMD operations (RaBitQuantizer-specific)
popcount_aligned_dim = ((d + 7) / 8) * 8;
size_t offset = (d + 7) / 8;
@@ -476,17 +369,6 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
bit ? (1 << (idim % 8)) : 0;
}
}
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
query_fac.qr_norm_L2sqr = fvec_norm_L2sqr(x, d);
}
}
FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
faiss/impl/index_read.cpp
+22
View File
@@ -43,6 +43,7 @@
#include <faiss/IndexPQFastScan.h>
#include <faiss/IndexPreTransform.h>
#include <faiss/IndexRaBitQ.h>
#include <faiss/IndexRaBitQFastScan.h>
#include <faiss/IndexRefine.h>
#include <faiss/IndexRowwiseMinMax.h>
#include <faiss/IndexScalarQuantizer.h>
@@ -1224,6 +1225,27 @@ Index* read_index(IOReader* f, int io_flags) {
imm->own_fields = true;
idx = imm;
} else if (h == fourcc("Irfs")) {
IndexRaBitQFastScan* idxqfs = new IndexRaBitQFastScan();
read_index_header(idxqfs, f);
read_RaBitQuantizer(&idxqfs->rabitq, f);
READVECTOR(idxqfs->center);
READ1(idxqfs->qb);
READVECTOR(idxqfs->factors_storage);
READ1(idxqfs->bbs);
READ1(idxqfs->ntotal2);
READ1(idxqfs->M2);
READ1(idxqfs->code_size);
// Need to initialize the FastScan base class fields
const size_t M_fastscan = (idxqfs->d + 3) / 4;
constexpr size_t nbits_fastscan = 4;
idxqfs->M = M_fastscan;
idxqfs->nbits = nbits_fastscan;
idxqfs->ksub = (1 << nbits_fastscan);
READVECTOR(idxqfs->codes);
idx = idxqfs;
} else if (h == fourcc("Ixrq")) {
IndexRaBitQ* idxq = new IndexRaBitQ();
read_index_header(idxq, f);
faiss/impl/index_write.cpp
+16
View File
@@ -41,6 +41,7 @@
#include <faiss/IndexPQFastScan.h>
#include <faiss/IndexPreTransform.h>
#include <faiss/IndexRaBitQ.h>
#include <faiss/IndexRaBitQFastScan.h>
#include <faiss/IndexRefine.h>
#include <faiss/IndexRowwiseMinMax.h>
#include <faiss/IndexScalarQuantizer.h>
@@ -861,6 +862,21 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
WRITE1(h);
write_index_header(imm_2, f);
write_index(imm_2->index, f);
} else if (
const IndexRaBitQFastScan* idxqfs =
dynamic_cast<const IndexRaBitQFastScan*>(idx)) {
uint32_t h = fourcc("Irfs");
WRITE1(h);
write_index_header(idx, f);
write_RaBitQuantizer(&idxqfs->rabitq, f);
WRITEVECTOR(idxqfs->center);
WRITE1(idxqfs->qb);
WRITEVECTOR(idxqfs->factors_storage);
WRITE1(idxqfs->bbs);
WRITE1(idxqfs->ntotal2);
WRITE1(idxqfs->M2);
WRITE1(idxqfs->code_size);
WRITEVECTOR(idxqfs->codes);
} else if (
const IndexRaBitQ* idxq = dynamic_cast<const IndexRaBitQ*>(idx)) {
uint32_t h = fourcc("Ixrq");
faiss/impl/pq4_fast_scan.cpp
+30 -4
View File
@@ -49,7 +49,19 @@ void pq4_pack_codes(
size_t nb,
size_t bbs,
size_t nsq,
uint8_t* blocks) {
uint8_t* blocks,
size_t code_stride) {
// Determine stride: use custom if provided, otherwise use legacy
// calculation
size_t actual_stride = (code_stride == 0) ? (M + 1) / 2 : code_stride;
// Input validation for custom stride
if (code_stride != 0) {
FAISS_THROW_IF_NOT_MSG(
code_stride >= (M + 1) / 2,
"Custom stride must be >= minimum code size");
}
FAISS_THROW_IF_NOT(bbs % 32 == 0);
FAISS_THROW_IF_NOT(nb % bbs == 0);
FAISS_THROW_IF_NOT(nsq % 2 == 0);
@@ -72,7 +84,8 @@ void pq4_pack_codes(
for (size_t i = 0; i < bbs; i += 32) {
std::array<uint8_t, 32> c, c0, c1;
get_matrix_column(
codes, ntotal, (M + 1) / 2, i0 + i, sq / 2, c);
codes, ntotal, actual_stride, i0 + i, sq / 2, c);
for (int j = 0; j < 32; j++) {
c0[j] = c[j] & 15;
c1[j] = c[j] >> 4;
@@ -97,7 +110,19 @@ void pq4_pack_codes_range(
size_t i1,
size_t bbs,
size_t nsq,
uint8_t* blocks) {
uint8_t* blocks,
size_t code_stride) {
// Determine stride: use custom if provided, otherwise use legacy
// calculation
size_t actual_stride = (code_stride == 0) ? (M + 1) / 2 : code_stride;
// Input validation for custom stride
if (code_stride != 0) {
FAISS_THROW_IF_NOT_MSG(
code_stride >= (M + 1) / 2,
"Custom stride must be >= minimum code size");
}
#ifdef FAISS_BIG_ENDIAN
const uint8_t perm0[16] = {
8, 0, 9, 1, 10, 2, 11, 3, 12, 4, 13, 5, 14, 6, 15, 7};
@@ -117,7 +142,8 @@ void pq4_pack_codes_range(
for (size_t i = 0; i < bbs; i += 32) {
std::array<uint8_t, 32> c, c0, c1;
get_matrix_column(
codes, i1 - i0, (M + 1) / 2, i_base + i, sq / 2, c);
codes, i1 - i0, actual_stride, i_base + i, sq / 2, c);
for (int j = 0; j < 32; j++) {
c0[j] = c[j] & 15;
c1[j] = c[j] >> 4;
faiss/impl/pq4_fast_scan.h
+8 -2
View File
@@ -37,6 +37,8 @@ struct SIMDResultHandler;
* @param nsq number of sub-quantizers (=M rounded up to a muliple of 2)
* @param bbs size of database blocks (multiple of 32)
* @param blocks output array, size nb * nsq / 2.
* @param code_stride optional stride between consecutive codes (0 = use
default (M + 1) / 2)
*/
void pq4_pack_codes(
const uint8_t* codes,
@@ -45,7 +47,8 @@ void pq4_pack_codes(
size_t nb,
size_t bbs,
size_t nsq,
uint8_t* blocks);
uint8_t* blocks,
size_t code_stride = 0);
/** Same as pack_codes but write in a given range of the output,
* leaving the rest untouched. Assumes allocated entries are 0 on input.
@@ -54,6 +57,8 @@ void pq4_pack_codes(
* @param i0 first output code to write
* @param i1 last output code to write
* @param blocks output array, size at least ceil(i1 / bbs) * bbs * nsq / 2
* @param code_stride optional stride between consecutive codes (0 = use
* default (M + 1) / 2)
*/
void pq4_pack_codes_range(
const uint8_t* codes,
@@ -62,7 +67,8 @@ void pq4_pack_codes_range(
size_t i1,
size_t bbs,
size_t nsq,
uint8_t* blocks);
uint8_t* blocks,
size_t code_stride = 0);
/** get a single element from a packed codes table
*
faiss/python/swigfaiss.swig
+7
View File
@@ -104,6 +104,7 @@ typedef uint64_t size_t;
#include <faiss/IndexFastScan.h>
#include <faiss/IndexAdditiveQuantizerFastScan.h>
#include <faiss/IndexPQFastScan.h>
#include <faiss/impl/FastScanDistancePostProcessing.h>
#include <faiss/utils/simdlib.h>
#include <faiss/impl/simd_result_handlers.h>
@@ -191,7 +192,9 @@ typedef uint64_t size_t;
#include <faiss/IndexNeuralNetCodec.h>
#include <faiss/impl/RaBitQuantizer.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/IndexRaBitQ.h>
#include <faiss/IndexRaBitQFastScan.h>
#include <faiss/IndexIVFRaBitQ.h>
%}
@@ -564,6 +567,7 @@ void gpu_sync_all_devices()
%include <faiss/IndexIVFPQR.h>
%include <faiss/Index2Layer.h>
%include <faiss/impl/FastScanDistancePostProcessing.h>
%include <faiss/IndexFastScan.h>
%include <faiss/IndexAdditiveQuantizerFastScan.h>
%include <faiss/IndexPQFastScan.h>
@@ -610,7 +614,9 @@ struct faiss::simd16uint16 {};
%include <faiss/IndexNeuralNetCodec.h>
%include <faiss/impl/RaBitQuantizer.h>
%include <faiss/impl/RaBitQUtils.h>
%include <faiss/IndexRaBitQ.h>
%include <faiss/IndexRaBitQFastScan.h>
%include <faiss/IndexIVFRaBitQ.h>
%ignore faiss::BufferList::Buffer;
@@ -704,6 +710,7 @@ struct faiss::simd16uint16 {};
DOWNCAST ( IndexShardsIVF )
DOWNCAST2 ( IndexShards, IndexShardsTemplateT_faiss__Index_t )
DOWNCAST2 ( IndexReplicas, IndexReplicasTemplateT_faiss__Index_t )
DOWNCAST ( IndexRaBitQFastScan )
DOWNCAST ( IndexRaBitQ )
DOWNCAST ( IndexIVFRaBitQ )
DOWNCAST ( IndexIVFIndependentQuantizer)
tests/test_rabitq.py
+350 -26
View File
@@ -5,9 +5,8 @@
import unittest
import faiss
import numpy as np
import faiss
from faiss.contrib import datasets
@@ -208,7 +207,7 @@ class TestRaBitQ(unittest.TestCase):
index_pq.train(ds.get_train())
index_pq.add(ds.get_database())
D_pq, I_pq = index_pq.search(ds.get_queries(), k)
_, I_pq = index_pq.search(ds.get_queries(), k)
loss_pq = 1 - eval_I(I_pq)
print(f"{random_rotate=:1}, {loss_pq=:5.3f}")
np.testing.assert_(loss_pq < 0.25, f"{loss_pq}")
@@ -223,24 +222,34 @@ class TestRaBitQ(unittest.TestCase):
index_rbq.train(ds.get_train())
index_rbq.add(ds.get_database())
def test(params):
_, I_rbq = index_rbq.search(ds.get_queries(), k, params=params)
def test(params, ratio_threshold, index, rotate, loss_pq_val):
_, I_rbq = index.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=}, {params.centered=:1}: "
f"{loss_rbq=:5.3f} = loss_pq * {loss_rbq / loss_pq:5.3f}"
f"{rotate=:1}, {params.qb=}, {params.centered=:1}: "
f"{loss_rbq=:5.3f} = loss_pq * "
f"{loss_rbq / loss_pq_val:5.3f}"
f" < {ratio_threshold=:.2f}"
)
np.testing.assert_(loss_rbq < loss_pq * ratio_threshold)
np.testing.assert_(loss_rbq < loss_pq_val * ratio_threshold)
for qb in [1, 2, 3, 4, 8]:
print()
for centered in [False, True]:
test(faiss.RaBitQSearchParameters(qb=qb, centered=centered))
params = faiss.RaBitQSearchParameters(
qb=qb, centered=centered
)
ratio_threshold = 2 ** (1 / qb)
test(
params,
ratio_threshold,
index_rbq,
random_rotate,
loss_pq,
)
def test_comparison_vs_pq_L2(self):
self.do_comparison_vs_pq_test(faiss.METRIC_L2)
@@ -346,7 +355,7 @@ class TestIVFRaBitQ(unittest.TestCase):
index_flat = faiss.IndexIVFFlat(quantizer, d, nlist, metric_type)
index_flat.train(xt)
index_flat.add(xb)
D_f, I_f = index_flat.search(
_, I_f = index_flat.search(
xq, k, params=faiss.IVFSearchParameters(nprobe=nprobe)
)
@@ -383,7 +392,7 @@ class TestIVFRaBitQ(unittest.TestCase):
index_pq.train(xt)
index_pq.add(xb)
D_pq, I_pq = index_pq.search(
_, I_pq = index_pq.search(
xq, k, params=faiss.IVFPQSearchParameters(nprobe=nprobe)
)
loss_pq = 1 - eval_I(I_pq)
@@ -391,27 +400,33 @@ class TestIVFRaBitQ(unittest.TestCase):
print(f"{random_rotate=:1}, {loss_pq=:5.3f}")
np.testing.assert_(loss_pq < 0.25, f"{loss_pq}")
def test(params):
D_rbq, I_rbq = index_rbq.search(xq, k, params=params)
def test(params, ratio_threshold, index, rotate, loss_pq_val):
_, I_rbq = index.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=}, {params.centered=:1}: "
f"{loss_rbq=:5.3f} = loss_pq * {loss_rbq / loss_pq:5.3f}"
f"{rotate=:1}, {params.qb=}, {params.centered=:1}: "
f"{loss_rbq=:5.3f} = loss_pq * "
f"{loss_rbq / loss_pq_val:5.3f}"
f" < {ratio_threshold=:.2f}"
)
np.testing.assert_(loss_rbq < loss_pq * ratio_threshold)
np.testing.assert_(loss_rbq < loss_pq_val * ratio_threshold)
for qb in [1, 2, 3, 4, 8]:
print()
for centered in [False, True]:
params = faiss.IVFRaBitQSearchParameters(
nprobe=nprobe, qb=qb, centered=centered
)
ratio_threshold = 2 ** (1 / qb)
test(
faiss.IVFRaBitQSearchParameters(
nprobe=nprobe, qb=qb, centered=centered
)
params,
ratio_threshold,
index_rbq,
random_rotate,
loss_pq,
)
def test_comparison_vs_pq_L2(self):
@@ -430,14 +445,16 @@ class TestIVFRaBitQ(unittest.TestCase):
ref_rbq.add(ds.get_database())
index_flat = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
index_rbq = faiss.IndexIVFRaBitQ(index_flat, ds.d, nlist, faiss.METRIC_L2)
index_rbq = faiss.IndexIVFRaBitQ(
index_flat, ds.d, nlist, faiss.METRIC_L2
)
index_rbq.qb = 4
index_rbq.train(ds.get_train())
index_rbq.add(ds.get_database())
for nprobe in 1, 4, 16:
ref_rbq.nprobe = nprobe
Dref, Iref = ref_rbq.search(ds.get_queries(), k)
_, Iref = ref_rbq.search(ds.get_queries(), k)
r_ref_k = faiss.eval_intersection(
Iref[:, :k], ds.get_groundtruth()[:, :k]
) / (ds.nq * k)
@@ -468,7 +485,9 @@ class TestIVFRaBitQ(unittest.TestCase):
ref_rbq.add(ds.get_database())
index_flat = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
index_rbq = faiss.IndexIVFRaBitQ(index_flat, ds.d, nlist, faiss.METRIC_L2)
index_rbq = faiss.IndexIVFRaBitQ(
index_flat, ds.d, nlist, faiss.METRIC_L2
)
index_rbq.qb = 4
# wrap with random rotations
@@ -481,7 +500,7 @@ class TestIVFRaBitQ(unittest.TestCase):
for nprobe in 1, 4, 16:
ref_rbq.nprobe = nprobe
Dref, Iref = ref_rbq.search(ds.get_queries(), k)
_, Iref = ref_rbq.search(ds.get_queries(), k)
r_ref_k = faiss.eval_intersection(
Iref[:, :k], ds.get_groundtruth()[:, :k]
) / (ds.nq * k)
@@ -490,7 +509,7 @@ class TestIVFRaBitQ(unittest.TestCase):
params = faiss.IVFRaBitQSearchParameters()
params.qb = index_rbq.qb
params.nprobe = nprobe
Dnew, Inew, stats2 = faiss.search_with_parameters(
_, Inew, _ = faiss.search_with_parameters(
index_cand, ds.get_queries(), k, params, output_stats=True
)
r_new_k = faiss.eval_intersection(
@@ -524,6 +543,311 @@ class TestIVFRaBitQ(unittest.TestCase):
self.do_test_serde("IVF16,RaBitQ")
class TestRaBitQFastScan(unittest.TestCase):
def do_comparison_vs_rabitq_test(
self, metric_type=faiss.METRIC_L2, bbs=32
):
"""Test IndexRaBitQFastScan produces similar results to IndexRaBitQ"""
ds = datasets.SyntheticDataset(128, 4096, 4096, 100)
k = 10
# IndexRaBitQ baseline
index_rbq = faiss.IndexRaBitQ(ds.d, metric_type)
index_rbq.train(ds.get_train())
index_rbq.add(ds.get_database())
_, I_rbq = index_rbq.search(ds.get_queries(), k)
# IndexRaBitQFastScan
index_rbq_fs = faiss.IndexRaBitQFastScan(ds.d, metric_type, bbs)
index_rbq_fs.train(ds.get_train())
index_rbq_fs.add(ds.get_database())
_, I_rbq_fs = index_rbq_fs.search(ds.get_queries(), k)
index_flat = faiss.IndexFlat(ds.d, metric_type)
index_flat.train(ds.get_train())
index_flat.add(ds.get_database())
_, I_f = index_flat.search(ds.get_queries(), k)
# Evaluate against ground truth
eval_rbq = faiss.eval_intersection(I_rbq[:, :k], I_f[:, :k])
eval_rbq /= ds.nq * k
eval_rbq_fs = faiss.eval_intersection(I_rbq_fs[:, :k], I_f[:, :k])
eval_rbq_fs /= ds.nq * k
print(
f"RaBitQ baseline is {eval_rbq}, "
f"RaBitQFastScan is {eval_rbq_fs}"
)
# FastScan should be similar to baseline
np.testing.assert_(abs(eval_rbq - eval_rbq_fs) < 0.05)
def test_comparison_vs_rabitq_L2(self):
self.do_comparison_vs_rabitq_test(faiss.METRIC_L2)
def test_comparison_vs_rabitq_IP(self):
self.do_comparison_vs_rabitq_test(faiss.METRIC_INNER_PRODUCT)
def test_encode_decode_consistency(self):
"""Test that encoding and decoding operations are consistent"""
ds = datasets.SyntheticDataset(128, 1000, 0, 0) # No queries/db needed
# Test with IndexRaBitQFastScan
index_rbq_fs = faiss.IndexRaBitQFastScan(ds.d, faiss.METRIC_L2)
index_rbq_fs.train(ds.get_train())
# Test encode/decode on a subset of training data
test_vectors = ds.get_train()[:100]
# Test compute_codes and sa_decode
# This tests that factors are properly embedded in codes
codes = np.empty(
(len(test_vectors), index_rbq_fs.code_size), dtype=np.uint8
)
index_rbq_fs.compute_codes(
faiss.swig_ptr(codes),
len(test_vectors),
faiss.swig_ptr(test_vectors)
)
# sa_decode should work directly with embedded codes
decoded_fs = index_rbq_fs.sa_decode(codes)
# Check reconstruction error for FastScan
distances_fs = np.sum((test_vectors - decoded_fs) ** 2, axis=1)
avg_distance_fs = np.mean(distances_fs)
print(f"Average FastScan reconstruction error: {avg_distance_fs}")
# Compare with original IndexRaBitQ on the SAME dataset
index_rbq = faiss.IndexRaBitQ(ds.d, faiss.METRIC_L2)
index_rbq.train(ds.get_train())
# Encode with original RaBitQ (correct API - returns encoded array)
codes_orig = index_rbq.sa_encode(test_vectors)
# Decode with original RaBitQ
decoded_orig = index_rbq.sa_decode(codes_orig)
# Check reconstruction error for original
distances_orig = np.sum((test_vectors - decoded_orig) ** 2, axis=1)
avg_distance_orig = np.mean(distances_orig)
print(
f"Average original RaBitQ reconstruction error: "
f"{avg_distance_orig}"
)
# Print comparison
print(
f"Error difference (FastScan - Original): "
f"{avg_distance_fs - avg_distance_orig}"
)
# FastScan should have similar reconstruction error to original RaBitQ
np.testing.assert_(
abs(avg_distance_fs - avg_distance_orig) < 0.01
) # Should be nearly identical
def test_query_quantization_bits(self):
"""Test different query quantization bit settings"""
ds = datasets.SyntheticDataset(64, 2000, 2000, 50)
k = 10
index_rbq_fs = faiss.IndexRaBitQFastScan(ds.d, faiss.METRIC_L2)
index_rbq_fs.train(ds.get_train())
index_rbq_fs.add(ds.get_database())
# Test different qb values
results = {}
for qb in [4, 6, 8]:
index_rbq_fs.qb = qb
_, I = index_rbq_fs.search(ds.get_queries(), k)
results[qb] = I
# All should produce reasonable results
index_flat = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
index_flat.train(ds.get_train())
index_flat.add(ds.get_database())
_, I_f = index_flat.search(ds.get_queries(), k)
for qb in [4, 6, 8]:
eval_qb = faiss.eval_intersection(results[qb][:, :k], I_f[:, :k])
eval_qb /= ds.nq * k
print(f"Query quantization qb={qb} recall: {eval_qb}")
np.testing.assert_(eval_qb > 0.4) # Should be reasonable
def test_small_dataset(self):
"""Test on a small dataset to ensure basic functionality"""
d = 32
n = 100
nq = 10
rs = np.random.RandomState(123)
xb = rs.rand(n, d).astype(np.float32)
xq = rs.rand(nq, d).astype(np.float32)
index_rbq_fs = faiss.IndexRaBitQFastScan(d, faiss.METRIC_L2)
index_rbq_fs.train(xb)
index_rbq_fs.add(xb)
k = 5
distances, labels = index_rbq_fs.search(xq, k)
# Check output shapes and validity
np.testing.assert_equal(distances.shape, (nq, k))
np.testing.assert_equal(labels.shape, (nq, k))
# Check that labels are valid indices
np.testing.assert_(np.all(labels >= 0))
np.testing.assert_(np.all(labels < n))
# Check that distances are non-negative (for L2)
np.testing.assert_(np.all(distances >= 0))
# Quick recall check against exact search
index_flat = faiss.IndexFlat(d, faiss.METRIC_L2)
index_flat.train(xb)
index_flat.add(xb)
_, I_f = index_flat.search(xq, k)
# Evaluate recall
recall = faiss.eval_intersection(labels[:, :k], I_f[:, :k])
recall /= (nq * k)
print(f"Small dataset recall: {recall:.3f}")
np.testing.assert_(
recall > 0.4
) # Should be reasonable for small dataset
def test_comparison_vs_pq_fastscan(self):
"""Compare RaBitQFastScan to PQFastScan as a performance baseline"""
ds = datasets.SyntheticDataset(128, 4096, 4096, 100)
k = 10
# PQFastScan baseline
index_pq_fs = faiss.IndexPQFastScan(ds.d, 16, 4, faiss.METRIC_L2)
index_pq_fs.train(ds.get_train())
index_pq_fs.add(ds.get_database())
_, I_pq_fs = index_pq_fs.search(ds.get_queries(), k)
# RaBitQFastScan
index_rbq_fs = faiss.IndexRaBitQFastScan(ds.d, faiss.METRIC_L2)
index_rbq_fs.train(ds.get_train())
index_rbq_fs.add(ds.get_database())
_, I_rbq_fs = index_rbq_fs.search(ds.get_queries(), k)
index_flat = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
index_flat.train(ds.get_train())
index_flat.add(ds.get_database())
_, I_f = index_flat.search(ds.get_queries(), k)
# Evaluate both against ground truth
eval_pq_fs = faiss.eval_intersection(I_pq_fs[:, :k], I_f[:, :k])
eval_pq_fs /= ds.nq * k
eval_rbq_fs = faiss.eval_intersection(I_rbq_fs[:, :k], I_f[:, :k])
eval_rbq_fs /= ds.nq * k
print(
f"PQFastScan is {eval_pq_fs}, "
f"RaBitQFastScan is {eval_rbq_fs}"
)
# RaBitQFastScan should have reasonable performance similar to regular
# RaBitQ
np.testing.assert_(eval_rbq_fs > 0.55)
def test_serialization(self):
"""Test serialization and deserialization of RaBitQFastScan"""
ds = datasets.SyntheticDataset(64, 1000, 100, 20)
index_rbq_fs = faiss.IndexRaBitQFastScan(ds.d, faiss.METRIC_L2)
index_rbq_fs.train(ds.get_train())
index_rbq_fs.add(ds.get_database())
Dref, Iref = index_rbq_fs.search(ds.get_queries(), 10)
# Serialize and deserialize
b = faiss.serialize_index(index_rbq_fs)
index2 = faiss.deserialize_index(b)
Dnew, Inew = index2.search(ds.get_queries(), 10)
# Results should be identical
np.testing.assert_array_equal(Dref, Dnew)
np.testing.assert_array_equal(Iref, Inew)
def test_memory_management(self):
"""Test that memory is managed correctly during operations"""
ds = datasets.SyntheticDataset(128, 2000, 2000, 50)
index_rbq_fs = faiss.IndexRaBitQFastScan(ds.d, faiss.METRIC_L2)
index_rbq_fs.train(ds.get_train())
# Add data in chunks to test memory management
chunk_size = 500
for i in range(0, ds.nb, chunk_size):
end_idx = min(i + chunk_size, ds.nb)
chunk_data = ds.get_database()[i:end_idx]
index_rbq_fs.add(chunk_data)
# Verify total count
np.testing.assert_equal(index_rbq_fs.ntotal, ds.nb)
# Test search still works and produces reasonable recall
_, I = index_rbq_fs.search(ds.get_queries(), 5)
np.testing.assert_equal(I.shape, (ds.nq, 5))
# Compare against ground truth to verify recall
index_flat = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
index_flat.train(ds.get_train())
index_flat.add(ds.get_database())
_, I_f = index_flat.search(ds.get_queries(), 5)
# Calculate recall - should be reasonable despite multiple add() calls
recall = faiss.eval_intersection(I[:, :5], I_f[:, :5])
recall /= (ds.nq * 5)
# If embedded factors are corrupted by multiple add() calls,
# recall will be very low
np.testing.assert_(
recall > 0.1,
f"Recall too low: {recall:.3f} - suggests multiple "
f"add() calls corrupted embedded factors"
)
def test_invalid_parameters(self):
"""Test proper error handling for invalid parameters"""
# Invalid dimension
with np.testing.assert_raises(Exception):
faiss.IndexRaBitQFastScan(0, faiss.METRIC_L2)
# Invalid metric (should only support L2 and IP)
try:
faiss.IndexRaBitQFastScan(64, faiss.METRIC_Lp)
np.testing.assert_(
False, "Should have raised exception for invalid metric"
)
except RuntimeError:
pass # Expected
def test_thread_safety(self):
"""Test that parallel operations work correctly"""
ds = datasets.SyntheticDataset(64, 2000, 2000, 100)
index_rbq_fs = faiss.IndexRaBitQFastScan(ds.d, faiss.METRIC_L2)
index_rbq_fs.train(ds.get_train())
index_rbq_fs.add(ds.get_database())
# Search with multiple threads (implicitly tested through OpenMP)
k = 10
distances, labels = index_rbq_fs.search(ds.get_queries(), k)
# Basic sanity checks
np.testing.assert_equal(distances.shape, (ds.nq, k))
np.testing.assert_equal(labels.shape, (ds.nq, k))
np.testing.assert_(np.all(distances >= 0))
np.testing.assert_(np.all(labels >= 0))
np.testing.assert_(np.all(labels < ds.nb))
class TestRaBitQuantizerEncodeDecode(unittest.TestCase):
def do_test_encode_decode(self, d, metric):
# rabitq must precisely reconstruct a vector,