Move KleidiAI support detection to CMake logic. (microsoft#26178)

edgchen1 · web-flow · commit 2e8a45ab059e · 2025-09-30T06:55:45.000-07:00
### Description
&lt;!-- Describe your changes. --&gt;

Use of the KleidiAI library is not supported on all platforms. Previously, this support detection was done in both
`tools/ci_build/build.py` and `cmake/CMakeLists.txt`. This change consolidates and simplifies the checks. They are now done in one place in the CMake logic.

Move setting of onnxruntime_target_platform into a separate file.

### Motivation and Context
&lt;!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. --&gt;

Clean up and simplify.
diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt
@@ -91,11 +91,6 @@ option(onnxruntime_USE_SVE "Build with SVE support in MLAS" OFF)
 option(onnxruntime_USE_ARM_NEON_NCHWC "Build with ARM Neon NCHWc kernels in MLAS" OFF)
 
 option(onnxruntime_USE_KLEIDIAI "Build with KleidiAI integration in MLAS" OFF)
-# iOS simulator build explicitly builds targets with USE_KLEIDIAI=ON so attempting to force override if so
-if(APPLE AND CMAKE_OSX_ARCHITECTURES MATCHES "x86_64")
-    message(WARNING "Disabling KleidiAI: not supported on Apple x86_64 platforms")
-  set(onnxruntime_USE_KLEIDIAI OFF CACHE BOOL "" FORCE)
-endif()
 option(onnxruntime_BUILD_UNIT_TESTS "Build ONNXRuntime unit tests" ON)
 option(onnxruntime_BUILD_CSHARP "Build C# library" OFF)
 option(onnxruntime_BUILD_OBJC "Build Objective-C library" OFF)
@@ -259,6 +254,8 @@ if("${CMAKE_C_COMPILER_ID}" STREQUAL "GNU" AND CMAKE_C_COMPILER_VERSION VERSION_
   message(FATAL_ERROR  "GCC version must be greater than or equal to 11.1")
 endif()
 
+include(detect_onnxruntime_target_platform.cmake)
+
 # ENABLE_TRAINING includes all training functionality
 # The following 2 entry points
 # 1. ORTModule
@@ -435,14 +432,6 @@ set(ORTTRAINING_SOURCE_DIR ${ORTTRAINING_ROOT}/orttraining)
 
 include(adjust_global_compile_flags.cmake)
 
-if (APPLE)
-  if (NOT CMAKE_OSX_ARCHITECTURES)
-    message("Building ONNX Runtime for ${CMAKE_HOST_SYSTEM_PROCESSOR} CPU ARCH")
-  endif()
-elseif (NOT WIN32 AND NOT APPLE)
-  message("Building ONNX Runtime for ${onnxruntime_target_platform} CPU ARCH")
-endif()
-
 # We need to link with libatomic on systems that do not have built-in atomics, or
 # don't have built-in support for 8 byte atomics
 # Derived from https://github.com/protocolbuffers/protobuf/blob/master/cmake/CMakeLists.txt
@@ -514,6 +503,66 @@ if (onnxruntime_BUILD_SHARED_LIB OR onnxruntime_ENABLE_PYTHON)
   endif()
 endif()
 
+if (onnxruntime_USE_ARM_NEON_NCHWC)
+  message(STATUS "Building MLAS with ARM Neon NCHWc kernels")
+endif()
+
+if(onnxruntime_USE_SVE)
+  if(LINUX AND onnxruntime_target_platform STREQUAL "aarch64")
+    check_cxx_compiler_flag("-march=armv8.2-a+sve" HAS_ARM64_SVE)
+    if(HAS_ARM64_SVE)
+      message(STATUS "Compiler supports SVE!")
+    else()
+      message(WARNING "onnxruntime_USE_SVE was set but compiler does not support SVE. It will be disabled.")
+      set(onnxruntime_USE_SVE OFF)
+    endif()
+  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)
+  endif()
+endif()
+
+if(onnxruntime_USE_KLEIDIAI)
+  function(is_kleidiai_supported is_supported_var)
+    # check for supported target platforms
+    if(NOT (onnxruntime_target_platform STREQUAL "aarch64" OR
+            onnxruntime_target_platform STREQUAL "ARM64" OR
+            onnxruntime_target_platform STREQUAL "arm64"))
+      message(WARNING "KleidiAI is not supported on this platform.")
+
+      set(${is_supported_var} FALSE PARENT_SCOPE)
+      return()
+    endif()
+
+    # check for compiler support
+    if(MSVC)
+      # TODO detect on MSVC
+    else()
+      check_cxx_compiler_flag(-march=armv8.2-a+dotprod HAS_ARM64_DOTPROD)
+      check_cxx_compiler_flag(-march=armv8.2-a+i8mm HAS_ARM64_I8MM)
+      if(NOT HAS_ARM64_DOTPROD)
+        message(WARNING "The compiler doesn't support dotprod instructions.")
+      endif()
+      if(NOT HAS_ARM64_I8MM)
+        message(WARNING "The compiler doesn't support i8mm instructions.")
+      endif()
+      if(NOT HAS_ARM64_DOTPROD OR NOT HAS_ARM64_I8MM)
+        set(${is_supported_var} FALSE PARENT_SCOPE)
+        return()
+      endif()
+    endif()
+
+    set(${is_supported_var} TRUE PARENT_SCOPE)
+  endfunction()
+
+  is_kleidiai_supported(is_kleidiai_supported_result)
+
+  if(NOT is_kleidiai_supported_result)
+    message(WARNING "onnxruntime_USE_KLEIDIAI was set but it is not supported. It will be disabled.")
+    set(onnxruntime_USE_KLEIDIAI OFF)
+  endif()
+endif()
+
 #Dependencies begin
 get_filename_component(ONNXRUNTIME_ROOT "${ONNXRUNTIME_ROOT}" ABSOLUTE)
 get_filename_component(ORTTRAINING_ROOT "${ORTTRAINING_ROOT}" ABSOLUTE)
@@ -664,47 +713,6 @@ else()
   endif()
 endif()
 
-if (onnxruntime_USE_ARM_NEON_NCHWC)
-  message(STATUS "Building MLAS with ARM Neon NCHWc kernels")
-endif()
-
-if(onnxruntime_USE_SVE)
-  if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64" AND CMAKE_SYSTEM_NAME STREQUAL "Linux")
-    check_cxx_compiler_flag("-march=armv8.2-a+sve" HAS_ARM64_SVE)
-    if(HAS_ARM64_SVE)
-      message(STATUS "Compiler supports SVE!")
-    else()
-      message(WARNING "onnxruntime_USE_SVE was set but compiler does not support SVE. It will be disabled.")
-      set(onnxruntime_USE_SVE OFF)
-    endif()
-  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)
-  endif()
-endif()
-
-if (onnxruntime_USE_KLEIDIAI AND (
-      (onnxruntime_target_platform STREQUAL "aarch64") OR
-      (onnxruntime_target_platform STREQUAL "ARM64") OR
-      (APPLE AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")))
-
-  # TODO Add checks for MSVC Compilation
-  if(NOT MSVC)
-    check_cxx_compiler_flag(-march=armv8.2-a+dotprod HAS_ARM64_DOTPROD)
-    check_cxx_compiler_flag(-march=armv8.2-a+i8mm HAS_ARM64_I8MM)
-    if (NOT HAS_ARM64_DOTPROD)
-      message(FATAL_ERROR "The compiler doesn't support dotprod")
-    endif()
-    if (NOT HAS_ARM64_I8MM)
-      message(FATAL_ERROR "The compiler doesn't support i8mm")
-    endif()
-  else()
-    message(STATUS "Skipping -march= checks on MSVC (not supported), assuming dotprod/i8mm support manually.")
-    set(HAS_ARM64_DOTPROD TRUE)
-    set(HAS_ARM64_I8MM TRUE)
-  endif()
-endif()
-
 #names in this var must match the directory names under onnxruntime/core/providers
 #ONNXRUNTIME_PROVIDER_NAMES is the list of providers that needs to export additional symbols in the global namespace.
 #For example CUDA EP exports "OrtSessionOptionsAppendExecutionProvider_CUDA", which is a global function.
diff --git a/cmake/adjust_global_compile_flags.cmake b/cmake/adjust_global_compile_flags.cmake
@@ -217,30 +217,20 @@ endmacro()
 #Set global compile flags for all the source code(including third_party code like protobuf)
 #This section must be before any add_subdirectory, otherwise build may fail because /MD,/MT mismatch
 if (MSVC)
-  if (CMAKE_VS_PLATFORM_NAME)
-    # Multi-platform generator
-    set(onnxruntime_target_platform ${CMAKE_VS_PLATFORM_NAME})
-  else()
-    set(onnxruntime_target_platform ${CMAKE_SYSTEM_PROCESSOR})
-  endif()
-  if (onnxruntime_target_platform STREQUAL "ARM64")
-    set(onnxruntime_target_platform "ARM64")
-    enable_language(ASM_MARMASM)
-  elseif (onnxruntime_target_platform STREQUAL "ARM64EC")
+  if (onnxruntime_target_platform STREQUAL "ARM64" OR
+      onnxruntime_target_platform STREQUAL "ARM64EC" OR
+      onnxruntime_target_platform STREQUAL "ARM")
     enable_language(ASM_MARMASM)
-  elseif (onnxruntime_target_platform STREQUAL "ARM" OR CMAKE_GENERATOR MATCHES "ARM")
-    set(onnxruntime_target_platform "ARM")
-    enable_language(ASM_MARMASM)
-  elseif (onnxruntime_target_platform STREQUAL "x64" OR onnxruntime_target_platform STREQUAL "x86_64" OR onnxruntime_target_platform STREQUAL "AMD64" OR CMAKE_GENERATOR MATCHES "Win64")
-    set(onnxruntime_target_platform "x64")
-    enable_language(ASM_MASM)
-  elseif (onnxruntime_target_platform STREQUAL "Win32" OR onnxruntime_target_platform STREQUAL "x86" OR onnxruntime_target_platform STREQUAL "i386" OR onnxruntime_target_platform STREQUAL "i686")
-    set(onnxruntime_target_platform "x86")
+  elseif (onnxruntime_target_platform STREQUAL "x64" OR
+          onnxruntime_target_platform STREQUAL "x86")
     enable_language(ASM_MASM)
-    message("Enabling SAFESEH for x86 build")
-    set(CMAKE_ASM_MASM_FLAGS "${CMAKE_ASM_MASM_FLAGS} /safeseh")
+
+    if (onnxruntime_target_platform STREQUAL "x86")
+      message("Enabling SAFESEH for x86 build")
+      set(CMAKE_ASM_MASM_FLAGS "${CMAKE_ASM_MASM_FLAGS} /safeseh")
+    endif()
   else()
-    message(FATAL_ERROR "Unknown CMAKE_SYSTEM_PROCESSOR: ${CMAKE_SYSTEM_PROCESSOR}")
+    message(FATAL_ERROR "Unsupported onnxruntime_target_platform value: ${onnxruntime_target_platform}")
   endif()
 
   #Always enable exception handling, even for Windows ARM
@@ -269,34 +259,6 @@ if (MSVC)
     set(CMAKE_CXX_FLAGS_MINSIZEREL "${CMAKE_CXX_FLAGS_MINSIZEREL} /Gw /GL")
   endif()
 else()
-  if (NOT APPLE)
-    #XXX: Sometimes the value of CMAKE_SYSTEM_PROCESSOR is set but it's wrong. For example, if you run an armv7 docker
-    #image on an aarch64 machine with an aarch64 Ubuntu host OS, in the docker instance cmake may still report
-    # CMAKE_SYSTEM_PROCESSOR as aarch64 by default. Given compiling this code may need more than 2GB memory, we do not
-    # support compiling for ARM32 natively(only support cross-compiling), we will ignore this issue for now.
-    if(NOT CMAKE_SYSTEM_PROCESSOR)
-      message(WARNING "CMAKE_SYSTEM_PROCESSOR is not set. Please set it in your toolchain cmake file.")
-      # Try to detect it
-      if("${CMAKE_C_COMPILER_ID}" STREQUAL "GNU" OR "${CMAKE_C_COMPILER_ID}" STREQUAL "Clang")
-        execute_process(
-		COMMAND "${CMAKE_C_COMPILER}" -dumpmachine
-		OUTPUT_VARIABLE GCC_DUMP_MACHINE_OUT OUTPUT_STRIP_TRAILING_WHITESPACE
-		ERROR_VARIABLE _err
-		RESULT_VARIABLE _res
-		)
-		if(NOT _res EQUAL 0)
-			message(SEND_ERROR "Failed to run 'gcc -dumpmachine':\n ${_res}")
-		endif()
-		string(REPLACE "-" ";" GCC_DUMP_MACHINE_OUT_LIST "${GCC_DUMP_MACHINE_OUT}")
-		list(LENGTH GCC_DUMP_MACHINE_OUT_LIST GCC_TRIPLET_LEN)
-		if(GCC_TRIPLET_LEN EQUAL 4)
-		  list(GET GCC_DUMP_MACHINE_OUT_LIST 0 CMAKE_SYSTEM_PROCESSOR)
-          message("Setting CMAKE_SYSTEM_PROCESSOR to ${CMAKE_SYSTEM_PROCESSOR}")
-        endif()
-      endif()
-    endif()
-    set(onnxruntime_target_platform ${CMAKE_SYSTEM_PROCESSOR})
-  endif()
   if (onnxruntime_BUILD_FOR_NATIVE_MACHINE)
     string(APPEND CMAKE_CXX_FLAGS " -march=native -mtune=native")
     string(APPEND CMAKE_C_FLAGS " -march=native -mtune=native")
diff --git a/cmake/detect_onnxruntime_target_platform.cmake b/cmake/detect_onnxruntime_target_platform.cmake
@@ -0,0 +1,80 @@
+# This file will set the onnxruntime_target_platform variable, if applicable.
+# onnxruntime_target_platform identifies the platform to compile for.
+block(PROPAGATE onnxruntime_target_platform)
+
+unset(onnxruntime_target_platform)
+
+if (MSVC)
+  if (CMAKE_VS_PLATFORM_NAME)
+    # Multi-platform generator
+    set(onnxruntime_target_platform ${CMAKE_VS_PLATFORM_NAME})
+  else()
+    set(onnxruntime_target_platform ${CMAKE_SYSTEM_PROCESSOR})
+  endif()
+
+  if (onnxruntime_target_platform STREQUAL "ARM64" OR
+      onnxruntime_target_platform STREQUAL "ARM64EC")
+    # Do nothing. We'll just use the current value of onnxruntime_target_platform.
+  elseif (onnxruntime_target_platform STREQUAL "ARM" OR
+          CMAKE_GENERATOR MATCHES "ARM")
+    set(onnxruntime_target_platform "ARM")
+  elseif (onnxruntime_target_platform STREQUAL "x64" OR
+          onnxruntime_target_platform STREQUAL "x86_64" OR
+          onnxruntime_target_platform STREQUAL "AMD64" OR
+          CMAKE_GENERATOR MATCHES "Win64")
+    set(onnxruntime_target_platform "x64")
+  elseif (onnxruntime_target_platform STREQUAL "Win32" OR
+          onnxruntime_target_platform STREQUAL "x86" OR
+          onnxruntime_target_platform STREQUAL "i386" OR
+          onnxruntime_target_platform STREQUAL "i686")
+    set(onnxruntime_target_platform "x86")
+  else()
+    message(FATAL_ERROR "Unknown target platform: ${onnxruntime_target_platform}")
+  endif()
+elseif(APPLE)
+  if(DEFINED CMAKE_OSX_ARCHITECTURES)
+    # We'll only set onnxruntime_target_platform when CMAKE_OSX_ARCHITECTURES specifies a single architecture.
+    list(LENGTH CMAKE_OSX_ARCHITECTURES CMAKE_OSX_ARCHITECTURES_LEN)
+    if(CMAKE_OSX_ARCHITECTURES_LEN EQUAL 1)
+      set(onnxruntime_target_platform ${CMAKE_OSX_ARCHITECTURES})
+    endif()
+  else()
+    set(onnxruntime_target_platform ${CMAKE_SYSTEM_PROCESSOR})
+  endif()
+else()
+  #XXX: Sometimes the value of CMAKE_SYSTEM_PROCESSOR is set but it's wrong. For example, if you run an armv7 docker
+  #image on an aarch64 machine with an aarch64 Ubuntu host OS, in the docker instance cmake may still report
+  # CMAKE_SYSTEM_PROCESSOR as aarch64 by default. Given compiling this code may need more than 2GB memory, we do not
+  # support compiling for ARM32 natively(only support cross-compiling), we will ignore this issue for now.
+  if(NOT CMAKE_SYSTEM_PROCESSOR)
+    message(WARNING "CMAKE_SYSTEM_PROCESSOR is not set. Please set it in your toolchain cmake file.")
+    # Try to detect it
+    if("${CMAKE_C_COMPILER_ID}" STREQUAL "GNU" OR "${CMAKE_C_COMPILER_ID}" STREQUAL "Clang")
+      execute_process(
+          COMMAND "${CMAKE_C_COMPILER}" -dumpmachine
+          OUTPUT_VARIABLE GCC_DUMP_MACHINE_OUT
+          OUTPUT_STRIP_TRAILING_WHITESPACE
+          ERROR_VARIABLE _err
+          RESULT_VARIABLE _res
+      )
+      if(NOT _res EQUAL 0)
+        message(SEND_ERROR "Failed to run 'gcc -dumpmachine':\n ${_res}")
+      endif()
+      string(REPLACE "-" ";" GCC_DUMP_MACHINE_OUT_LIST "${GCC_DUMP_MACHINE_OUT}")
+      list(LENGTH GCC_DUMP_MACHINE_OUT_LIST GCC_TRIPLET_LEN)
+      if(GCC_TRIPLET_LEN EQUAL 4)
+        list(GET GCC_DUMP_MACHINE_OUT_LIST 0 CMAKE_SYSTEM_PROCESSOR)
+        message("Setting CMAKE_SYSTEM_PROCESSOR to ${CMAKE_SYSTEM_PROCESSOR}")
+      endif()
+    endif()
+  endif()
+  set(onnxruntime_target_platform ${CMAKE_SYSTEM_PROCESSOR})
+endif()
+
+if(DEFINED onnxruntime_target_platform)
+  message(STATUS "onnxruntime_target_platform = ${onnxruntime_target_platform}")
+else()
+  message(WARNING "onnxruntime_target_platform is not set")
+endif()
+
+endblock()
diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py
@@ -879,22 +879,8 @@ def generate_build_tree(
     if args.use_snpe:
         cmake_args += ["-Donnxruntime_USE_SNPE=ON"]
 
-    # Set onnxruntime_USE_KLEIDIAI based on:
-    # * Default value above is NO.
-    # * Leave disabled if "no_kleidiai" argument was specified.
-    # * Enable if the target is Android and args.android_abi contains arm64*
-    # * Enable for a Windows cross compile build if compile target is an Arm one.
-    # * Finally enable if platform.machine contains "arm64" and not a WebAssembly build. This should cover the following cases:
-    #     *  Linux on Arm
-    #     *  MacOs (case must be ignored)
-    # * TODO Delegate responsibility for Onnxruntime_USE_KLEIDIAI = ON to CMake logic
     if not args.no_kleidiai:
-        if (
-            (args.android and "arm64" in args.android_abi.lower())
-            or (is_windows() and (args.arm64 or args.arm64ec or args.arm) and platform.architecture()[0] != "AMD64")
-            or ("arm64" in platform.machine().lower() and not args.build_wasm)
-        ):
-            cmake_args += ["-Donnxruntime_USE_KLEIDIAI=ON"]
+        cmake_args += ["-Donnxruntime_USE_KLEIDIAI=ON"]
 
     if args.enable_arm_neon_nchwc:
         cmake_args += ["-Donnxruntime_USE_ARM_NEON_NCHWC=ON"]
