V2 enable specifying a subset of SIMD levels that are implemented at dispatching time. (#4959)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4959

The link with the original diff D96943554 was lost for some reason, so reviving

So far the with_simd_level function would assume that all  implementations are available for each function. This is not the case, and  is effectively not the case when we enable AVX512_SPR.
This diff has a specialized with_selected_simd_levels that takes a bitmask of the available implementations as template parameter.
The previous with_simd_level remains, it assumes by default the AVX2, NEON and AVX512 implementations are available.

Reviewed By: algoriddle

Differential Revision: D97247328

fbshipit-source-id: 47367cda5ed4e4c02b095343878049caa47c9d38

faiss/CMakeLists.txt

@@ -333,6 +333,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/exhaustive_L2sqr_blas_cmax.h
   utils/simd_impl/IVFFlatScanner-inl.h
   utils/simd_impl/distances_sse-inl.h
 )
@@ -404,7 +405,7 @@ else()
   add_compile_options(/bigobj)
 endif()
 target_sources(faiss_avx512_spr PRIVATE ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC})
-target_compile_definitions(faiss_avx512_spr PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512)
+target_compile_definitions(faiss_avx512_spr PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR )
 
 add_library(faiss_sve ${FAISS_SRC})
 if(NOT FAISS_OPT_LEVEL STREQUAL "sve")

faiss/IndexAdditiveQuantizer.cpp

@@ -13,6 +13,7 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 namespace faiss {

faiss/IndexFlatCodes.cpp

@@ -13,6 +13,7 @@
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 namespace faiss {

faiss/IndexIVFFlat.cpp

@@ -21,6 +21,7 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/expanded_scanners.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 #define THE_SIMD_LEVEL SIMDLevel::NONE

faiss/IndexIVFFlatPanorama.cpp

@@ -20,6 +20,7 @@
 #include <faiss/impl/ResultHandler.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 #include <faiss/utils/utils.h>
 

faiss/IndexIVFPQFastScan.cpp

@@ -18,6 +18,7 @@
 #include <faiss/impl/simdlib/simdlib_dispatch.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 #include <faiss/invlists/BlockInvertedLists.h>

faiss/IndexPQ.cpp

@@ -152,7 +152,9 @@ FlatCodesDistanceComputer* get_FlatCodesDistanceComputer1(
 } // namespace
 
 FlatCodesDistanceComputer* IndexPQ::get_FlatCodesDistanceComputer() const {
-    DISPATCH_SIMDLevel(get_FlatCodesDistanceComputer1, *this);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return get_FlatCodesDistanceComputer1<SL>(*this);
+    });
 }
 
 /*****************************************

faiss/impl/fast_scan/fast_scan.cpp

@@ -426,17 +426,10 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
         int64_t* ids,
         const IDSelector* sel,
         bool with_id_map) {
-    DISPATCH_SIMDLevel(
-            make_fast_scan_scanner_impl,
-            is_max,
-            impl,
-            nq,
-            ntotal,
-            k,
-            distances,
-            ids,
-            sel,
-            with_id_map);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return make_fast_scan_scanner_impl<SL>(
+                is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map);
+    });
 }
 
 std::unique_ptr<FastScanCodeScanner> make_range_scanner(
@@ -445,8 +438,9 @@ std::unique_ptr<FastScanCodeScanner> make_range_scanner(
         float radius,
         size_t ntotal,
         const IDSelector* sel) {
-    DISPATCH_SIMDLevel(
-            make_range_scanner_impl, is_max, rres, radius, ntotal, sel);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return make_range_scanner_impl<SL>(is_max, rres, radius, ntotal, sel);
+    });
 }
 
 std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner(
@@ -457,15 +451,10 @@ std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner(
         size_t q0,
         size_t q1,
         const IDSelector* sel) {
-    DISPATCH_SIMDLevel(
-            make_partial_range_scanner_impl,
-            is_max,
-            pres,
-            radius,
-            ntotal,
-            q0,
-            q1,
-            sel);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return make_partial_range_scanner_impl<SL>(
+                is_max, pres, radius, ntotal, q0, q1, sel);
+    });
 }
 
 std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner(
@@ -478,17 +467,18 @@ std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner(
         const IDSelector* sel,
         const FastScanDistancePostProcessing& context,
         bool is_multi_bit) {
-    DISPATCH_SIMDLevel(
-            rabitq_make_knn_scanner_impl,
-            index,
-            is_max,
-            nq,
-            k,
-            distances,
-            ids,
-            sel,
-            context,
-            is_multi_bit);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return rabitq_make_knn_scanner_impl<SL>(
+                index,
+                is_max,
+                nq,
+                k,
+                distances,
+                ids,
+                sel,
+                context,
+                is_multi_bit);
+    });
 }
 
 std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner(
@@ -500,16 +490,10 @@ std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner(
         int64_t* ids,
         const FastScanDistancePostProcessing* context,
         bool multi_bit) {
-    DISPATCH_SIMDLevel(
-            rabitq_ivf_make_knn_scanner_impl,
-            is_max,
-            index,
-            nq,
-            k,
-            distances,
-            ids,
-            context,
-            multi_bit);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return rabitq_ivf_make_knn_scanner_impl<SL>(
+                is_max, index, nq, k, distances, ids, context, multi_bit);
+    });
 }
 
 } // namespace faiss

faiss/impl/fast_scan/fast_scan.h

@@ -219,7 +219,7 @@ std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl(
         bool with_id_map);
 
 /// Runtime dispatch wrapper: selects the best available SIMD level
-/// (via DISPATCH_SIMDLevel) and delegates to the corresponding
+/// (via with_simd_level) and delegates to the corresponding
 /// make_fast_scan_scanner_impl<SL> specialization.
 std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
         bool is_max,

faiss/impl/pq_code_distance/pq_code_distance-avx512.cpp

@@ -159,50 +159,6 @@ void pq_code_distance_four_impl<SIMDLevel::AVX512>(
     }
 }
 
-#ifdef COMPILE_SIMD_AVX512_SPR
-// AVX512_SPR: Sapphire Rapids is a superset of AVX512. Reuse the
-// AVX512 implementation until a dedicated SPR specialization is written.
-
-// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
-template <>
-float pq_code_distance_single_impl<SIMDLevel::AVX512_SPR>(
-        size_t M,
-        size_t nbits,
-        const float* sim_table,
-        const uint8_t* code) {
-    return pq_code_distance_single_impl<SIMDLevel::AVX512>(
-            M, nbits, sim_table, code);
-}
-
-// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
-template <>
-void pq_code_distance_four_impl<SIMDLevel::AVX512_SPR>(
-        size_t M,
-        size_t nbits,
-        const float* sim_table,
-        const uint8_t* __restrict code0,
-        const uint8_t* __restrict code1,
-        const uint8_t* __restrict code2,
-        const uint8_t* __restrict code3,
-        float& result0,
-        float& result1,
-        float& result2,
-        float& result3) {
-    pq_code_distance_four_impl<SIMDLevel::AVX512>(
-            M,
-            nbits,
-            sim_table,
-            code0,
-            code1,
-            code2,
-            code3,
-            result0,
-            result1,
-            result2,
-            result3);
-}
-#endif // COMPILE_SIMD_AVX512_SPR
-
 } // namespace pq_code_distance
 } // namespace faiss
 

faiss/impl/pq_code_distance/pq_code_distance-generic.cpp

@@ -9,7 +9,7 @@
 // 1. _impl specializations for NONE (and ARM_NEON), using scalar code.
 // 2. Non-templated PQ code distance dispatch wrappers
 //    (pq_code_distance_single, pq_code_distance_four) declared in
-//    pq_code_distance.h. These use DISPATCH_SIMDLevel to route to the
+//    pq_code_distance.h. These use with_simd_level to route to the
 //    best available SIMD implementation via pq_code_distance_*_impl
 //    function template specializations defined in the per-SIMD .cpp files.
 
@@ -107,7 +107,9 @@ float pq_code_distance_single(
         size_t nbits,
         const float* sim_table,
         const uint8_t* code) {
-    DISPATCH_SIMDLevel(pq_code_distance_single_impl, M, nbits, sim_table, code);
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return pq_code_distance_single_impl<SL>(M, nbits, sim_table, code);
+    });
 }
 
 void pq_code_distance_four(
@@ -122,19 +124,20 @@ void pq_code_distance_four(
         float& result1,
         float& result2,
         float& result3) {
-    DISPATCH_SIMDLevel(
-            pq_code_distance_four_impl,
-            M,
-            nbits,
-            sim_table,
-            code0,
-            code1,
-            code2,
-            code3,
-            result0,
-            result1,
-            result2,
-            result3);
+    with_simd_level([&]<SIMDLevel SL>() {
+        pq_code_distance_four_impl<SL>(
+                M,
+                nbits,
+                sim_table,
+                code0,
+                code1,
+                code2,
+                code3,
+                result0,
+                result1,
+                result2,
+                result3);
+    });
 }
 
 } // namespace pq_code_distance

faiss/impl/simd_dispatch.h

@@ -23,86 +23,99 @@
 
 namespace faiss {
 
-/*********************** x86 SIMD dispatch cases */
+/** Defining which SIMD levels are available for a given function is via a
+ * binary mask. Here we predefine the most common masks.
+ *  */
 
