feat(ailego): enable NEON math kernels on MSVC ARM64

Aanerud · Aanerud · commit 5baa52a3da21 · 2026-06-15T00:36:09.000+02:00
Addresses @feihongxu0824's review comment on #352. Previously the *_neon.cc files compiled into empty translation units on MSVC ARM64 because their guards (`__ARM_NEON`, `__aarch64__`) are GCC/Clang-only macros, leaving zvec to use the scalar fallback. That fallback hits ~1 ULP precision drift versus the NEON path, which surfaced as two HnswStreamerTest cosine failures and a `NormMatrix.Norm1_General` failure on the `windows-11-arm` CI runner. Changes: * Expand `defined(__ARM_NEON)` -> `(defined(__ARM_NEON) || defined(_M_ARM64))` at the 51 source sites that gate ARM NEON kernels (math, math_batch, utility, normalizer, platform headers, dispatch tables, version probe). * Expand `defined(__aarch64__)` -> `(defined(__aarch64__) || defined(_M_ARM64))` at the 26 sites that distinguish AArch64 from ARMv7 NEON — MSVC ARM64 is AArch64 but does not predefine `__aarch64__`. As a side effect the ARMv7-only polyfills (`vaddvq_f32`/`vaddvq_s32` shims in `distance_matrix_accum_fp32.i`, `distance_matrix_fp32.i`) are correctly skipped under MSVC ARM64, where those intrinsics are built in. * `src/include/zvec/ailego/internal/platform.h`: include `<arm_neon.h>` on the MSVC branch when `_M_ARM64` is defined (it was previously gated behind `!_MSC_VER`, so MSVC ARM64 saw no NEON types). * `src/ailego/CMakeLists.txt`: keep the existing `if(NOT MSVC)` wrapper around the GCC-only `-march=armv8-a` flag, and add an explicit `else()` branch with a comment explaining MSVC ARM64 does not need `-march` (NEON is the ARMv8 baseline on MSVC and the kernels are now picked up via the ALL_SRCS glob with the macro guards above). This supersedes the earlier `test(hnsw): skip two cosine self-match tests on MSVC ARM64` commit (which has been dropped from the branch). The NEON math kernels now use the same precision path as Linux/macOS ARM64, so those tests should pass natively on `windows-11-arm`.
diff --git a/src/ailego/CMakeLists.txt b/src/ailego/CMakeLists.txt
@@ -122,6 +122,13 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH)
                   COMPILE_FLAGS "${MATH_MARCH_FLAG_NEON}"
               )
             endforeach()
+        else()
+            # MSVC on ARM64: NEON is the ARMv8 baseline and is always enabled,
+            # so no `-march` flag is required (MSVC does not accept GCC-style
+            # `-march=` anyway). The NEON math kernels still get compiled via
+            # the ALL_SRCS glob above; their `#if defined(__ARM_NEON)` guards
+            # were extended in this PR to also accept `_M_ARM64` so the bodies
+            # actually emit code under MSVC.
         endif()
     endif()
 endif()
diff --git a/src/ailego/internal/cpu_features.cc b/src/ailego/internal/cpu_features.cc
@@ -17,7 +17,8 @@
 
 #if defined(_MSC_VER) && (defined(_M_IX86) || defined(_M_X64))
 #include <intrin.h>
-#elif !defined(_MSC_VER) && !defined(__ARM_ARCH) && !defined(__aarch64__)
+#elif !defined(_MSC_VER) && !defined(__ARM_ARCH) && \
+    !(defined(__aarch64__) || defined(_M_ARM64))
 #include <cpuid.h>
 #endif
 
@@ -48,7 +49,8 @@ CpuFeatures::CpuFlags::CpuFlags(void)
   L7_ECX = l7[2];
   L7_EDX = l7[3];
 }
