Wire RaBitQ search through FastScanCodeScanner dispatch

faiss/CMakeLists.txt

@@ -9,16 +9,19 @@
 # Architecture-specific: only include files for the current build architecture
 # =============================================================================
 set(FAISS_SIMD_AVX2_SRC
+  impl/fast_scan/impl-avx2.cpp
   impl/pq_code_distance/pq_code_distance-avx2.cpp
   impl/scalar_quantizer/sq-avx2.cpp
   utils/simd_impl/distances_avx2.cpp
 )
 set(FAISS_SIMD_AVX512_SRC
+  impl/fast_scan/impl-avx512.cpp
   impl/pq_code_distance/pq_code_distance-avx512.cpp
   impl/scalar_quantizer/sq-avx512.cpp
   utils/simd_impl/distances_avx512.cpp
 )
 set(FAISS_SIMD_NEON_SRC
+  impl/fast_scan/impl-neon.cpp
   impl/scalar_quantizer/sq-neon.cpp
   utils/simd_impl/distances_aarch64.cpp
 )
@@ -117,7 +120,7 @@ set(FAISS_SRC
   impl/kmeans1d.cpp
   impl/lattice_Zn.cpp
   impl/mapped_io.cpp
-  impl/fast_scan/pq4_fast_scan.cpp
+  impl/fast_scan/fast_scan.cpp
   impl/fast_scan/pq4_fast_scan_search_1.cpp
   impl/fast_scan/pq4_fast_scan_search_qbs.cpp
   impl/residual_quantizer_encode_steps.cpp
@@ -262,10 +265,13 @@ set(FAISS_HEADERS
   impl/kmeans1d.h
   impl/lattice_Zn.h
   impl/platform_macros.h
-  impl/fast_scan/pq4_fast_scan.h
+  impl/fast_scan/accumulate_loops.h
+  impl/fast_scan/dispatching.h
+  impl/fast_scan/fast_scan.h
   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/IndexAdditiveQuantizerFastScan.cpp

@@ -14,7 +14,7 @@
 #include <faiss/impl/LocalSearchQuantizer.h>
 #include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
-#include <faiss/impl/fast_scan/pq4_fast_scan.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/utils/quantize_lut.h>
 #include <faiss/utils/utils.h>
 

faiss/IndexFastScan.cpp

@@ -16,7 +16,7 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
-#include <faiss/impl/fast_scan/pq4_fast_scan.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/impl/fast_scan/simd_result_handlers.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/quantize_lut.h>
@@ -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

@@ -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/IndexIVFAdditiveQuantizerFastScan.cpp

@@ -15,7 +15,7 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
-#include <faiss/impl/fast_scan/pq4_fast_scan.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/invlists/BlockInvertedLists.h>
 #include <faiss/utils/distances.h>

faiss/IndexIVFFastScan.cpp

@@ -19,7 +19,7 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
-#include <faiss/impl/fast_scan/pq4_fast_scan.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/hamming.h>
@@ -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

@@ -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/IndexIVFPQFastScan.cpp

@@ -22,7 +22,7 @@
 
 #include <faiss/invlists/BlockInvertedLists.h>
 
-#include <faiss/impl/fast_scan/pq4_fast_scan.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/impl/fast_scan/simd_result_handlers.h>
 
 namespace faiss {