-#ifdef COMPILE_SIMD_AVX2
-#define DISPATCH_SIMDLevel_AVX2(f, ...) \
-    case SIMDLevel::AVX2:               \
-        return f<SIMDLevel::AVX2>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_AVX2(f, ...)
+constexpr int AVAILABLE_SIMD_LEVELS_NONE = (1 << int(SIMDLevel::NONE));
+
+constexpr int AVAILABLE_SIMD_LEVELS_AVX2_NEON = AVAILABLE_SIMD_LEVELS_NONE |
+        (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::ARM_NEON));
+
+// A0: same + AVX512
+constexpr int AVAILABLE_SIMD_LEVELS_A0 =
+        AVAILABLE_SIMD_LEVELS_AVX2_NEON | (1 << int(SIMDLevel::AVX512));
+
+// A1: same + ARM_SVE (for functions with dedicated SVE implementations)
+constexpr int AVAILABLE_SIMD_LEVELS_A1 =
+        AVAILABLE_SIMD_LEVELS_A0 | (1 << int(SIMDLevel::ARM_SVE));
+
+// A2: NONE + AVX2 + ARM_SVE only (for functions with only these
+// implementations)
+constexpr int AVAILABLE_SIMD_LEVELS_A2 = AVAILABLE_SIMD_LEVELS_NONE |
+        (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::ARM_SVE));
+
+constexpr int AVAILABLE_SIMD_LEVELS_ALL = -1;
+
+/** The complete dispatching function. It takes into account:
+ * - the currently selected SIMD level
+ * - the compiled in SIMD levels (given by COMPILE_SIMD_XXX)
+ * - the available SIMD implementations for that particular function (given by
+ * available_levels)
+ */
+
+template <int available_levels, typename LambdaType>
+inline auto with_selected_simd_levels(LambdaType&& action) {
+#ifdef FAISS_ENABLE_DD
+    switch (SIMDConfig::level) {
+        // For x86 -- try from highest to lowest level
+
+#ifdef COMPILE_SIMD_AVX512_SPR
+        case SIMDLevel::AVX512_SPR:
+            if constexpr (
+                    available_levels & (1 << int(SIMDLevel::AVX512_SPR))) {
+                return action.template operator()<SIMDLevel::AVX512_SPR>();
+            }
+            [[fallthrough]];
 #endif
 
 #ifdef COMPILE_SIMD_AVX512
-#define DISPATCH_SIMDLevel_AVX512(f, ...) \
-    case SIMDLevel::AVX512:               \
-        return f<SIMDLevel::AVX512>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_AVX512(f, ...)
+        case SIMDLevel::AVX512:
+            if constexpr (available_levels & (1 << int(SIMDLevel::AVX512))) {
+                return action.template operator()<SIMDLevel::AVX512>();
+            }
+            [[fallthrough]];
 #endif
 
-#ifdef COMPILE_SIMD_AVX512_SPR
-#define DISPATCH_SIMDLevel_AVX512_SPR(f, ...) \
-    case SIMDLevel::AVX512_SPR:               \
-        return f<SIMDLevel::AVX512_SPR>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_AVX512_SPR(f, ...)
+#ifdef COMPILE_SIMD_AVX2
+        case SIMDLevel::AVX2:
+            if constexpr (available_levels & (1 << int(SIMDLevel::AVX2))) {
+                return action.template operator()<SIMDLevel::AVX2>();
+            }
+            [[fallthrough]];
 #endif
 
-/*********************** ARM SIMD dispatch cases */
+            // For ARM, try from highest to lowest level
+#ifdef COMPILE_SIMD_ARM_SVE
+        case SIMDLevel::ARM_SVE:
+            if constexpr (available_levels & (1 << int(SIMDLevel::ARM_SVE))) {
+                return action.template operator()<SIMDLevel::ARM_SVE>();
+            }
+            [[fallthrough]];
+#endif
 
 #ifdef COMPILE_SIMD_ARM_NEON
-#define DISPATCH_SIMDLevel_ARM_NEON(f, ...) \
-    case SIMDLevel::ARM_NEON:               \
-        return f<SIMDLevel::ARM_NEON>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_ARM_NEON(f, ...)
+        case SIMDLevel::ARM_NEON:
+            if constexpr (available_levels & (1 << int(SIMDLevel::ARM_NEON))) {
+                return action.template operator()<SIMDLevel::ARM_NEON>();
+            }
+            [[fallthrough]];
 #endif
-
-#ifdef COMPILE_SIMD_ARM_SVE
-#define DISPATCH_SIMDLevel_ARM_SVE(f, ...) \
-    case SIMDLevel::ARM_SVE:               \
-        return f<SIMDLevel::ARM_SVE>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_ARM_SVE(f, ...)
-#endif
-
-/*********************** Main dispatch macro */
-
-#ifdef FAISS_ENABLE_DD
-
-// DD mode: runtime dispatch based on SIMDConfig::level
-#define DISPATCH_SIMDLevel(f, ...)                         \
-    switch (SIMDConfig::level) {                           \
-        case SIMDLevel::NONE:                              \
-            return f<SIMDLevel::NONE>(__VA_ARGS__);        \
-            DISPATCH_SIMDLevel_AVX2(f, __VA_ARGS__);       \
-            DISPATCH_SIMDLevel_AVX512(f, __VA_ARGS__);     \
-            DISPATCH_SIMDLevel_AVX512_SPR(f, __VA_ARGS__); \
-            DISPATCH_SIMDLevel_ARM_NEON(f, __VA_ARGS__);   \
-            DISPATCH_SIMDLevel_ARM_SVE(f, __VA_ARGS__);    \
-        default:                                           \
-            FAISS_THROW_MSG("Invalid SIMD level");         \
+        default:
+            return action.template operator()<SIMDLevel::NONE>();
+    }
+#else // static dispatch
+    // In static mode, SINGLE_SIMD_LEVEL is a constexpr resolved at compile
+    // time. If the compiled level is not in the available set, fall through
+    // to NONE (mirroring the DD fallthrough behavior). Only SINGLE_SIMD_LEVEL
+    // and NONE have compiled specializations.
+    if constexpr (available_levels & (1 << int(SINGLE_SIMD_LEVEL))) {
+        return action.template operator()<SINGLE_SIMD_LEVEL>();
+    } else {
+        return action.template operator()<SIMDLevel::NONE>();
     }
-
-#else // Static mode
-
-// Static mode: direct call to compiled-in SIMD level (no runtime switch)
-#if defined(COMPILE_SIMD_AVX512_SPR)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::AVX512_SPR>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_AVX512)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::AVX512>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_AVX2)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::AVX2>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_ARM_SVE)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::ARM_SVE>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_ARM_NEON)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::ARM_NEON>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::NONE>(__VA_ARGS__)
 #endif
-
-#endif // FAISS_ENABLE_DD
+}
 
 /**
  * Dispatch to a lambda with SIMDLevel as a compile-time constant.
@@ -126,6 +139,8 @@ namespace faiss {
  *   });
  *
  * The lambda must be a generic lambda with a SIMDLevel template parameter.
+ * By default, the lambda uses levels AVX2 + AVX512 + NEON, since these are the
+ * most common cases.
  *
  * @param action A generic lambda with signature `template<SIMDLevel> T
  * operator()()`
@@ -133,24 +148,18 @@ namespace faiss {
  */
 template <typename LambdaType>
 inline auto with_simd_level(LambdaType&& action) {
-    DISPATCH_SIMDLevel(action.template operator());
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+            std::forward<LambdaType>(action));
 }
 
 /**
- * Like with_simd_level, but maps to the 256-bit SIMD equivalent:
- *   AVX512, AVX512_SPR -> AVX2
- *   ARM_SVE -> ARM_NEON
- *   AVX2, ARM_NEON, NONE -> unchanged
- *
- * Use for functions implemented with simd8float32 (256-bit) operations
+ * Use for functions implemented with simdXintY (256-bit) operations
  * that don't have dedicated AVX512 or SVE implementations.
  */
 template <typename LambdaType>
 inline auto with_simd_level_256bit(LambdaType&& action) {
-    return with_simd_level([&]<SIMDLevel level>() {
-        constexpr SIMDLevel level256 = simd256_level_selector<level>::value;
-        return action.template operator()<level256>();
-    });
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_AVX2_NEON>(
+            std::forward<LambdaType>(action));
 }
 
 } // namespace faiss

faiss/utils/distances.cpp

@@ -16,19 +16,15 @@
 
 #include <omp.h>
 
-#ifdef __AVX2__
-#include <immintrin.h>
-#elif defined(__ARM_FEATURE_SVE)
-#include <arm_sve.h>
-#endif
-
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
 
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/distances_fused/distances_fused.h>
+#include <faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h>
 
 #ifndef FINTEGER
 #define FINTEGER long
@@ -206,10 +202,12 @@ void fvec_norms_L2(
         const float* __restrict x,
         size_t d,
         size_t nx) {
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (nx > 10000)
-    for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-        nr[i] = sqrtf(fvec_norm_L2sqr_dispatch(x + i * d, d));
-    }
+        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+            nr[i] = sqrtf(fvec_norm_L2sqr<SL>(x + i * d, d));
+        }
+    });
 }
 
 void fvec_norms_L2sqr(
@@ -217,10 +215,12 @@ void fvec_norms_L2sqr(
         const float* __restrict x,
         size_t d,
         size_t nx) {
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (nx > 10000)
-    for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-        nr[i] = fvec_norm_L2sqr_dispatch(x + i * d, d);
-    }
+        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+            nr[i] = fvec_norm_L2sqr<SL>(x + i * d, d);
+        }
+    });
 }
 
 // The following is a workaround to a problem
