move distances implementation into SIMD specific compile units (#4906)

Summary:
Pull Request resolved: #4906

The main objective is to move the IVFFlat scanning code inside SIMD compile units, so that the loop over codes runs inside SIMD code.

Unfortunately, I could not instant
iate all the implementations with with_MetricType, hence the macro to define them.

Reviewed By: algoriddle

Differential Revision: D95936389

fbshipit-source-id: ae0515fcdcd7c880f329014b264b9dd05065f1c6

Author: mdouze

faiss/CMakeLists.txt

@@ -289,7 +289,6 @@ set(FAISS_HEADERS
   utils/WorkerThread.h
   utils/distances.h
   utils/distances_dispatch.h
-  utils/extra_distances-inl.h
   utils/extra_distances.h
   utils/fp16-fp16c.h
   utils/fp16-inl.h
@@ -331,6 +330,7 @@ set(FAISS_HEADERS
   utils/hamming_distance/avx512-inl.h
   utils/simd_impl/distances_autovec-inl.h
   utils/simd_impl/distances_simdlib256.h
+  utils/simd_impl/IVFFlatScanner-inl.h
   utils/simd_impl/distances_sse-inl.h
 )
 

faiss/IndexIVFFlat.cpp

@@ -13,7 +13,6 @@
 
 #include <cinttypes>
 #include <cstdio>
-#include <numeric>
 
 #include <faiss/IndexFlat.h>
 
@@ -23,6 +22,11 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/expanded_scanners.h>
 #include <faiss/utils/extra_distances.h>
+
+#define THE_SIMD_LEVEL SIMDLevel::NONE
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
+
 #include <faiss/utils/utils.h>
 
 namespace faiss {
@@ -145,48 +149,6 @@ void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
     }
 }
 
-namespace {
-
-template <typename VectorDistance>
-struct IVFFlatScanner : InvertedListScanner {
-    VectorDistance vd;
-    using C = typename VectorDistance::C;
-
-    IVFFlatScanner(
-            const VectorDistance& vd,
-            bool store_pairs,
-            const IDSelector* sel)
-            : InvertedListScanner(store_pairs, sel), vd(vd) {
-        keep_max = vd.is_similarity;
-        code_size = vd.d * sizeof(float);
-    }
-
-    const float* xi;
-    void set_query(const float* query) override {
-        this->xi = query;
-    }
-
-    void set_list(idx_t list_no, float /* coarse_dis */) override {
-        this->list_no = list_no;
-    }
-
-    float distance_to_code(const uint8_t* code) const final {
-        const float* yj = (float*)code;
-        return vd(xi, yj);
-    }
-
-    // redefining the scan_codes allows to inline the distance_to_code
-    size_t scan_codes(
-            size_t list_size,
-            const uint8_t* codes,
-            const idx_t* ids,
-            ResultHandler& handler) const {
-        return run_scan_codes_fix_C<C>(*this, list_size, codes, ids, handler);
-    }
-};
-
-} // anonymous namespace
-
 InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,

faiss/IndexIVFFlat.h

@@ -62,6 +62,38 @@ struct IndexIVFFlat : IndexIVF {
     IndexIVFFlat();
 };
 
+template <typename VectorDistance>
+struct IVFFlatScanner : InvertedListScanner {
+    VectorDistance vd;
+    using C = typename VectorDistance::C;
+
+    IVFFlatScanner(
+            const VectorDistance& vd,
+            bool store_pairs,
+            const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel), vd(vd) {
+        keep_max = vd.is_similarity;
+        code_size = vd.d * sizeof(float);
+    }
+
+    const float* xi;
+    void set_query(const float* query) override {
+        this->xi = query;
+    }
+
+    void set_list(idx_t list_no, float /* coarse_dis */) override {
+        this->list_no = list_no;
+    }
+
+    float distance_to_code(const uint8_t* code) const final;
+
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const;
+};
+
 struct IndexIVFFlatDedup : IndexIVFFlat {
     /** Maps ids stored in the index to the ids of vectors that are
      *  the same. When a vector is unique, it does not appear in the

faiss/utils/distances_simd.cpp

@@ -18,11 +18,10 @@
 #include <faiss/impl/platform_macros.h>
 #include <faiss/impl/simdlib/simdlib_dispatch.h>
 
-#define AUTOVEC_LEVEL SIMDLevel::NONE
+#define THE_SIMD_LEVEL SIMDLevel::NONE
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
 
-#define THE_SIMDLEVEL SIMDLevel::NONE
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/simd_impl/distances_simdlib256.h>