Batch the IVF ScalarQuantizer scan four codes at a time

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 <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.
@@ -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 <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>>(
-                *this, list_size, codes, ids, handler);
+        if constexpr (sq_batch_4_is_exact<DCClass>()) {
+            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>>(
-                *this, list_size, codes, ids, handler);
+        if constexpr (sq_batch_4_is_exact<DCClass>()) {
+            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);
+        }
     }
 };