@@ -234,29 +234,35 @@ void fvec_norms_L2sqr(
 // The workaround below is explicitly branching
 // off to a codepath without omp.
 
-#define FVEC_RENORM_L2_IMPL                     \
-    float* __restrict xi = x + i * d;           \
-                                                \
-    float nr = fvec_norm_L2sqr_dispatch(xi, d); \
-                                                \
-    if (nr > 0) {                               \
-        size_t j;                               \
-        const float inv_nr = 1.0 / sqrtf(nr);   \
-        for (j = 0; j < d; j++)                 \
-            xi[j] *= inv_nr;                    \
-    }
-
 void fvec_renorm_L2_noomp(size_t d, size_t nx, float* __restrict x) {
-    for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-        FVEC_RENORM_L2_IMPL
-    }
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+            float* __restrict xi = x + i * d;
+            float nr = fvec_norm_L2sqr<SL>(xi, d);
+            if (nr > 0) {
+                const float inv_nr = 1.0 / sqrtf(nr);
+                for (size_t j = 0; j < d; j++) {
+                    xi[j] *= inv_nr;
+                }
+            }
+        }
+    });
 }
 
 void fvec_renorm_L2_omp(size_t d, size_t nx, float* __restrict x) {
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (nx > 10000)
-    for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-        FVEC_RENORM_L2_IMPL
-    }
+        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+            float* __restrict xi = x + i * d;
+            float nr = fvec_norm_L2sqr<SL>(xi, d);
+            if (nr > 0) {
+                const float inv_nr = 1.0 / sqrtf(nr);
+                for (size_t j = 0; j < d; j++) {
+                    xi[j] *= inv_nr;
+                }
+            }
+        }
+    });
 }
 
 void fvec_renorm_L2(size_t d, size_t nx, float* __restrict x) {
@@ -289,22 +295,24 @@ void exhaustive_inner_product_seq(
 #pragma omp parallel num_threads(nt)
     {
         SingleResultHandler resi(res);
+        with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp for
-        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-            const float* x_i = x + i * d;
-            const float* y_j = y;
+            for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+                const float* x_i = x + i * d;
+                const float* y_j = y;
 
-            resi.begin(i);
+                resi.begin(i);
 
-            for (size_t j = 0; j < ny; j++, y_j += d) {
-                if (!res.is_in_selection(j)) {
-                    continue;
+                for (size_t j = 0; j < ny; j++, y_j += d) {
+                    if (!res.is_in_selection(j)) {
+                        continue;
+                    }
+                    float ip = fvec_inner_product<SL>(x_i, y_j, d);
+                    resi.add_result(ip, j);
                 }
-                float ip = fvec_inner_product_dispatch(x_i, y_j, d);
-                resi.add_result(ip, j);
+                resi.end();
             }
-            resi.end();
-        }
+        });
     }
 }
 
@@ -323,20 +331,22 @@ void exhaustive_L2sqr_seq(
 #pragma omp parallel num_threads(nt)
     {
         SingleResultHandler resi(res);
+        with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp for
-        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-            const float* x_i = x + i * d;
-            const float* y_j = y;
-            resi.begin(i);
-            for (size_t j = 0; j < ny; j++, y_j += d) {
-                if (!res.is_in_selection(j)) {
-                    continue;
+            for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+                const float* x_i = x + i * d;
+                const float* y_j = y;
+                resi.begin(i);
+                for (size_t j = 0; j < ny; j++, y_j += d) {
+                    if (!res.is_in_selection(j)) {
+                        continue;
+                    }
+                    float disij = fvec_L2sqr<SL>(x_i, y_j, d);
+                    resi.add_result(disij, j);
                 }
-                float disij = fvec_L2sqr_dispatch(x_i, y_j, d);
-                resi.add_result(disij, j);
+                resi.end();
             }
-            resi.end();
-        }
+        });
     }
 }
 
@@ -497,397 +507,13 @@ void exhaustive_L2sqr_blas(
         size_t nx,
         size_t ny,
         BlockResultHandler& res,
-        const float* /* y_norms */ = nullptr) {
-    exhaustive_L2sqr_blas_default_impl(x, y, d, nx, ny, res);
+        const float* y_norms = nullptr) {
+    exhaustive_L2sqr_blas_default_impl(x, y, d, nx, ny, res, y_norms);
 }
 
-#ifdef __AVX2__
-void exhaustive_L2sqr_blas_cmax_avx2(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        Top1BlockResultHandler<CMax<float, int64_t>>& res,
-        const float* y_norms) {
-    // BLAS does not like empty matrices
-    if (nx == 0 || ny == 0) {
-        return;
-    }
+} // anonymous namespace
 
