Add NEON support and move fvec_madd for NEON

Summary:
Main changes
--

1. simdlib_neon.h has been updated to use the templates properly. This is the largest file change in the diff.

2. simd_levels.h now has 2 separate macros for DISPATCH_SIMDLevel and DISPATCH_SIMDLevel_AND_RETURN. The first one just executes the given function. The second one will return a value.

3. xplat.bzl now reads the architecture and decides the COMPILE_SIMD_<> flags based on that. After we move from static to dynamic dispatch, this will be enabled.

4. Adds distances_neon.cpp and distances_sve.cpp just for fvec_madd. However, SVE code is not invoked yet. The import for detection does not quite work. It can be separate work.

5. updates RQ and PQ to accomodate simdlib_neon.h changes. Following https://docs.google.com/document/d/1hVwoW-SrPxbBxJ3Nvl58tOaKKgkzDgVxPJbx92TSTvs/edit?tab=t.0#bookmark=id.ujcwvusaoqp8, we move some functions to headers and `extern template` some things.

Note 1
--
Please yell at me if we want to break the above up a bit and I am very happy to. Items 1 and 2 could be in a separate diff IMO. They are here because we needed it for dynamic dispatch, then we decided on static dispatch.

Note 2
--
These two tests still fail to build. I didn't fix them yet. We can take it up later in the stack, as per the google doc?
1. test_code_distance.cpp (distance_single_code_generic no longer exists)
2. test_simdlib.cpp (needs templatization)

Note 3
--
autodeps failing due to `<asm/hwcap.h>` in the SVE check. SVE does not quite work yet and it isn't enabled, so we can come back to this (or I can comment it out). Feel free to comment if you have a preference.

Reviewed By: mdouze

Differential Revision: D74275821

Author: Subhadeep Karan

faiss/CMakeLists.txt

