[ARM] MLAS: portable machine-code SVE elementwise kernels, FEXPA exp

cmake/CMakeLists.txt

@@ -544,6 +544,11 @@ if(onnxruntime_USE_SVE)
       message(WARNING "onnxruntime_USE_SVE was set but compiler does not support SVE. It will be disabled.")
       set(onnxruntime_USE_SVE OFF)
     endif()
+  elseif(WIN32 AND onnxruntime_target_platform STREQUAL "ARM64")
+    # The SVE elementwise kernels ship as portable machine-code assembly
+    # (hex-encoded instruction words assembled by armasm64), so no compiler
+    # SVE support is needed; the kernels are runtime-gated on HasArmSve().
+    message(STATUS "SVE enabled via portable machine-code kernels.")
   else()
     message(WARNING "onnxruntime_USE_SVE was set but it is not supported on this platform. It will be disabled.")
     set(onnxruntime_USE_SVE OFF)

cmake/onnxruntime_mlas.cmake

@@ -170,6 +170,21 @@ function(setup_mlas_source_for_windows)
 	  if (onnxruntime_USE_KLEIDIAI)
         setup_kleidiai()
       endif()
+
+      if (onnxruntime_USE_SVE)
+        # Portable machine-code SVE elementwise kernels: a plain C++ dispatch
+        # layer plus hex-encoded assembly. kai_asm_macros.h resolves to
+        # armasm64 directives (AREA/PROC/DCD) during the cl.exe preprocessing
+        # step applied to mlas_platform_preprocess_srcs below, so no SVE
+        # toolchain support is required.
+        target_sources(onnxruntime_mlas PRIVATE
+          ${MLAS_SRC_DIR}/sve/mlasi_sve.h
+          ${MLAS_SRC_DIR}/sve/elementwise_sve_dispatch.cpp
+        )
+        list(APPEND mlas_platform_preprocess_srcs ${MLAS_SRC_DIR}/aarch64/elementwise_sve_asm.S)
+        list(APPEND mlas_private_compile_definitions MLAS_USE_SVE)
+        set(mlas_private_compile_definitions ${mlas_private_compile_definitions} PARENT_SCOPE)
+      endif()
     else()
       target_sources(onnxruntime_mlas PRIVATE
         ${MLAS_SRC_DIR}/qgemm_kernel_neon.cpp
@@ -552,11 +567,20 @@ else()
         # Conditionally add the SVE implementation if compiler supports it
         if (onnxruntime_USE_SVE)
           list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/mlasi_sve.h)
-          list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/elementwise_sve.cpp)
-          list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/elementwise_sve_fp16.cpp)
-          list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/mlas_sve_fp16.h)
-          set_source_files_properties(${MLAS_SRC_DIR}/sve/elementwise_sve.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+sve+fp16 ")
-          set_source_files_properties(${MLAS_SRC_DIR}/sve/elementwise_sve_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+sve+fp16 ")
+          list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/elementwise_sve_dispatch.cpp)
+          # The portable machine-code variant is the production default; set
+          # to OFF to build the SVE intrinsics reference implementation (the
+          # regeneration source for elementwise_sve_asm.S).
+          set(onnxruntime_SVE_ELEMENTWISE_ASM ON)
+          if (onnxruntime_SVE_ELEMENTWISE_ASM)
+            # Portable machine-code variant (same symbols as the intrinsics TUs).
+            list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/aarch64/elementwise_sve_asm.S)
+          else()
+            list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/elementwise_sve.cpp)
+            list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/elementwise_sve_fp16.cpp)
+            set_source_files_properties(${MLAS_SRC_DIR}/sve/elementwise_sve.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+sve+fp16 -DMLAS_SVE_SUMEXP_FEXPA=1 ")
+            set_source_files_properties(${MLAS_SRC_DIR}/sve/elementwise_sve_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+sve+fp16 ")
+          endif()
           list(APPEND mlas_private_compile_definitions MLAS_USE_SVE)
         endif()
 

onnxruntime/core/common/cpuid_info.cc

@@ -317,6 +317,12 @@ void CPUIDInfo::ArmWindowsInit() {
     has_arm_sme2_ = cpuinfo_has_arm_sme2();
   }
 #endif  // defined(CPUINFO_SUPPORTED)
+
+#if defined(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE)
+  // Available from Windows 11 24H2 SDKs; older SDKs lack the constant, in
+  // which case SVE stays disabled at runtime.
+  has_arm_sve_ = IsProcessorFeaturePresent(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) != 0;
+#endif
 }
 
 #elif defined(__APPLE__)  // ^ defined(_WIN32)

onnxruntime/core/mlas/lib/aarch64/kai_asm_macros.h

@@ -0,0 +1,63 @@
+/*++
+
+Copyright (c) Microsoft Corporation. All rights reserved.
+
+Licensed under the MIT License.
+
+Module Name:
+
+    kai_asm_macros.h
+
+Abstract:
+
+    Portable assembly macros for AArch64 kernels that must assemble under
+    both the GNU assembler (Linux, macOS) and Microsoft armasm64 (Windows).
+
+    The macro set and its names are adopted verbatim from Arm's KleidiAI
+    library (kai/kai_common_sve_asm.S, SPDX-License-Identifier: Apache-2.0),
+    which inlines this block at the top of each of its portable *_asm.S
+    files; keeping the same names lets kernels move between the two trees
+    unchanged. This header exists so ONNX Runtime kernels can share a single
+    copy.
+
+    Instructions that a given assembler may not accept as mnemonics (all SVE
+    instructions, in the case of armasm64) are emitted as raw instruction
+    words via KAI_ASM_INST, which expands to ".inst" under GAS and "DCD"
+    under armasm64 — the same bytes either way.
+
+--*/
+
+#if defined(_MSC_VER)
+
+    #define KAI_ASM_GLOBAL(name) GLOBAL name
+    #define KAI_ASM_FUNCTION_TYPE(name)
+    #define KAI_ASM_FUNCTION_LABEL(name) name PROC
+    #define KAI_ASM_FUNCTION_END(name) ENDP
+
+    #define KAI_ASM_CODE(name) AREA name, CODE, READONLY
+    #define KAI_ASM_ALIGN
+    #define KAI_ASM_LABEL(name) name
+    #define KAI_ASM_INST(hex) DCD hex
+    #define KAI_ASM_END END
+
+#else
+
+    #if defined(__APPLE__)
+        #define KAI_ASM_GLOBAL(name) .globl _##name
+        #define KAI_ASM_FUNCTION_TYPE(name)
+        #define KAI_ASM_FUNCTION_LABEL(name) _##name:
+        #define KAI_ASM_FUNCTION_END(name)
+    #else
+        #define KAI_ASM_GLOBAL(name) .global name
+        #define KAI_ASM_FUNCTION_TYPE(name) .type name, %function
+        #define KAI_ASM_FUNCTION_LABEL(name) name:
+        #define KAI_ASM_FUNCTION_END(name) .size name, .-name
+    #endif
+
+    #define KAI_ASM_CODE(name) .text
+    #define KAI_ASM_ALIGN .p2align 4,,11
+    #define KAI_ASM_LABEL(name) name:
+    #define KAI_ASM_INST(hex) .inst hex
+    #define KAI_ASM_END
+
+#endif

onnxruntime/core/mlas/lib/compute.cpp

@@ -277,7 +277,7 @@ Return Value:
 
 --*/
 {
-#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_RISCV64)
+#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_RISCV64) || defined(MLAS_USE_SVE)
     GetMlasPlatform().ComputeExpF32Kernel(Input, Output, N);
 #else
     MlasComputeExpF32Kernel(Input, Output, N);

onnxruntime/core/mlas/lib/mlasi.h

@@ -1640,15 +1640,16 @@ struct MLAS_PLATFORM {
     MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL* ComputeSumExpF32Kernel;
     MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL* ComputeLogSoftmaxOutputF32Kernel;
     MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL* ComputeSoftmaxOutputF32Kernel;
+    MLAS_COMPUTE_UNARY_FLOAT_KERNEL* ComputeExpF32Kernel;
+    MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL* ReduceMinimumMaximumF32Kernel;
+    MLAS_COMPUTE_UNARY_FLOAT_KERNEL* TanhKernelRoutine;
 #endif
 #if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_RISCV64)
     // Hoisted under combined guard so future "shared between AMD64 and RISCV64"
     // changes won't produce duplicate-member errors. Each platform's init code
     // assigns these fields independently in platform.cpp.
     MLAS_COMPUTE_UNARY_FLOAT_KERNEL* GeluErfKernelRoutine;
     MLAS_COMPUTE_UNARY_FLOAT_KERNEL* SiluKernelRoutine;
-    MLAS_COMPUTE_UNARY_FLOAT_KERNEL* TanhKernelRoutine;
-    MLAS_COMPUTE_UNARY_FLOAT_KERNEL* ComputeExpF32Kernel;
 #endif
 
 #if defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV)
@@ -1683,7 +1684,6 @@ MLAS_COMPUTE_TANH_FP16_KERNEL* TanhFP16KernelRoutine = nullptr;
     MLAS_POOL_FLOAT_KERNEL* PoolFloatKernel[MlasPoolingKindCount];
     MLAS_QLINEAR_BINARY_OP_S8_KERNEL* QLinearAddS8Kernel;
     MLAS_QLINEAR_BINARY_OP_U8_KERNEL* QLinearAddU8Kernel;
-    MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL* ReduceMinimumMaximumF32Kernel;
     MLAS_QUANTIZE_LINEAR_S8_KERNEL* QuantizeLinearS8Kernel;
     MLAS_QUANTIZE_LINEAR_U8_KERNEL* QuantizeLinearU8Kernel;
     MLAS_QUANTIZE_LINEAR_S16_KERNEL* QuantizeLinearS16Kernel;

onnxruntime/core/mlas/lib/platform.cpp

@@ -729,33 +729,47 @@ Return Value:
     if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
         this->ErfKernelRoutine = MlasSveErfKernel;
         this->LogisticKernelRoutine = MlasSveLogisticKernel;
+        this->ComputeExpF32Kernel = MlasSveComputeExpF32Kernel;
+        this->TanhKernelRoutine = MlasSveTanhKernel;
         this->ReduceMaximumF32Kernel = MlasSveReduceMaximumF32Kernel;
+        this->ReduceMinimumMaximumF32Kernel = MlasSveReduceMinimumMaximumF32Kernel;
         this->ComputeSumExpF32Kernel = MlasSveComputeSumExpF32Kernel;
         this->ComputeLogSoftmaxOutputF32Kernel = MlasSveComputeLogSoftmaxOutputF32Kernel;
         this->ComputeSoftmaxOutputF32Kernel = MlasSveComputeSoftmaxOutputF32Kernel;
     }
     else{
         this->ErfKernelRoutine = MlasErfKernel;
         this->LogisticKernelRoutine = MlasLogisticKernel;
+        this->ComputeExpF32Kernel = MlasComputeExpF32Kernel;
+        this->TanhKernelRoutine = MlasTanhKernel;
         this->ReduceMaximumF32Kernel = MlasReduceMaximumF32Kernel;
+        this->ReduceMinimumMaximumF32Kernel = MlasReduceMinimumMaximumF32Kernel;
         this->ComputeSumExpF32Kernel = MlasComputeSumExpF32Kernel;
         this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32Kernel;
         this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32Kernel;
     }
 #endif
 
-#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && !defined(_WIN32)
+#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED)
+    //
+    // The SVE fp16 kernels are portable machine code and available on every
+    // platform; the NEON fp16 fallbacks are only built on non-Windows
+    // targets. On Windows without SVE hardware the routines stay null and
+    // the callers use their scalar fallbacks.
+    //
     #if defined(MLAS_USE_SVE)
         if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
             this->ErfFP16KernelRoutine = MlasSveErfFP16Kernel;
             this->GeluFP16KernelRoutine = MlasSveGeluFP16Kernel;
             this->TanhFP16KernelRoutine = MlasSveTanhFP16Kernel;
         }
         else{
+        #if !defined(_WIN32)
             this->ErfFP16KernelRoutine = MlasNeonErfFP16Kernel;
             this->GeluFP16KernelRoutine = MlasNeonGeluFP16Kernel;
+        #endif
         }
-    #else
+    #elif !defined(_WIN32)
         this->ErfFP16KernelRoutine = MlasNeonErfFP16Kernel;
         this->GeluFP16KernelRoutine = MlasNeonGeluFP16Kernel;
     #endif

onnxruntime/core/mlas/lib/quantize.cpp

@@ -2311,7 +2311,7 @@ Return Value:
 
 --*/
 {
-#if defined(MLAS_TARGET_AMD64)
+#if defined(MLAS_TARGET_AMD64) || defined(MLAS_USE_SVE)
     GetMlasPlatform().ReduceMinimumMaximumF32Kernel(Input, Min, Max, N);
 #else
     MlasReduceMinimumMaximumF32Kernel(Input, Min, Max, N);