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

Wire RaBitQ search through FastScanCodeScanner dispatch (#4903)

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

Add make_knn_scanner() virtual to IndexFastScan and IndexIVFFastScan
that returns a SIMD-dispatched FastScanCodeScanner. Override in
RaBitQ classes to produce rabitq-specific scanners via per-SIMD TU
factories (rabitq_dispatching.h). Move RaBitQ handler method bodies
into headers, add context pointer, and update rabitq_result_handler.h
for SL template param. Wire IVFRaBitQFastScanScanner::scan_codes to
use the new scanner. Update build files and SWIG ignores.

Callers (search_implem_12/14, search_dispatch_implem) are unchanged
and still use make_knn_handler.

Reviewed By: mdouze

Differential Revision: D96116557

fbshipit-source-id: 9ead718ab863f9bc63953c8f9e7958d32ed896f2
This commit is contained in:
Gergely Szilvasy
2026-03-12 02:15:42 -07:00
committed by meta-codesync[bot]
parent 4a0fba1fe2
commit 57190f91cd
17 changed files with 782 additions and 553 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+1
View File
@@ -271,6 +271,7 @@ set(FAISS_HEADERS
impl/fast_scan/decompose_qbs.h
impl/fast_scan/kernels_simd256.h
impl/fast_scan/kernels_simd512.h
impl/fast_scan/rabitq_dispatching.h
impl/fast_scan/rabitq_result_handler.h
impl/residual_quantizer_encode_steps.h
impl/simd_dispatch.h
faiss/IndexFastScan.cpp
+13
View File
@@ -211,6 +211,19 @@ void estimators_from_tables_generic(
} // anonymous namespace
std::unique_ptr<FastScanCodeScanner> IndexFastScan::make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
size_t ntotal,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing&) const {
return make_fast_scan_knn_scanner(
is_max, 0, n, ntotal, k, distances, labels, sel);
}
// Default implementation of make_knn_handler with centralized fallback logic
SIMDResultHandlerToFloat* IndexFastScan::make_knn_handler(
bool is_max,
faiss/IndexFastScan.h
+19
View File
@@ -7,8 +7,11 @@
#pragma once
#include <memory>
#include <faiss/Index.h>
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
#include <faiss/impl/fast_scan/fast_scan.h>
#include <faiss/utils/AlignedTable.h>
namespace faiss {
@@ -121,6 +124,22 @@ struct IndexFastScan : Index {
const float* x,
const FastScanDistancePostProcessing& context) const = 0;
/** Create a SIMD-dispatched scanner for knn search.
*
* Returns a FastScanCodeScanner that bundles handler + accumulation
* kernel behind the SIMD dispatch boundary.
* The scanner's accumulate methods dispatch to the optimal SIMD level.
*/
virtual std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
size_t ntotal,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context = {}) const;
/** Create a KNN handler for this index type
*
* This method can be overridden by derived classes to provide
faiss/IndexIVFFastScan.cpp
+20
View File
@@ -495,6 +495,26 @@ int compute_search_nslice(
} // namespace
std::unique_ptr<FastScanCodeScanner> IndexIVFFastScan::make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing&) const {
return make_fast_scan_knn_scanner(
is_max,
0,
n,
0,
k,
distances,
labels,
sel,
/*with_id_map=*/true);
}
SIMDResultHandlerToFloat* IndexIVFFastScan::make_knn_handler(
bool is_max,
int impl,
faiss/IndexIVFFastScan.h
+26
View File
@@ -9,6 +9,7 @@
#include <faiss/IndexIVF.h>
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
#include <faiss/impl/fast_scan/fast_scan.h>
#include <faiss/utils/AlignedTable.h>
namespace faiss {
@@ -221,6 +222,31 @@ struct IndexIVFFastScan : IndexIVF {
RangeSearchResult* result,
const SearchParameters* params = nullptr) const override;
/** Create a SIMD-dispatched scanner for knn search (IVF variant).
*
* Returns a FastScanCodeScanner that bundles handler + accumulation
* kernel behind the SIMD dispatch boundary. ntotal is not passed
* because IVF sets it per-list via handler->ntotal.
* Derived classes that need custom handlers (e.g. RaBitQ) override
* this to return nullptr, falling back to make_knn_handler.
*
* @param is_max whether to use CMax comparator (true) or CMin
* @param n number of queries
* @param k number of neighbors to find
* @param distances output distances array
* @param labels output labels array
* @param sel optional ID selector
* @return scanner, or nullptr if unsupported
*/
virtual std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context = {}) const;
/** Create a KNN handler for this index type
*
* This method can be overridden by derived classes to provide
faiss/IndexIVFRaBitQFastScan.cpp
+23 -273
View File
@@ -539,6 +539,19 @@ void IndexIVFRaBitQFastScan::decode_fastscan_to_residual(
}
}
std::unique_ptr<FastScanCodeScanner> IndexIVFRaBitQFastScan::make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* /*sel*/,
const FastScanDistancePostProcessing& context) const {
const bool is_multibit = (rabitq.nb_bits - 1) > 0;
return rabitq_ivf_make_knn_scanner(
is_max, this, n, k, distances, labels, &context, is_multibit);
}
// Implementation of virtual make_knn_handler method
SIMDResultHandlerToFloat* IndexIVFRaBitQFastScan::make_knn_handler(
bool is_max,
@@ -562,266 +575,6 @@ SIMDResultHandlerToFloat* IndexIVFRaBitQFastScan::make_knn_handler(
}
}
/*********************************************************
* simd_result_handlers::IVFRaBitQHeapHandler implementation
*********************************************************/
template <class C, SIMDLevel SL>
simd_result_handlers::IVFRaBitQHeapHandler<C, SL>::IVFRaBitQHeapHandler(
const IndexIVFRaBitQFastScan* idx,
size_t nq_val,
size_t k_val,
float* distances,
int64_t* labels,
const FastScanDistancePostProcessing* ctx,
bool multibit)
: simd_result_handlers::ResultHandlerCompare<C, true>(
nq_val,
0,
nullptr),
index(idx),
heap_distances(distances),
heap_labels(labels),
nq(nq_val),
k(k_val),
context(ctx),
is_multibit(multibit),
storage_size(idx->compute_per_vector_storage_size()),
packed_block_size(((idx->M2 + 1) / 2) * idx->bbs),
full_block_size(idx->get_block_stride()),
packer(idx->get_CodePacker()) {
current_list_no = 0;
probe_indices.clear();
// Initialize heaps in constructor (standard pattern from HeapHandler)
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, SIMDLevel SL>
void simd_result_handlers::IVFRaBitQHeapHandler<C, SL>::handle(
size_t q,
size_t b,
simd16uint16 d0,
simd16uint16 d1) {
// Store the original local query index before adjust_with_origin changes it
size_t local_q = q;
this->adjust_with_origin(q, d0, d1);
ALIGNED(32) uint16_t d32tab[32];
d0.store(d32tab);
d1.store(d32tab + 16);
float* const heap_dis = heap_distances + q * k;
int64_t* const heap_ids = heap_labels + q * k;
FAISS_THROW_IF_NOT_FMT(
!probe_indices.empty() && local_q < probe_indices.size(),
"set_list_context() must be called before handle() - probe_indices size: %zu, local_q: %zu, global_q: %zu",
probe_indices.size(),
local_q,
q);
// Access query factors directly from array via ProcessingContext
if (!context || !context->query_factors) {
FAISS_THROW_MSG(
"Query factors not available: FastScanDistancePostProcessing with query_factors required");
}
// Use probe_rank from probe_indices for compact storage indexing
size_t probe_rank = probe_indices[local_q];
size_t nprobe = context->nprobe > 0 ? context->nprobe : index->nprobe;
size_t storage_idx = q * nprobe + probe_rank;
const auto& query_factors = context->query_factors[storage_idx];
const float one_a =
this->normalizers ? (1.0f / this->normalizers[2 * q]) : 1.0f;
const float bias = this->normalizers ? this->normalizers[2 * q + 1] : 0.0f;
uint64_t idx_base = this->j0 + b * 32;
if (idx_base >= this->ntotal) {
return;
}
size_t max_positions = std::min<size_t>(32, this->ntotal - idx_base);
// Stats tracking for two-stage search
// n_1bit_evaluations: candidates evaluated using 1-bit lower bound
// n_multibit_evaluations: candidates requiring full multi-bit distance
size_t local_1bit_evaluations = 0;
size_t local_multibit_evaluations = 0;
// Process each candidate vector in the SIMD batch
for (size_t j = 0; j < max_positions; j++) {
const int64_t result_id = this->adjust_id(b, j);
if (result_id < 0) {
continue;
}
const float normalized_distance = d32tab[j] * one_a + bias;
const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
list_codes_ptr,
idx_base + j,
index->bbs,
packed_block_size,
full_block_size,
storage_size);
if (is_multibit) {
// Track candidates actually considered for two-stage filtering
local_1bit_evaluations++;
// Multi-bit: use SignBitFactorsWithError and two-stage search
const SignBitFactorsWithError& full_factors =
*reinterpret_cast<const SignBitFactorsWithError*>(base_ptr);
// Compute 1-bit adjusted distance using shared helper
float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
full_factors,
query_factors,
index->centered,
index->qb,
index->d);
// Adaptive filtering: decide whether to compute full distance
const bool is_similarity =
index->metric_type == MetricType::METRIC_INNER_PRODUCT;
float g_error = query_factors.g_error;
bool should_refine = rabitq_utils::should_refine_candidate(
dist_1bit,
full_factors.f_error,
g_error,
heap_dis[0],
is_similarity);
if (should_refine) {
local_multibit_evaluations++;
// Compute local_offset: position within current inverted list
size_t local_offset = this->j0 + b * 32 + j;
// Compute full multi-bit distance
float dist_full = compute_full_multibit_distance(
result_id, local_q, q, local_offset);
// Update heap if this distance is better
if (Cfloat::cmp(heap_dis[0], dist_full)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, dist_full, result_id);
nup++;
}
}
} else {
const auto& db_factors =
*reinterpret_cast<const SignBitFactors*>(base_ptr);
// Compute adjusted distance using shared helper
float adjusted_distance =
rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
db_factors,
query_factors,
index->centered,
index->qb,
index->d);
if (Cfloat::cmp(heap_dis[0], adjusted_distance)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, adjusted_distance, result_id);
nup++;
}
}
}
// Update global stats atomically
#pragma omp atomic
rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
#pragma omp atomic
rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
}
template <class C, SIMDLevel SL>
void simd_result_handlers::IVFRaBitQHeapHandler<C, SL>::set_list_context(
size_t list_no,
const std::vector<int>& probe_map) {
current_list_no = list_no;
probe_indices = probe_map;
list_codes_ptr = index->invlists->get_codes(list_no);
}
template <class C, SIMDLevel SL>
void simd_result_handlers::IVFRaBitQHeapHandler<C, SL>::begin(
const float* norms) {
this->normalizers = norms;
}
template <class C, SIMDLevel SL>
void simd_result_handlers::IVFRaBitQHeapHandler<C, SL>::end() {
#pragma omp parallel for
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);
}
}
template <class C, SIMDLevel SL>
float simd_result_handlers::IVFRaBitQHeapHandler<C, SL>::
compute_full_multibit_distance(
size_t /*db_idx*/,
size_t local_q,
size_t global_q,
size_t local_offset) const {
const size_t ex_bits = index->rabitq.nb_bits - 1;
const size_t dim = index->d;
const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
list_codes_ptr,
local_offset,
index->bbs,
packed_block_size,
full_block_size,
storage_size);
const size_t ex_code_size = (dim * ex_bits + 7) / 8;
const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError);
const ExtraBitsFactors& ex_fac = *reinterpret_cast<const ExtraBitsFactors*>(
base_ptr + sizeof(SignBitFactorsWithError) + ex_code_size);
// Use local_q to access probe_indices (batch-local), global_q for storage
size_t probe_rank = probe_indices[local_q];
size_t nprobe = context->nprobe > 0 ? context->nprobe : index->nprobe;
size_t storage_idx = global_q * nprobe + probe_rank;
const auto& query_factors = context->query_factors[storage_idx];
size_t list_no = current_list_no;
InvertedLists::ScopedCodes list_codes(index->invlists, list_no);
std::vector<uint8_t> unpacked_code(index->code_size);
packer->unpack_1(list_codes.get(), local_offset, unpacked_code.data());
const uint8_t* sign_bits = unpacked_code.data();
return rabitq_utils::compute_full_multibit_distance(
sign_bits,
ex_code,
ex_fac,
query_factors.rotated_q.data(),
(index->metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_factors.q_dot_c
: query_factors.qr_to_c_L2sqr,
dim,
ex_bits,
index->metric_type);
}
/*********************************************************
* IVFRaBitQFastScanScanner implementation
*********************************************************/
@@ -905,17 +658,15 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
std::vector<float> curr_dists(k, distances[0]);
std::vector<idx_t> curr_labels(k, labels[0]);
std::unique_ptr<SIMDResultHandlerToFloat> handler(
index.make_knn_handler(
!keep_max,
impl,
nq,
k,
curr_dists.data(),
curr_labels.data(),
sel,
context,
&normalizers[0]));
auto scanner = index.make_knn_scanner(
!keep_max,
nq,
k,
curr_dists.data(),
curr_labels.data(),
sel,
context);
auto* handler = scanner->handler();
int qmap1[1] = {0};
handler->q_map = qmap1;
@@ -930,14 +681,13 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
std::vector<int> probe_map = {0};
handler->set_list_context(list_no, probe_map);
pq4_accumulate_loop(
scanner->accumulate_loop(
1,
roundup(ntotal, index.bbs),
index.bbs,
static_cast<int>(index.M2),
codes,
LUT,
*handler,
0,
index.get_block_stride());
faiss/IndexIVFRaBitQFastScan.h
+232
View File
@@ -148,6 +148,16 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
const IVFSearchParameters* params = nullptr,
IndexIVFStats* stats = nullptr) const override;
/// RaBitQ scanner via rabitq_ivf_make_knn_scanner
std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context = {}) const override;
/// Override to create RaBitQ-specific handlers
SIMDResultHandlerToFloat* make_knn_handler(
bool is_max,
@@ -172,4 +182,226 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
using IVFRaBitQHeapHandler = simd_result_handlers::IVFRaBitQHeapHandler<C>;
};
/*********************************************************
* IVFRaBitQHeapHandler method bodies — defined here (after
* IndexIVFRaBitQFastScan is complete) to break the circular
* dependency between rabitq_result_handler.h and this header.
*********************************************************/
namespace simd_result_handlers {
template <class C, SIMDLevel SL>
IVFRaBitQHeapHandler<C, SL>::IVFRaBitQHeapHandler(
const IndexIVFRaBitQFastScan* idx,
size_t nq_val,
size_t k_val,
float* distances,
int64_t* labels,
const FastScanDistancePostProcessing* ctx,
bool multibit)
: ResultHandlerCompare<C, true, SL>(nq_val, 0, nullptr),
index(idx),
heap_distances(distances),
heap_labels(labels),
nq(nq_val),
k(k_val),
context(ctx),
is_multibit(multibit),
storage_size(idx->compute_per_vector_storage_size()),
packed_block_size(((idx->M2 + 1) / 2) * idx->bbs),
full_block_size(idx->get_block_stride()),
packer(idx->get_CodePacker()) {
current_list_no = 0;
probe_indices.clear();
for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
heap_heapify<Cfloat>(k, heap_distances + q * k, heap_labels + q * k);
}
}
template <class C, SIMDLevel SL>
void IVFRaBitQHeapHandler<C, SL>::handle(
size_t q,
size_t b,
simd16uint16 d0,
simd16uint16 d1) {
size_t local_q = q;
this->adjust_with_origin(q, d0, d1);
ALIGNED(32) uint16_t d32tab[32];
d0.store(d32tab);
d1.store(d32tab + 16);
float* const heap_dis = heap_distances + q * k;
int64_t* const heap_ids = heap_labels + q * k;
FAISS_THROW_IF_NOT_FMT(
!probe_indices.empty() && local_q < probe_indices.size(),
"set_list_context() must be called before handle() - probe_indices size: %zu, local_q: %zu, global_q: %zu",
probe_indices.size(),
local_q,
q);
if (!context || !context->query_factors) {
FAISS_THROW_MSG(
"Query factors not available: FastScanDistancePostProcessing with query_factors required");
}
size_t probe_rank = probe_indices[local_q];
size_t nprobe = context->nprobe > 0 ? context->nprobe : index->nprobe;
size_t storage_idx = q * nprobe + probe_rank;
const auto& query_factors = context->query_factors[storage_idx];
const float one_a =
this->normalizers ? (1.0f / this->normalizers[2 * q]) : 1.0f;
const float bias = this->normalizers ? this->normalizers[2 * q + 1] : 0.0f;
uint64_t idx_base = this->j0 + b * 32;
if (idx_base >= this->ntotal) {
return;
}
size_t max_positions = std::min<size_t>(32, this->ntotal - idx_base);
size_t local_1bit_evaluations = 0;
size_t local_multibit_evaluations = 0;
for (size_t j = 0; j < max_positions; j++) {
const int64_t result_id = this->adjust_id(b, j);
if (result_id < 0) {
continue;
}
const float normalized_distance = d32tab[j] * one_a + bias;
const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
list_codes_ptr,
idx_base + j,
index->bbs,
packed_block_size,
full_block_size,
storage_size);
if (is_multibit) {
local_1bit_evaluations++;
const SignBitFactorsWithError& full_factors =
*reinterpret_cast<const SignBitFactorsWithError*>(base_ptr);
float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
full_factors,
query_factors,
index->centered,
index->qb,
index->d);
const bool is_similarity =
index->metric_type == MetricType::METRIC_INNER_PRODUCT;
bool should_refine = rabitq_utils::should_refine_candidate(
dist_1bit,
full_factors.f_error,
query_factors.g_error,
heap_dis[0],
is_similarity);
if (should_refine) {
local_multibit_evaluations++;
size_t local_offset = this->j0 + b * 32 + j;
float dist_full = compute_full_multibit_distance(
result_id, local_q, q, local_offset);
if (Cfloat::cmp(heap_dis[0], dist_full)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, dist_full, result_id);
nup++;
}
}
} else {
const auto& db_factors =
*reinterpret_cast<const SignBitFactors*>(base_ptr);
float adjusted_distance =
rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
db_factors,
query_factors,
index->centered,
index->qb,
index->d);
if (Cfloat::cmp(heap_dis[0], adjusted_distance)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, adjusted_distance, result_id);
nup++;
}
}
}
#pragma omp atomic
rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
#pragma omp atomic
rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
}
template <class C, SIMDLevel SL>
void IVFRaBitQHeapHandler<C, SL>::set_list_context(
size_t list_no,
const std::vector<int>& probe_map) {
current_list_no = list_no;
probe_indices = probe_map;
list_codes_ptr = index->invlists->get_codes(list_no);
}
template <class C, SIMDLevel SL>
void IVFRaBitQHeapHandler<C, SL>::begin(const float* norms) {
this->normalizers = norms;
}
template <class C, SIMDLevel SL>
void IVFRaBitQHeapHandler<C, SL>::end() {
#pragma omp parallel for
for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
heap_reorder<Cfloat>(k, heap_distances + q * k, heap_labels + q * k);
}
}
template <class C, SIMDLevel SL>
float IVFRaBitQHeapHandler<C, SL>::compute_full_multibit_distance(
size_t /*db_idx*/,
size_t local_q,
size_t global_q,
size_t local_offset) const {
const size_t ex_bits = index->rabitq.nb_bits - 1;
const size_t dim = index->d;
const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
list_codes_ptr,
local_offset,
index->bbs,
packed_block_size,
full_block_size,
storage_size);
const size_t ex_code_size = (dim * ex_bits + 7) / 8;
const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError);
const ExtraBitsFactors& ex_fac = *reinterpret_cast<const ExtraBitsFactors*>(
base_ptr + sizeof(SignBitFactorsWithError) + ex_code_size);
size_t probe_rank = probe_indices[local_q];
size_t nprobe_val = context->nprobe > 0 ? context->nprobe : index->nprobe;
size_t storage_idx_val = global_q * nprobe_val + probe_rank;
const auto& query_factors = context->query_factors[storage_idx_val];
InvertedLists::ScopedCodes list_codes(index->invlists, current_list_no);
std::vector<uint8_t> unpacked_code(index->code_size);
packer->unpack_1(list_codes.get(), local_offset, unpacked_code.data());
return rabitq_utils::compute_full_multibit_distance(
unpacked_code.data(),
ex_code,
ex_fac,
query_factors.rotated_q.data(),
(index->metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_factors.q_dot_c
: query_factors.qr_to_c_L2sqr,
dim,
ex_bits,
index->metric_type);
}
} // namespace simd_result_handlers
} // namespace faiss
faiss/IndexRaBitQFastScan.cpp
+13 -237
View File
@@ -522,241 +522,18 @@ void IndexRaBitQFastScan::search(
}
}
// 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,
std::unique_ptr<FastScanCodeScanner> IndexRaBitQFastScan::make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
size_t /*ntotal*/,
float* distances,
int64_t* labels,
const IDSelector* sel_in,
const FastScanDistancePostProcessing& ctx,
bool multi_bit)
: RHC(nq_val, index->ntotal, sel_in),
rabitq_index(index),
heap_distances(distances),
heap_labels(labels),
nq(nq_val),
k(k_val),
context(ctx),
is_multi_bit(multi_bit),
storage_size(index->compute_per_vector_storage_size()),
packed_block_size(((index->M2 + 1) / 2) * index->bbs),
full_block_size(index->get_block_stride()),
packer(index->get_CodePacker()) {
// 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 != nullptr) {
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;
// Compute block auxiliary region base pointer once per batch.
// Since bbs=32, each batch of 32 vectors aligns to one block.
const size_t block_idx = base_db_idx / rabitq_index->bbs;
const uint8_t* aux_base = rabitq_index->codes.get() +
block_idx * full_block_size + packed_block_size;
// Stats tracking for multi-bit two-stage search only
// n_1bit_evaluations: candidates evaluated using 1-bit lower bound
// n_multibit_evaluations: candidates requiring full multi-bit distance
size_t local_1bit_evaluations = 0;
size_t local_multibit_evaluations = 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 block auxiliary region
const uint8_t* base_ptr = aux_base + i * storage_size;
if (is_multi_bit) {
// Track candidates actually considered for two-stage filtering
local_1bit_evaluations++;
const SignBitFactorsWithError& full_factors =
*reinterpret_cast<const SignBitFactorsWithError*>(base_ptr);
float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
full_factors,
query_factors_data,
rabitq_index->centered,
rabitq_index->qb,
rabitq_index->d);
// Adaptive filtering: decide whether to compute full distance
const bool is_similarity = rabitq_index->metric_type ==
MetricType::METRIC_INNER_PRODUCT;
bool should_refine = rabitq_utils::should_refine_candidate(
dist_1bit,
full_factors.f_error,
context.query_factors ? context.query_factors[q].g_error
: 0.0f,
heap_dis[0],
is_similarity);
if (should_refine) {
local_multibit_evaluations++;
float dist_full = compute_full_multibit_distance(db_idx, q);
if (Cfloat::cmp(heap_dis[0], dist_full)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, dist_full, db_idx);
}
}
} else {
const rabitq_utils::SignBitFactors& db_factors =
*reinterpret_cast<const rabitq_utils::SignBitFactors*>(
base_ptr);
float adjusted_distance =
rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
db_factors,
query_factors_data,
rabitq_index->centered,
rabitq_index->qb,
rabitq_index->d);
// 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);
}
}
}
// Update global stats atomically
#pragma omp atomic
rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
#pragma omp atomic
rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
}
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);
}
}
template <class C, bool with_id_map>
float RaBitQHeapHandler<C, with_id_map>::compute_lower_bound(
float dist_1bit,
size_t db_idx,
size_t q) const {
// Access f_error from block auxiliary region
const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
rabitq_index->codes.get(),
db_idx,
rabitq_index->bbs,
packed_block_size,
full_block_size,
storage_size);
const SignBitFactorsWithError& db_factors =
*reinterpret_cast<const SignBitFactorsWithError*>(base_ptr);
float f_error = db_factors.f_error;
// Get g_error from query factors (query-dependent error term)
float g_error = 0.0f;
if (context.query_factors != nullptr) {
g_error = context.query_factors[q].g_error;
}
// Compute error adjustment: f_error * g_error
float error_adjustment = f_error * g_error;
return dist_1bit - error_adjustment;
}
template <class C, bool with_id_map>
float RaBitQHeapHandler<C, with_id_map>::compute_full_multibit_distance(
size_t db_idx,
size_t q) const {
const size_t ex_bits = rabitq_index->rabitq.nb_bits - 1;
const size_t dim = rabitq_index->d;
const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
rabitq_index->codes.get(),
db_idx,
rabitq_index->bbs,
packed_block_size,
full_block_size,
storage_size);
const size_t ex_code_size = (dim * ex_bits + 7) / 8;
const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError);
const ExtraBitsFactors& ex_fac = *reinterpret_cast<const ExtraBitsFactors*>(
base_ptr + sizeof(SignBitFactorsWithError) + ex_code_size);
// Get query factors reference (avoid copying)
const rabitq_utils::QueryFactorsData& query_factors =
context.query_factors[q];
// Get sign bits from FastScan packed format
std::vector<uint8_t> unpacked_code(rabitq_index->code_size);
packer->unpack_1(rabitq_index->codes.get(), db_idx, unpacked_code.data());
const uint8_t* sign_bits = unpacked_code.data();
return rabitq_utils::compute_full_multibit_distance(
sign_bits,
ex_code,
ex_fac,
query_factors.rotated_q.data(),
(rabitq_index->metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_factors.q_dot_c
: query_factors.qr_to_c_L2sqr,
dim,
ex_bits,
rabitq_index->metric_type);
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context) const {
const bool is_multi_bit = rabitq.nb_bits > 1;
return rabitq_make_knn_scanner(
this, is_max, n, k, distances, labels, sel, context, is_multi_bit);
}
// Implementation of virtual make_knn_handler method
@@ -770,15 +547,14 @@ SIMDResultHandlerToFloat* IndexRaBitQFastScan::make_knn_handler(
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context) const {
// Use runtime boolean for multi-bit mode
const bool multi_bit = rabitq.nb_bits > 1;
if (is_max) {
return new RaBitQHeapHandler<CMax<uint16_t, int>, false>(
this, n, k, distances, labels, sel, context, multi_bit);
this, n, k, distances, labels, sel, &context, multi_bit);
} else {
return new RaBitQHeapHandler<CMin<uint16_t, int>, false>(
this, n, k, distances, labels, sel, context, multi_bit);
this, n, k, distances, labels, sel, &context, multi_bit);
}
}
faiss/IndexRaBitQFastScan.h
+185 -14
View File
@@ -92,6 +92,17 @@ struct IndexRaBitQFastScan : IndexFastScan {
idx_t* labels,
const SearchParameters* params = nullptr) const override;
/// RaBitQ scanner wired through rabitq_make_knn_scanner
std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
bool is_max,
idx_t n,
idx_t k,
size_t ntotal,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context = {}) const override;
/// Override to create RaBitQ-specific handlers
SIMDResultHandlerToFloat* make_knn_handler(
bool is_max,
@@ -122,18 +133,24 @@ struct IndexRaBitQFastScan : IndexFastScan {
*
* @tparam C Comparator type (CMin/CMax) for heap operations
* @tparam with_id_map Whether to use id mapping (similar to HeapHandler)
* @tparam SL SIMD level for per-TU template instantiation
*/
template <class C, bool with_id_map = false>
template <
class C,
bool with_id_map = false,
SIMDLevel SL = SINGLE_SIMD_LEVEL_256>
struct RaBitQHeapHandler
: simd_result_handlers::ResultHandlerCompare<C, with_id_map> {
using RHC = simd_result_handlers::ResultHandlerCompare<C, with_id_map>;
: simd_result_handlers::ResultHandlerCompare<C, with_id_map, SL> {
using RHC = simd_result_handlers::ResultHandlerCompare<C, with_id_map, SL>;
using RHC::normalizers;
static constexpr SIMDLevel SL256 = simd256_level_selector<SL>::value;
using simd16uint16 = simd16uint16_tpl<SL256>;
const IndexRaBitQFastScan* rabitq_index;
float* heap_distances; // [nq * k]
int64_t* heap_labels; // [nq * k]
const size_t nq, k;
const FastScanDistancePostProcessing&
const FastScanDistancePostProcessing*
context; // Processing context with query offset
const bool is_multi_bit; // Runtime flag for multi-bit mode
@@ -156,22 +173,176 @@ struct RaBitQHeapHandler
float* distances,
int64_t* labels,
const IDSelector* sel_in,
const FastScanDistancePostProcessing& context,
bool multi_bit);
const FastScanDistancePostProcessing* ctx,
bool multi_bit)
: RHC(nq_val, index->ntotal, sel_in),
rabitq_index(index),
heap_distances(distances),
heap_labels(labels),
nq(nq_val),
k(k_val),
context(ctx),
is_multi_bit(multi_bit),
storage_size(index->compute_per_vector_storage_size()),
packed_block_size(((index->M2 + 1) / 2) * index->bbs),
full_block_size(index->get_block_stride()),
packer(index->get_CodePacker()) {
#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);
}
}
void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override;
void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override {
ALIGNED(32) uint16_t d32tab[32];
d0.store(d32tab);
d1.store(d32tab + 16);
void begin(const float* norms);
float* const heap_dis = heap_distances + q * k;
int64_t* const heap_ids = heap_labels + q * k;
void end();
rabitq_utils::QueryFactorsData query_factors_data = {};
if (context && context->query_factors != nullptr) {
query_factors_data = context->query_factors[q];
}
const float one_a = normalizers ? (1.0f / normalizers[2 * q]) : 1.0f;
const float bias = normalizers ? normalizers[2 * q + 1] : 0.0f;
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;
const size_t block_idx = base_db_idx / rabitq_index->bbs;
const uint8_t* aux_base = rabitq_index->codes.get() +
block_idx * full_block_size + packed_block_size;
size_t local_1bit_evaluations = 0;
size_t local_multibit_evaluations = 0;
for (size_t i = 0; i < max_vectors; i++) {
const size_t db_idx = base_db_idx + i;
const float normalized_distance = d32tab[i] * one_a + bias;
const uint8_t* base_ptr = aux_base + i * storage_size;
if (is_multi_bit) {
local_1bit_evaluations++;
const SignBitFactorsWithError& full_factors =
*reinterpret_cast<const SignBitFactorsWithError*>(
base_ptr);
float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
full_factors,
query_factors_data,
rabitq_index->centered,
rabitq_index->qb,
rabitq_index->d);
const bool is_similarity = rabitq_index->metric_type ==
MetricType::METRIC_INNER_PRODUCT;
bool should_refine = rabitq_utils::should_refine_candidate(
dist_1bit,
full_factors.f_error,
context && context->query_factors
? context->query_factors[q].g_error
: 0.0f,
heap_dis[0],
is_similarity);
if (should_refine) {
local_multibit_evaluations++;
float dist_full = compute_full_multibit_distance(db_idx, q);
if (Cfloat::cmp(heap_dis[0], dist_full)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, dist_full, db_idx);
}
}
} else {
const rabitq_utils::SignBitFactors& db_factors =
*reinterpret_cast<const rabitq_utils::SignBitFactors*>(
base_ptr);
float adjusted_distance =
rabitq_utils::compute_1bit_adjusted_distance(
normalized_distance,
db_factors,
query_factors_data,
rabitq_index->centered,
rabitq_index->qb,
rabitq_index->d);
if (Cfloat::cmp(heap_dis[0], adjusted_distance)) {
heap_replace_top<Cfloat>(
k, heap_dis, heap_ids, adjusted_distance, db_idx);
}
}
}
#pragma omp atomic
rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
#pragma omp atomic
rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
}
void begin(const float* norms) {
normalizers = norms;
}
void end() {
#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);
}
}
private:
/// Compute full multi-bit distance for a candidate vector (multi-bit only)
float compute_full_multibit_distance(size_t db_idx, size_t q) const;
float compute_full_multibit_distance(size_t db_idx, size_t q) const {
const size_t ex_bits = rabitq_index->rabitq.nb_bits - 1;
const size_t dim = rabitq_index->d;
/// Compute lower bound using 1-bit distance and error bound (multi-bit
/// only)
float compute_lower_bound(float dist_1bit, size_t db_idx, size_t q) const;
const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
rabitq_index->codes.get(),
db_idx,
rabitq_index->bbs,
packed_block_size,
full_block_size,
storage_size);
const size_t ex_code_size = (dim * ex_bits + 7) / 8;
const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError);
const ExtraBitsFactors& ex_fac =
*reinterpret_cast<const ExtraBitsFactors*>(
base_ptr + sizeof(SignBitFactorsWithError) +
ex_code_size);
const rabitq_utils::QueryFactorsData& query_factors =
context->query_factors[q];
std::vector<uint8_t> unpacked_code(rabitq_index->code_size);
packer->unpack_1(
rabitq_index->codes.get(), db_idx, unpacked_code.data());
const uint8_t* sign_bits = unpacked_code.data();
return rabitq_utils::compute_full_multibit_distance(
sign_bits,
ex_code,
ex_fac,
query_factors.rotated_q.data(),
(rabitq_index->metric_type == MetricType::METRIC_INNER_PRODUCT)
? query_factors.q_dot_c
: query_factors.qr_to_c_L2sqr,
dim,
ex_bits,
rabitq_index->metric_type);
}
};
} // namespace faiss
faiss/impl/fast_scan/fast_scan.cpp
+46 -1
View File
@@ -362,7 +362,8 @@ int pq4_pack_LUT_qbs_q_map(
***************************************************************/
#define THE_LEVEL_TO_DISPATCH SIMDLevel::NONE
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/rabitq_dispatching.h> // IWYU pragma: keep
#undef THE_LEVEL_TO_DISPATCH
namespace faiss {
@@ -390,4 +391,48 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
with_id_map);
}
std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner(
const IndexRaBitQFastScan* index,
bool is_max,
size_t nq,
int64_t k,
float* distances,
int64_t* ids,
const IDSelector* sel,
const FastScanDistancePostProcessing& context,
bool is_multi_bit) {
DISPATCH_SIMDLevel(
rabitq_make_knn_scanner_impl,
index,
is_max,
nq,
k,
distances,
ids,
sel,
context,
is_multi_bit);
}
std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner(
bool is_max,
const IndexIVFRaBitQFastScan* index,
size_t nq,
size_t k,
float* distances,
int64_t* ids,
const FastScanDistancePostProcessing* context,
bool multi_bit) {
DISPATCH_SIMDLevel(
rabitq_ivf_make_knn_scanner_impl,
is_max,
index,
nq,
k,
distances,
ids,
context,
multi_bit);
}
} // namespace faiss
faiss/impl/fast_scan/fast_scan.h
+58 -1
View File
@@ -26,9 +26,9 @@
namespace faiss {
struct IDSelector;
struct SIMDResultHandler;
struct SIMDResultHandlerToFloat;
struct IDSelector;
/** Pack codes for consumption by the SIMD kernels.
* The unused bytes are set to 0.
@@ -288,4 +288,61 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
const IDSelector* sel,
bool with_id_map = false);
/***************************************************************
* RaBitQ scanner factory: per-SIMD specializations live in
* rabitq_dispatching.h, included by each per-SIMD TU.
***************************************************************/
struct IndexRaBitQFastScan;
struct IndexIVFRaBitQFastScan;
struct FastScanDistancePostProcessing;
/// Per-SIMD factory (primary template; specializations in rabitq_dispatching.h)
template <SIMDLevel SL>
std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl(
const IndexRaBitQFastScan* index,
bool is_max,
size_t nq,
int64_t k,
float* distances,
int64_t* ids,
const IDSelector* sel,
const FastScanDistancePostProcessing& context,
bool is_multi_bit);
/// Runtime dispatch wrapper for rabitq_make_knn_scanner_impl
std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner(
const IndexRaBitQFastScan* index,
bool is_max,
size_t nq,
int64_t k,
float* distances,
int64_t* ids,
const IDSelector* sel,
const FastScanDistancePostProcessing& context,
bool is_multi_bit);
/// Per-SIMD IVF RaBitQ scanner factory.
template <SIMDLevel SL>
std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner_impl(
bool is_max,
const IndexIVFRaBitQFastScan* index,
size_t nq,
size_t k,
float* distances,
int64_t* ids,
const FastScanDistancePostProcessing* context,
bool multi_bit);
/// Runtime dispatch wrapper for IVF RaBitQ scanner.
std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner(
bool is_max,
const IndexIVFRaBitQFastScan* index,
size_t nq,
size_t k,
float* distances,
int64_t* ids,
const FastScanDistancePostProcessing* context,
bool multi_bit);
} // namespace faiss
faiss/impl/fast_scan/impl-avx2.cpp
+2 -1
View File
@@ -8,6 +8,7 @@
#ifdef COMPILE_SIMD_AVX2
#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/rabitq_dispatching.h> // IWYU pragma: keep
#endif // COMPILE_SIMD_AVX2
faiss/impl/fast_scan/impl-avx512.cpp
+2 -1
View File
@@ -8,6 +8,7 @@
#ifdef COMPILE_SIMD_AVX512
#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX512
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/rabitq_dispatching.h> // IWYU pragma: keep
#endif // COMPILE_SIMD_AVX512
faiss/impl/fast_scan/impl-neon.cpp
+33 -1
View File
@@ -8,7 +8,8 @@
#ifdef COMPILE_SIMD_ARM_NEON
#define THE_LEVEL_TO_DISPATCH SIMDLevel::ARM_NEON
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
#include <faiss/impl/fast_scan/rabitq_dispatching.h> // IWYU pragma: keep
// ARM_SVE: forward to ARM_NEON implementation until a dedicated SVE
// specialization is written (same pattern as scalar_quantizer/sq-neon.cpp).
@@ -32,6 +33,37 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl<
is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map);
}
template <>
std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl<
SIMDLevel::ARM_SVE>(
const IndexRaBitQFastScan* index,
bool is_max,
size_t nq,
int64_t k,
float* distances,
int64_t* ids,
const IDSelector* sel,
const FastScanDistancePostProcessing& context,
bool is_multi_bit) {
return rabitq_make_knn_scanner_impl<SIMDLevel::ARM_NEON>(
index, is_max, nq, k, distances, ids, sel, context, is_multi_bit);
}
template <>
std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner_impl<
SIMDLevel::ARM_SVE>(
bool is_max,
const IndexIVFRaBitQFastScan* index,
size_t nq,
size_t k,
float* distances,
int64_t* ids,
const FastScanDistancePostProcessing* context,
bool multi_bit) {
return rabitq_ivf_make_knn_scanner_impl<SIMDLevel::ARM_NEON>(
is_max, index, nq, k, distances, ids, context, multi_bit);
}
} // namespace faiss
#endif // COMPILE_SIMD_ARM_SVE
faiss/impl/fast_scan/rabitq_dispatching.h
+81
View File
@@ -0,0 +1,81 @@
/*
* 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
/**
* @file rabitq_dispatching.h
* @brief Per-SIMD TU dispatch for RaBitQ flat scanner.
*
* Included after dispatching.h in each per-SIMD TU, so that
* ScannerMixIn from dispatching.h is visible (same TU).
*
* Provides the rabitq_make_knn_scanner_impl<SL> specialization
* that wraps RaBitQHeapHandler in ScannerMixIn.
*/
#ifndef THE_LEVEL_TO_DISPATCH
#error "Define THE_LEVEL_TO_DISPATCH before including this header"
#endif
#include <faiss/IndexIVFRaBitQFastScan.h>
#include <faiss/IndexRaBitQFastScan.h>
#include <faiss/impl/fast_scan/fast_scan.h>
// ScannerMixIn is visible from dispatching.h (same TU)
namespace faiss {
template <>
std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl<
THE_LEVEL_TO_DISPATCH>(
const IndexRaBitQFastScan* index,
bool is_max,
size_t nq,
int64_t k,
float* distances,
int64_t* ids,
const IDSelector* sel,
const FastScanDistancePostProcessing& context,
bool is_multi_bit) {
if (is_max) {
using H = RaBitQHeapHandler<CMax<uint16_t, int>, false>;
return std::make_unique<ScannerMixIn<H>>(
index, nq, k, distances, ids, sel, &context, is_multi_bit);
} else {
using H = RaBitQHeapHandler<CMin<uint16_t, int>, false>;
return std::make_unique<ScannerMixIn<H>>(
index, nq, k, distances, ids, sel, &context, is_multi_bit);
}
}
// IVF RaBitQ scanner factory
template <>
std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner_impl<
THE_LEVEL_TO_DISPATCH>(
bool is_max,
const IndexIVFRaBitQFastScan* index,
size_t nq,
size_t k,
float* distances,
int64_t* ids,
const FastScanDistancePostProcessing* context,
bool multi_bit) {
if (is_max) {
using C = CMax<uint16_t, int64_t>;
using H = simd_result_handlers::IVFRaBitQHeapHandler<C>;
return std::make_unique<ScannerMixIn<H>>(
index, nq, k, distances, ids, context, multi_bit);
} else {
using C = CMin<uint16_t, int64_t>;
using H = simd_result_handlers::IVFRaBitQHeapHandler<C>;
return std::make_unique<ScannerMixIn<H>>(
index, nq, k, distances, ids, context, multi_bit);
}
}
} // namespace faiss
faiss/impl/fast_scan/rabitq_result_handler.h
+14 -22
View File
@@ -11,33 +11,29 @@
#include <vector>
#include <faiss/impl/CodePacker.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/RaBitQStats.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
#include <faiss/impl/fast_scan/simd_result_handlers.h>
#include <faiss/invlists/InvertedLists.h>
#include <faiss/utils/Heap.h>
namespace faiss {
// Forward declaration — full definition needed only in implementation
// Forward declaration — full definition in IndexIVFRaBitQFastScan.h
struct IndexIVFRaBitQFastScan;
namespace simd_result_handlers {
// Import shared utilities from RaBitQUtils
using rabitq_utils::ExtraBitsFactors;
using rabitq_utils::SignBitFactors;
using rabitq_utils::SignBitFactorsWithError;
/** SIMD result handler for IndexIVFRaBitQFastScan that applies
* RaBitQ-specific distance corrections during batch processing.
*
* This handler processes batches of 32 distance computations from SIMD
* kernels, applies RaBitQ distance formula adjustments (factors and
* normalizers), and immediately updates result heaps. This eliminates the
* need for post-processing and provides significant performance benefits.
*
* Key optimizations:
* - Direct heap integration with no intermediate result storage
* - Batch-level computation of normalizers and query factors
* - Specialized handling for both centered and non-centered quantization
* modes
* - Efficient inner product metric corrections
* - Uses runtime boolean for multi-bit mode
*
* @tparam C Comparator type (CMin/CMax) for heap operations
* @tparam SL SIMD level for dynamic dispatch
*/
@@ -71,6 +67,9 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare<C, true, SL> {
CMax<float, int64_t>,
CMin<float, int64_t>>::type;
// Constructor and method bodies are defined at the bottom of
// IndexIVFRaBitQFastScan.h (after the full struct definition is
// available) to break the circular header dependency.
IVFRaBitQHeapHandler(
const IndexIVFRaBitQFastScan* idx,
size_t nq_val,
@@ -82,7 +81,6 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare<C, true, SL> {
void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) override;
/// Override base class virtual method to receive context information
void set_list_context(size_t list_no, const std::vector<int>& probe_map)
override;
@@ -95,14 +93,8 @@ struct IVFRaBitQHeapHandler : ResultHandlerCompare<C, true, SL> {
}
private:
/// Compute full multi-bit distance for a candidate vector (multi-bit
/// only)
/// @param db_idx Global database vector index
/// @param local_q Batch-local query index (for probe_indices access)
/// @param global_q Global query index (for storage indexing)
/// @param local_offset Offset within the current inverted list
float compute_full_multibit_distance(
size_t /*db_idx*/,
size_t db_idx,
size_t local_q,
size_t global_q,
size_t local_offset) const;
faiss/python/swigfaiss.swig
+14 -2
View File
@@ -599,11 +599,23 @@ void gpu_sync_all_devices()
%include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
// Ignore make_knn_handler methods — they return SIMDResultHandlerToFloat*
// which is an internal SIMD type not exposed to Python.
// Ignore FastScanCodeScanner and related factory functions — internal SIMD types.
%ignore faiss::FastScanCodeScanner;
%ignore faiss::make_fast_scan_scanner_impl;
%ignore faiss::make_fast_scan_knn_scanner;
%ignore faiss::rabitq_make_knn_scanner_impl;
%ignore faiss::rabitq_make_knn_scanner;
%ignore faiss::rabitq_ivf_make_knn_scanner_impl;
%ignore faiss::rabitq_ivf_make_knn_scanner;
// Ignore make_knn_scanner / make_knn_handler — they return internal SIMD types.
%ignore faiss::IndexFastScan::make_knn_scanner;
%ignore faiss::IndexFastScan::make_knn_handler;
%ignore faiss::IndexIVFFastScan::make_knn_scanner;
%ignore faiss::IndexIVFFastScan::make_knn_handler;
%ignore faiss::IndexRaBitQFastScan::make_knn_scanner;
%ignore faiss::IndexRaBitQFastScan::make_knn_handler;
%ignore faiss::IndexIVFRaBitQFastScan::make_knn_scanner;
%ignore faiss::IndexIVFRaBitQFastScan::make_knn_handler;
%include <faiss/IndexFastScan.h>