Skip to content

Commit 957a351

Browse files
algoriddlefacebook-github-bot
authored andcommitted
Wire RaBitQ search through FastScanCodeScanner dispatch
Summary: 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. Differential Revision: D96116557
1 parent f783bed commit 957a351

17 files changed

Lines changed: 786 additions & 558 deletions

faiss/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ set(FAISS_HEADERS
271271
impl/fast_scan/decompose_qbs.h
272272
impl/fast_scan/kernels_simd256.h
273273
impl/fast_scan/kernels_simd512.h
274-
impl/fast_scan/rabitq_result_handler.h
274+
impl/fast_scan/rabitq_dispatching.h
275275
impl/residual_quantizer_encode_steps.h
276276
impl/simd_dispatch.h
277277
impl/fast_scan/simd_result_handlers.h

faiss/IndexFastScan.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,19 @@ void estimators_from_tables_generic(
211211

212212
} // anonymous namespace
213213

214+
std::unique_ptr<FastScanCodeScanner> IndexFastScan::make_knn_scanner(
215+
bool is_max,
216+
idx_t n,
217+
idx_t k,
218+
size_t ntotal,
219+
float* distances,
220+
idx_t* labels,
221+
const IDSelector* sel,
222+
const FastScanDistancePostProcessing&) const {
223+
return make_fast_scan_knn_scanner(
224+
is_max, 0, n, ntotal, k, distances, labels, sel);
225+
}
226+
214227
// Default implementation of make_knn_handler with centralized fallback logic
215228
SIMDResultHandlerToFloat* IndexFastScan::make_knn_handler(
216229
bool is_max,

faiss/IndexFastScan.h

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,11 @@
77

88
#pragma once
99

10+
#include <memory>
11+
1012
#include <faiss/Index.h>
1113
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
14+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
1215
#include <faiss/utils/AlignedTable.h>
1316

1417
namespace faiss {
@@ -121,6 +124,22 @@ struct IndexFastScan : Index {
121124
const float* x,
122125
const FastScanDistancePostProcessing& context) const = 0;
123126

127+
/** Create a SIMD-dispatched scanner for knn search.
128+
*
129+
* Returns a FastScanCodeScanner that bundles handler + accumulation
130+
* kernel behind the SIMD dispatch boundary.
131+
* The scanner's accumulate methods dispatch to the optimal SIMD level.
132+
*/
133+
virtual std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
134+
bool is_max,
135+
idx_t n,
136+
idx_t k,
137+
size_t ntotal,
138+
float* distances,
139+
idx_t* labels,
140+
const IDSelector* sel,
141+
const FastScanDistancePostProcessing& context = {}) const;
142+
124143
/** Create a KNN handler for this index type
125144
*
126145
* This method can be overridden by derived classes to provide

faiss/IndexIVFFastScan.cpp

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -495,6 +495,26 @@ int compute_search_nslice(
495495

496496
} // namespace
497497

498+
std::unique_ptr<FastScanCodeScanner> IndexIVFFastScan::make_knn_scanner(
499+
bool is_max,
500+
idx_t n,
501+
idx_t k,
502+
float* distances,
503+
idx_t* labels,
504+
const IDSelector* sel,
505+
const FastScanDistancePostProcessing&) const {
506+
return make_fast_scan_knn_scanner(
507+
is_max,
508+
0,
509+
n,
510+
0,
511+
k,
512+
distances,
513+
labels,
514+
sel,
515+
/*with_id_map=*/true);
516+
}
517+
498518
SIMDResultHandlerToFloat* IndexIVFFastScan::make_knn_handler(
499519
bool is_max,
500520
int impl,

faiss/IndexIVFFastScan.h

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
#include <faiss/IndexIVF.h>
1111
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
12+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
1213
#include <faiss/utils/AlignedTable.h>
1314

1415
namespace faiss {
@@ -221,6 +222,31 @@ struct IndexIVFFastScan : IndexIVF {
221222
RangeSearchResult* result,
222223
const SearchParameters* params = nullptr) const override;
223224

225+
/** Create a SIMD-dispatched scanner for knn search (IVF variant).
226+
*
227+
* Returns a FastScanCodeScanner that bundles handler + accumulation
228+
* kernel behind the SIMD dispatch boundary. ntotal is not passed
229+
* because IVF sets it per-list via handler->ntotal.
230+
* Derived classes that need custom handlers (e.g. RaBitQ) override
231+
* this to return nullptr, falling back to make_knn_handler.
232+
*
233+
* @param is_max whether to use CMax comparator (true) or CMin
234+
* @param n number of queries
235+
* @param k number of neighbors to find
236+
* @param distances output distances array
237+
* @param labels output labels array
238+
* @param sel optional ID selector
239+
* @return scanner, or nullptr if unsupported
240+
*/
241+
virtual std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
242+
bool is_max,
243+
idx_t n,
244+
idx_t k,
245+
float* distances,
246+
idx_t* labels,
247+
const IDSelector* sel,
248+
const FastScanDistancePostProcessing& context = {}) const;
249+
224250
/** Create a KNN handler for this index type
225251
*
226252
* This method can be overridden by derived classes to provide

0 commit comments

Comments
 (0)