From dcf8ebb34623c1b3dde08e495ef974c92a179c4e Mon Sep 17 00:00:00 2001 From: Michael Norris Date: Fri, 25 Sep 2026 13:47:19 -0700 Subject: [PATCH] Support IDSelector in flat fast-scan search (#5660) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5660 `IndexFastScan::search` rejected every `SearchParameters` object, so a flat fast-scan index could not take an `IDSelector` at all. `IndexSQFastScan` avoided this by unpacking all `ntotal` codes into a fresh buffer on every call, then running the generic scanner. That allocates and rewrites 12.8 MB per query at 200k vectors and d=128. The kernel already supports a selector. `make_fast_scan_knn_scanner` takes an `IDSelector*`, and the result handlers test it after computing a 32-code SIMD block. The flat path passed `nullptr` in that slot while the IVF path passed the real selector. This forwards it instead, and deletes the unpack fallback. `adjust_id` returns `j0 + 32 * b + j` and applies `id_map` when one is set, so the id the selector tests is correct for the flat layout with no change. Two paths still refuse a selector, both deliberately: - `implem` 2, 3 and 4 have no scanner and cannot apply one. The default `implem` of 0 never selects them. - `IndexRaBitQFastScan` keeps its refusal here. The next diff adds the selector test to its kernel and removes the refusal. Reviewed By: junjieqi Differential Revision: D121036637 fbshipit-source-id: d199debb81244275b96a540bdfb8d8dd5fc31c18 --- faiss/IndexFastScan.cpp | 48 +++++++++++++++++++++++------------ faiss/IndexFastScan.h | 9 ++++--- faiss/IndexRaBitQFastScan.cpp | 1 + faiss/IndexSQFastScan.cpp | 46 ++------------------------------- 4 files changed, 41 insertions(+), 63 deletions(-) diff --git a/faiss/IndexFastScan.cpp b/faiss/IndexFastScan.cpp index bd7ed4838..72f4fec0a 100644 --- a/faiss/IndexFastScan.cpp +++ b/faiss/IndexFastScan.cpp @@ -267,15 +267,16 @@ void IndexFastScan::search( float* distances, idx_t* labels, const SearchParameters* params) const { - FAISS_THROW_IF_MSG(params, "search params not supported for this index"); FAISS_THROW_IF_NOT(k > 0); + const IDSelector* sel = params ? params->sel : nullptr; FastScanDistancePostProcessing empty_context{}; if (metric_type == METRIC_L2) { - search_dispatch_implem(n, x, k, distances, labels, empty_context); + search_dispatch_implem( + n, x, k, distances, labels, empty_context, sel); } else { search_dispatch_implem( - n, x, k, distances, labels, empty_context); + n, x, k, distances, labels, empty_context, sel); } } @@ -286,7 +287,8 @@ void IndexFastScan::search_dispatch_implem( idx_t k, float* distances, idx_t* labels, - const FastScanDistancePostProcessing& context) const { + const FastScanDistancePostProcessing& context, + const IDSelector* sel) const { using Cfloat = typename std::conditional< is_max, CMax, @@ -317,6 +319,10 @@ void IndexFastScan::search_dispatch_implem( FAISS_THROW_MSG("not implemented"); } else if (implem == 2 || implem == 3 || implem == 4) { FAISS_THROW_IF_NOT(orig_codes); + FAISS_THROW_IF_MSG( + sel, + "IDSelector is not supported by implem 2/3/4; leave implem at " + "0 so a selector-aware kernel is chosen"); search_implem_234(n, x, k, distances, labels, context); } else if (impl >= 12 && impl <= 15) { FAISS_THROW_IF_NOT(ntotal < INT_MAX); @@ -328,9 +334,11 @@ void IndexFastScan::search_dispatch_implem( // - OpenMP disabled (omp_get_max_threads() = 1) if (nt < 2) { if (impl == 12 || impl == 13) { - search_implem_12(n, x, k, distances, labels, impl, context); + search_implem_12( + n, x, k, distances, labels, impl, context, sel); } else { - search_implem_14(n, x, k, distances, labels, impl, context); + search_implem_14( + n, x, k, distances, labels, impl, context, sel); } } else { // explicitly slice over threads @@ -355,7 +363,8 @@ void IndexFastScan::search_dispatch_implem( dis_i, lab_i, impl, - thread_context); + thread_context, + sel); } else { search_implem_14( i1 - i0, @@ -364,7 +373,8 @@ void IndexFastScan::search_dispatch_implem( dis_i, lab_i, impl, - thread_context); + thread_context, + sel); } } } @@ -440,7 +450,8 @@ void IndexFastScan::search_implem_12( float* distances, idx_t* labels, int impl, - const FastScanDistancePostProcessing& context) const { + const FastScanDistancePostProcessing& context, + const IDSelector* sel) const { FAISS_THROW_IF_NOT(bbs == 32); // handle qbs2 blocking by recursive call @@ -460,7 +471,8 @@ void IndexFastScan::search_implem_12( distances + i0 * k, labels + i0 * k, impl, - sub_context); + sub_context, + sel); } return; } @@ -492,7 +504,7 @@ void IndexFastScan::search_implem_12( FAISS_THROW_IF_NOT(LUT_nq == n); auto scanner = make_knn_scanner( - C::is_max, n, k, ntotal, distances, labels, nullptr, impl, context); + C::is_max, n, k, ntotal, distances, labels, sel, impl, context); auto* rh = scanner->handler(); rh->normalizers = normalizers.get(); // Note: skip & 2 previously set handler->disable (run kernel, @@ -525,7 +537,8 @@ void IndexFastScan::search_implem_14( float* distances, idx_t* labels, int impl, - const FastScanDistancePostProcessing& context) const { + const FastScanDistancePostProcessing& context, + const IDSelector* sel) const { FAISS_THROW_IF_NOT(bbs % 32 == 0); // The accumulate loop dispatch table only instantiates certain @@ -553,7 +566,8 @@ void IndexFastScan::search_implem_14( distances + i0 * k, labels + i0 * k, impl, - sub_context); + sub_context, + sel); } return; } @@ -577,7 +591,7 @@ void IndexFastScan::search_implem_14( LUT.get()); auto scanner = make_knn_scanner( - C::is_max, n, k, ntotal, distances, labels, nullptr, impl, context); + C::is_max, n, k, ntotal, distances, labels, sel, impl, context); auto* rh = scanner->handler(); rh->normalizers = normalizers.get(); // Note: skip & 2 previously set handler->disable (run kernel, @@ -606,7 +620,8 @@ template void IndexFastScan::search_dispatch_implem( idx_t k, float* distances, idx_t* labels, - const FastScanDistancePostProcessing& context) const; + const FastScanDistancePostProcessing& context, + const IDSelector* sel) const; template void IndexFastScan::search_dispatch_implem( idx_t n, @@ -614,7 +629,8 @@ template void IndexFastScan::search_dispatch_implem( idx_t k, float* distances, idx_t* labels, - const FastScanDistancePostProcessing& context) const; + const FastScanDistancePostProcessing& context, + const IDSelector* sel) const; void IndexFastScan::reconstruct(idx_t key, float* recons) const { FAISS_THROW_IF_NOT_FMT( diff --git a/faiss/IndexFastScan.h b/faiss/IndexFastScan.h index 8f73e4322..e323cd1a6 100644 --- a/faiss/IndexFastScan.h +++ b/faiss/IndexFastScan.h @@ -155,7 +155,8 @@ struct IndexFastScan : Index { idx_t k, float* distances, idx_t* labels, - const FastScanDistancePostProcessing& context) const; + const FastScanDistancePostProcessing& context, + const IDSelector* sel = nullptr) const; template void search_implem_234( @@ -174,7 +175,8 @@ struct IndexFastScan : Index { float* distances, idx_t* labels, int impl, - const FastScanDistancePostProcessing& context) const; + const FastScanDistancePostProcessing& context, + const IDSelector* sel = nullptr) const; template void search_implem_14( @@ -184,7 +186,8 @@ struct IndexFastScan : Index { float* distances, idx_t* labels, int impl, - const FastScanDistancePostProcessing& context) const; + const FastScanDistancePostProcessing& context, + const IDSelector* sel = nullptr) const; /** Reconstruct a vector from its code * diff --git a/faiss/IndexRaBitQFastScan.cpp b/faiss/IndexRaBitQFastScan.cpp index 8dacbd6a5..3bd6f1111 100644 --- a/faiss/IndexRaBitQFastScan.cpp +++ b/faiss/IndexRaBitQFastScan.cpp @@ -512,6 +512,7 @@ void IndexRaBitQFastScan::search( float* distances, idx_t* labels, const SearchParameters* params) const { + // The kernel gains its selector test in the next diff. FAISS_THROW_IF_MSG(params, "search params not supported for this index"); // Create query factors array on stack - memory managed by caller diff --git a/faiss/IndexSQFastScan.cpp b/faiss/IndexSQFastScan.cpp index 0fbf66e3c..4bf99674f 100644 --- a/faiss/IndexSQFastScan.cpp +++ b/faiss/IndexSQFastScan.cpp @@ -14,7 +14,6 @@ #include #include -#include #include #include #include @@ -24,7 +23,6 @@ #include #include #include -#include #include #include @@ -113,48 +111,8 @@ void IndexSQFastScan::search( const SearchParameters* params) const { FAISS_THROW_IF_NOT(k > 0); FAISS_THROW_IF_NOT(is_trained); - - const IDSelector* sel = params ? params->sel : nullptr; - - if (!sel) { - IndexFastScan::search(n, x, k, distances, labels, nullptr); - return; - } - - // IDSelector path: unpack codes and use SQ scanner - std::vector flat_codes(ntotal * sq.code_size); - { - std::unique_ptr packer(get_CodePacker()); - for (idx_t i = 0; i < ntotal; i++) { - packer->unpack_1( - codes.data(), i, flat_codes.data() + i * sq.code_size); - } - } - -#pragma omp parallel - { - std::unique_ptr scanner( - sq.select_InvertedListScanner(metric_type, nullptr, true, sel)); - scanner->list_no = 0; - -#pragma omp for - for (idx_t i = 0; i < n; i++) { - float* D = distances + k * i; - idx_t* I = labels + k * i; - if (metric_type == METRIC_L2) { - maxheap_heapify(k, D, I); - } else { - minheap_heapify(k, D, I); - } - scanner->set_query(x + i * d); - scanner->scan_codes(ntotal, flat_codes.data(), nullptr, D, I, k); - if (metric_type == METRIC_L2) { - maxheap_reorder(k, D, I); - } else { - minheap_reorder(k, D, I); - } - } - } + // The kernel applies the selector itself, so no generic-scanner fallback. + IndexFastScan::search(n, x, k, distances, labels, params); } // -----------------------------------------------------------------------