Batch the IVF ScalarQuantizer scan four codes at a time (#5456)

Summary:
The IVF ScalarQuantizer inverted-list scan distances one code at a time through `run_scan_codes1`, even though the SQ distance computers already expose a four-at-a-time kernel, `query_to_codes_batch_4` (the one HNSW uses through `distances_batch_4`). This adds `run_scan_codes4`, a batched inner loop that distances four consecutive codes per step and then applies the same threshold compare, `add_result`, and threshold refresh as `run_scan_codes1`, in id order, so the sequence of heap updates is identical. The two SQ scanners route their no-selector path through it; the selector path stays on `run_scan_codes1`.

Four independent accumulator chains let the out-of-order core overlap the latency-bound reduction that a single code cannot hide. Results are unchanged: 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, so a compile-time check keeps those on the scalar scan. Verified bit-identical distances and ids against the scalar scan for QT_8bit, QT_8bit_uniform, QT_6bit, QT_4bit, QT_fp16, and QT_8bit_direct, in both metrics, at d in {64, 128, 256}.

Measured on AVX2, end-to-end `IndexIVFScalarQuantizer` search of QT_8bit, single thread, batched against the scalar scan:

| metric | d | scalar | batched | speedup |
| --- | --- | --- | --- | --- |
| L2 | 128 | 60.4 ms | 40.3 ms | 1.50x |
| L2 | 256 | 170.9 ms | 105.7 ms | 1.62x |
| L2 | 256, out of cache | 403.5 ms | 262.3 ms | 1.54x |
| IP | 128 | 46.5 ms | 38.4 ms | 1.21x |
| IP | 256 | 108.7 ms | 83.8 ms | 1.30x |

The L2 gain is the larger one, around 1.5-1.6x, and it holds when the index spills out of cache. Inner product gains less since its shorter chain is already partly overlapped by the core. When batching does not help it matches the scalar scan, so there is no regression.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5456

Reviewed By: junjieqi

Differential Revision: D118857126

Pulled By: mnorris11

fbshipit-source-id: df7b02c2ede8699c599956370b87a949f1b7db9d

faiss/impl/expanded_scanners.h

@@ -67,6 +67,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 <class ScannerType, typename C, bool store_pairs>
+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.
@@ -136,6 +196,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 <class C, class ScannerType>
+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<ScannerType, C, true, true>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes1<ScannerType, C, false, true>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    } else {
+        if (scanner.store_pairs) {
+            return run_scan_codes4<ScannerType, C, true>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes4<ScannerType, C, false>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    }
+}
+
 template <class ScannerType>
 size_t run_scan_codes(
         const ScannerType& scanner,

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 <class DC>
+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<CMin<float, idx_t>>(
+        if constexpr (sq_batch_4_is_exact<DCClass>()) {
-                *this, list_size, codes, ids, handler);
+            return run_scan_codes4_fix_C<CMin<float, idx_t>>(
+                    *this, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_C<CMin<float, idx_t>>(
+                    *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<CMax<float, idx_t>>(
+        if constexpr (sq_batch_4_is_exact<DCClass>()) {
-                *this, list_size, codes, ids, handler);
+            return run_scan_codes4_fix_C<CMax<float, idx_t>>(
+                    *this, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_C<CMax<float, idx_t>>(
+                    *this, list_size, codes, ids, handler);
+        }
     }
 };