@@ -76,7 +76,9 @@ set(FAISS_SRC
   impl/pq4_fast_scan.cpp
   impl/pq4_fast_scan_search_1.cpp
   impl/pq4_fast_scan_search_qbs.cpp
-  impl/residual_quantizer_encode_steps.cpp
+  impl/residual_quantizer_encode_steps/residual_quantizer_encode_steps.cpp
+  impl/residual_quantizer_encode_steps/residual_quantizer_encode_steps-avx2.cpp
+  impl/residual_quantizer_encode_steps/residual_quantizer_encode_steps-neon.cpp
   impl/zerocopy_io.cpp
   impl/NNDescent.cpp
   invlists/BlockInvertedLists.cpp
@@ -183,7 +185,7 @@ set(FAISS_HEADERS
   impl/lattice_Zn.h
   impl/platform_macros.h
   impl/pq4_fast_scan.h
-  impl/residual_quantizer_encode_steps.h
+  impl/residual_quantizer_encode_steps/residual_quantizer_encode_steps.h
   impl/simd_result_handlers.h
   impl/zerocopy_io.h
   impl/pq_code_distance/code_distance.h

faiss/IndexIVFPQ.cpp

@@ -438,7 +438,8 @@ void initialize_IVFPQ_precomputed_table(
 
             float* tab = &precomputed_table[i * pq.M * pq.ksub];
             pq.compute_inner_prod_table(centroid.data(), tab);
-            fvec_madd(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            fvec_madd<SIMDLevel::NONE>(
+                    pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
         }
     } else if (use_precomputed_table == 2) {
         const MultiIndexQuantizer* miq =
@@ -465,7 +466,8 @@ 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(pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
+            fvec_madd<SIMDLevel::NONE>(
+                    pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
         }
     }
 }
@@ -657,7 +659,7 @@ struct QueryTables {
         } else if (use_precomputed_table == 1) {
             dis0 = coarse_dis;
 
-            fvec_madd(
+            fvec_madd<SIMDLevel::NONE>(
                     pq.M * pq.ksub,
                     ivfpq.precomputed_table.data() + key * pq.ksub * pq.M,
                     -2.0,
@@ -693,7 +695,8 @@ struct QueryTables {
 
                 if (polysemous_ht == 0) {
                     // sum up with query-specific table
-                    fvec_madd(Mf * pq.ksub, pc, -2.0, qtab, ltab);
+                    fvec_madd<SIMDLevel::NONE>(
+                            Mf * pq.ksub, pc, -2.0, qtab, ltab);
                     ltab += Mf * pq.ksub;
                     qtab += Mf * pq.ksub;
                 } else {
@@ -1355,7 +1358,6 @@ InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
         const IDSelector* sel,
         const IVFSearchParameters*) const {
     DISPATCH_SIMDLevel(get_InvertedListScanner3, *this, store_pairs, sel);
-    return nullptr;
 }
 
 IndexIVFPQStats indexIVFPQ_stats;

faiss/IndexIVFPQFastScan.cpp

@@ -7,11 +7,6 @@
 
 #include <faiss/IndexIVFPQFastScan.h>
 
-#include <cassert>
-#include <cstdio>
-
-#include <memory>
-
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
@@ -181,7 +176,8 @@ bool IndexIVFPQFastScan::lookup_table_is_3d() const {
     return by_residual && metric_type == METRIC_L2;
 }
 
-void IndexIVFPQFastScan::compute_LUT(
+template <SIMDLevel SL>
+void IndexIVFPQFastScan::compute_LUT_helper(
         size_t n,
         const float* x,
         const CoarseQuantized& cq,
@@ -209,8 +205,7 @@ void IndexIVFPQFastScan::compute_LUT(
                     idx_t cij = cq.ids[ij];
 
                     if (cij >= 0) {
-                        // TODO avoid dynamic dispatch
-                        fvec_madd(
+                        fvec_madd<SL>(
                                 dim12,
                                 precomputed_table.get() + cij * dim12,
                                 -2,
@@ -269,4 +264,13 @@ void IndexIVFPQFastScan::compute_LUT(
     }
 }
 
+void IndexIVFPQFastScan::compute_LUT(
+        size_t n,
+        const float* x,
+        const CoarseQuantized& cq,
+        AlignedTable<float>& dis_tables,
+        AlignedTable<float>& biases) const {
+    DISPATCH_SIMDLevel(compute_LUT_helper, n, x, cq, dis_tables, biases);
+}
+
 } // namespace faiss

faiss/IndexIVFPQFastScan.h

@@ -13,6 +13,7 @@
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/ProductQuantizer.h>
 #include <faiss/utils/AlignedTable.h>
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
@@ -81,6 +82,15 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
             const CoarseQuantized& cq,
             AlignedTable<float>& dis_tables,
             AlignedTable<float>& biases) const override;
+
+   protected:
+    template <SIMDLevel SL>
+    void compute_LUT_helper(
+            size_t n,
+            const float* x,
+            const CoarseQuantized& cq,
+            AlignedTable<float>& dis_tables,
+            AlignedTable<float>& biases) const;
 };
 
 } // namespace faiss

faiss/impl/ResidualQuantizer.cpp

@@ -17,7 +17,7 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/residual_quantizer_encode_steps.h>
+#include <faiss/impl/residual_quantizer_encode_steps/residual_quantizer_encode_steps.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
@@ -456,7 +456,7 @@ void ResidualQuantizer::compute_codes_add_centroids(
                     cent,
                     pool0);
         } else if (use_beam_LUT == 1) {
-            compute_codes_add_centroids_mp_lut1(
+            compute_codes_add_centroids_mp_lut1<SIMDLevel::NONE>(
                     *this,
                     x + i0 * d,
                     codes_out + i0 * code_size,
@@ -500,7 +500,8 @@ void ResidualQuantizer::refine_beam_LUT(
         int32_t* out_codes,
         float* out_distances) const {
     RefineBeamLUTMemoryPool pool;
-    refine_beam_LUT_mp(
+
+    refine_beam_LUT_mp<SIMDLevel::NONE>(
             *this,
             n,
             query_norms,

faiss/impl/pq_4bit/pq4_fast_scan.cpp

@@ -9,12 +9,6 @@
 #include <faiss/impl/pq_4bit/pq4_fast_scan.h>
 #include <faiss/impl/pq_4bit/simd_result_handlers.h>
 
-#ifdef __x86_64__
-#ifdef __AVX2__
-#error "this should not be compiled with AVX2"
-#endif
-#endif
-
 #include <faiss/impl/pq_4bit/kernels_simd256.h>
 
 #include <faiss/impl/pq_4bit/dispatching.h>
@@ -376,6 +370,8 @@ int pq4_preferred_qbs(int n) {
 
 /**************************** Dispatching  */
 
+#ifdef COMPILE_SIMD_NONE
+
 template <>
 PQ4CodeScanner* make_pq4_scanner<SIMDLevel::NONE, false>(KNN_ARGS_LIST) {
     return make_pq4_scanner_1<SIMDLevel::NONE, false>(KNN_ARGS_LIST_2);
@@ -395,6 +391,7 @@ template <>
 PQ4CodeScanner* make_pq4_scanner<SIMDLevel::NONE, true>(PRES_ARGS_LIST) {
     return make_pq4_scanner_1<SIMDLevel::NONE, true>(PRES_ARGS_LIST_2);
 }
+#endif // COMPILE_SIMD_NONE
 
 template <bool with_id_map>
 PQ4CodeScanner* make_knn_scanner(
@@ -407,22 +404,31 @@ PQ4CodeScanner* make_knn_scanner(
         float* dis,
         idx_t* ids,
         const IDSelector* sel) {
-#ifdef COMPILE_SIMD_AVX512
-    if (SIMDConfig::level == SIMDLevel::AVX512) {
-        return make_pq4_scanner<SIMDLevel::AVX512, with_id_map>(
-                is_max, ns, ur, nq, ntotal, k, dis, ids, sel);
-    } else
+// code for dynamic dispatching is commented out
+#ifdef COMPILE_SIMD_AVX512F
+    //  if (SIMDConfig::level == SIMDLevel::AVX512F) {
+    return make_pq4_scanner<SIMDLevel::AVX512F, with_id_map>(
+            is_max, ns, ur, nq, ntotal, k, dis, ids, sel);
+    //  } else
 #endif
 #ifdef COMPILE_SIMD_AVX2
-            if (SIMDConfig::level == SIMDLevel::AVX2) {
-        return make_pq4_scanner<SIMDLevel::AVX2, with_id_map>(
-                is_max, ns, ur, nq, ntotal, k, dis, ids, sel);
-    } else
+    //         if (SIMDConfig::level == SIMDLevel::AVX2) {
+    return make_pq4_scanner<SIMDLevel::AVX2, with_id_map>(
+            is_max, ns, ur, nq, ntotal, k, dis, ids, sel);
+    // } else
 #endif
+#ifdef COMPILE_SIMD_NEON
+    //         if (SIMDConfig::level == SIMDLevel::ARM_NEON) {
+    return make_pq4_scanner<SIMDLevel::ARM_NEON, with_id_map>(
+            is_max, ns, ur, nq, ntotal, k, dis, ids, sel);
+    // } else
+#endif
+#ifdef COMPILE_SIMD_NONE
     {
         return make_pq4_scanner<SIMDLevel::NONE, with_id_map>(
                 is_max, ns, ur, nq, ntotal, k, dis, ids, sel);
     }
+#endif
 }
 
 PQ4CodeScanner* pq4_make_flat_knn_handler(
@@ -462,22 +468,30 @@ PQ4CodeScanner* pq4_make_ivf_range_handler(
         float radius,
         int norm_scale,
         const IDSelector* sel) {
-#ifdef COMPILE_SIMD_AVX512
-    if (SIMDConfig::level == SIMDLevel::AVX512) {
-        return make_pq4_scanner<SIMDLevel::AVX512, true>(
-                is_max, norm_scale, &rres, radius, 0, sel);
-    } else
+#ifdef COMPILE_SIMD_AVX512F
+    // if (SIMDConfig::level == SIMDLevel::AVX512F) {
+    return make_pq4_scanner<SIMDLevel::AVX512F, true>(
+            is_max, norm_scale, &rres, radius, 0, sel);
+    // } else
 #endif
 #ifdef COMPILE_SIMD_AVX2
-            if (SIMDConfig::level == SIMDLevel::AVX2) {
-        return make_pq4_scanner<SIMDLevel::AVX2, true>(
-                is_max, norm_scale, &rres, radius, 0, sel);
-    } else
+    //         if (SIMDConfig::level == SIMDLevel::AVX2) {
+    return make_pq4_scanner<SIMDLevel::AVX2, true>(
+            is_max, norm_scale, &rres, radius, 0, sel);
+    // } else
 #endif
+#ifdef COMPILE_SIMD_NEON
+    //         if (SIMDConfig::level == SIMDLevel::ARM_NEON) {
+    return make_pq4_scanner<SIMDLevel::ARM_NEON, true>(
+            is_max, norm_scale, &rres, radius, 0, sel);
+    // } else
+#endif
+#ifdef COMPILE_SIMD_NONE
     {
         return make_pq4_scanner<SIMDLevel::NONE, true>(
                 is_max, norm_scale, &rres, radius, 0, sel);
     }
+#endif
 }
 
 PQ4CodeScanner* pq4_make_ivf_partial_range_handler(
@@ -488,22 +502,30 @@ PQ4CodeScanner* pq4_make_ivf_partial_range_handler(
         idx_t i1,
         int norm_scale,
         const IDSelector* sel) {
-#ifdef COMPILE_SIMD_AVX512
-    if (SIMDConfig::level == SIMDLevel::AVX512) {
-        return make_pq4_scanner<SIMDLevel::AVX512, true>(
-                is_max, norm_scale, &pres, radius, 0, i0, i1, sel);
-    } else
+#ifdef COMPILE_SIMD_AVX512F
+    // if (SIMDConfig::level == SIMDLevel::AVX512F) {
+    return make_pq4_scanner<SIMDLevel::AVX512F, true>(
+            is_max, norm_scale, &pres, radius, 0, i0, i1, sel);
+    // } else
 #endif
 #ifdef COMPILE_SIMD_AVX2
-            if (SIMDConfig::level == SIMDLevel::AVX2) {
-        return make_pq4_scanner<SIMDLevel::AVX2, true>(
-                is_max, norm_scale, &pres, radius, 0, i0, i1, sel);
-    } else
+    //        if (SIMDConfig::level == SIMDLevel::AVX2) {
+    return make_pq4_scanner<SIMDLevel::AVX2, true>(
+            is_max, norm_scale, &pres, radius, 0, i0, i1, sel);
+    // } else
 #endif
+#ifdef COMPILE_SIMD_NEON
+    //        if (SIMDConfig::level == SIMDLevel::ARM_NEON) {
+    return make_pq4_scanner<SIMDLevel::ARM_NEON, true>(
+            is_max, norm_scale, &pres, radius, 0, i0, i1, sel);
+    // } else
+#endif
+#ifdef COMPILE_SIMD_NONE
     {
         return make_pq4_scanner<SIMDLevel::NONE, true>(
                 is_max, norm_scale, &pres, radius, 0, i0, i1, sel);
     }
+#endif
 }
 
 } // namespace faiss