Dynamic Dispatch Task 7 (distances.cpp) (#4559)

Summary: Pull Request resolved: #4559

Reviewed By: mnorris11

Differential Revision: D80500327

Author: Subhadeep Karan

demos/demo_simd_levels.py

@@ -6,9 +6,11 @@
 import os
 import time
 from collections import defaultdict
+import numpy as np
+import os
+from collections import defaultdict
 
 import faiss
-import numpy as np
 from faiss.contrib.datasets import SyntheticDataset
 
 
@@ -26,14 +28,13 @@
 index.train(ds.get_train())
 index.add(ds.get_database())
 
+simd_levels = [faiss.SIMDLevel_NONE, faiss.SIMDLevel_AVX2, faiss.SIMDLevel_AVX512]
 
 if False:
     faiss.omp_set_num_threads(1)
     print("PID=", os.getpid())
     input("press enter to continue")
-    # for simd_level in faiss.NONE, faiss.AVX2, faiss.AVX512F:
-    for simd_level in faiss.AVX2, faiss.AVX512F:
-
+    for simd_level in simd_levels:
         faiss.SIMDConfig.set_level(simd_level)
         print("simd_level=", faiss.SIMDConfig.get_level_name())
         for _ in range(1000):
@@ -42,8 +43,7 @@
 times = defaultdict(list)
 
 for _ in range(10):
-    for simd_level in (faiss.SIMDLevel_NONE, faiss.SIMDLevel_AVX2,
-                       faiss.SIMDLevel_AVX512F):
+    for simd_level in simd_levels:
         faiss.SIMDConfig.set_level(simd_level)
 
         t0 = time.time()

faiss/CMakeLists.txt

@@ -99,7 +99,6 @@ set(FAISS_SRC
   utils/utils.cpp
   utils/distances_fused/avx512.cpp
   utils/distances_fused/distances_fused.cpp
-  utils/distances_fused/simdlib_based.cpp
 )
 
 set(FAISS_HEADERS
@@ -227,7 +226,6 @@ set(FAISS_HEADERS
   utils/utils.h
   utils/distances_fused/avx512.h
   utils/distances_fused/distances_fused.h
-  utils/distances_fused/simdlib_based.h
   utils/approx_topk/approx_topk.h
   utils/approx_topk/avx2-inl.h
   utils/approx_topk/generic.h

faiss/IndexIVFPQ.cpp

@@ -373,6 +373,24 @@ size_t precomputed_table_max_bytes = ((size_t)1) << 31;
  * is faster when the length of the lists is > ksub * M.
  */
 
+void initialize_IVFPQ_precomputed_table(
+        int& use_precomputed_table,
+        const Index* quantizer,
+        const ProductQuantizer& pq,
+        AlignedTable<float>& precomputed_table,
+        bool by_residual,
+        bool verbose) {
+    DISPATCH_SIMDLevel(
+            initialize_IVFPQ_precomputed_table,
+            use_precomputed_table,
+            quantizer,
+            pq,
+            precomputed_table,
+            by_residual,
+            verbose);
+}
+
+template <SIMDLevel SL>
 void initialize_IVFPQ_precomputed_table(
         int& use_precomputed_table,
         const Index* quantizer,
@@ -438,8 +456,7 @@ void initialize_IVFPQ_precomputed_table(
 
             float* tab = &precomputed_table[i * pq.M * pq.ksub];
             pq.compute_inner_prod_table(centroid.data(), tab);
-            fvec_madd<SIMDLevel::NONE>(
-                    pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            fvec_madd<SL>(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
         }
     } else if (use_precomputed_table == 2) {
         const MultiIndexQuantizer* miq =
@@ -466,8 +483,7 @@ void initialize_IVFPQ_precomputed_table(
 
         for (size_t i = 0; i < cpq.ksub; i++) {
             float* tab = &precomputed_table[i * pq.M * pq.ksub];
-            fvec_madd<SIMDLevel::NONE>(
-                    pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            fvec_madd<SL>(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
         }
     }
 }
@@ -645,6 +661,11 @@ struct QueryTables {
      * compute tables for L2 distance
      *****************************************************/
 
+    float precompute_list_tables_L2() {
+        DISPATCH_SIMDLevel(precompute_list_tables_L2);
+    }
+
+    template <SIMDLevel SL>
     float precompute_list_tables_L2() {
         float dis0 = 0;
 
@@ -659,7 +680,7 @@ struct QueryTables {
         } else if (use_precomputed_table == 1) {
             dis0 = coarse_dis;
 
-            fvec_madd<SIMDLevel::NONE>(
+            fvec_madd<SL>(
                     pq.M * pq.ksub,
                     ivfpq.precomputed_table.data() + key * pq.ksub * pq.M,
                     -2.0,
@@ -695,13 +716,12 @@ struct QueryTables {
 
                 if (polysemous_ht == 0) {
                     // sum up with query-specific table
-                    fvec_madd<SIMDLevel::NONE>(
-                            Mf * pq.ksub, pc, -2.0, qtab, ltab);
+                    fvec_madd<SL>(Mf * pq.ksub, pc, -2.0, qtab, ltab);
                     ltab += Mf * pq.ksub;
                     qtab += Mf * pq.ksub;
                 } else {
                     for (int m = cm * Mf; m < (cm + 1) * Mf; m++) {
-                        q_code[m] = fvec_madd_and_argmin(
+                        q_code[m] = fvec_madd_and_argmin<SL>(
                                 pq.ksub, pc, -2, qtab, ltab);
                         pc += pq.ksub;
                         ltab += pq.ksub;

faiss/IndexIVFPQ.h

@@ -16,6 +16,7 @@
 #include <faiss/IndexPQ.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/AlignedTable.h>
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
@@ -172,6 +173,15 @@ void initialize_IVFPQ_precomputed_table(
         bool by_residual,
         bool verbose);
 
+template <SIMDLevel SL>
+void initialize_IVFPQ_precomputed_table(
+        int& use_precomputed_table,
+        const Index* quantizer,
+        const ProductQuantizer& pq,
+        AlignedTable<float>& precomputed_table,
+        bool by_residual,
+        bool verbose);
+
 /// statistics are robust to internal threading, but not if
 /// IndexIVFPQ::search_preassigned is called by multiple threads
 struct IndexIVFPQStats {

faiss/impl/FaissAssert.h

@@ -75,7 +75,7 @@
 
 #define FAISS_THROW_FMT(FMT, ...)                              \
     do {                                                       \
-        std::string __s;                                       \
+        ::std::string __s;                                     \
         int __size = snprintf(nullptr, 0, FMT, __VA_ARGS__);   \
         __s.resize(__size + 1);                                \
         snprintf(&__s[0], __s.size(), FMT, __VA_ARGS__);       \

faiss/impl/ResidualQuantizer.cpp

@@ -456,7 +456,7 @@ void ResidualQuantizer::compute_codes_add_centroids(
                     cent,
                     pool0);
         } else if (use_beam_LUT == 1) {
-            compute_codes_add_centroids_mp_lut1<SIMDLevel::NONE>(
+            compute_codes_add_centroids_mp_lut1(
                     *this,
                     x + i0 * d,
                     codes_out + i0 * code_size,
@@ -501,7 +501,8 @@ void ResidualQuantizer::refine_beam_LUT(
         float* out_distances) const {
     RefineBeamLUTMemoryPool pool;
 
-    refine_beam_LUT_mp<SIMDLevel::NONE>(
+    DISPATCH_SIMDLevel(
+            refine_beam_LUT_mp,
             *this,
             n,
             query_norms,

faiss/impl/pq_4bit/LookupTableScaler.h

@@ -7,9 +7,7 @@
 
 #pragma once
 
-#include <cstdint>
-#include <cstdlib>
-
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/simdlib.h>
 
 /*******************************************
@@ -43,7 +41,7 @@ struct DummyScaler {
         return simd16uint16(0);
     }
 
-#ifdef __AVX512F__
+#if defined(COMPILE_SIMD_AVX512) && defined(__AVX512F__)
 
     using simd64uint8 = simd64uint8<SIMDLevel::AVX512>;
     using simd32uint16 = simd32uint16<SIMDLevel::AVX512>;
@@ -96,7 +94,7 @@ struct Scaler2x4bit {
         return (simd16uint16(res) >> 8) * scale_simd;
     }
 
-#ifdef __AVX512F__
+#if defined(COMPILE_SIMD_AVX512) && defined(__AVX512F__)
     using simd64uint8 = simd64uint8<SIMDLevel::AVX512>;
     using simd32uint16 = simd32uint16<SIMDLevel::AVX512>;
 

faiss/impl/pq_4bit/decompose_qbs.h

@@ -10,6 +10,12 @@
 // This is not standalone code, it is intended to be included in the kernels-X.h
 // files.
 
+#include <cstddef>
+#include <cstdint>
+
+#include <faiss/impl/pq_4bit/simd_result_handlers.h>
+#include <faiss/utils/simd_levels.h>
+
 // handle at most 4 blocks of queries
 template <int QBS, class ResultHandler, class Scaler>
 void accumulate_q_4step(
@@ -19,15 +25,15 @@ void accumulate_q_4step(
         const uint8_t* LUT0,
         ResultHandler& res,
         const Scaler& scaler) {
-    constexpr SIMDLevel SL = ResultHandler::SL;
+    constexpr faiss::SIMDLevel SL = ResultHandler::SL;
     constexpr int Q1 = QBS & 15;
     constexpr int Q2 = (QBS >> 4) & 15;
     constexpr int Q3 = (QBS >> 8) & 15;
     constexpr int Q4 = (QBS >> 12) & 15;
     constexpr int SQ = Q1 + Q2 + Q3 + Q4;
 
     for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
-        FixedStorageHandler<SQ, 2, SL> res2;
+        faiss::simd_result_handlers::FixedStorageHandler<SQ, 2, SL> res2;
         const uint8_t* LUT = LUT0;
         kernel_accumulate_block<Q1>(nsq, codes, LUT, res2, scaler);
         LUT += Q1 * nsq * 16;
@@ -77,8 +83,8 @@ void accumulate(
         ResultHandler& res,
         const Scaler& scaler) {
     assert(nsq % 2 == 0);
-    assert(is_aligned_pointer(codes));
-    assert(is_aligned_pointer(LUT));
+    assert(faiss::is_aligned_pointer(codes));
+    assert(faiss::is_aligned_pointer(LUT));
 
 #define DISPATCH(NQ)                                     \
     case NQ:                                             \
@@ -107,8 +113,8 @@ void pq4_accumulate_loop_qbs_fixed_scaler(
         ResultHandler& res,
         const Scaler& scaler) {
     assert(nsq % 2 == 0);
-    assert(is_aligned_pointer(codes));
-    assert(is_aligned_pointer(LUT0));
+    assert(faiss::is_aligned_pointer(codes));
+    assert(faiss::is_aligned_pointer(LUT0));
 
     // try out optimized versions
     switch (qbs) {

faiss/impl/pq_4bit/dispatching.h

@@ -12,12 +12,14 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/platform_macros.h>
+#include <faiss/impl/pq_4bit/LookupTableScaler.h>
 #include <faiss/impl/pq_4bit/pq4_fast_scan.h>
 #include <faiss/impl/pq_4bit/simd_result_handlers.h>
 
-#include <faiss/impl/pq_4bit/kernels_simd256.h>
+#include <faiss/impl/pq_4bit/kernels_common.h>
 
 namespace faiss {
+using namespace simd_result_handlers;
 
 /** Mix-in class that manages both an SIMD result hander and offers the actual
  * scanning routines. */
@@ -69,14 +71,18 @@ PQ4CodeScanner* make_handler_2(
         const IDSelector* sel) {
     if (k == 1) {
         return new ScannerMixIn<
-                SingleResultHandler<C, with_id_map, SL>,
+                faiss::simd_result_handlers::
+                        SingleResultHandler<C, with_id_map, SL>,
                 Scaler>(norm_scale, nq, ntotal, dis, ids, sel);
     } else if (use_reservoir) {
-        return new ScannerMixIn<ReservoirHandler<C, with_id_map, SL>, Scaler>(
-                norm_scale, nq, ntotal, k, 2 * k, dis, ids, sel);
+        return new ScannerMixIn<
+                faiss::simd_result_handlers::
+                        ReservoirHandler<C, with_id_map, SL>,
+                Scaler>(norm_scale, nq, ntotal, k, 2 * k, dis, ids, sel);
     } else {
-        return new ScannerMixIn<HeapHandler<C, with_id_map, SL>, Scaler>(
-                norm_scale, nq, ntotal, k, dis, ids, sel);
+        return new ScannerMixIn<
+                faiss::simd_result_handlers::HeapHandler<C, with_id_map, SL>,
+                Scaler>(norm_scale, nq, ntotal, k, dis, ids, sel);
     }
 }
 
@@ -87,8 +93,9 @@ PQ4CodeScanner* make_handler_2(
         float radius,
         size_t ntotal,
         const IDSelector* sel) {
-    return new ScannerMixIn<RangeHandler<C, with_id_map, SL>, Scaler>(
-            norm_scale, rres, radius, ntotal, sel);
+    return new ScannerMixIn<
+            faiss::simd_result_handlers::RangeHandler<C, with_id_map, SL>,
+            Scaler>(norm_scale, rres, radius, ntotal, sel);
 }
 
 template <SIMDLevel SL, bool with_id_map, class C, class Scaler>
@@ -100,8 +107,10 @@ PQ4CodeScanner* make_handler_2(
         size_t q0,
         size_t q1,
         const IDSelector* sel) {
-    return new ScannerMixIn<PartialRangeHandler<C, with_id_map, SL>, Scaler>(
-            norm_scale, pres, radius, ntotal, q0, q1, sel);
+    return new ScannerMixIn<
+            faiss::simd_result_handlers::
+                    PartialRangeHandler<C, with_id_map, SL>,
+            Scaler>(norm_scale, pres, radius, ntotal, q0, q1, sel);
 }
 
 // this function dispatches runtime -> template parameters. It is generic for

faiss/impl/pq_4bit/impl-neon.cpp

@@ -7,7 +7,7 @@
 
 #include <faiss/utils/simd_levels.h>
 
-#ifdef COMPILE_SIMD_NEON
+#ifdef COMPILE_SIMD_ARM_NEON
 
 #ifndef __aarch64__
 #error "this can only run on aarch64"
@@ -42,4 +42,4 @@ PQ4CodeScanner* make_pq4_scanner<SIMDLevel::ARM_NEON, true>(PRES_ARGS_LIST) {
 
 } // namespace faiss
 
-#endif // COMPILE_SIMD_NEON
+#endif // COMPILE_SIMD_ARM_NEON