use NEON distance function if running on aarch64

ssciwr · Nov 10, 2023 · bc6c2a4 · bc6c2a4
1 parent 78d661c
commit bc6c2a4
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/include/hamming/hamming_impl.hh b/include/hamming/hamming_impl.hh
@@ -2,9 +2,7 @@
 #define _HAMMING_IMPL_HH
 
 #include <array>
-#if defined(__aarch64__) || defined(_M_ARM64)
-#include <cpuinfo_aarch64.h>
-#else
+#if !(defined(__aarch64__) || defined(_M_ARM64))
 #include <cpuinfo_x86.h>
 #endif
 #include <cstdint>
@@ -14,6 +12,9 @@
 #ifdef HAMMING_WITH_OPENMP
 #include <omp.h>
 #endif
+#ifdef HAMMING_WITH_NEON
+#include "hamming/distance_neon.hh"
+#endif
 #ifdef HAMMING_WITH_SSE2
 #include "hamming/distance_sse2.hh"
 #endif
@@ -101,11 +102,11 @@ std::vector<DistIntType> distances(std::vector<std::string> &data,
                        const std::vector<GeneBlock> &b) = distance_cpp;
 
 #if defined(__aarch64__) || defined(_M_ARM64)
-  const auto features = cpu_features::GetAarch64Info().features;
+#ifdef HAMMING_WITH_NEON
+  distance_func = distance_neon;
+#endif
 #else
   const auto features = cpu_features::GetX86Info().features;
-#endif
-
 #ifdef HAMMING_WITH_SSE2
   if (features.sse2) {
     distance_func = distance_sse2;
@@ -121,6 +122,7 @@ std::vector<DistIntType> distances(std::vector<std::string> &data,
     distance_func = distance_avx512;
   }
 #endif
+#endif
 
 #ifdef HAMMING_WITH_OPENMP
 #pragma omp parallel for schedule(static, 1)