Add defaulted SIMDLevel template parameter to handler and scaler types (#4867)

algoriddle · facebook-github-bot · commit d28354ee61b3 · 2026-03-09T03:01:38.000-07:00
Summary:

Templatize the result handler hierarchy and scaler types on SIMDLevel SL,
defaulted to SINGLE_SIMD_LEVEL_256. This allows per-SIMD TUs to instantiate
handlers and scalers with explicit SIMD levels (e.g., AVX2) for native
dispatch.

Result handlers: ResultHandlerCompare, SingleResultHandler, HeapHandler,
ReservoirHandler, RangeHandler, PartialRangeHandler — all gain SL parameter.

Scalers: DummyScaler templatized on SL. 512-bit methods use SL directly
(removing #ifdef __AVX512F__ guard — safe because template bodies only
instantiated when called). NormTableScaler stays non-template (public API).

FixedStorageHandler: add SL parameter, remove SIMDResultHandler base class
(never used polymorphically), remove final/virtual.

Pure refactor. All existing callers use defaults and compile unchanged.

Differential Revision: D95392149
diff --git a/faiss/IndexAdditiveQuantizerFastScan.cpp b/faiss/IndexAdditiveQuantizerFastScan.cpp
@@ -13,7 +13,6 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/impl/LocalSearchQuantizer.h>
-#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/utils/quantize_lut.h>
@@ -202,9 +201,8 @@ void IndexAdditiveQuantizerFastScan::search(
         return;
     }
 
-    NormTableScaler scaler(norm_scale);
     FastScanDistancePostProcessing context;
-    context.norm_scaler = &scaler;
+    context.pq2x4_scale = norm_scale;
     if (metric_type == METRIC_L2) {
         search_dispatch_implem<true>(n, x, k, distances, labels, context);
     } else {
diff --git a/faiss/IndexFastScan.cpp b/faiss/IndexFastScan.cpp
@@ -15,7 +15,6 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/impl/IDSelector.h>
-#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>
@@ -187,18 +186,18 @@ void estimators_from_tables_generic(
         BitstringReader bsr(codes + j * index.code_size, index.code_size);
         accu_t dis = 0;
         const dis_t* dt = dis_table;
-        int nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
+        int nscale = context.pq2x4_scale ? 2 : 0;
 
         for (size_t m = 0; m < index.M - nscale; m++) {
             uint64_t c = bsr.read(index.nbits);
             dis += dt[c];
             dt += index.ksub;
         }
 
-        if (nscale && context.norm_scaler) {
+        if (nscale) {
             for (size_t m = 0; m < nscale; m++) {
                 uint64_t c = bsr.read(index.nbits);
-                dis += context.norm_scaler->scale_one(dt[c]);
+                dis += dt[c] * context.pq2x4_scale;
                 dt += index.ksub;
             }
         }
@@ -545,7 +544,7 @@ void IndexFastScan::search_implem_12(
                 codes.get(),
                 LUT.get(),
                 *handler.get(),
-                context.norm_scaler,
+                context.pq2x4_scale,
                 get_block_stride());
     }
     if (!(skip & 8)) {
@@ -629,7 +628,7 @@ void IndexFastScan::search_implem_14(
                 codes.get(),
                 LUT.get(),
                 *handler.get(),
-                context.norm_scaler,
+                context.pq2x4_scale,
                 get_block_stride());
     }
     if (!(skip & 8)) {
diff --git a/faiss/IndexFastScan.h b/faiss/IndexFastScan.h
@@ -14,7 +14,6 @@
 namespace faiss {
 
 struct CodePacker;
-struct NormTableScaler;
 struct IDSelector;
 struct SIMDResultHandlerToFloat;
 
diff --git a/faiss/IndexIVFAdditiveQuantizerFastScan.cpp b/faiss/IndexIVFAdditiveQuantizerFastScan.cpp
@@ -15,7 +15,6 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/FastScanDistancePostProcessing.h>
-#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/invlists/BlockInvertedLists.h>
@@ -317,9 +316,8 @@ void IndexIVFAdditiveQuantizerFastScan::search(
         return;
     }
 
-    NormTableScaler scaler(norm_scale);
     FastScanDistancePostProcessing context;
-    context.norm_scaler = &scaler;
+    context.pq2x4_scale = norm_scale;
     IndexIVFFastScan::CoarseQuantized cq{nprobe};
     search_dispatch_implem(n, x, k, distances, labels, cq, context);
 }
diff --git a/faiss/IndexIVFFastScan.cpp b/faiss/IndexIVFFastScan.cpp
@@ -18,7 +18,6 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/FastScanDistancePostProcessing.h>
-#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>
@@ -239,7 +238,7 @@ void estimators_from_tables_generic(
         int64_t* heap_ids,
         const FastScanDistancePostProcessing& context) {
     using accu_t = typename C::T;
-    size_t nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
+    size_t nscale = context.pq2x4_scale ? 2 : 0;
     for (size_t j = 0; j < ncodes; ++j) {
         BitstringReader bsr(codes + j * index.code_size, index.code_size);
         accu_t dis = bias;
@@ -251,10 +250,10 @@ void estimators_from_tables_generic(
             dt += index.ksub;
         }
 
-        if (context.norm_scaler) {
+        if (nscale) {
             for (size_t m = 0; m < nscale; m++) {
                 uint64_t c = bsr.read(index.nbits);
-                dis += context.norm_scaler->scale_one(dt[c]);
+                dis += dt[c] * context.pq2x4_scale;
                 dt += index.ksub;
             }
         }
@@ -1031,7 +1030,7 @@ void IndexIVFFastScan::search_implem_10(
                     codes.get(),
                     LUT,
                     handler,
-                    context.norm_scaler,
+                    context.pq2x4_scale,
                     get_block_stride());
 
             ndis += ls;
@@ -1183,7 +1182,7 @@ void IndexIVFFastScan::search_implem_12(
                 codes.get(),
                 LUT.get(),
                 handler,
-                context.norm_scaler,
+                context.pq2x4_scale,
                 get_block_stride());
         // prepare for next loop
         i0 = i1;
@@ -1407,7 +1406,7 @@ void IndexIVFFastScan::search_implem_14(
                     codes.get(),
                     LUT.get(),
                     *handler.get(),
-                    context.norm_scaler,
+                    context.pq2x4_scale,
                     get_block_stride());
         }
 
diff --git a/faiss/IndexIVFFastScan.h b/faiss/IndexIVFFastScan.h
@@ -13,7 +13,6 @@
 
 namespace faiss {
 
-struct NormTableScaler;
 struct SIMDResultHandlerToFloat;
 struct Quantizer;
 
diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp
@@ -419,7 +419,7 @@ struct IVFPQFastScanScanner : InvertedListScanner {
                 codes,
                 LUT,
                 *handler,
-                nullptr,
+                0,
                 index.get_block_stride());
 
         // The handler is for the results of this iteration.
diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp
@@ -938,7 +938,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
                 codes,
                 LUT,
                 *handler,
-                nullptr,
+                0,
                 index.get_block_stride());
 
         // Combine results across iterations
diff --git a/faiss/impl/FastScanDistancePostProcessing.h b/faiss/impl/FastScanDistancePostProcessing.h
@@ -11,9 +11,6 @@
 
 namespace faiss {
 
-// Forward declarations
-struct NormTableScaler;
-
 namespace rabitq_utils {
 struct QueryFactorsData;
 }
@@ -22,8 +19,10 @@ struct QueryFactorsData;
  * Simple context object that holds processors for FastScan operations.
  * */
 struct FastScanDistancePostProcessing {
-    /// Norm scaling processor for Additive Quantizers (nullptr if not needed)
-    const NormTableScaler* norm_scaler = nullptr;
+    /// Norm scaling processor for Additive Quantizers.
+    /// The scale is encoded in a 2x4 bit PQ table, then scaled by this int.
+    /// Set to 0 if unused.
+    int pq2x4_scale = 0;
 
     /// Query factors data pointer for RaBitQ (nullptr if not needed)
     /// This pointer should point to the beginning of the relevant
@@ -41,7 +40,7 @@ struct FastScanDistancePostProcessing {
 
     /// Check if norm scaling is enabled
     bool has_norm_scaling() const {
-        return norm_scaler != nullptr;
+        return pq2x4_scale != 0;
     }
 
     /// Check if query factors processing is enabled
diff --git a/faiss/impl/LookupTableScaler.h b/faiss/impl/LookupTableScaler.h
@@ -20,8 +20,14 @@
 namespace faiss {
 
 /// no-op handler
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
 struct DummyScaler {
     static constexpr int nscale = 0;
+    static constexpr SIMDLevel SL256 = simd256_level_selector<SL>::value;
+    using simd32uint8 = simd32uint8_tpl<SL256>;
+    using simd16uint16 = simd16uint16_tpl<SL256>;
+    using simd64uint8 = simd64uint8_tpl<SL>;
+    using simd32uint16 = simd32uint16_tpl<SL>;
 
     inline simd32uint8 lookup(const simd32uint8&, const simd32uint8&) const {
         FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
@@ -38,7 +44,6 @@ struct DummyScaler {
         return simd16uint16(0);
     }
 
-#ifdef __AVX512F__
     inline simd64uint8 lookup(const simd64uint8&, const simd64uint8&) const {
         FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
         return simd64uint8(0);
@@ -53,7 +58,6 @@ struct DummyScaler {
         FAISS_THROW_MSG("DummyScaler::scale_hi should not be called.");
         return simd32uint16(0);
     }
-#endif
 
     template <class dist_t>
     inline dist_t scale_one(const dist_t&) const {
diff --git a/faiss/impl/pq4_fast_scan.h b/faiss/impl/pq4_fast_scan.h
@@ -24,7 +24,6 @@
 
 namespace faiss {
 
-struct NormTableScaler;
 struct SIMDResultHandler;
 
 /** Pack codes for consumption by the SIMD kernels.
@@ -139,7 +138,7 @@ void pq4_accumulate_loop(
         const uint8_t* codes,
         const uint8_t* LUT,
         SIMDResultHandler& res,
-        const NormTableScaler* scaler,
+        int pq2x4_scale,
         size_t block_stride);
 
 /* qbs versions, supported only for bbs=32.
@@ -190,7 +189,7 @@ int pq4_pack_LUT_qbs_q_map(
  * @param codes   encoded database vectors (packed)
  * @param LUT     look-up table (packed)
  * @param res     call-back for the results
- * @param scaler  scaler to scale the encoded norm
+ * @param pq2x4_scale  scaler to scale the encoded norm
  * @param block_stride  stride in bytes between consecutive blocks.
  */
 void pq4_accumulate_loop_qbs(
@@ -200,7 +199,7 @@ void pq4_accumulate_loop_qbs(
         const uint8_t* codes,
         const uint8_t* LUT,
         SIMDResultHandler& res,
-        const NormTableScaler* scaler,
+        int pq2x4_scale,
         size_t block_stride);
 
 /** Wrapper of pq4_accumulate_loop_qbs using simple StoreResultHandler
diff --git a/faiss/impl/pq4_fast_scan_search_1.cpp b/faiss/impl/pq4_fast_scan_search_1.cpp
@@ -182,34 +182,19 @@ void pq4_accumulate_loop_fixed_handler(
         const uint8_t* codes,
         const uint8_t* LUT,
         ResultHandler& res,
-        const NormTableScaler* scaler,
+        int pq2x4_scale,
         size_t block_stride) {
-    if (scaler) {
+    if (pq2x4_scale) {
+        NormTableScaler scaler(pq2x4_scale);
         pq4_accumulate_loop_fixed_scaler(
-                nq, nb, bbs, nsq, codes, LUT, res, *scaler, block_stride);
+                nq, nb, bbs, nsq, codes, LUT, res, scaler, block_stride);
     } else {
-        DummyScaler dscaler;
+        DummyScaler<> dscaler;
         pq4_accumulate_loop_fixed_scaler(
                 nq, nb, bbs, nsq, codes, LUT, res, dscaler, block_stride);
     }
 }
 
-struct Run_pq4_accumulate_loop {
-    template <class ResultHandler>
-    void f(ResultHandler& res,
-           int nq,
-           size_t nb,
-           int bbs,
-           int nsq,
-           const uint8_t* codes,
-           const uint8_t* LUT,
-           const NormTableScaler* scaler,
-           size_t block_stride) {
-        pq4_accumulate_loop_fixed_handler(
-                nq, nb, bbs, nsq, codes, LUT, res, scaler, block_stride);
-    }
-};
-
 } // anonymous namespace
 
 void pq4_accumulate_loop(
@@ -220,11 +205,20 @@ void pq4_accumulate_loop(
         const uint8_t* codes,
         const uint8_t* LUT,
         SIMDResultHandler& res,
-        const NormTableScaler* scaler,
+        int pq2x4_scale,
         size_t block_stride) {
-    Run_pq4_accumulate_loop consumer;
-    dispatch_SIMDResultHandler(
-            res, consumer, nq, nb, bbs, nsq, codes, LUT, scaler, block_stride);
+    with_SIMDResultHandler(res, [&](auto& handler) {
+        pq4_accumulate_loop_fixed_handler(
+                nq,
+                nb,
+                bbs,
+                nsq,
+                codes,
+                LUT,
+                handler,
+                pq2x4_scale,
+                block_stride);
+    });
 }
 
 } // namespace faiss
diff --git a/faiss/impl/pq4_fast_scan_search_qbs.cpp b/faiss/impl/pq4_fast_scan_search_qbs.cpp
diff --git a/faiss/impl/simd_result_handlers.h b/faiss/impl/simd_result_handlers.h
