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

Switch all search paths to FastScanCodeScanner and remove make_knn_handler (#4904)

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

Now that every make_knn_scanner() returns a real scanner (never nullptr),
switch all callers to use scanners directly and remove make_knn_handler.

Add impl param to make_knn_scanner and its factory so heap vs reservoir
selection respects the implem setting. Switch IndexFastScan
search_implem_12/14, IndexIVFFastScan search_dispatch_implem/range_search,
and IndexIVFPQFastScan scan_codes to scanner-only paths. Add range scanner
factories (make_range_scanner, make_partial_range_scanner). Remove
make_knn_handler from all classes. Change search_implem_10/12 scanner
param from pointer to reference.

Reviewed By: mdouze

Differential Revision: D96116555

fbshipit-source-id: 977c24180d67bf90f7c151ac4ca6fe4e2d5e1f92
This commit is contained in:
Gergely Szilvasy
2026-03-12 02:15:42 -07:00
committed by meta-codesync[bot]
parent 57190f91cd
commit ea11f0ef9c
15 changed files with 315 additions and 409 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/IndexFastScan.cpp
+24 -85
View File
@@ -24,8 +24,6 @@
namespace faiss {
using namespace simd_result_handlers;
inline size_t roundup(size_t a, size_t b) {
return (a + b - 1) / b * b;
}
@@ -219,48 +217,10 @@ std::unique_ptr<FastScanCodeScanner> IndexFastScan::make_knn_scanner(
float* distances,
idx_t* labels,
const IDSelector* sel,
int impl,
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,
int impl,
idx_t n,
idx_t k,
size_t ntotal,
float* distances,
idx_t* labels,
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 {
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);
}
}
is_max, impl, n, ntotal, k, distances, labels, sel);
}
using namespace quantize_lut;
@@ -481,7 +441,6 @@ void IndexFastScan::search_implem_12(
idx_t* labels,
int impl,
const FastScanDistancePostProcessing& context) const {
using RH = ResultHandlerCompare<C, false>;
FAISS_THROW_IF_NOT(bbs == 32);
// handle qbs2 blocking by recursive call
@@ -532,36 +491,26 @@ 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(
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();
if (skip & 4) {
// pass
} else {
pq4_accumulate_loop_qbs(
auto scanner = make_knn_scanner(
C::is_max, n, k, ntotal, distances, labels, nullptr, impl, context);
auto* rh = scanner->handler();
rh->normalizers = normalizers.get();
// Note: skip & 2 previously set handler->disable (run kernel,
// discard results). Through the scanner path, skip & 2 now skips
// the kernel entirely (same as skip & 4), since disable is not
// accessible through the SIMDResultHandlerToFloat* interface.
if (!(skip & (2 | 4))) {
scanner->accumulate_loop_qbs(
qbs,
ntotal2,
M2,
codes.get(),
LUT.get(),
*handler.get(),
context.pq2x4_scale,
get_block_stride());
}
if (!(skip & 8)) {
handler->end();
rh->end();
}
}
@@ -576,7 +525,6 @@ void IndexFastScan::search_implem_14(
idx_t* labels,
int impl,
const FastScanDistancePostProcessing& context) const {
using RH = ResultHandlerCompare<C, false>;
FAISS_THROW_IF_NOT(bbs % 32 == 0);
int qbs2 = qbs == 0 ? 4 : qbs;
@@ -616,36 +564,27 @@ void IndexFastScan::search_implem_14(
AlignedTable<uint8_t> LUT(n * dim12);
pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get());
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();
if (skip & 4) {
// pass
} else {
pq4_accumulate_loop(
auto scanner = make_knn_scanner(
C::is_max, n, k, ntotal, distances, labels, nullptr, impl, context);
auto* rh = scanner->handler();
rh->normalizers = normalizers.get();
// Note: skip & 2 previously set handler->disable (run kernel,
// discard results). Through the scanner path, skip & 2 now skips
// the kernel entirely (same as skip & 4), since disable is not
// accessible through the SIMDResultHandlerToFloat* interface.
if (!(skip & (2 | 4))) {
scanner->accumulate_loop(
n,
ntotal2,
bbs,
M2,
codes.get(),
LUT.get(),
*handler.get(),
context.pq2x4_scale,
get_block_stride());
}
if (!(skip & 8)) {
handler->end();
rh->end();
}
}
faiss/IndexFastScan.h
+1 -29
View File
@@ -18,7 +18,6 @@ namespace faiss {
struct CodePacker;
struct IDSelector;
struct SIMDResultHandlerToFloat;
/** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now.
*
@@ -138,36 +137,9 @@ struct IndexFastScan : Index {
float* distances,
idx_t* labels,
const IDSelector* sel,
int impl = 0,
const FastScanDistancePostProcessing& context = {}) const;
/** 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 context processing context for distance post-processing
* @return pointer to created handler (never returns nullptr)
*/
virtual SIMDResultHandlerToFloat* 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,
faiss/IndexIVFFastScan.cpp
+113 -145
View File
@@ -28,8 +28,6 @@
namespace faiss {
using namespace simd_result_handlers;
inline size_t roundup(size_t a, size_t b) {
return (a + b - 1) / b * b;
}
@@ -397,28 +395,6 @@ void IndexIVFFastScan::range_search(
namespace {
template <class C>
ResultHandlerCompare<C, true>* make_knn_handler_fixC(
int impl,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const float* normalizers) {
using HeapHC = HeapHandler<C, true>;
using ReservoirHC = ReservoirHandler<C, true>;
using SingleResultHC = SingleResultHandler<C, true>;
if (k == 1) {
return new SingleResultHC(n, 0, distances, labels, sel);
} else if (impl % 2 == 0) {
return new HeapHC(n, 0, k, distances, labels, sel, normalizers);
} else /* if (impl % 2 == 1) */ {
return new ReservoirHC(n, 0, k, 2 * k, distances, labels, sel);
}
}
using CoarseQuantized = IndexIVFFastScan::CoarseQuantized;
struct CoarseQuantizedWithBuffer : CoarseQuantized {
@@ -502,10 +478,11 @@ std::unique_ptr<FastScanCodeScanner> IndexIVFFastScan::make_knn_scanner(
float* distances,
idx_t* labels,
const IDSelector* sel,
int impl,
const FastScanDistancePostProcessing&) const {
return make_fast_scan_knn_scanner(
is_max,
0,
impl,
n,
0,
k,
@@ -515,25 +492,6 @@ std::unique_ptr<FastScanCodeScanner> IndexIVFFastScan::make_knn_scanner(
/*with_id_map=*/true);
}
SIMDResultHandlerToFloat* IndexIVFFastScan::make_knn_handler(
bool is_max,
int impl,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing&,
const float* normalizers) const {
if (is_max) {
return make_knn_handler_fixC<CMax<uint16_t, int64_t>>(
impl, n, k, distances, labels, sel, normalizers);
} else {
return make_knn_handler_fixC<CMin<uint16_t, int64_t>>(
impl, n, k, distances, labels, sel, normalizers);
}
}
void IndexIVFFastScan::search_dispatch_implem(
idx_t n,
const float* x,
@@ -549,7 +507,6 @@ void IndexIVFFastScan::search_dispatch_implem(
params ? params->quantizer_params : nullptr;
bool is_max = !is_similarity_metric(metric_type);
using RH = SIMDResultHandlerToFloat;
if (n == 0) {
return;
@@ -607,43 +564,37 @@ void IndexIVFFastScan::search_dispatch_implem(
size_t ndis = 0, nlist_visited = 0;
if (!multiple_threads) {
// clang-format off
if (impl == 12 || impl == 13) {
std::unique_ptr<RH> handler(
static_cast<RH*>(this->make_knn_handler(
is_max,
impl,
n,
k,
distances,
labels,
sel,
context))
);
search_implem_12(
n, x, *handler.get(),
cq, &ndis, &nlist_visited, context, params);
} else if (impl == 14 || impl == 15) {
if (impl == 14 || impl == 15) {
search_implem_14(
n, x, k, distances, labels,
cq, impl, context, params);
n, x, k, distances, labels, cq, impl, context, params);
} else {
std::unique_ptr<RH> handler(
static_cast<RH*>(this->make_knn_handler(
is_max,
impl,
n,
k,
distances,
labels,
sel,
context))
);
search_implem_10(
n, x, *handler.get(), cq,
&ndis, &nlist_visited, context, params);
auto scanner = make_knn_scanner(
is_max, n, k, distances, labels, sel, impl, context);
auto* handler = scanner->handler();
if (impl == 12 || impl == 13) {
search_implem_12(
n,
x,
*handler,
cq,
&ndis,
&nlist_visited,
context,
params,
*scanner);
} else {
search_implem_10(
n,
x,
*handler,
cq,
&ndis,
&nlist_visited,
context,
params,
*scanner);
}
}
// clang-format on
} else {
// explicitly slice over threads
int nslice = compute_search_nslice(this, n, cq.nprobe);
@@ -671,27 +622,39 @@ void IndexIVFFastScan::search_dispatch_implem(
thread_context.query_factors += i0 * nprobe;
}
std::unique_ptr<RH> handler(
static_cast<RH*>(this->make_knn_handler(
is_max,
impl,
i1 - i0,
k,
dis_i,
lab_i,
sel,
thread_context)));
// clang-format off
auto scanner = make_knn_scanner(
is_max,
i1 - i0,
k,
dis_i,
lab_i,
sel,
impl,
thread_context);
auto* handler = scanner->handler();
if (impl == 12 || impl == 13) {
search_implem_12(
i1 - i0, x + i0 * d, *handler.get(),
cq_i, &ndis, &nlist_visited, thread_context, params);
i1 - i0,
x + i0 * d,
*handler,
cq_i,
&ndis,
&nlist_visited,
thread_context,
params,
*scanner);
} else {
search_implem_10(
i1 - i0, x + i0 * d, *handler.get(),
cq_i, &ndis, &nlist_visited, thread_context, params);
i1 - i0,
x + i0 * d,
*handler,
cq_i,
&ndis,
&nlist_visited,
thread_context,
params,
*scanner);
}
// clang-format on
}
}
}
@@ -749,20 +712,30 @@ void IndexIVFFastScan::range_search_dispatch_implem(
size_t ndis = 0, nlist_visited = 0;
if (!multiple_threads) { // single thread
std::unique_ptr<SIMDResultHandlerToFloat> handler;
if (is_max) {
handler.reset(new RangeHandler<CMax<uint16_t, int64_t>, true>(
rres, radius, 0, sel));
} else {
handler.reset(new RangeHandler<CMin<uint16_t, int64_t>, true>(
rres, radius, 0, sel));
}
auto scanner = make_range_scanner(is_max, rres, radius, 0, sel);
auto* handler = scanner->handler();
if (impl == 12) {
search_implem_12(
n, x, *handler.get(), cq, &ndis, &nlist_visited, context);
n,
x,
*handler,
cq,
&ndis,
&nlist_visited,
context,
nullptr,
*scanner);
} else if (impl == 10) {
search_implem_10(
n, x, *handler.get(), cq, &ndis, &nlist_visited, context);
n,
x,
*handler,
cq,
&ndis,
&nlist_visited,
context,
nullptr,
*scanner);
} else {
FAISS_THROW_FMT("Range search implem %d not implemented", impl);
}
@@ -781,35 +754,32 @@ void IndexIVFFastScan::range_search_dispatch_implem(
if (!cq_i.done()) {
cq_i.quantize_slice(quantizer, x, quantizer_params);
}
std::unique_ptr<SIMDResultHandlerToFloat> handler;
if (is_max) {
handler.reset(new PartialRangeHandler<
CMax<uint16_t, int64_t>,
true>(pres, radius, 0, i0, i1, sel));
} else {
handler.reset(new PartialRangeHandler<
CMin<uint16_t, int64_t>,
true>(pres, radius, 0, i0, i1, sel));
}
auto scanner = make_partial_range_scanner(
is_max, pres, radius, 0, i0, i1, sel);
auto* handler = scanner->handler();
if (impl == 12 || impl == 13) {
search_implem_12(
i1 - i0,
x + i0 * d,
*handler.get(),
*handler,
cq_i,
&ndis,
&nlist_visited,
context);
context,
nullptr,
*scanner);
} else {
search_implem_10(
i1 - i0,
x + i0 * d,
*handler.get(),
*handler,
cq_i,
&ndis,
&nlist_visited,
context);
context,
nullptr,
*scanner);
}
}
pres.finalize();
@@ -986,7 +956,8 @@ void IndexIVFFastScan::search_implem_10(
size_t* ndis_out,
size_t* nlist_out,
const FastScanDistancePostProcessing& context,
const IVFSearchParameters* /* params */) const {
const IVFSearchParameters* /* params */,
FastScanCodeScanner& scanner) const {
size_t dim12 = ksub * M2;
AlignedTable<uint8_t> dis_tables;
AlignedTable<uint16_t> biases;
@@ -1042,14 +1013,13 @@ void IndexIVFFastScan::search_implem_10(
probe_map[0] = static_cast<int>(j);
handler.set_list_context(list_no, probe_map);
pq4_accumulate_loop(
scanner.accumulate_loop(
1,
roundup(ls, bbs),
bbs,
M2,
codes.get(),
LUT,
handler,
context.pq2x4_scale,
get_block_stride());
@@ -1071,7 +1041,8 @@ void IndexIVFFastScan::search_implem_12(
size_t* ndis_out,
size_t* nlist_out,
const FastScanDistancePostProcessing& context,
const IVFSearchParameters* /* params */) const {
const IVFSearchParameters* /* params */,
FastScanCodeScanner& scanner) const {
if (n == 0) { // does not work well with reservoir
return;
}
@@ -1195,13 +1166,12 @@ void IndexIVFFastScan::search_implem_12(
}
handler.set_list_context(list_no, probe_map);
pq4_accumulate_loop_qbs(
scanner.accumulate_loop_qbs(
qbs_for_list,
list_size,
M2,
codes.get(),
LUT.get(),
handler,
context.pq2x4_scale,
get_block_stride());
// prepare for next loop
@@ -1335,25 +1305,24 @@ void IndexIVFFastScan::search_implem_14(
std::vector<idx_t> local_idx(k * n);
std::vector<float> local_dis(k * n);
// prepare the result handlers
std::unique_ptr<SIMDResultHandlerToFloat> handler(
this->make_knn_handler(
is_max,
impl,
n,
k,
local_dis.data(),
local_idx.data(),
sel,
context));
handler->begin(normalizers.get());
auto scanner = make_knn_scanner(
is_max,
n,
k,
local_dis.data(),
local_idx.data(),
sel,
impl,
context);
SIMDResultHandlerToFloat* handler_ptr = scanner->handler();
handler_ptr->begin(normalizers.get());
int actual_qbs2 = this->qbs2 ? this->qbs2 : 11;
std::vector<uint16_t> tmp_bias;
if (biases.get()) {
tmp_bias.resize(actual_qbs2);
handler->dbias = tmp_bias.data();
handler_ptr->dbias = tmp_bias.data();
}
std::set<int> q_set;
@@ -1405,9 +1374,9 @@ void IndexIVFFastScan::search_implem_14(
// prepare the handler
handler->ntotal = list_size;
handler->q_map = q_map.data();
handler->id_map = ids.get();
handler_ptr->ntotal = list_size;
handler_ptr->q_map = q_map.data();
handler_ptr->id_map = ids.get();
// Set context information for handlers that need additional data
// All queries in this batch access the same list_no, but each
@@ -1417,21 +1386,20 @@ void IndexIVFFastScan::search_implem_14(
const QC& qc = qcs[i];
probe_map[i - i0] = qc.rank;
}
handler->set_list_context(list_no, probe_map);
handler_ptr->set_list_context(list_no, probe_map);
pq4_accumulate_loop_qbs(
scanner->accumulate_loop_qbs(
qbs_for_list,
list_size,
M2,
codes.get(),
LUT.get(),
*handler.get(),
context.pq2x4_scale,
get_block_stride());
}
// labels is in-place for HeapHC
handler->end();
handler_ptr->end();
// merge per-thread results
#pragma omp single
faiss/IndexIVFFastScan.h
+8 -45
View File
@@ -227,8 +227,8 @@ struct IndexIVFFastScan : IndexIVF {
* 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.
* Derived classes override this to provide custom handlers
* (e.g. RaBitQ).
*
* @param is_max whether to use CMax comparator (true) or CMin
* @param n number of queries
@@ -236,7 +236,7 @@ struct IndexIVFFastScan : IndexIVF {
* @param distances output distances array
* @param labels output labels array
* @param sel optional ID selector
* @return scanner, or nullptr if unsupported
* @return scanner
*/
virtual std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
bool is_max,
@@ -245,48 +245,9 @@ struct IndexIVFFastScan : IndexIVF {
float* distances,
idx_t* labels,
const IDSelector* sel,
int impl = 0,
const FastScanDistancePostProcessing& context = {}) const;
/** Create a KNN handler for this index type
*
* This method can be overridden by derived classes to provide
* specialized handlers (e.g., IVFRaBitQHeapHandler for RaBitQ indexes).
* Base implementation creates standard handlers based on k and impl.
*
* @param is_max true for max-heap (inner product), false for
* min-heap (L2 distance)
* @param impl implementation number:
* - even (10, 12, 14): use heap for top-k
* - odd (11, 13, 15): use reservoir sampling
* @param n number of queries
* @param k number of neighbors to find per query
* @param distances output array for distances (n * k), will be
* populated by handler
* @param labels output array for result IDs (n * k), will be
* populated by handler
* @param sel optional ID selector to filter results (nullptr =
* no filtering)
* @param context processing context containing additional data
* @param normalizers optional array of size 2*n for converting quantized
* uint16 distances to float.
*
* @return Allocated result handler (caller owns and must delete).
* Handler processes SIMD batches and populates distances/labels.
*
* @note The returned handler must be deleted by caller after use.
* Typical usage: handler->begin() → process batches → handler->end()
*/
virtual SIMDResultHandlerToFloat* make_knn_handler(
bool is_max,
int impl,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context,
const float* normalizers = nullptr) const;
// dispatch to implementations and parallelize
void search_dispatch_implem(
idx_t n,
@@ -340,7 +301,8 @@ struct IndexIVFFastScan : IndexIVF {
size_t* ndis_out,
size_t* nlist_out,
const FastScanDistancePostProcessing& context,
const IVFSearchParameters* params = nullptr) const;
const IVFSearchParameters* params,
FastScanCodeScanner& scanner) const;
void search_implem_12(
idx_t n,
@@ -350,7 +312,8 @@ struct IndexIVFFastScan : IndexIVF {
size_t* ndis_out,
size_t* nlist_out,
const FastScanDistancePostProcessing& context,
const IVFSearchParameters* params = nullptr) const;
const IVFSearchParameters* params,
FastScanCodeScanner& scanner) const;
// implem 14 is multithreaded internally across nprobes and queries
void search_implem_14(
faiss/IndexIVFPQFastScan.cpp
+16 -26
View File
@@ -22,13 +22,12 @@
#include <faiss/invlists/BlockInvertedLists.h>
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
#include <faiss/impl/fast_scan/fast_scan.h>
#include <faiss/impl/fast_scan/simd_result_handlers.h>
namespace faiss {
using namespace simd_result_handlers;
inline size_t roundup(size_t a, size_t b) {
return (a + b - 1) / b * b;
}
@@ -384,47 +383,38 @@ struct IVFPQFastScanScanner : InvertedListScanner {
// the prior loop
std::vector<float> curr_dists(k, distances[0]);
std::vector<idx_t> curr_labels(k, labels[0]);
FastScanDistancePostProcessing empty_context{};
std::unique_ptr<SIMDResultHandlerToFloat> handler(
index.make_knn_handler(
!keep_max,
impl,
nq,
k,
curr_dists.data(),
curr_labels.data(),
sel,
empty_context,
&normalizers[0]));
auto scanner = index.make_knn_scanner(
!keep_max, nq, k, curr_dists.data(), curr_labels.data(), sel);
SIMDResultHandlerToFloat* rh = scanner->handler();
// This does not quite match search_implem_10, but it is fine because
// the scanner operates on a single query at a time, and this value is
// used as the query index. For a single query, the value is always 0.
int qmap1[1] = {0};
handler->q_map = qmap1;
handler->begin(&normalizers[0]);
rh->q_map = qmap1;
rh->begin(&normalizers[0]);
const uint8_t* LUT = dis_tables.get();
handler->dbias = biases.get();
rh->dbias = biases.get();
rh->ntotal = ntotal;
rh->id_map = ids;
handler->ntotal = ntotal;
handler->id_map = ids;
pq4_accumulate_loop(
scanner->accumulate_loop(
1,
roundup(ntotal, index.bbs),
index.bbs,
static_cast<int>(index.M2),
codes,
LUT,
*handler,
dis_tables.get(),
0,
index.get_block_stride());
// The handler is for the results of this iteration.
// Then we need a second heap to combine across iterations.
handler->end();
rh->end();
if (keep_max) {
minheap_addn(
k,
@@ -443,7 +433,7 @@ struct IVFPQFastScanScanner : InvertedListScanner {
k);
}
return handler->num_updates();
return rh->num_updates();
}
};
faiss/IndexIVFRaBitQFastScan.cpp
+2 -24
View File
@@ -18,7 +18,6 @@
#include <faiss/impl/RaBitQuantizerMultiBit.h>
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
#include <faiss/impl/fast_scan/fast_scan.h>
#include <faiss/impl/fast_scan/simd_result_handlers.h>
#include <faiss/invlists/BlockInvertedLists.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/utils.h>
@@ -546,35 +545,13 @@ std::unique_ptr<FastScanCodeScanner> IndexIVFRaBitQFastScan::make_knn_scanner(
float* distances,
idx_t* labels,
const IDSelector* /*sel*/,
int /*impl*/,
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,
int /* impl */,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* /* sel */,
const FastScanDistancePostProcessing& context,
const float* /* normalizers */) const {
const size_t ex_bits = rabitq.nb_bits - 1;
const bool is_multibit = ex_bits > 0;
if (is_max) {
return new IVFRaBitQHeapHandler<CMax<uint16_t, int64_t>>(
this, n, k, distances, labels, &context, is_multibit);
} else {
return new IVFRaBitQHeapHandler<CMin<uint16_t, int64_t>>(
this, n, k, distances, labels, &context, is_multibit);
}
}
/*********************************************************
* IVFRaBitQFastScanScanner implementation
*********************************************************/
@@ -665,6 +642,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
curr_dists.data(),
curr_labels.data(),
sel,
0,
context);
auto* handler = scanner->handler();
faiss/IndexIVFRaBitQFastScan.h
+1 -12
View File
@@ -156,20 +156,9 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
float* distances,
idx_t* labels,
const IDSelector* sel,
int impl = 0,
const FastScanDistancePostProcessing& context = {}) const override;
/// Override to create RaBitQ-specific handlers
SIMDResultHandlerToFloat* make_knn_handler(
bool is_max,
int /* impl */,
idx_t n,
idx_t k,
float* distances,
idx_t* labels,
const IDSelector* sel,
const FastScanDistancePostProcessing& context,
const float* normalizers = nullptr) const override;
/// Get an InvertedListScanner for single-query scanning.
/// This provides compatibility with the standard IVF search interface
InvertedListScanner* get_InvertedListScanner(
faiss/IndexRaBitQFastScan.cpp
+1 -22
View File
@@ -530,32 +530,11 @@ std::unique_ptr<FastScanCodeScanner> IndexRaBitQFastScan::make_knn_scanner(
float* distances,
idx_t* labels,
const IDSelector* sel,
int /*impl*/,
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
SIMDResultHandlerToFloat* 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 {
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);
} else {
return new RaBitQHeapHandler<CMin<uint16_t, int>, false>(
this, n, k, distances, labels, sel, &context, multi_bit);
}
}
} // namespace faiss
faiss/IndexRaBitQFastScan.h
+1 -12
View File
@@ -101,19 +101,8 @@ struct IndexRaBitQFastScan : IndexFastScan {
float* distances,
idx_t* labels,
const IDSelector* sel,
int impl = 0,
const FastScanDistancePostProcessing& context = {}) const override;
/// Override to create RaBitQ-specific handlers
SIMDResultHandlerToFloat* 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
faiss/impl/fast_scan/dispatching.h
+44
View File
@@ -166,4 +166,48 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl<
}
}
/***************************************************************
* Range search scanner factories.
***************************************************************/
template <>
std::unique_ptr<FastScanCodeScanner> make_range_scanner_impl<
THE_LEVEL_TO_DISPATCH>(
bool is_max,
RangeSearchResult& rres,
float radius,
size_t ntotal,
const IDSelector* sel) {
if (is_max) {
using C = CMax<uint16_t, int64_t>;
return std::make_unique<ScannerMixIn<RangeHandler<C, true>>>(
rres, radius, ntotal, sel);
} else {
using C = CMin<uint16_t, int64_t>;
return std::make_unique<ScannerMixIn<RangeHandler<C, true>>>(
rres, radius, ntotal, sel);
}
}
template <>
std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner_impl<
THE_LEVEL_TO_DISPATCH>(
bool is_max,
RangeSearchPartialResult& pres,
float radius,
size_t ntotal,
size_t q0,
size_t q1,
const IDSelector* sel) {
if (is_max) {
using C = CMax<uint16_t, int64_t>;
return std::make_unique<ScannerMixIn<PartialRangeHandler<C, true>>>(
pres, radius, ntotal, q0, q1, sel);
} else {
using C = CMin<uint16_t, int64_t>;
return std::make_unique<ScannerMixIn<PartialRangeHandler<C, true>>>(
pres, radius, ntotal, q0, q1, sel);
}
}
} // namespace faiss
faiss/impl/fast_scan/fast_scan.cpp
+29
View File
@@ -391,6 +391,35 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
with_id_map);
}
std::unique_ptr<FastScanCodeScanner> make_range_scanner(
bool is_max,
RangeSearchResult& rres,
float radius,
size_t ntotal,
const IDSelector* sel) {
DISPATCH_SIMDLevel(
make_range_scanner_impl, is_max, rres, radius, ntotal, sel);
}
std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner(
bool is_max,
RangeSearchPartialResult& pres,
float radius,
size_t ntotal,
size_t q0,
size_t q1,
const IDSelector* sel) {
DISPATCH_SIMDLevel(
make_partial_range_scanner_impl,
is_max,
pres,
radius,
ntotal,
q0,
q1,
sel);
}
std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner(
const IndexRaBitQFastScan* index,
bool is_max,
faiss/impl/fast_scan/fast_scan.h
+39
View File
@@ -27,6 +27,8 @@
namespace faiss {
struct IDSelector;
struct RangeSearchResult;
struct RangeSearchPartialResult;
struct SIMDResultHandler;
struct SIMDResultHandlerToFloat;
@@ -288,6 +290,43 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
const IDSelector* sel,
bool with_id_map = false);
/// Per-SIMD range scanner factories (defined in per-SIMD TUs via dispatching.h)
template <SIMDLevel SL>
std::unique_ptr<FastScanCodeScanner> make_range_scanner_impl(
bool is_max,
RangeSearchResult& rres,
float radius,
size_t ntotal,
const IDSelector* sel);
template <SIMDLevel SL>
std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner_impl(
bool is_max,
RangeSearchPartialResult& pres,
float radius,
size_t ntotal,
size_t q0,
size_t q1,
const IDSelector* sel);
/// Runtime dispatch: range search scanner.
std::unique_ptr<FastScanCodeScanner> make_range_scanner(
bool is_max,
RangeSearchResult& rres,
float radius,
size_t ntotal,
const IDSelector* sel);
/// Runtime dispatch: partial range search scanner (per-thread).
std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner(
bool is_max,
RangeSearchPartialResult& pres,
float radius,
size_t ntotal,
size_t q0,
size_t q1,
const IDSelector* sel);
/***************************************************************
* RaBitQ scanner factory: per-SIMD specializations live in
* rabitq_dispatching.h, included by each per-SIMD TU.
faiss/impl/fast_scan/impl-neon.cpp
+26
View File
@@ -33,6 +33,32 @@ 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> make_range_scanner_impl<
SIMDLevel::ARM_SVE>(
bool is_max,
RangeSearchResult& rres,
float radius,
size_t ntotal,
const IDSelector* sel) {
return make_range_scanner_impl<SIMDLevel::ARM_NEON>(
is_max, rres, radius, ntotal, sel);
}
template <>
std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner_impl<
SIMDLevel::ARM_SVE>(
bool is_max,
RangeSearchPartialResult& pres,
float radius,
size_t ntotal,
size_t q0,
size_t q1,
const IDSelector* sel) {
return make_partial_range_scanner_impl<SIMDLevel::ARM_NEON>(
is_max, pres, radius, ntotal, q0, q1, sel);
}
template <>
std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl<
SIMDLevel::ARM_SVE>(
faiss/python/swigfaiss.swig
+7 -5
View File
@@ -603,20 +603,22 @@ void gpu_sync_all_devices()
%ignore faiss::FastScanCodeScanner;
%ignore faiss::make_fast_scan_scanner_impl;
%ignore faiss::make_fast_scan_knn_scanner;
%ignore faiss::make_range_scanner_impl;
%ignore faiss::make_range_scanner;
%ignore faiss::make_partial_range_scanner_impl;
%ignore faiss::make_partial_range_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 make_knn_scanner — 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::IndexIVFFastScan::search_implem_10;
%ignore faiss::IndexIVFFastScan::search_implem_12;
%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>
%include <faiss/IndexAdditiveQuantizerFastScan.h>
tests/test_fast_scan.py
+3 -4
View File
@@ -41,16 +41,15 @@ class TestSearch(unittest.TestCase):
# hopefully the jitter in executtion time will not produce
# too many spurious test failures. Unoptimized timings are
# not exploitable, hence the flag test on that as well.
# TODO(DD): Add DD dispatch to fast_scan (pq4_fast_scan.cpp) and remove
# the "DD" exclusion below. Currently fast_scan is compiled with SSE4-only
# flags in DD mode, so it doesn't benefit from AVX2/AVX512 SIMD.
# DD mode uses virtual dispatch which, combined with ASAN overhead,
# makes the 4x speed threshold unreliable on dev machines.
@unittest.skipUnless(
('AVX2' in faiss.get_compile_options() or
'AVX512' in faiss.get_compile_options() or
'NEON' in faiss.get_compile_options()) and
"OPTIMIZE" in faiss.get_compile_options() and
"DD" not in faiss.get_compile_options(),
"only test in static mode with avx2 or neon (DD lacks fast_scan dispatch)")
"only test in static optimized mode with avx2/avx512/neon")
def test_PQ4_speed(self):
ds = datasets.SyntheticDataset(32, 2000, 5000, 1000)
xt = ds.get_train()