-#elif !defined(_MSC_VER) && !defined(__ARM_ARCH) && !defined(__aarch64__)
+#elif !defined(_MSC_VER) && !defined(__ARM_ARCH) && \
+    !(defined(__aarch64__) || defined(_M_ARM64))
 CpuFeatures::CpuFlags::CpuFlags(void)
     : L1_ECX(0), L1_EDX(0), L7_EBX(0), L7_ECX(0), L7_EDX(0) {
   uint32_t eax, ebx, ecx, edx;
@@ -336,7 +338,7 @@ bool CpuFeatures::HYPERVISOR(void) {
 
 const char *CpuFeatures::Intrinsics(void) {
   return ""
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
          "Neon"
 #if defined(__ARM_FEATURE_CRC32)
          "+CRC"
diff --git a/src/ailego/math/distance_matrix_accum_fp32.i b/src/ailego/math/distance_matrix_accum_fp32.i
@@ -30,7 +30,7 @@
                                             _mm512_castps_si512(b)))
 #endif  // __AVX512DQ__
 
-#if defined(__ARM_NEON) && !defined(__aarch64__)
+#if (defined(__ARM_NEON) || defined(_M_ARM64)) && !(defined(__aarch64__) || defined(_M_ARM64))
 static inline float32_t vaddvq_f32(float32x4_t v) {
   float32x2_t s = vadd_f32(vget_low_f32(v), vget_high_f32(v));
   return vget_lane_f32(vpadd_f32(s, s), 0);
@@ -42,7 +42,7 @@ static inline int32_t vaddvq_s32(int32x4_t v) {
 }
 #endif  //__ARM_NEON && !__aarch64__
 
-#if defined(__aarch64__)
+#if (defined(__aarch64__) || defined(_M_ARM64))
 #define ACCUM_FP32_2X1_NEON ACCUM_FP32_2X1_NEON_A64
 #else
 #define ACCUM_FP32_2X1_NEON ACCUM_FP32_2X1_NEON_A32
diff --git a/src/ailego/math/distance_matrix_fp32.i b/src/ailego/math/distance_matrix_fp32.i
@@ -26,7 +26,7 @@
   _mm256_insertf128_ps(_mm256_castps128_ps256(b), (a), 1)
 #endif  // __AVX__
 
-#if defined(__ARM_NEON) && !defined(__aarch64__)
+#if (defined(__ARM_NEON) || defined(_M_ARM64)) && !(defined(__aarch64__) || defined(_M_ARM64))
 #define vdupq_laneq_f32(a, b) vdupq_n_f32(vgetq_lane_f32(a, b))
 #endif  // __ARM_NEON && __aarch64__
 
diff --git a/src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc b/src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc
@@ -18,7 +18,7 @@
 namespace zvec {
 namespace ailego {
 
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 void SquaredEuclideanDistanceFp32NEON(const float *lhs, const float *rhs,
                                       size_t size, float *out);
 #endif
@@ -49,7 +49,7 @@ void SquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(const ValueType *m,
                                                           const ValueType *q,
                                                           size_t dim,
                                                           float *out) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   SquaredEuclideanDistanceFp32NEON(m, q, dim, out);
 #else
 #if defined(__AVX512F__)
diff --git a/src/ailego/math/euclidean_distance_matrix_fp32_neon.cc b/src/ailego/math/euclidean_distance_matrix_fp32_neon.cc
@@ -19,7 +19,7 @@
 namespace zvec {
 namespace ailego {
 
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 //! Squared Euclidean Distance
 void SquaredEuclideanDistanceFp32NEON(const float *lhs, const float *rhs,
                                       size_t size, float *out) {
diff --git a/src/ailego/math/inner_product_matrix_fp32_dispatch.cc b/src/ailego/math/inner_product_matrix_fp32_dispatch.cc
@@ -20,7 +20,7 @@ namespace ailego {
 //--------------------------------------------------
 // Dense
 //--------------------------------------------------
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 float InnerProductFp32NEON(const float *lhs, const float *rhs, size_t size);
 float MinusInnerProductFp32NEON(const float *lhs, const float *rhs,
                                 size_t size);
@@ -49,7 +49,7 @@ float MinusInnerProductFp32Scalar(const float *lhs, const float *rhs,
 //! Compute the distance between matrix and query (FP32, M=1, N=1)
 void InnerProductMatrix<float, 1, 1>::Compute(const float *m, const float *q,
                                               size_t dim, float *out) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   *out = InnerProductFp32NEON(m, q, dim);
 #else
 #if defined(__AVX512F__)
@@ -80,7 +80,7 @@ void InnerProductMatrix<float, 1, 1>::Compute(const float *m, const float *q,
 void MinusInnerProductMatrix<float, 1, 1>::Compute(const float *m,
                                                    const float *q, size_t dim,
                                                    float *out) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   *out = MinusInnerProductFp32NEON(m, q, dim);
 #else
 #if defined(__AVX512F__)
diff --git a/src/ailego/math/inner_product_matrix_fp32_neon.cc b/src/ailego/math/inner_product_matrix_fp32_neon.cc
@@ -22,7 +22,7 @@ namespace ailego {
 //--------------------------------------------------
 // Dense
 //--------------------------------------------------
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 float InnerProductFp32NEON(const float *lhs, const float *rhs, size_t size) {
   const float *last = lhs + size;
   const float *last_aligned = lhs + ((size >> 3) << 3);
diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc
@@ -18,7 +18,7 @@
 namespace zvec {
 namespace ailego {
 
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 float InnerProductAndSquaredNormFp32NEON(const float *lhs, const float *rhs,
                                          size_t size, float *sql, float *sqr);
 #endif
@@ -98,7 +98,7 @@ void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
 void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
     const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,
     float *out) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   float u2{0.0f};
   float v2{0.0f};
   float sum = InnerProductAndSquaredNormFp32NEON(p, q, dim, &u2, &v2);
diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc
@@ -19,7 +19,7 @@
 namespace zvec {
 namespace ailego {
 
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 //! Compute the Inner Product between p and q, and each Squared L2-Norm value
 float InnerProductAndSquaredNormFp32NEON(const float *lhs, const float *rhs,
                                          size_t size, float *sql, float *sqr) {
diff --git a/src/ailego/math/norm1_matrix.h b/src/ailego/math/norm1_matrix.h
@@ -116,7 +116,8 @@ struct Norm1Matrix<
   }
 };
 
-#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))
+#if defined(__SSE__) || ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+                         (defined(__aarch64__) || defined(_M_ARM64)))
 /*! L1-Norm Matrix (FP32, M=1)
  */
 template <>
@@ -129,6 +130,8 @@ struct Norm1Matrix<float, 1> {
 };
 #endif  // __SSE__ || (__ARM_NEON && __aarch64__)
 
+// MSVC ARM64 lacks `float16_t` without ARMv8.2 FP16; gate FP16 NEON
+// specialization to gcc/clang aarch64.
 #if (defined(__F16C__) && defined(__AVX__)) || \
     (defined(__ARM_NEON) && defined(__aarch64__))
 /*! L1-Norm Matrix (FP16, M=1)
diff --git a/src/ailego/math/norm1_matrix_fp16.cc b/src/ailego/math/norm1_matrix_fp16.cc
@@ -67,6 +67,9 @@ static const __m512 ABS_MASK_FP32_AVX512 =
 //! Calculate sum of absolute (NEON)
 #define SA_FP16_NEON(v_m, v_sum) v_sum = vaddq_f16(vabsq_f16(v_m), v_sum);
 
+// MSVC ARM64 lacks `float16_t` without ARMv8.2 FP16, so the NEON FP16
+// kernel is gated to gcc/clang aarch64. The generic Float16 path is used
+// on MSVC ARM64.
 #if (defined(__F16C__) && defined(__AVX__)) || \
     (defined(__ARM_NEON) && defined(__aarch64__))
 //! Compute the L1-norm of vectors (FP16, M=1)
diff --git a/src/ailego/math/norm1_matrix_fp32.cc b/src/ailego/math/norm1_matrix_fp32.cc
@@ -56,11 +56,12 @@ namespace ailego {
 //! Calculate sum of absolute (NEON)
 #define SA_FP32_NEON(v_m, v_sum) v_sum = vaddq_f32(vabsq_f32(v_m), v_sum);
 
-#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))
+#if defined(__SSE__) || ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+                         (defined(__aarch64__) || defined(_M_ARM64)))
 //! Compute the L1-norm of vectors (FP32, M=1)
 void Norm1Matrix<float, 1>::Compute(const ValueType *m, size_t dim,
                                     float *out) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   NORM_FP32_1_NEON(m, dim, out, )
 #else
 #if defined(__AVX512F__)
diff --git a/src/ailego/math/norm2_matrix.h b/src/ailego/math/norm2_matrix.h
@@ -371,7 +371,8 @@ struct SquaredNorm2Matrix<uint8_t, M, typename std::enable_if<M >= 2>::type> {
   }
 };
 
-#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))
+#if defined(__SSE__) || ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+                         (defined(__aarch64__) || defined(_M_ARM64)))
 /*! L2-Norm Matrix (FP32, M=1)
  */
 template <>
@@ -395,6 +396,8 @@ struct SquaredNorm2Matrix<float, 1> {
 };
 #endif  // __SSE__ || (__ARM_NEON && __aarch64__)
 
+// MSVC ARM64 lacks `float16_t` without ARMv8.2 FP16; gate FP16 NEON
+// specialization to gcc/clang aarch64.
 #if (defined(__F16C__) && defined(__AVX__)) || \
     (defined(__ARM_NEON) && defined(__aarch64__))
 /*! L2-Norm Matrix (FP16, M=1)
diff --git a/src/ailego/math/norm2_matrix_fp16.cc b/src/ailego/math/norm2_matrix_fp16.cc
@@ -52,6 +52,8 @@ namespace ailego {
 //! Calculate sum of squared (NEON)
 #define SS_FP16_NEON(v_m, v_sum) v_sum = vfmaq_f16(v_sum, v_m, v_m);
 
+// MSVC ARM64 lacks `float16_t` without ARMv8.2 FP16, so the NEON FP16
+// kernel is gated to gcc/clang aarch64.
 #if (defined(__F16C__) && defined(__AVX__)) || \
     (defined(__ARM_NEON) && defined(__aarch64__))
 //! Compute the L2-norm of vectors (FP16, M=1)
diff --git a/src/ailego/math/norm2_matrix_fp32.cc b/src/ailego/math/norm2_matrix_fp32.cc
@@ -43,11 +43,12 @@ namespace ailego {
 //! Calculate sum of squared (NEON)
 #define SS_FP32_NEON(v_m, v_sum) v_sum = vfmaq_f32(v_sum, v_m, v_m);
 
-#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))
+#if defined(__SSE__) || ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+                         (defined(__aarch64__) || defined(_M_ARM64)))
 //! Compute the L2-norm of vectors (FP32, M=1)
 void Norm2Matrix<float, 1>::Compute(const ValueType *m, size_t dim,
                                     float *out) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   NORM_FP32_1_NEON(m, dim, out, std::sqrt)
 #else
 #if defined(__AVX512F__)
@@ -69,7 +70,7 @@ void Norm2Matrix<float, 1>::Compute(const ValueType *m, size_t dim,
 //! Compute the squared L2-norm of vectors (FP32, M=1)
 void SquaredNorm2Matrix<float, 1>::Compute(const ValueType *m, size_t dim,
                                            float *out) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   NORM_FP32_1_NEON(m, dim, out, )
 #else
 #if defined(__AVX512F__)
diff --git a/src/ailego/math/normalizer.cc b/src/ailego/math/normalizer.cc
@@ -17,7 +17,8 @@
 namespace zvec {
 namespace ailego {
 
-#if (defined(__ARM_NEON) && defined(__aarch64__))
+#if ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+     (defined(__aarch64__) || defined(_M_ARM64)))
 static inline void NormalizeNEON(float *arr, size_t dim, float norm) {
   float *last = arr + dim;
   float *last_aligned = arr + ((dim >> 3) << 3);
@@ -43,6 +44,11 @@ static inline void NormalizeNEON(float *arr, size_t dim, float norm) {
   }
 }
 
+// MSVC ARM64 does not declare `float16_t` unless ARMv8.2 FP16 is enabled
+// (`_M_ARM64FP16` / `/arch:armv8.2`). Skip the FP16 NEON variants on MSVC
+// ARM64 — the FP16 `Normalizer<Float16>::Compute` dispatch below falls
+// through to the scalar conversion path.
+#if !defined(_MSC_VER)
 #if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)
 static inline void NormalizeNEON(float16_t *arr, size_t dim, float norm) {
   float16_t *last = arr + dim;
@@ -114,6 +120,7 @@ static inline void NormalizeNEON(float16_t *arr, size_t dim, float norm) {
   }
 }
 #endif  // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+#endif  // !_MSC_VER (FP16 NEON: gcc/clang aarch64 only)
 #endif  // __ARM_NEON && __aarch64__
 
 #if defined(__AVX__)
@@ -392,10 +399,11 @@ static inline void NormalizeSSE(float *arr, size_t dim, float norm) {
 }
 #endif  // __SSE__
 
-#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))
+#if defined(__SSE__) || ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+                         (defined(__aarch64__) || defined(_M_ARM64)))
 //! Compute the norm of vector
 void Normalizer<float>::Compute(ValueType *arr, size_t dim, float norm) {
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
   NormalizeNEON(arr, dim, norm);
 #else
 #if defined(__AVX512F__)
@@ -415,6 +423,10 @@ void Normalizer<float>::Compute(ValueType *arr, size_t dim, float norm) {
 }
 #endif  // __SSE__ || (__ARM_NEON && __aarch64__)
 
+// MSVC ARM64 lacks `float16_t` without ARMv8.2 FP16, so the FP16
+// `NormalizeNEON` variants are not defined there (see top of file).
+// Exclude MSVC from this specialization so MSVC ARM64 falls back to the
+// generic uint16_t-based Float16 normalize path.
 #if (defined(__F16C__) && defined(__AVX__)) || \
     (defined(__ARM_NEON) && defined(__aarch64__))
 //! Compute the norm of vector
diff --git a/src/ailego/math/normalizer.h b/src/ailego/math/normalizer.h
@@ -51,7 +51,8 @@ struct Normalizer {
   }
 };
 
-#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__))
+#if defined(__SSE__) || ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+                         (defined(__aarch64__) || defined(_M_ARM64)))
 /*! Normalizer (FP32)
  */
 template <>
@@ -80,6 +81,8 @@ struct Normalizer<float> {
 };
 #endif  // __SSE__ || (__ARM_NEON && __aarch64__)
 
+// MSVC ARM64 lacks `float16_t` without ARMv8.2 FP16, so the FP16 NEON
+// `Normalizer<Float16>` specialization is gated to gcc/clang aarch64.
 #if (defined(__F16C__) && defined(__AVX__)) || \
     (defined(__ARM_NEON) && defined(__aarch64__))
 /*! Normalizer (FP16)
diff --git a/src/ailego/utility/bitset_helper.cc b/src/ailego/utility/bitset_helper.cc
@@ -23,7 +23,7 @@
 #define bitset_popcount64 _mm_popcnt_u64
 #endif  // !__SSE4_2__
 
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 static inline void bitset_and(uint32_t *lhs, const uint32_t *rhs, size_t size) {
   uint32_t *last = lhs + size;
   uint32_t *last_aligned = lhs + ((size >> 2) << 2);
@@ -1480,7 +1480,8 @@ static inline bool bitset_test_none(const uint32_t *lhs, size_t size) {
 #endif  // AILEGO_M64
 #endif  // __AVX2__
 
-#if (defined(__ARM_NEON) && defined(__aarch64__))
+#if ((defined(__ARM_NEON) || defined(_M_ARM64)) && \
+     (defined(__aarch64__) || defined(_M_ARM64)))
 static inline size_t bitset_cardinality(const uint32_t *lhs, size_t size) {
   const uint32_t *last = lhs + size;
   const uint32_t *last_aligned = lhs + ((size >> 2) << 2);
diff --git a/src/ailego/utility/float_helper.cc b/src/ailego/utility/float_helper.cc
@@ -22,7 +22,10 @@
 // #define float32(x) _cvtsh_ss(x)
 // #endif  // __F16C__ && __AVX__
 
-#if defined(__aarch64__)
+// MSVC ARM64 lacks the GCC/Clang `__fp16` extension type, so keep this
+// path gated on `__aarch64__` (predefined only by GCC/Clang on AArch64).
+// MSVC ARM64 falls through to the F16C/scalar paths below.
+#if defined(__aarch64__) && !defined(_MSC_VER)
 static inline float float32(uint16_t val) {
   __fp16 f;
   memcpy(&f, &val, sizeof(val));
diff --git a/src/ailego/version.i b/src/ailego/version.i
@@ -225,7 +225,7 @@
 #define AILEGO_VERSION_OPENMP ""
 #endif
 
-#if defined(__ARM_NEON)
+#if (defined(__ARM_NEON) || defined(_M_ARM64))
 #define AILEGO_VERSION_SIMD "  Arm Neon Instruction Set\n"
 #elif defined(__AVX512FP16__)
 #define AILEGO_VERSION_SIMD "  AVX-512FP16 Instruction Set\n"
diff --git a/src/include/zvec/ailego/buffer/concurrentqueue.h b/src/include/zvec/ailego/buffer/concurrentqueue.h
diff --git a/src/include/zvec/ailego/internal/platform.h b/src/include/zvec/ailego/internal/platform.h
diff --git a/src/include/zvec/ailego/utility/float_helper.h b/src/include/zvec/ailego/utility/float_helper.h