-    /* block sizes */
-    const size_t bs_x = distance_compute_blas_query_bs;
-    const size_t bs_y = distance_compute_blas_database_bs;
-    // const size_t bs_x = 16, bs_y = 16;
-    std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]);
-    std::unique_ptr<float[]> x_norms(new float[nx]);
-    std::unique_ptr<float[]> del2;
-
-    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
-
-    if (!y_norms) {
-        float* y_norms2 = new float[ny];
-        del2.reset(y_norms2);
-        fvec_norms_L2sqr(y_norms2, y, d, ny);
-        y_norms = y_norms2;
-    }
-
-    for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
-        size_t i1 = i0 + bs_x;
-        if (i1 > nx) {
-            i1 = nx;
-        }
-
-        res.begin_multiple(i0, i1);
-
-        for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
-            size_t j1 = j0 + bs_y;
-            if (j1 > ny) {
-                j1 = ny;
-            }
-            /* compute the actual dot products */
-            {
-                float one = 1, zero = 0;
-                FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
-                sgemm_("Transpose",
-                       "Not transpose",
-                       &nyi,
-                       &nxi,
-                       &di,
-                       &one,
-                       y + j0 * d,
-                       &di,
-                       x + i0 * d,
-                       &di,
-                       &zero,
-                       ip_block.get(),
-                       &nyi);
-            }
-            for (size_t i = i0; i < i1; i++) {
-                float* ip_line = ip_block.get() + (i - i0) * (j1 - j0);
-
-                _mm_prefetch((const char*)ip_line, _MM_HINT_NTA);
-                _mm_prefetch((const char*)(ip_line + 16), _MM_HINT_NTA);
-
-                // constant
-                const __m256 mul_minus2 = _mm256_set1_ps(-2);
-
-                // Track 8 min distances + 8 min indices.
-                // All the distances tracked do not take x_norms[i]
-                //   into account in order to get rid of extra
-                //   _mm256_add_ps(x_norms[i], ...) instructions
-                //   is distance computations.
-                __m256 min_distances =
-                        _mm256_set1_ps(res.dis_tab[i] - x_norms[i]);
-
-                // these indices are local and are relative to j0.
-                // so, value 0 means j0.
-                __m256i min_indices = _mm256_set1_epi32(0);
-
-                __m256i current_indices =
-                        _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
-                const __m256i indices_delta = _mm256_set1_epi32(8);
-
-                // current j index
-                size_t idx_j = 0;
-                size_t count = j1 - j0;
-
-                // process 16 elements per loop
-                for (; idx_j < (count / 16) * 16; idx_j += 16, ip_line += 16) {
-                    _mm_prefetch((const char*)(ip_line + 32), _MM_HINT_NTA);
-                    _mm_prefetch((const char*)(ip_line + 48), _MM_HINT_NTA);
-
-                    // load values for norms
-                    const __m256 y_norm_0 =
-                            _mm256_loadu_ps(y_norms + idx_j + j0 + 0);
-                    const __m256 y_norm_1 =
-                            _mm256_loadu_ps(y_norms + idx_j + j0 + 8);
-
-                    // load values for dot products
-                    const __m256 ip_0 = _mm256_loadu_ps(ip_line + 0);
-                    const __m256 ip_1 = _mm256_loadu_ps(ip_line + 8);
-
-                    // compute dis = y_norm[j] - 2 * dot(x_norm[i], y_norm[j]).
-                    // x_norm[i] was dropped off because it is a constant for a
-                    // given i. We'll deal with it later.
-                    __m256 distances_0 =
-                            _mm256_fmadd_ps(ip_0, mul_minus2, y_norm_0);
-                    __m256 distances_1 =
-                            _mm256_fmadd_ps(ip_1, mul_minus2, y_norm_1);
-
-                    // compare the new distances to the min distances
-                    // for each of the first group of 8 AVX2 components.
-                    const __m256 comparison_0 = _mm256_cmp_ps(
-                            min_distances, distances_0, _CMP_LE_OS);
-
-                    // update min distances and indices with closest vectors if
-                    // needed.
-                    min_distances = _mm256_blendv_ps(
-                            distances_0, min_distances, comparison_0);
-                    min_indices = _mm256_castps_si256(_mm256_blendv_ps(
-                            _mm256_castsi256_ps(current_indices),
-                            _mm256_castsi256_ps(min_indices),
-                            comparison_0));
-                    current_indices =
-                            _mm256_add_epi32(current_indices, indices_delta);
-
-                    // compare the new distances to the min distances
-                    // for each of the second group of 8 AVX2 components.
-                    const __m256 comparison_1 = _mm256_cmp_ps(
-                            min_distances, distances_1, _CMP_LE_OS);
-
-                    // update min distances and indices with closest vectors if
-                    // needed.
-                    min_distances = _mm256_blendv_ps(
-                            distances_1, min_distances, comparison_1);
-                    min_indices = _mm256_castps_si256(_mm256_blendv_ps(
-                            _mm256_castsi256_ps(current_indices),
-                            _mm256_castsi256_ps(min_indices),
-                            comparison_1));
-                    current_indices =
-                            _mm256_add_epi32(current_indices, indices_delta);
-                }
-
-                // dump values and find the minimum distance / minimum index
-                float min_distances_scalar[8];
-                uint32_t min_indices_scalar[8];
-                _mm256_storeu_ps(min_distances_scalar, min_distances);
-                _mm256_storeu_si256(
-                        (__m256i*)(min_indices_scalar), min_indices);
-
-                float current_min_distance = res.dis_tab[i];
-                uint32_t current_min_index = res.ids_tab[i];
-
-                // This unusual comparison is needed to maintain the behavior
-                // of the original implementation: if two indices are
-                // represented with equal distance values, then
-                // the index with the min value is returned.
-                for (size_t jv = 0; jv < 8; jv++) {
-                    // add missing x_norms[i]
-                    float distance_candidate =
-                            min_distances_scalar[jv] + x_norms[i];
-
-                    // negative values can occur for identical vectors
-                    //    due to roundoff errors.
-                    if (distance_candidate < 0) {
-                        distance_candidate = 0;
-                    }
-
-                    int64_t index_candidate = min_indices_scalar[jv] + j0;
-
-                    if (current_min_distance > distance_candidate) {
-                        current_min_distance = distance_candidate;
-                        current_min_index = index_candidate;
-                    } else if (
-                            current_min_distance == distance_candidate &&
-                            current_min_index > index_candidate) {
-                        current_min_index = index_candidate;
-                    }
-                }
-
-                // process leftovers
-                for (; idx_j < count; idx_j++, ip_line++) {
-                    float ip = *ip_line;
-                    float dis = x_norms[i] + y_norms[idx_j + j0] - 2 * ip;
-                    // negative values can occur for identical vectors
-                    //    due to roundoff errors.
-                    if (dis < 0) {
-                        dis = 0;
-                    }
-
-                    if (current_min_distance > dis) {
-                        current_min_distance = dis;
-                        current_min_index = idx_j + j0;
-                    }
-                }
-
-                //
-                res.add_result(i, current_min_distance, current_min_index);
-            }
-        }
-        // Does nothing for SingleBestResultHandler, but
-        // keeping the call for the consistency.
-        res.end_multiple();
-        InterruptCallback::check();
-    }
-}
-#elif defined(__ARM_FEATURE_SVE)
-void exhaustive_L2sqr_blas_cmax_sve(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        Top1BlockResultHandler<CMax<float, int64_t>>& res,
-        const float* y_norms) {
-    // BLAS does not like empty matrices
-    if (nx == 0 || ny == 0)
-        return;
-
-    /* block sizes */
-    const size_t bs_x = distance_compute_blas_query_bs;
-    const size_t bs_y = distance_compute_blas_database_bs;
-    // const size_t bs_x = 16, bs_y = 16;
-    std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]);
-    std::unique_ptr<float[]> x_norms(new float[nx]);
-    std::unique_ptr<float[]> del2;
-
-    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
-
-    const size_t lanes = svcntw();
-
-    if (!y_norms) {
-        float* y_norms2 = new float[ny];
-        del2.reset(y_norms2);
-        fvec_norms_L2sqr(y_norms2, y, d, ny);
-        y_norms = y_norms2;
-    }
-
-    for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
-        size_t i1 = i0 + bs_x;
-        if (i1 > nx)
-            i1 = nx;
-
-        res.begin_multiple(i0, i1);
-
-        for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
-            size_t j1 = j0 + bs_y;
-            if (j1 > ny)
-                j1 = ny;
-            /* compute the actual dot products */
-            {
-                float one = 1, zero = 0;
-                FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
-                sgemm_("Transpose",
-                       "Not transpose",
-                       &nyi,
-                       &nxi,
-                       &di,
-                       &one,
-                       y + j0 * d,
-                       &di,
-                       x + i0 * d,
-                       &di,
-                       &zero,
-                       ip_block.get(),
-                       &nyi);
-            }
-            for (size_t i = i0; i < i1; i++) {
-                const size_t count = j1 - j0;
-                float* ip_line = ip_block.get() + (i - i0) * count;
-
-                svprfw(svwhilelt_b32_u64(0, count), ip_line, SV_PLDL1KEEP);
-                svprfw(svwhilelt_b32_u64(lanes, count),
-                       ip_line + lanes,
-                       SV_PLDL1KEEP);
-
-                // Track lanes min distances + lanes min indices.
-                // All the distances tracked do not take x_norms[i]
-                //   into account in order to get rid of extra
-                //   vaddq_f32(x_norms[i], ...) instructions
-                //   is distance computations.
-                auto min_distances = svdup_n_f32(res.dis_tab[i] - x_norms[i]);
-
-                // these indices are local and are relative to j0.
-                // so, value 0 means j0.
-                auto min_indices = svdup_n_u32(0u);
-
-                auto current_indices = svindex_u32(0u, 1u);
-
-                // process lanes * 2 elements per loop
-                for (size_t idx_j = 0; idx_j < count;
-                     idx_j += lanes * 2, ip_line += lanes * 2) {
-                    svprfw(svwhilelt_b32_u64(idx_j + lanes * 2, count),
-                           ip_line + lanes * 2,
-                           SV_PLDL1KEEP);
-                    svprfw(svwhilelt_b32_u64(idx_j + lanes * 3, count),
-                           ip_line + lanes * 3,
-                           SV_PLDL1KEEP);
-
-                    // mask
-                    const auto mask_0 = svwhilelt_b32_u64(idx_j, count);
-                    const auto mask_1 = svwhilelt_b32_u64(idx_j + lanes, count);
-
-                    // load values for norms
-                    const auto y_norm_0 =
-                            svld1_f32(mask_0, y_norms + idx_j + j0 + 0);
-                    const auto y_norm_1 =
-                            svld1_f32(mask_1, y_norms + idx_j + j0 + lanes);
-
-                    // load values for dot products
-                    const auto ip_0 = svld1_f32(mask_0, ip_line + 0);
-                    const auto ip_1 = svld1_f32(mask_1, ip_line + lanes);
-
-                    // compute dis = y_norm[j] - 2 * dot(x_norm[i], y_norm[j]).
-                    // x_norm[i] was dropped off because it is a constant for a
-                    // given i. We'll deal with it later.
-                    const auto distances_0 =
-                            svmla_n_f32_z(mask_0, y_norm_0, ip_0, -2.f);
-                    const auto distances_1 =
-                            svmla_n_f32_z(mask_1, y_norm_1, ip_1, -2.f);
-
-                    // compare the new distances to the min distances
-                    // for each of the first group of 4 ARM SIMD components.
-                    auto comparison =
-                            svcmpgt_f32(mask_0, min_distances, distances_0);
-
-                    // update min distances and indices with closest vectors if
-                    // needed.
-                    min_distances =
-                            svsel_f32(comparison, distances_0, min_distances);
-                    min_indices =
-                            svsel_u32(comparison, current_indices, min_indices);
-                    current_indices = svadd_n_u32_x(
-                            mask_0,
-                            current_indices,
-                            static_cast<uint32_t>(lanes));
-
-                    // compare the new distances to the min distances
-                    // for each of the second group of 4 ARM SIMD components.
-                    comparison =
-                            svcmpgt_f32(mask_1, min_distances, distances_1);
-
-                    // update min distances and indices with closest vectors if
-                    // needed.
-                    min_distances =
-                            svsel_f32(comparison, distances_1, min_distances);
-                    min_indices =
-                            svsel_u32(comparison, current_indices, min_indices);
-                    current_indices = svadd_n_u32_x(
-                            mask_1,
-                            current_indices,
-                            static_cast<uint32_t>(lanes));
-                }
-
-                // add missing x_norms[i]
-                // negative values can occur for identical vectors
-                //    due to roundoff errors.
-                auto mask = svwhilelt_b32_u64(0, count);
-                min_distances = svadd_n_f32_z(
-                        svcmpge_n_f32(mask, min_distances, -x_norms[i]),
-                        min_distances,
-                        x_norms[i]);
-                min_indices = svadd_n_u32_x(
-                        mask, min_indices, static_cast<uint32_t>(j0));
-                mask = svcmple_n_f32(mask, min_distances, res.dis_tab[i]);
-                if (svcntp_b32(svptrue_b32(), mask) == 0)
-                    res.add_result(i, res.dis_tab[i], res.ids_tab[i]);
-                else {
-                    const auto min_distance = svminv_f32(mask, min_distances);
-                    const auto min_index = svminv_u32(
-                            svcmpeq_n_f32(mask, min_distances, min_distance),
-                            min_indices);
-                    res.add_result(i, min_distance, min_index);
-                }
-            }
-        }
-        // Does nothing for SingleBestResultHandler, but
-        // keeping the call for the consistency.
-        res.end_multiple();
-        InterruptCallback::check();
-    }
-}
-#endif
+namespace {
 
 // an override if only a single closest point is needed
 template <>
@@ -899,43 +525,20 @@ void exhaustive_L2sqr_blas<Top1BlockResultHandler<CMax<float, int64_t>>>(
         size_t ny,
         Top1BlockResultHandler<CMax<float, int64_t>>& res,
         const float* y_norms) {
-#if defined(__AVX2__)
     // use a faster fused kernel if available
     if (exhaustive_L2sqr_fused_cmax(x, y, d, nx, ny, res, y_norms)) {
-        // the kernel is available and it is complete, we're done.
         return;
     }
 
-    // run the specialized AVX2 implementation
-    exhaustive_L2sqr_blas_cmax_avx2(x, y, d, nx, ny, res, y_norms);
-
-#elif defined(__ARM_FEATURE_SVE)
-    // use a faster fused kernel if available
-    if (exhaustive_L2sqr_fused_cmax(x, y, d, nx, ny, res, y_norms)) {
-        // the kernel is available and it is complete, we're done.
-        return;
-    }
-
-    // run the specialized SVE implementation
-    exhaustive_L2sqr_blas_cmax_sve(x, y, d, nx, ny, res, y_norms);
-
-#elif defined(__aarch64__)
-    // use a faster fused kernel if available
-    if (exhaustive_L2sqr_fused_cmax(x, y, d, nx, ny, res, y_norms)) {
-        // the kernel is available and it is complete, we're done.
-        return;
-    }
-
-    // run the default implementation
-    exhaustive_L2sqr_blas_default_impl<
-            Top1BlockResultHandler<CMax<float, int64_t>>>(
-            x, y, d, nx, ny, res, y_norms);
-#else
-    // run the default implementation
-    exhaustive_L2sqr_blas_default_impl<
-            Top1BlockResultHandler<CMax<float, int64_t>>>(
-            x, y, d, nx, ny, res, y_norms);
-#endif
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A2>([&]<SIMDLevel SL>() {
+        if constexpr (SL == SIMDLevel::AVX2 || SL == SIMDLevel::ARM_SVE) {
+            exhaustive_L2sqr_blas_cmax<SL>(x, y, d, nx, ny, res, y_norms);
+        } else {
+            exhaustive_L2sqr_blas_default_impl<
+                    Top1BlockResultHandler<CMax<float, int64_t>>>(
+                    x, y, d, nx, ny, res, y_norms);
+        }
+    });
 }
 
 struct Run_search_inner_product {
@@ -1132,19 +735,21 @@ void fvec_inner_products_by_idx(
         size_t d,
         size_t nx,
         size_t ny) {
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-    for (int64_t j = 0; j < static_cast<int64_t>(nx); j++) {
-        const int64_t* __restrict idsj = ids + j * ny;
-        const float* xj = x + j * d;
-        float* __restrict ipj = ip + j * ny;
-        for (size_t i = 0; i < ny; i++) {
-            if (idsj[i] < 0) {
-                ipj[i] = -INFINITY;
-            } else {
-                ipj[i] = fvec_inner_product_dispatch(xj, y + d * idsj[i], d);
+        for (int64_t j = 0; j < static_cast<int64_t>(nx); j++) {
+            const int64_t* __restrict idsj = ids + j * ny;
+            const float* xj = x + j * d;
+            float* __restrict ipj = ip + j * ny;
+            for (size_t i = 0; i < ny; i++) {
+                if (idsj[i] < 0) {
+                    ipj[i] = -INFINITY;
+                } else {
+                    ipj[i] = fvec_inner_product<SL>(xj, y + d * idsj[i], d);
+                }
             }
         }
-    }
+    });
 }
 
 /* compute the inner product between x and a subset y of ny vectors,
@@ -1157,19 +762,21 @@ void fvec_L2sqr_by_idx(
         size_t d,
         size_t nx,
         size_t ny) {
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-    for (int64_t j = 0; j < static_cast<int64_t>(nx); j++) {
-        const int64_t* __restrict idsj = ids + j * ny;
-        const float* xj = x + j * d;
-        float* __restrict disj = dis + j * ny;
-        for (size_t i = 0; i < ny; i++) {
-            if (idsj[i] < 0) {
-                disj[i] = INFINITY;
-            } else {
-                disj[i] = fvec_L2sqr_dispatch(xj, y + d * idsj[i], d);
+        for (int64_t j = 0; j < static_cast<int64_t>(nx); j++) {
+            const int64_t* __restrict idsj = ids + j * ny;
+            const float* xj = x + j * d;
+            float* __restrict disj = dis + j * ny;
+            for (size_t i = 0; i < ny; i++) {
+                if (idsj[i] < 0) {
+                    disj[i] = INFINITY;
+                } else {
+                    disj[i] = fvec_L2sqr<SL>(xj, y + d * idsj[i], d);
+                }
             }
         }
-    }
+    });
 }
 
 void pairwise_indexed_L2sqr(
@@ -1180,14 +787,16 @@ void pairwise_indexed_L2sqr(
         const float* y,
         const int64_t* iy,
         float* dis) {
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (n > 1)
-    for (int64_t j = 0; j < static_cast<int64_t>(n); j++) {
-        if (ix[j] >= 0 && iy[j] >= 0) {
-            dis[j] = fvec_L2sqr_dispatch(x + d * ix[j], y + d * iy[j], d);
-        } else {
-            dis[j] = INFINITY;
+        for (int64_t j = 0; j < static_cast<int64_t>(n); j++) {
+            if (ix[j] >= 0 && iy[j] >= 0) {
+                dis[j] = fvec_L2sqr<SL>(x + d * ix[j], y + d * iy[j], d);
+            } else {
+                dis[j] = INFINITY;
+            }
         }
-    }
+    });
 }
 
 void pairwise_indexed_inner_product(
@@ -1198,15 +807,17 @@ void pairwise_indexed_inner_product(
         const float* y,
         const int64_t* iy,
         float* dis) {
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (n > 1)
-    for (int64_t j = 0; j < static_cast<int64_t>(n); j++) {
-        if (ix[j] >= 0 && iy[j] >= 0) {
-            dis[j] = fvec_inner_product_dispatch(
-                    x + d * ix[j], y + d * iy[j], d);
-        } else {
-            dis[j] = -INFINITY;
+        for (int64_t j = 0; j < static_cast<int64_t>(n); j++) {
+            if (ix[j] >= 0 && iy[j] >= 0) {
+                dis[j] =
+                        fvec_inner_product<SL>(x + d * ix[j], y + d * iy[j], d);
+            } else {
+                dis[j] = -INFINITY;
+            }
         }
-    }
+    });
 }
 
 /* Find the nearest neighbors for nx queries in a set of ny vectors
@@ -1236,16 +847,18 @@ void knn_inner_products_by_idx(
         int64_t* __restrict idxi = res_ids + i * k;
         minheap_heapify(k, simi, idxi);
 
-        for (j = 0; j < nsubset; j++) {
-            if (idsi[j] < 0 || static_cast<size_t>(idsi[j]) >= ny) {
-                break;
-            }
-            float ip = fvec_inner_product_dispatch(x_, y + d * idsi[j], d);
+        with_simd_level([&]<SIMDLevel SL>() {
+            for (j = 0; j < nsubset; j++) {
+                if (idsi[j] < 0 || static_cast<size_t>(idsi[j]) >= ny) {
+                    break;
+                }
+                float ip = fvec_inner_product<SL>(x_, y + d * idsi[j], d);
 
-            if (ip > simi[0]) {
-                minheap_replace_top(k, simi, idxi, ip, idsi[j]);
+                if (ip > simi[0]) {
+                    minheap_replace_top(k, simi, idxi, ip, idsi[j]);
+                }
             }
-        }
+        });
         minheap_reorder(k, simi, idxi);
     }
 }
@@ -1272,16 +885,18 @@ void knn_L2sqr_by_idx(
         float* __restrict simi = res_vals + i * k;
         int64_t* __restrict idxi = res_ids + i * k;
         maxheap_heapify(k, simi, idxi);
-        for (size_t j = 0; j < nsubset; j++) {
-            if (idsi[j] < 0 || static_cast<size_t>(idsi[j]) >= ny) {
-                break;
-            }
-            float disij = fvec_L2sqr_dispatch(x_, y + d * idsi[j], d);
+        with_simd_level([&]<SIMDLevel SL>() {
+            for (size_t j = 0; j < nsubset; j++) {
+                if (idsi[j] < 0 || static_cast<size_t>(idsi[j]) >= ny) {
+                    break;
+                }
+                float disij = fvec_L2sqr<SL>(x_, y + d * idsi[j], d);
 
-            if (disij < simi[0]) {
-                maxheap_replace_top(k, simi, idxi, disij, idsi[j]);
+                if (disij < simi[0]) {
+                    maxheap_replace_top(k, simi, idxi, disij, idsi[j]);
+                }
             }
-        }
+        });
         maxheap_reorder(k, simi, idxi);
     }
 }
@@ -1312,25 +927,27 @@ void pairwise_L2sqr(
     // store in beginning of distance matrix to avoid malloc
     float* b_norms = dis;
 
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (nb > 1)
-    for (int64_t i = 0; i < nb; i++) {
-        b_norms[i] = fvec_norm_L2sqr_dispatch(xb + i * ldb, d);
-    }
+        for (int64_t i = 0; i < nb; i++) {
+            b_norms[i] = fvec_norm_L2sqr<SL>(xb + i * ldb, d);
+        }
 
 #pragma omp parallel for
-    for (int64_t i = 1; i < nq; i++) {
-        float q_norm = fvec_norm_L2sqr_dispatch(xq + i * ldq, d);
-        for (int64_t j = 0; j < nb; j++) {
-            dis[i * ldd + j] = q_norm + b_norms[j];
+        for (int64_t i = 1; i < nq; i++) {
+            float q_norm = fvec_norm_L2sqr<SL>(xq + i * ldq, d);
+            for (int64_t j = 0; j < nb; j++) {
+                dis[i * ldd + j] = q_norm + b_norms[j];
+            }
         }
-    }
 
-    {
-        float q_norm = fvec_norm_L2sqr_dispatch(xq, d);
-        for (int64_t j = 0; j < nb; j++) {
-            dis[j] += q_norm;
+        {
+            float q_norm = fvec_norm_L2sqr<SL>(xq, d);
+            for (int64_t j = 0; j < nb; j++) {
+                dis[j] += q_norm;
+            }
         }
-    }
+    });
 
     {
         FINTEGER nbi = nb, nqi = nq, di = d, ldqi = ldq, ldbi = ldb, lddi = ldd;

faiss/utils/distances_dispatch.h

@@ -13,7 +13,7 @@
  *
  * This is a PRIVATE header. Do not include in public APIs or user code.
  *
- * These wrappers call DISPATCH_SIMDLevel to route to the correct SIMD
+ * These wrappers call with_simd_level to route to the correct SIMD
  * implementation. They are plain inline functions with a _dispatch suffix
  * (e.g. fvec_L2sqr_dispatch). Internal callers that want inlining include
  * this header and call the _dispatch variants directly.
@@ -24,30 +24,36 @@
 
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/extra_distances.h>
 
 namespace faiss {
 
 inline float fvec_L1_dispatch(const float* x, const float* y, size_t d) {
-    DISPATCH_SIMDLevel(fvec_L1, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_L1<SL>(x, y, d); });
 }
 
 inline float fvec_Linf_dispatch(const float* x, const float* y, size_t d) {
-    DISPATCH_SIMDLevel(fvec_Linf, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_Linf<SL>(x, y, d); });
 }
 
 inline float fvec_norm_L2sqr_dispatch(const float* x, size_t d) {
-    DISPATCH_SIMDLevel(fvec_norm_L2sqr, x, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_norm_L2sqr<SL>(x, d); });
 }
 
 inline float fvec_L2sqr_dispatch(const float* x, const float* y, size_t d) {
-    DISPATCH_SIMDLevel(fvec_L2sqr, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_L2sqr<SL>(x, y, d); });
 }
 
 inline float fvec_inner_product_dispatch(
         const float* x,
         const float* y,
         size_t d) {
-    DISPATCH_SIMDLevel(fvec_inner_product, x, y, d);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { return fvec_inner_product<SL>(x, y, d); });
 }
 
 inline void fvec_inner_product_batch_4_dispatch(
@@ -61,18 +67,10 @@ inline void fvec_inner_product_batch_4_dispatch(
         float& dis1,
         float& dis2,
         float& dis3) {
-    DISPATCH_SIMDLevel(
-            fvec_inner_product_batch_4,
-            x,
-            y0,
-            y1,
-            y2,
-            y3,
-            d,
-            dis0,
-            dis1,
-            dis2,
-            dis3);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_inner_product_batch_4<SL>(
+                x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    });
 }
 
 inline void fvec_L2sqr_batch_4_dispatch(
@@ -86,8 +84,9 @@ inline void fvec_L2sqr_batch_4_dispatch(
         float& dis1,
         float& dis2,
         float& dis3) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_batch_4, x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_L2sqr_batch_4<SL>(x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
+    });
 }
 
 inline void fvec_L2sqr_ny_transposed_dispatch(
@@ -98,8 +97,9 @@ inline void fvec_L2sqr_ny_transposed_dispatch(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_ny_transposed, dis, x, y, y_sqlen, d, d_offset, ny);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_L2sqr_ny_transposed<SL>(dis, x, y, y_sqlen, d, d_offset, ny);
+    });
 }
 
 inline void fvec_inner_products_ny_dispatch(
@@ -108,7 +108,9 @@ inline void fvec_inner_products_ny_dispatch(
         const float* y,
         size_t d,
         size_t ny) {
-    DISPATCH_SIMDLevel(fvec_inner_products_ny, ip, x, y, d, ny);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
+        fvec_inner_products_ny<SL>(ip, x, y, d, ny);
+    });
 }
 
 inline void fvec_L2sqr_ny_dispatch(
@@ -117,7 +119,8 @@ inline void fvec_L2sqr_ny_dispatch(
         const float* y,
         size_t d,
         size_t ny) {
-    DISPATCH_SIMDLevel(fvec_L2sqr_ny, dis, x, y, d, ny);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { fvec_L2sqr_ny<SL>(dis, x, y, d, ny); });
 }
 
 inline size_t fvec_L2sqr_ny_nearest_dispatch(
@@ -126,8 +129,11 @@ inline size_t fvec_L2sqr_ny_nearest_dispatch(
         const float* y,
         size_t d,
         size_t ny) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_ny_nearest, distances_tmp_buffer, x, y, d, ny);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() {
+                return fvec_L2sqr_ny_nearest<SL>(
+                        distances_tmp_buffer, x, y, d, ny);
+            });
 }
 
 inline size_t fvec_L2sqr_ny_nearest_y_transposed_dispatch(
@@ -138,15 +144,11 @@ inline size_t fvec_L2sqr_ny_nearest_y_transposed_dispatch(
         size_t d,
         size_t d_offset,
         size_t ny) {
-    DISPATCH_SIMDLevel(
-            fvec_L2sqr_ny_nearest_y_transposed,
-            distances_tmp_buffer,
-            x,
-            y,
-            y_sqlen,
-            d,
-            d_offset,
-            ny);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() {
+                return fvec_L2sqr_ny_nearest_y_transposed<SL>(
+                        distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
+            });
 }
 
 inline void fvec_madd_dispatch(
@@ -155,7 +157,8 @@ inline void fvec_madd_dispatch(
         float bf,
         const float* b,
         float* c) {
-    DISPATCH_SIMDLevel(fvec_madd, n, a, bf, b, c);
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() { fvec_madd<SL>(n, a, bf, b, c); });
 }
 
 inline int fvec_madd_and_argmin_dispatch(
@@ -164,7 +167,10 @@ inline int fvec_madd_and_argmin_dispatch(
         float bf,
         const float* b,
         float* c) {
-    DISPATCH_SIMDLevel(fvec_madd_and_argmin, n, a, bf, b, c);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>(
+            [&]<SIMDLevel SL>() {
+                return fvec_madd_and_argmin<SL>(n, a, bf, b, c);
+            });
 }
 
 inline void fvec_sub_dispatch(
@@ -208,4 +214,37 @@ inline void compute_PQ_dis_tables_dsub2_dispatch(
     });
 }
 
+/***************************************************************************
+ * Dispatching function that takes a lambda directly.
+ * The lambda should be templated on VectorDistance, eg.:
+ *
+ *   auto result = with_VectorDistance(
+ *       metric, metric_arg, [&]<class VD>(VD vd) {
+ *           return vd(x, y);
+ *       });
+ **************************************************************************/
+
+template <typename LambdaType>
+auto with_VectorDistance(
+        size_t d,
+        MetricType metric,
+        float metric_arg,
+        LambdaType&& action) {
+    auto dispatch_metric = [&]<MetricType mt>() {
+        auto call = [&]<SIMDLevel level>() {
+            VectorDistance<mt, level> vd = {d, metric_arg};
+            return action(vd);
+        };
+
+        constexpr bool has_simd = mt == METRIC_INNER_PRODUCT ||
+                mt == METRIC_L2 || mt == METRIC_L1 || mt == METRIC_Linf;
+        if constexpr (!has_simd) {
+            return call.template operator()<SIMDLevel::NONE>();
+        } else {
+            return with_simd_level(call);
+        }
+    };
+    return with_metric_type(metric, dispatch_metric);
+}
+
 } // namespace faiss

faiss/utils/extra_distances.cpp

@@ -7,6 +7,7 @@
 
 // -*- c++ -*-
 
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
 
 #include <omp.h>

faiss/utils/extra_distances.h

@@ -13,9 +13,10 @@
 #include <cstdint>
 
 #include <faiss/MetricType.h>
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
-#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/ordered_key_value.h>
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
@@ -125,39 +126,6 @@ struct VectorDistance {
     float operator()(const float* x, const float* y) const;
 };
 
-/***************************************************************************
- * Dispatching function that takes a lambda directly.
- * The lambda should be templated on VectorDistance, eg.:
- *
- *   auto result = with_VectorDistance(
- *       metric, metric_arg, [&]<class VD>(VD vd) {
- *           return vd(x, y);
- *       });
- **************************************************************************/
-
-template <typename LambdaType>
-auto with_VectorDistance(
-        size_t d,
-        MetricType metric,
-        float metric_arg,
-        LambdaType&& action) {
-    auto dispatch_metric = [&]<MetricType mt>() {
-        auto call = [&]<SIMDLevel level>() {
-            VectorDistance<mt, level> vd = {d, metric_arg};
-            return action(vd);
-        };
-
-        constexpr bool has_simd = mt == METRIC_INNER_PRODUCT ||
-                mt == METRIC_L2 || mt == METRIC_L1 || mt == METRIC_Linf;
-        if constexpr (!has_simd) {
-            return call.template operator()<SIMDLevel::NONE>();
-        } else {
-            DISPATCH_SIMDLevel(call.template operator());
-        }
-    };
-    return with_metric_type(metric, dispatch_metric);
-}
-
 #endif // SWIG
 
 } // namespace faiss

faiss/utils/pq_code_distance.h

@@ -29,7 +29,7 @@ namespace pq_code_distance {
  *
  * PQCodeDistance<PQDecoderT, SL> computes PQ code distances at a given
  * SIMD level. The dispatch site (IndexIVFPQ.cpp, IndexPQ.cpp) uses
- * DISPATCH_SIMDLevel to select SL at runtime, which instantiates
+ * with_simd_level to select SL at runtime, which instantiates
  * PQCodeDistance for ALL decoder types (PQDecoder8, PQDecoder16,
  * PQDecoderGeneric) at the chosen level.
  *
@@ -211,7 +211,7 @@ struct PQCodeDistance {
  * Non-templated PQ code distance dispatch (PQDecoder8 only).
  *
  * These follow the same pattern as distances.h: the caller does not
- * name a SIMDLevel. Internally they dispatch via DISPATCH_SIMDLevel
+ * name a SIMDLevel. Internally they dispatch via with_simd_level
  * to the best available SIMD implementation (DD: runtime detection,
  * static: compile-time selection). Definitions are in
  * pq_code_distance-generic.cpp.

faiss/utils/simd_impl/distances_arm_sve.cpp

@@ -9,6 +9,33 @@
 
 #include <faiss/utils/distances.h>
 
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances_fused/distances_fused.h>
+#include <faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h>
+
+#ifndef FINTEGER
+#define FINTEGER long
+#endif
+
+extern "C" {
+
+int sgemm_(
+        const char* transa,
+        const char* transb,
+        FINTEGER* m,
+        FINTEGER* n,
+        FINTEGER* k,
+        const float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* b,
+        FINTEGER* ldb,
+        float* beta,
+        float* c,
+        FINTEGER* ldc);
+}
+
 #define THE_SIMD_LEVEL SIMDLevel::ARM_SVE
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
@@ -567,4 +594,181 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
     return nearest_idx;
 }
 
+template <>
+void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms) {
+    // BLAS does not like empty matrices
+    if (nx == 0 || ny == 0)
+        return;
+
+    /* block sizes */
+    const size_t bs_x = distance_compute_blas_query_bs;
+    const size_t bs_y = distance_compute_blas_database_bs;
+    // const size_t bs_x = 16, bs_y = 16;
+    std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]);
+    std::unique_ptr<float[]> x_norms(new float[nx]);
+    std::unique_ptr<float[]> del2;
+
+    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
+
+    const size_t lanes = svcntw();
+
+    if (!y_norms) {
+        float* y_norms2 = new float[ny];
+        del2.reset(y_norms2);
+        fvec_norms_L2sqr(y_norms2, y, d, ny);
+        y_norms = y_norms2;
+    }
+
+    for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
+        size_t i1 = i0 + bs_x;
+        if (i1 > nx)
+            i1 = nx;
+
+        res.begin_multiple(i0, i1);
+
+        for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
+            size_t j1 = j0 + bs_y;
+            if (j1 > ny)
+                j1 = ny;
+            /* compute the actual dot products */
+            {
+                float one = 1, zero = 0;
+                FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
+                sgemm_("Transpose",
+                       "Not transpose",
+                       &nyi,
+                       &nxi,
+                       &di,
+                       &one,
+                       y + j0 * d,
+                       &di,
+                       x + i0 * d,
+                       &di,
+                       &zero,
+                       ip_block.get(),
+                       &nyi);
+            }
+            for (int64_t i = i0; i < i1; i++) {
+                const size_t count = j1 - j0;
+                float* ip_line = ip_block.get() + (i - i0) * count;
+
+                svprfw(svwhilelt_b32_u64(0, count), ip_line, SV_PLDL1KEEP);
+                svprfw(svwhilelt_b32_u64(lanes, count),
+                       ip_line + lanes,
+                       SV_PLDL1KEEP);
+
+                // Track lanes min distances + lanes min indices.
+                // All the distances tracked do not take x_norms[i]
+                //   into account in order to get rid of extra
+                //   vaddq_f32(x_norms[i], ...) instructions
+                //   is distance computations.
+                auto min_distances = svdup_n_f32(res.dis_tab[i] - x_norms[i]);
+
+                // these indices are local and are relative to j0.
+                // so, value 0 means j0.
+                auto min_indices = svdup_n_u32(0u);
+
+                auto current_indices = svindex_u32(0u, 1u);
+
+                // process lanes * 2 elements per loop
+                for (size_t idx_j = 0; idx_j < count;
+                     idx_j += lanes * 2, ip_line += lanes * 2) {
+                    svprfw(svwhilelt_b32_u64(idx_j + lanes * 2, count),
+                           ip_line + lanes * 2,
+                           SV_PLDL1KEEP);
+                    svprfw(svwhilelt_b32_u64(idx_j + lanes * 3, count),
+                           ip_line + lanes * 3,
+                           SV_PLDL1KEEP);
+
+                    // mask
+                    const auto mask_0 = svwhilelt_b32_u64(idx_j, count);
+                    const auto mask_1 = svwhilelt_b32_u64(idx_j + lanes, count);
+
+                    // load values for norms
+                    const auto y_norm_0 =
+                            svld1_f32(mask_0, y_norms + idx_j + j0 + 0);
+                    const auto y_norm_1 =
+                            svld1_f32(mask_1, y_norms + idx_j + j0 + lanes);
+
+                    // load values for dot products
+                    const auto ip_0 = svld1_f32(mask_0, ip_line + 0);
+                    const auto ip_1 = svld1_f32(mask_1, ip_line + lanes);
+
+                    // compute dis = y_norm[j] - 2 * dot(x_norm[i], y_norm[j]).
+                    // x_norm[i] was dropped off because it is a constant for a
+                    // given i. We'll deal with it later.
+                    const auto distances_0 =
+                            svmla_n_f32_z(mask_0, y_norm_0, ip_0, -2.f);
+                    const auto distances_1 =
+                            svmla_n_f32_z(mask_1, y_norm_1, ip_1, -2.f);
+
+                    // compare the new distances to the min distances
+                    // for each of the first group of 4 ARM SIMD components.
+                    auto comparison =
+                            svcmpgt_f32(mask_0, min_distances, distances_0);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances =
+                            svsel_f32(comparison, distances_0, min_distances);
+                    min_indices =
+                            svsel_u32(comparison, current_indices, min_indices);
+                    current_indices = svadd_n_u32_x(
+                            mask_0,
+                            current_indices,
+                            static_cast<uint32_t>(lanes));
+
+                    // compare the new distances to the min distances
+                    // for each of the second group of 4 ARM SIMD components.
+                    comparison =
+                            svcmpgt_f32(mask_1, min_distances, distances_1);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances =
+                            svsel_f32(comparison, distances_1, min_distances);
+                    min_indices =
+                            svsel_u32(comparison, current_indices, min_indices);
+                    current_indices = svadd_n_u32_x(
+                            mask_1,
+                            current_indices,
+                            static_cast<uint32_t>(lanes));
+                }
+
+                // add missing x_norms[i]
+                // negative values can occur for identical vectors
+                //    due to roundoff errors.
+                auto mask = svwhilelt_b32_u64(0, count);
+                min_distances = svadd_n_f32_z(
+                        svcmpge_n_f32(mask, min_distances, -x_norms[i]),
+                        min_distances,
+                        x_norms[i]);
+                min_indices = svadd_n_u32_x(
+                        mask, min_indices, static_cast<uint32_t>(j0));
+                mask = svcmple_n_f32(mask, min_distances, res.dis_tab[i]);
+                if (svcntp_b32(svptrue_b32(), mask) == 0)
+                    res.add_result(i, res.dis_tab[i], res.ids_tab[i]);
+                else {
+                    const auto min_distance = svminv_f32(mask, min_distances);
+                    const auto min_index = svminv_u32(
+                            svcmpeq_n_f32(mask, min_distances, min_distance),
+                            min_indices);
+                    res.add_result(i, min_distance, min_index);
+                }
+            }
+        }
+        // Does nothing for SingleBestResultHandler, but
+        // keeping the call for the consistency.
+        res.end_multiple();
+        InterruptCallback::check();
+    }
+}
+
 } // namespace faiss

faiss/utils/simd_impl/distances_avx2.cpp

@@ -9,6 +9,33 @@
 
 #include <immintrin.h>
 
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances_fused/distances_fused.h>
+#include <faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h>
+
+#ifndef FINTEGER
+#define FINTEGER long
+#endif
+
+extern "C" {
+
+int sgemm_(
+        const char* transa,
+        const char* transb,
+        FINTEGER* m,
+        FINTEGER* n,
+        FINTEGER* k,
+        const float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* b,
+        FINTEGER* ldb,
+        float* beta,
+        float* c,
+        FINTEGER* ldc);
+}
+
 #define THE_SIMD_LEVEL SIMDLevel::AVX2
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
@@ -1187,4 +1214,215 @@ int fvec_madd_and_argmin<SIMDLevel::AVX2>(
     return fvec_madd_and_argmin_sse(n, a, bf, b, c);
 }
 
+template <>
+void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms) {
+    // BLAS does not like empty matrices
+    if (nx == 0 || ny == 0) {
+        return;
+    }
+
+    /* block sizes */
+    const size_t bs_x = distance_compute_blas_query_bs;
+    const size_t bs_y = distance_compute_blas_database_bs;
+    // const size_t bs_x = 16, bs_y = 16;
+    std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]);
+    std::unique_ptr<float[]> x_norms(new float[nx]);
+    std::unique_ptr<float[]> del2;
+
+    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
+
+    if (!y_norms) {
+        float* y_norms2 = new float[ny];
+        del2.reset(y_norms2);
+        fvec_norms_L2sqr(y_norms2, y, d, ny);
+        y_norms = y_norms2;
+    }
+
+    for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
+        size_t i1 = i0 + bs_x;
+        if (i1 > nx) {
+            i1 = nx;
+        }
+
+        res.begin_multiple(i0, i1);
+
+        for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
+            size_t j1 = j0 + bs_y;
+            if (j1 > ny) {
+                j1 = ny;
+            }
+            /* compute the actual dot products */
+            {
+                float one = 1, zero = 0;
+                FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
+                sgemm_("Transpose",
+                       "Not transpose",
+                       &nyi,
+                       &nxi,
+                       &di,
+                       &one,
+                       y + j0 * d,
+                       &di,
+                       x + i0 * d,
+                       &di,
+                       &zero,
+                       ip_block.get(),
+                       &nyi);
+            }
+            for (int64_t i = i0; i < i1; i++) {
+                float* ip_line = ip_block.get() + (i - i0) * (j1 - j0);
+
+                _mm_prefetch((const char*)ip_line, _MM_HINT_NTA);
+                _mm_prefetch((const char*)(ip_line + 16), _MM_HINT_NTA);
+
+                // constant
+                const __m256 mul_minus2 = _mm256_set1_ps(-2);
+
+                // Track 8 min distances + 8 min indices.
+                // All the distances tracked do not take x_norms[i]
+                //   into account in order to get rid of extra
+                //   _mm256_add_ps(x_norms[i], ...) instructions
+                //   is distance computations.
+                __m256 min_distances =
+                        _mm256_set1_ps(res.dis_tab[i] - x_norms[i]);
+
+                // these indices are local and are relative to j0.
+                // so, value 0 means j0.
+                __m256i min_indices = _mm256_set1_epi32(0);
+
+                __m256i current_indices =
+                        _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+                const __m256i indices_delta = _mm256_set1_epi32(8);
+
+                // current j index
+                size_t idx_j = 0;
+                size_t count = j1 - j0;
+
+                // process 16 elements per loop
+                for (; idx_j < (count / 16) * 16; idx_j += 16, ip_line += 16) {
+                    _mm_prefetch((const char*)(ip_line + 32), _MM_HINT_NTA);
+                    _mm_prefetch((const char*)(ip_line + 48), _MM_HINT_NTA);
+
+                    // load values for norms
+                    const __m256 y_norm_0 =
+                            _mm256_loadu_ps(y_norms + idx_j + j0 + 0);
+                    const __m256 y_norm_1 =
+                            _mm256_loadu_ps(y_norms + idx_j + j0 + 8);
+
+                    // load values for dot products
+                    const __m256 ip_0 = _mm256_loadu_ps(ip_line + 0);
+                    const __m256 ip_1 = _mm256_loadu_ps(ip_line + 8);
+
+                    // compute dis = y_norm[j] - 2 * dot(x_norm[i], y_norm[j]).
+                    // x_norm[i] was dropped off because it is a constant for a
+                    // given i. We'll deal with it later.
+                    __m256 distances_0 =
+                            _mm256_fmadd_ps(ip_0, mul_minus2, y_norm_0);
+                    __m256 distances_1 =
+                            _mm256_fmadd_ps(ip_1, mul_minus2, y_norm_1);
+
+                    // compare the new distances to the min distances
+                    // for each of the first group of 8 AVX2 components.
+                    const __m256 comparison_0 = _mm256_cmp_ps(
+                            min_distances, distances_0, _CMP_LE_OS);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances = _mm256_blendv_ps(
+                            distances_0, min_distances, comparison_0);
+                    min_indices = _mm256_castps_si256(_mm256_blendv_ps(
+                            _mm256_castsi256_ps(current_indices),
+                            _mm256_castsi256_ps(min_indices),
+                            comparison_0));
+                    current_indices =
+                            _mm256_add_epi32(current_indices, indices_delta);
+
+                    // compare the new distances to the min distances
+                    // for each of the second group of 8 AVX2 components.
+                    const __m256 comparison_1 = _mm256_cmp_ps(
+                            min_distances, distances_1, _CMP_LE_OS);
+
+                    // update min distances and indices with closest vectors if
+                    // needed.
+                    min_distances = _mm256_blendv_ps(
+                            distances_1, min_distances, comparison_1);
+                    min_indices = _mm256_castps_si256(_mm256_blendv_ps(
+                            _mm256_castsi256_ps(current_indices),
+                            _mm256_castsi256_ps(min_indices),
+                            comparison_1));
+                    current_indices =
+                            _mm256_add_epi32(current_indices, indices_delta);
+                }
+
+                // dump values and find the minimum distance / minimum index
+                float min_distances_scalar[8];
+                uint32_t min_indices_scalar[8];
+                _mm256_storeu_ps(min_distances_scalar, min_distances);
+                _mm256_storeu_si256(
+                        (__m256i*)(min_indices_scalar), min_indices);
+
+                float current_min_distance = res.dis_tab[i];
+                uint32_t current_min_index = res.ids_tab[i];
+
+                // This unusual comparison is needed to maintain the behavior
+                // of the original implementation: if two indices are
+                // represented with equal distance values, then
+                // the index with the min value is returned.
+                for (size_t jv = 0; jv < 8; jv++) {
+                    // add missing x_norms[i]
+                    float distance_candidate =
+                            min_distances_scalar[jv] + x_norms[i];
+
+                    // negative values can occur for identical vectors
+                    //    due to roundoff errors.
+                    if (distance_candidate < 0) {
+                        distance_candidate = 0;
+                    }
+
+                    int64_t index_candidate = min_indices_scalar[jv] + j0;
+
+                    if (current_min_distance > distance_candidate) {
+                        current_min_distance = distance_candidate;
+                        current_min_index = index_candidate;
+                    } else if (
+                            current_min_distance == distance_candidate &&
+                            current_min_index > index_candidate) {
+                        current_min_index = index_candidate;
+                    }
+                }
+
+                // process leftovers
+                for (; idx_j < count; idx_j++, ip_line++) {
+                    float ip = *ip_line;
+                    float dis = x_norms[i] + y_norms[idx_j + j0] - 2 * ip;
+                    // negative values can occur for identical vectors
+                    //    due to roundoff errors.
+                    if (dis < 0) {
+                        dis = 0;
+                    }
+
+                    if (current_min_distance > dis) {
+                        current_min_distance = dis;
+                        current_min_index = idx_j + j0;
+                    }
+                }
+
+                //
+                res.add_result(i, current_min_distance, current_min_index);
+            }
+        }
+        // Does nothing for SingleBestResultHandler, but
+        // keeping the call for the consistency.
+        res.end_multiple();
+        InterruptCallback::check();
+    }
+}
+
 } // namespace faiss

faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h

@@ -0,0 +1,27 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+/// BLAS-accelerated exhaustive L2 search for the k=1 (top-1) case.
+/// Specializations live in the per-SIMD translation units under simd_impl/.
+template <SIMDLevel>
+void exhaustive_L2sqr_blas_cmax(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms);
+
+} // namespace faiss

faiss/utils/simd_levels.cpp

@@ -10,6 +10,7 @@
 #include <cstdlib>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 
 namespace faiss {
 
@@ -192,9 +193,6 @@ SIMDLevel SIMDConfig::auto_detect_simd_level() {
     return detected_level;
 }
 
-// Include private header for DISPATCH_SIMDLevel macro
-#include <faiss/impl/simd_dispatch.h>
-
 namespace {
 
 template <SIMDLevel Level>
@@ -205,7 +203,8 @@ SIMDLevel get_dispatched_level_impl() {
 } // namespace
 
 SIMDLevel SIMDConfig::get_dispatched_level() {
-    DISPATCH_SIMDLevel(get_dispatched_level_impl);
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_ALL>(
+            [&]<SIMDLevel SL>() { return get_dispatched_level_impl<SL>(); });
 }
 
 #else // Static mode

faiss/utils/simd_levels.h

@@ -87,8 +87,12 @@ inline constexpr SIMDLevel SINGLE_SIMD_LEVEL_256 =
         simd256_level_selector<SINGLE_SIMD_LEVEL>::value;
 
 /// Number of float32 lanes for a given SIMD level.
+/// ARM_SVE is variable-width (128–2048 bits); no single constant is correct.
 template <SIMDLevel SL>
 constexpr int simd_width() {
+    static_assert(
+            SL != SIMDLevel::ARM_SVE,
+            "simd_width<ARM_SVE> is not supported: SVE is variable-width");
     if constexpr (SL == SIMDLevel::AVX512 || SL == SIMDLevel::AVX512_SPR)
         return 16;
     else if constexpr (SL == SIMDLevel::AVX2 || SL == SIMDLevel::ARM_NEON)
@@ -139,7 +143,7 @@ struct FAISS_API SIMDConfig {
     static bool is_simd_level_available(SIMDLevel level);
 
     /// Returns the SIMD level via the dispatch mechanism.
-    /// In DD mode, uses DISPATCH_SIMDLevel internally.
+    /// In DD mode, uses with_simd_level internally.
     /// In static mode, returns the compiled-in level.
     /// Useful for verification: get_level() == get_dispatched_level()
     static SIMDLevel get_dispatched_level();

tests/test_fast_scan_ivf.py

@@ -278,7 +278,7 @@ class TestEquivPQ(unittest.TestCase):
         Dnew, Inew = index_pq.search(xq, 4)
 
         np.testing.assert_array_equal(Iref, Inew)
-        np.testing.assert_array_equal(Dref, Dnew)
+        np.testing.assert_allclose(Dref, Dnew, rtol=1e-5)
 
         index_pq2 = faiss.IndexPQFastScan(index_pq)
         index_pq2.implem = 12
@@ -288,7 +288,7 @@ class TestEquivPQ(unittest.TestCase):
         index2.implem = 12
         Dnew, Inew = index2.search(xq, 4)
         np.testing.assert_array_equal(Iref, Inew)
-        np.testing.assert_array_equal(Dref, Dnew)
+        np.testing.assert_allclose(Dref, Dnew, rtol=1e-5)
 
         # test encode and decode