Skip to content

Wire flat + IVF RaBitQ search through FastScanCodeScanner dispatch#4901

Open
algoriddle wants to merge 4 commits intofacebookresearch:mainfrom
algoriddle:export-D96018604
Open

Wire flat + IVF RaBitQ search through FastScanCodeScanner dispatch#4901
algoriddle wants to merge 4 commits intofacebookresearch:mainfrom
algoriddle:export-D96018604

Conversation

@algoriddle
Copy link
Contributor

@algoriddle algoriddle commented Mar 10, 2026

Summary:
The previous diff switched flat + IVF PQ search to FastScanCodeScanner
dispatch, but left both RaBitQ indexes' make_knn_scanner() returning
nullptr (falling back to make_knn_handler). This diff wires both flat
and IVF RaBitQ so they use the SIMD-dispatched scanner path.

Flat RaBitQ changes:

  • Add context parameter to make_knn_scanner on IndexFastScan,
    IndexIVFFastScan, and all derived classes so RaBitQ can access
    query_factors needed for distance correction
  • Add SL template parameter to RaBitQHeapHandler and move method
    bodies from .cpp to .h (required for per-TU template instantiation)
  • Change RaBitQHeapHandler::context from reference to pointer for
    safe lifetime management inside ScannerMixIn unique_ptr
  • Wire IndexRaBitQFastScan::make_knn_scanner() to rabitq_make_knn_scanner

IVF RaBitQ changes:

  • Move IVFRaBitQHeapHandler method bodies from IndexIVFRaBitQFastScan.cpp
    to IndexIVFRaBitQFastScan.h (after struct definition, breaking circular
    dependency with rabitq_result_handler.h)
  • Add rabitq_ivf_make_knn_scanner[_impl] factory + dispatch
  • Wire IndexIVFRaBitQFastScan::make_knn_scanner() to return scanner
  • Update IVFRaBitQFastScanScanner::scan_codes to use scanner path

Common changes:

  • Create rabitq_dispatching.h: per-SIMD TU header with flat + IVF
    scanner factories, included after dispatching.h in each SIMD TU
  • Add ARM_SVE forwarders for both flat and IVF RaBitQ scanners

All FastScan search paths are now fully DD-wired.

Differential Revision: D96018604

@algoriddle @meta-codesync
Extract RaBitQ result handler to impl/fast_scan/
7c4721e 
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
@algoriddle @meta-codesync
Add FastScanCodeScanner dispatch boundary with per-SIMD TUs
ccd8d02 
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
@algoriddle @meta-codesync
Switch flat + IVF search to FastScanCodeScanner dispatch
b017ccb 
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
@algoriddle @meta-codesync
Wire flat + IVF RaBitQ search through FastScanCodeScanner dispatch
7c779a8 
Summary:
The previous diff switched flat + IVF PQ search to FastScanCodeScanner
dispatch, but left both RaBitQ indexes' make_knn_scanner() returning
nullptr (falling back to make_knn_handler). This diff wires both flat
and IVF RaBitQ so they use the SIMD-dispatched scanner path.

Flat RaBitQ changes:
- Add `context` parameter to make_knn_scanner on IndexFastScan,
  IndexIVFFastScan, and all derived classes so RaBitQ can access
  query_factors needed for distance correction
- Add SL template parameter to RaBitQHeapHandler and move method
  bodies from .cpp to .h (required for per-TU template instantiation)
- Change RaBitQHeapHandler::context from reference to pointer for
  safe lifetime management inside ScannerMixIn unique_ptr
- Wire IndexRaBitQFastScan::make_knn_scanner() to rabitq_make_knn_scanner

IVF RaBitQ changes:
- Move IVFRaBitQHeapHandler method bodies from IndexIVFRaBitQFastScan.cpp
  to IndexIVFRaBitQFastScan.h (after struct definition, breaking circular
  dependency with rabitq_result_handler.h)
- Add rabitq_ivf_make_knn_scanner[_impl] factory + dispatch
- Wire IndexIVFRaBitQFastScan::make_knn_scanner() to return scanner
- Update IVFRaBitQFastScanScanner::scan_codes to use scanner path

Common changes:
- Create rabitq_dispatching.h: per-SIMD TU header with flat + IVF
  scanner factories, included after dispatching.h in each SIMD TU
- Add ARM_SVE forwarders for both flat and IVF RaBitQ scanners

All FastScan search paths are now fully DD-wired.

Differential Revision: D96018604
@meta-cla meta-cla bot added the CLA Signed label Mar 10, 2026
@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 D96018604.

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

Reviewers

No reviews

Assignees

No one assigned

Labels

CLA Signed fb-exported meta-exported

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@algoriddle