Stop aarch64 falling back to SIMDLevel::NONE in faiss dispatch (#5572)

faiss/IndexFlat.cpp

@@ -270,9 +270,11 @@ struct FlatIPDis : FlatCodesDistanceComputer {
 FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
     FlatCodesDistanceComputer* dc = nullptr;
     if (metric_type == METRIC_L2) {
-        with_simd_level([&]<SIMDLevel SL>() { dc = new FlatL2Dis<SL>(*this); });
+        with_simd_level_a1(
+                [&]<SIMDLevel SL>() { dc = new FlatL2Dis<SL>(*this); });
     } else if (metric_type == METRIC_INNER_PRODUCT) {
-        with_simd_level([&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
+        with_simd_level_a1(
+                [&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
     } else {
         dc = get_extra_distance_computer(d, metric_type, metric_arg, get_xb());
     }
@@ -404,7 +406,7 @@ FlatCodesDistanceComputer* IndexFlatL2::get_FlatCodesDistanceComputer() const {
     if (metric_type == METRIC_L2) {
         if (!cached_l2norms.empty()) {
             FlatCodesDistanceComputer* dc = nullptr;
-            with_simd_level([&]<SIMDLevel SL>() {
+            with_simd_level_a1([&]<SIMDLevel SL>() {
                 dc = new FlatL2WithNormsDis<SL>(*this);
             });
             return dc;
@@ -784,7 +786,7 @@ void IndexFlatPanorama::search_subset(
         idx_t k,
         float* distances,
         idx_t* labels) const {
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_a1([&]<SIMDLevel SL>() {
         with_metric_type(metric_type, [&]<MetricType M>() {
             constexpr bool is_sim = is_similarity_metric(M);
             using C = std::conditional_t<

faiss/SuperKMeans.cpp

@@ -621,7 +621,7 @@ void super_kmeans_assign_iteration(
             }
 
             // One SIMD dispatch per (xi, yj) tile.
-            with_simd_level([&]<SIMDLevel SL>() {
+            with_simd_level_a1([&]<SIMDLevel SL>() {
                 [[maybe_unused]] const int omp_chunk_local = cp.omp_chunk;
                 int64_t tile_total = 0;
                 int64_t tile_pruned = 0;

faiss/impl/AdditiveQuantizer.cpp

@@ -376,7 +376,7 @@ void AdditiveQuantizer::compute_centroid_norms(float* norms) const {
     size_t ntotal = (size_t)1 << tot_bits;
     int64_t ntotal_signed = ntotal;
     // TODO: make tree of partial sums
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_a1([&]<SIMDLevel SL>() {
 #pragma omp parallel
         {
             std::vector<float> tmp(d);

faiss/impl/ProductQuantizer.cpp

@@ -280,7 +280,6 @@ void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
 } // namespace
 
 void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
-    // a1: fvec_L2sqr_ny_nearest / _y_transposed have ARM_SVE specializations
     with_simd_level_a1([&]<SIMDLevel SL>() {
         switch (nbits) {
             case 8:
@@ -443,7 +442,6 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
 
 void ProductQuantizer::compute_distance_table(const float* x, float* dis_table)
         const {
-    // a1: fvec_L2sqr_ny / _transposed have ARM_SVE specializations
     with_simd_level_a1([&]<SIMDLevel SL>() {
         if (transposed_centroids.empty()) {
             // use regular version
@@ -826,7 +824,6 @@ void ProductQuantizer::compute_sdc_table() {
     sdc_table.resize(M * ksub * ksub);
 
     if (dsub < 4) {
-        // a1: fvec_L2sqr_ny has an ARM_SVE specialization
         with_simd_level_a1([&]<SIMDLevel SL>() {
 #pragma omp parallel for
             for (int64_t mk = 0; mk < static_cast<int64_t>(M * ksub); mk++) {

faiss/impl/pq_code_distance/pq_code_distance-generic.cpp

@@ -7,7 +7,7 @@
 
 // This TU provides non-templated PQ code distance dispatch wrappers
 // (pq_code_distance_8bit_single, pq_code_distance_8bit_four) declared
-// in pq_code_distance-inl.h. These use with_simd_level to route to the
+// in pq_code_distance-inl.h. These use with_simd_level_a1 to route to the
 // best available SIMD implementation via pq_code_distance_8bit_*_impl
 // function template specializations.
 //
@@ -34,7 +34,7 @@ void pq_scan_8bit(
         float* heap_dis,
         int64_t* heap_ids,
         bool max_heap) {
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_a1([&]<SIMDLevel SL>() {
         pq_scan_8bit_impl<SL>(
                 M, dis_table, codes, ncodes, k, heap_dis, heap_ids, max_heap);
     });
@@ -44,7 +44,7 @@ float pq_code_distance_8bit_single(
         size_t M,
         const float* sim_table,
         const uint8_t* code) {
-    return with_simd_level([&]<SIMDLevel SL>() {
+    return with_simd_level_a1([&]<SIMDLevel SL>() {
         return pq_code_distance_8bit_single_impl<SL>(M, sim_table, code);
     });
 }
@@ -60,7 +60,7 @@ void pq_code_distance_8bit_four(
         float& result1,
         float& result2,
         float& result3) {
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_a1([&]<SIMDLevel SL>() {
         pq_code_distance_8bit_four_impl<SL>(
                 M,
                 sim_table,

faiss/impl/scalar_quantizer/distance_computers.h

@@ -78,8 +78,8 @@ template <class Similarity, SIMDLevel SL>
 struct DistanceComputerByte : SQDistanceComputer {};
 
 // Byte-domain distance computer for QT_8bit_direct_signed (storage is
-// value+128). Only specialized for AVX512_SPR; other levels fall back to
-// the float-domain DCTemplate path via the dispatch logic.
+// value+128). A level added to the sq-dispatch.h chain without a
+// specialization here instantiates this empty template.
 template <class Similarity, SIMDLevel SL>
 struct DistanceComputerByteSigned : SQDistanceComputer {};
 

faiss/impl/scalar_quantizer/sq-dispatch.h

@@ -531,7 +531,8 @@ SQDistanceComputer* select_distance_computer_body(
                     return new DistanceComputerByte<Sim, SL2>(
                             static_cast<int>(d), trained);
                 }
-            } else if constexpr (SL2 == SIMDLevel::AVX2) {
+            } else if constexpr (
+                    SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
                 if (d % 16 == 0) {
                     return new DistanceComputerByte<Sim, SL2>(
                             static_cast<int>(d), trained);
@@ -551,7 +552,8 @@ SQDistanceComputer* select_distance_computer_body(
                     return new DistanceComputerByteSigned<Sim, SL2>(
                             static_cast<int>(d), trained);
                 }
-            } else if constexpr (SL2 == SIMDLevel::AVX2) {
+            } else if constexpr (
+                    SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
                 if (d % 16 == 0) {
                     return new DistanceComputerByteSigned<Sim, SL2>(
                             static_cast<int>(d), trained);
@@ -744,7 +746,8 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                         return scan.template
                         operator()<DistanceComputerByte<Similarity, SL2>>();
                     }
-                } else if constexpr (SL2 == SIMDLevel::AVX2) {
+                } else if constexpr (
+                        SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
                     if (d % 16 == 0) {
                         return scan.template
                         operator()<DistanceComputerByte<Similarity, SL2>>();
@@ -765,7 +768,8 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                         return scan.template operator()<
                                 DistanceComputerByteSigned<Similarity, SL2>>();
                     }
-                } else if constexpr (SL2 == SIMDLevel::AVX2) {
+                } else if constexpr (
+                        SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
                     if (d % 16 == 0) {
                         return scan.template operator()<
                                 DistanceComputerByteSigned<Similarity, SL2>>();

faiss/impl/scalar_quantizer/sq-neon.cpp

@@ -618,6 +618,64 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
     }
 };
 
+// Byte-domain kernels for QT_8bit_direct{,_signed}. The dispatch only
+// selects them when d % 16 == 0, so no loop needs a tail.
+
+namespace {
+
+// The accumulator stays unsigned: a vmull_u8 square reaches 255*255 = 65025,
+// which an int16 lane would read as negative.
+FAISS_ALWAYS_INLINE int neon_byte_l2sqr(
+        const uint8_t* code1,
+        const uint8_t* code2,
+        int d) {
+    uint32x4_t accu = vdupq_n_u32(0);
+    for (int i = 0; i < d; i += 16) {
+        const uint8x16_t diff =
+                vabdq_u8(vld1q_u8(code1 + i), vld1q_u8(code2 + i));
+        accu = vpadalq_u16(
+                accu, vmull_u8(vget_low_u8(diff), vget_low_u8(diff)));
+        accu = vpadalq_u16(
+                accu, vmull_u8(vget_high_u8(diff), vget_high_u8(diff)));
+    }
+    return static_cast<int>(vaddvq_u32(accu));
+}
+
+FAISS_ALWAYS_INLINE int neon_byte_ip(
+        const uint8_t* code1,
+        const uint8_t* code2,
+        int d) {
+    uint32x4_t accu = vdupq_n_u32(0);
+    for (int i = 0; i < d; i += 16) {
+        const uint8x16_t c1 = vld1q_u8(code1 + i);
+        const uint8x16_t c2 = vld1q_u8(code2 + i);
+        accu = vpadalq_u16(accu, vmull_u8(vget_low_u8(c1), vget_low_u8(c2)));
+        accu = vpadalq_u16(accu, vmull_u8(vget_high_u8(c1), vget_high_u8(c2)));
+    }
+    return static_cast<int>(vaddvq_u32(accu));
+}
+
+// The codes store value + 128. For x in 0 to 255, x ^ 0x80 read as int8 is
+// exactly x - 128, which is how the bias comes off before vmull_s8.
+FAISS_ALWAYS_INLINE int neon_byte_ip_unbias(
+        const uint8_t* code1,
+        const uint8_t* code2,
+        int d) {
+    const uint8x16_t bias = vdupq_n_u8(0x80);
+    int32x4_t accu = vdupq_n_s32(0);
+    for (int i = 0; i < d; i += 16) {
+        const int8x16_t c1 =
+                vreinterpretq_s8_u8(veorq_u8(vld1q_u8(code1 + i), bias));
+        const int8x16_t c2 =
+                vreinterpretq_s8_u8(veorq_u8(vld1q_u8(code2 + i), bias));
+        accu = vpadalq_s16(accu, vmull_s8(vget_low_s8(c1), vget_low_s8(c2)));
+        accu = vpadalq_s16(accu, vmull_s8(vget_high_s8(c1), vget_high_s8(c2)));
+    }
+    return static_cast<int>(vaddvq_s32(accu));
+}
+
+} // namespace
+
 template <class Similarity>
 struct DistanceComputerByte<Similarity, SIMDLevel::ARM_NEON>
         : SQDistanceComputer {
@@ -630,16 +688,11 @@ struct DistanceComputerByte<Similarity, SIMDLevel::ARM_NEON>
 
     int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
             const {
-        int accu = 0;
-        for (int i = 0; i < d; i++) {
-            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
-                accu += int(code1[i]) * code2[i];
-            } else {
-                int diff = int(code1[i]) - code2[i];
-                accu += diff * diff;
-            }
+        if constexpr (Sim::metric_type == METRIC_INNER_PRODUCT) {
+            return neon_byte_ip(code1, code2, d);
+        } else {
+            return neon_byte_l2sqr(code1, code2, d);
         }
-        return accu;
     }
 
     void set_query(const float* x) final {
@@ -663,6 +716,48 @@ struct DistanceComputerByte<Similarity, SIMDLevel::ARM_NEON>
     }
 };
 
+template <class Similarity>
+struct DistanceComputerByteSigned<Similarity, SIMDLevel::ARM_NEON>
+        : SQDistanceComputer {
+    using Sim = Similarity;
+
+    int d;
+    std::vector<uint8_t> tmp;
+
+    DistanceComputerByteSigned(int d, const std::vector<float>&)
+            : d(d), tmp(d) {}
+
+    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        if constexpr (Sim::metric_type == METRIC_INNER_PRODUCT) {
+            return neon_byte_ip_unbias(code1, code2, d);
+        } else {
+            // The bias cancels in the difference.
+            return neon_byte_l2sqr(code1, code2, d);
+        }
+    }
+
+    void set_query(const float* x) final {
+        for (int i = 0; i < d; i++) {
+            tmp[i] = uint8_t(int(x[i]) + 128);
+        }
+    }
+
+    int compute_distance(const float* x, const uint8_t* code) {
+        set_query(x);
+        return compute_code_distance(tmp.data(), code);
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return compute_code_distance(
+                codes + i * code_size, codes + j * code_size);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return compute_code_distance(tmp.data(), code);
+    }
+};
+
 /**********************************************************
  * TurboQuant masked_sum NEON specialization (scalar fallback)
  **********************************************************/

faiss/impl/simd_dispatch.h

@@ -46,11 +46,6 @@ constexpr int AVAILABLE_SIMD_LEVELS_A0_SPR =
 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;
 
 constexpr SIMDLevel get_simd_fallback(SIMDLevel level) {

faiss/utils/distances_dispatch.h

@@ -241,7 +241,7 @@ auto with_VectorDistance(
         if constexpr (!has_simd) {
             return call.template operator()<SIMDLevel::NONE>();
         } else {
-            return with_simd_level(call);
+            return with_simd_level_a1(call);
         }
     };
     return with_metric_type(metric, dispatch_metric);

faiss/utils/simd_impl/super_kmeans_dispatch.h

@@ -11,10 +11,6 @@
 // highest available SIMD specialization at runtime (DD mode) or the
 // compiled-in level (static mode).
 //
-// The A1 level mask is required: plain with_simd_level() omits the ARM_SVE bit,
-// so on an SVE host the SVE specialization would never be instantiated.
-// ARM_NEON has no specialization and falls through to the scalar primary
-// template.
 
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/simd_impl/super_kmeans_kernels.h>

faiss/utils/simd_levels.cpp

@@ -8,6 +8,7 @@
 #include <faiss/utils/simd_levels.h>
 
 #include <cstdint>
+#include <cstdio>
 #include <cstdlib>
 
 #if defined(_MSC_VER)
@@ -156,6 +157,35 @@ void detect_x86_uarch_flags() {}
 // NOLINTNEXTLINE(facebook-avoid-non-const-global-variables)
 static SIMDConfig simd_config_initializer;
 
+namespace {
+
+/// Levels this binary holds code for. Must mirror the case labels in
+/// with_selected_simd_levels.
+uint64_t compiled_simd_levels() {
+    uint64_t mask = uint64_t(1) << static_cast<int>(SIMDLevel::NONE);
+#ifdef COMPILE_SIMD_AVX2
+    mask |= uint64_t(1) << static_cast<int>(SIMDLevel::AVX2);
+#endif
+#ifdef COMPILE_SIMD_AVX512
+    mask |= uint64_t(1) << static_cast<int>(SIMDLevel::AVX512);
+#endif
+#ifdef COMPILE_SIMD_AVX512_SPR
+    mask |= uint64_t(1) << static_cast<int>(SIMDLevel::AVX512_SPR);
+#endif
+#ifdef COMPILE_SIMD_ARM_NEON
+    mask |= uint64_t(1) << static_cast<int>(SIMDLevel::ARM_NEON);
+#endif
+#ifdef COMPILE_SIMD_ARM_SVE
+    mask |= uint64_t(1) << static_cast<int>(SIMDLevel::ARM_SVE);
+#endif
+#ifdef COMPILE_SIMD_RISCV_RVV
+    mask |= uint64_t(1) << static_cast<int>(SIMDLevel::RISCV_RVV);
+#endif
+    return mask;
+}
+
+} // namespace
+
 SIMDConfig::SIMDConfig(const char** faiss_simd_level_env) {
     // Support dependency injection for testing
     const char* env_var = faiss_simd_level_env ? *faiss_simd_level_env
@@ -164,7 +194,22 @@ SIMDConfig::SIMDConfig(const char** faiss_simd_level_env) {
     if (!env_var) {
         level = auto_detect_simd_level();
     } else {
-        level = to_simd_level(env_var);
+        // Forcing a level the CPU lacks is allowed. Forcing one the binary
+        // does not hold is not: dispatch would fall to NONE and skip every
+        // level between. Walk down to the nearest compiled level instead.
+        const uint64_t compiled = compiled_simd_levels();
+        const SIMDLevel requested = to_simd_level(env_var);
+        level = requested;
+        while (((compiled >> static_cast<int>(level)) & 1) == 0) {
+            level = get_simd_fallback(level);
+        }
+        if (level != requested) {
+            fprintf(stderr,
+                    "faiss: FAISS_SIMD_LEVEL=%s is not compiled into this "
+                    "build, using %s instead\n",
+                    to_string(requested).c_str(),
+                    to_string(level).c_str());
+        }
         supported_simd_levels = (1 << static_cast<int>(level));
     }
     supported_simd_levels |= (1 << static_cast<int>(SIMDLevel::NONE));