Skip to content

Switch flat + IVF search to FastScanCodeScanner dispatch#4896

Open
algoriddle wants to merge 3 commits intofacebookresearch:mainfrom
algoriddle:export-D95950481
Open

Switch flat + IVF search to FastScanCodeScanner dispatch#4896
algoriddle wants to merge 3 commits intofacebookresearch:mainfrom
algoriddle:export-D95950481

Conversation

@algoriddle
Copy link
Contributor

Summary:
Switch the actual search paths in IndexFastScan and IndexIVFFastScan to
use the SIMD-dispatched FastScanCodeScanner instead of the old
make_knn_handler() + free-function accumulate pattern.

Each index gains a virtual make_knn_scanner() that returns a
FastScanCodeScanner (bundling handler + kernel behind the SIMD dispatch
boundary). Every search function follows the same pattern: try
make_knn_scanner() first; if it returns nullptr (RaBitQ with custom
handlers), fall back to the existing make_knn_handler() path.

Flat index (IndexFastScan):

  • search_implem_12 (QBS path): scanner->accumulate_loop_qbs(), fallback
  • search_implem_14 (multi-BB path): scanner->accumulate_loop(), fallback

IVF index (IndexIVFFastScan):

  • search_dispatch_implem: try scanner in both single-thread and
    multi-thread paths, pass to implem_10/12
  • search_implem_10: accept optional scanner, use scanner->accumulate_loop()
  • search_implem_12: accept optional scanner, use scanner->accumulate_loop_qbs()
  • search_implem_14: create per-thread scanner, use if available

IVFPQFastScan scanner (IVFPQFastScanScanner::scan_codes):

  • Try scanner->accumulate_loop(), fallback to handler

RaBitQ overrides (IndexRaBitQFastScan, IndexIVFRaBitQFastScan):

  • make_knn_scanner() returns nullptr (scanner wiring comes in next diff)

SWIG: ignore FastScanCodeScanner, make_fast_scan_knn_scanner,
make_fast_scan_scanner_impl, all make_knn_scanner overrides, and
search_implem_10/12 (which now take FastScanCodeScanner* parameter).

Differential Revision: D95950481

@meta-codesync
Copy link
Contributor

meta-codesync bot commented Mar 10, 2026

@algoriddle has exported this pull request. If you are a Meta employee, you can view the originating Diff in D95950481.

Summary:

Move IVFRaBitQHeapHandler from a nested class inside IndexIVFRaBitQFastScan
to a standalone template in impl/fast_scan/rabitq_result_handler.h. Add SL
template parameter for future DD support. A using alias in the index class
preserves source compatibility.

Pure refactor — no behavior change.

Differential Revision: D95950482
…research#4897)

Summary:

Add `FastScanCodeScanner`, a virtual base that bundles handler + kernel
behind the SIMD dispatch boundary. In DD mode, `SINGLE_SIMD_LEVEL = NONE`
so the existing fast scan code path uses emulated SIMD types. The new
scanner provides per-SIMD translation units (AVX2, AVX512, ARM_NEON)
compiled with the correct ISA flags, and a factory function
(`make_fast_scan_knn_scanner`) that uses `DISPATCH_SIMDLevel` to select
the right TU at runtime.

This follows the proven `THE_LEVEL_TO_DISPATCH` pattern from the scalar
quantizer per-SIMD TUs (`sq-dispatch.h`). Each per-SIMD TU includes
`dispatching.h` which provides:
- `ScannerMixIn<Handler>`: wraps a concrete handler and calls accumulation
  kernels (both search_1 multi-BB and QBS paths)
- Factory specialization `make_fast_scan_scanner_impl<SL>()` with
  combinatorial dispatch over `is_max × with_id_map × handler_type`
  (SingleResultHandler for k=1, HeapHandler for k≤20, ReservoirHandler
  for k>20)

New files:
- `impl/fast_scan/dispatching.h` — dispatch template header
- `impl/fast_scan/impl-avx2.cpp` — AVX2 per-SIMD TU
- `impl/fast_scan/impl-avx512.cpp` — AVX512 per-SIMD TU
- `impl/fast_scan/impl-neon.cpp` — ARM NEON TU (with ARM_SVE forwarding)

Modified files:
- `impl/fast_scan/pq4_fast_scan.h` — FastScanCodeScanner base + factory decl
- `impl/fast_scan/pq4_fast_scan.cpp` — NONE specialization + dispatch wrapper
- `xplat.bzl` / `CMakeLists.txt` — register SIMD files and header

Note: RaBitQ handler is not wired through FastScanCodeScanner in this
diff. That comes in later diffs when callers are switched.

Differential Revision: D95950483
…earch#4896)

Summary:

Switch the actual search paths in IndexFastScan and IndexIVFFastScan to
use the SIMD-dispatched FastScanCodeScanner instead of the old
make_knn_handler() + free-function accumulate pattern.

Each index gains a virtual make_knn_scanner() that returns a
FastScanCodeScanner (bundling handler + kernel behind the SIMD dispatch
boundary). Every search function follows the same pattern: try
make_knn_scanner() first; if it returns nullptr (RaBitQ with custom
handlers), fall back to the existing make_knn_handler() path.

Flat index (IndexFastScan):
- search_implem_12 (QBS path): scanner->accumulate_loop_qbs(), fallback
- search_implem_14 (multi-BB path): scanner->accumulate_loop(), fallback

IVF index (IndexIVFFastScan):
- search_dispatch_implem: try scanner in both single-thread and
  multi-thread paths, pass to implem_10/12
- search_implem_10: accept optional scanner, use scanner->accumulate_loop()
- search_implem_12: accept optional scanner, use scanner->accumulate_loop_qbs()
- search_implem_14: create per-thread scanner, use if available

IVFPQFastScan scanner (IVFPQFastScanScanner::scan_codes):
- Try scanner->accumulate_loop(), fallback to handler

RaBitQ overrides (IndexRaBitQFastScan, IndexIVFRaBitQFastScan):
- make_knn_scanner() returns nullptr (scanner wiring comes in next diff)

SWIG: ignore FastScanCodeScanner, make_fast_scan_knn_scanner,
make_fast_scan_scanner_impl, all make_knn_scanner overrides, and
search_implem_10/12 (which now take FastScanCodeScanner* parameter).

Differential Revision: D95950481
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant