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

@@ -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<true>(n, x, k, distances, labels, empty_context);
+        search_dispatch_implem<true>(
+                n, x, k, distances, labels, empty_context, sel);
     } else {
         search_dispatch_implem<false>(
-                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<float, int64_t>,
@@ -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<Cfloat>(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<C>(n, x, k, distances, labels, impl, context);
+                search_implem_12<C>(
+                        n, x, k, distances, labels, impl, context, sel);
             } else {
-                search_implem_14<C>(n, x, k, distances, labels, impl, context);
+                search_implem_14<C>(
+                        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<C>(
                             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<true>(
         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<false>(
         idx_t n,
@@ -614,7 +629,8 @@ template void IndexFastScan::search_dispatch_implem<false>(
         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(

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 <class Cfloat>
     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 <class C>
     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
      *

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

faiss/IndexSQFastScan.cpp

@@ -14,7 +14,6 @@
 #include <memory>
 #include <vector>
 
-#include <faiss/IndexIVF.h>
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/CodePacker.h>
@@ -24,7 +23,6 @@
 #include <faiss/impl/ScalarQuantizer.h>
 #include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
 #include <faiss/impl/fast_scan/fast_scan.h>
-#include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/utils.h>
 
@@ -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<uint8_t> flat_codes(ntotal * sq.code_size);
-    {
-        std::unique_ptr<CodePacker> 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<InvertedListScanner> 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);
 }
 
 // -----------------------------------------------------------------------