From 0cfa03a00a09436d140e4a4811e449567371e3ba Mon Sep 17 00:00:00 2001 From: Aaryan Mehta <73230976+blazingphoenix7@users.noreply.github.com> Date: Fri, 10 Jul 2026 21:15:21 -0400 Subject: [PATCH] Batch the IVF ScalarQuantizer scan four codes at a time The SQ inverted-list scanners call distance_to_code one code at a time via run_scan_codes1, even though the SQ distance computers already expose query_to_codes_batch_4 (the batched kernel HNSW uses). Add run_scan_codes4, which distances four consecutive codes per step and applies the threshold and heap bookkeeping in id order so the heap-update sequence matches run_scan_codes1 exactly, then route the SQ scanners' no-selector path through it. The batched kernel reconstructs each code, which is bit-identical to the single-code path for every quantizer except the uniform ones, whose single-code path predecodes the query. A compile-time check keeps the uniform quantizers on the scalar scan, so search results are unchanged for every quantizer type. The selector path and all non-SQ scanners are unchanged. --- faiss/impl/expanded_scanners.h | 88 ++++++++++++++++++++++++++ faiss/impl/scalar_quantizer/scanners.h | 61 ++++++++++++++++-- 2 files changed, 145 insertions(+), 4 deletions(-) diff --git a/faiss/impl/expanded_scanners.h b/faiss/impl/expanded_scanners.h index 8eab6ab2ab..1518ceb2f3 100644 --- a/faiss/impl/expanded_scanners.h +++ b/faiss/impl/expanded_scanners.h @@ -63,6 +63,66 @@ size_t run_scan_codes1( return nup; } +// Batched variant of run_scan_codes1 for the SQ scanners: distances four codes +// per step via distance_to_codes_batch_4, then applies the threshold and heap +// updates in id order so results match run_scan_codes1. No-selector path only. +template +size_t run_scan_codes4( + const ScannerType& scanner, + size_t list_size, + const uint8_t* codes, + const idx_t* ids, + ResultHandler& handler) { + size_t nup = 0; + size_t list_no = scanner.list_no; + size_t code_size = scanner.code_size; + float threshold = handler.threshold; + + size_t j = 0; + for (; j + 4 <= list_size; j += 4) { + float dis[4]; + scanner.distance_to_codes_batch_4( + codes, + codes + code_size, + codes + 2 * code_size, + codes + 3 * code_size, + dis[0], + dis[1], + dis[2], + dis[3]); + handler.stats.scan_cnt += 4; + for (size_t b = 0; b < 4; b++) { + if (C::cmp(threshold, dis[b])) { + int64_t id = + store_pairs ? lo_build(list_no, j + b) : ids[j + b]; + if (handler.add_result(dis[b], id)) { + handler.stats.nheap_updates++; + nup++; + threshold = handler.threshold; + } + } + } + codes += 4 * code_size; + } + + // tail: the final < 4 codes, one at a time + for (; j < list_size; j++) { + handler.stats.scan_cnt++; + float dis = scanner.distance_to_code(codes); + if (C::cmp(threshold, dis)) { + int64_t id = store_pairs ? lo_build(list_no, j) : ids[j]; + if (handler.add_result(dis, id)) { + handler.stats.nheap_updates++; + nup++; + threshold = handler.threshold; + } + } + codes += code_size; + } + + return nup; +} + /***************************************************************************** * The following functions dispatch runtime parameters to templates, with * possibly some already-fixed templates. @@ -132,6 +192,34 @@ size_t run_scan_codes_fix_C( } } +// Routing wrapper for the SQ scanners: the no-selector path takes the batched +// run_scan_codes4, the selector path stays on run_scan_codes1. +template +size_t run_scan_codes4_fix_C( + const ScannerType& scanner, + size_t list_size, + const uint8_t* codes, + const idx_t* ids, + ResultHandler& handler) { + if (scanner.sel) { + if (scanner.store_pairs) { + return run_scan_codes1( + scanner, list_size, codes, ids, handler); + } else { + return run_scan_codes1( + scanner, list_size, codes, ids, handler); + } + } else { + if (scanner.store_pairs) { + return run_scan_codes4( + scanner, list_size, codes, ids, handler); + } else { + return run_scan_codes4( + scanner, list_size, codes, ids, handler); + } + } +} + template size_t run_scan_codes( const ScannerType& scanner, diff --git a/faiss/impl/scalar_quantizer/scanners.h b/faiss/impl/scalar_quantizer/scanners.h index 6f63f6e249..cb4ec60bbd 100644 --- a/faiss/impl/scalar_quantizer/scanners.h +++ b/faiss/impl/scalar_quantizer/scanners.h @@ -35,6 +35,19 @@ namespace scalar_quantizer { using QuantizerType = ScalarQuantizer::QuantizerType; using SQDistanceComputer = ScalarQuantizer::SQDistanceComputer; +// True when query_to_codes_batch_4 matches four query_to_code calls. The SIMD +// DCTemplate batch differs from the single path only for the uniform +// quantizers, whose single path predecodes the query, so those stay on the +// scalar scan and batched results never differ from scalar. +template +constexpr bool sq_batch_4_is_exact() { + if constexpr (requires { DC::has_decode_raw(); }) { + return !DC::has_decode_raw(); + } else { + return true; + } +} + /******************************************************************* * IVFSQScannerIP / IVFSQScannerL2 — moved from anonymous namespace * in ScalarQuantizer.cpp @@ -74,13 +87,35 @@ struct IVFSQScannerIP : InvertedListScanner { return accu0 + dc.query_to_code(code); } + void distance_to_codes_batch_4( + const uint8_t* code_0, + const uint8_t* code_1, + const uint8_t* code_2, + const uint8_t* code_3, + float& dis0, + float& dis1, + float& dis2, + float& dis3) const { + dc.query_to_codes_batch_4( + code_0, code_1, code_2, code_3, dis0, dis1, dis2, dis3); + dis0 += accu0; + dis1 += accu0; + dis2 += accu0; + dis3 += accu0; + } + size_t scan_codes( size_t list_size, const uint8_t* codes, const idx_t* ids, ResultHandler& handler) const override { - return run_scan_codes_fix_C>( - *this, list_size, codes, ids, handler); + if constexpr (sq_batch_4_is_exact()) { + return run_scan_codes4_fix_C>( + *this, list_size, codes, ids, handler); + } else { + return run_scan_codes_fix_C>( + *this, list_size, codes, ids, handler); + } } }; @@ -133,13 +168,31 @@ struct IVFSQScannerL2 : InvertedListScanner { return dc.query_to_code(code); } + void distance_to_codes_batch_4( + const uint8_t* code_0, + const uint8_t* code_1, + const uint8_t* code_2, + const uint8_t* code_3, + float& dis0, + float& dis1, + float& dis2, + float& dis3) const { + dc.query_to_codes_batch_4( + code_0, code_1, code_2, code_3, dis0, dis1, dis2, dis3); + } + size_t scan_codes( size_t list_size, const uint8_t* codes, const idx_t* ids, ResultHandler& handler) const override { - return run_scan_codes_fix_C>( - *this, list_size, codes, ids, handler); + if constexpr (sq_batch_4_is_exact()) { + return run_scan_codes4_fix_C>( + *this, list_size, codes, ids, handler); + } else { + return run_scan_codes_fix_C>( + *this, list_size, codes, ids, handler); + } } };