Update KleidiAI (#2692)

* Update KleidiAI

* up

* up

* up

Author: metascroy

torchao/experimental/CMakeLists.txt

@@ -90,24 +90,40 @@ if(TORCHAO_BUILD_CPU_AARCH64)
     add_subdirectory(kernels/cpu/aarch64)
 endif()
 
-
-
 if (NOT TARGET cpuinfo)
     # For some reason cpuinfo package has unused functions/variables
     # TODO (T215533422): fix upstream
-    set(CPUINFO_BUILD_UNIT_TESTS OFF CACHE BOOL "Disable unit tests" FORCE)
-    set(CPUINFO_BUILD_MOCK_TESTS OFF CACHE BOOL "Disable mock tests" FORCE)
-    set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "Disable benchmarks" FORCE)
     add_compile_options(-Wno-unused-function -Wno-unused-variable)
     set(CMAKE_POLICY_VERSION_MINIMUM 3.5)
     include(FetchContent)
+    set(CPUINFO_BUILD_UNIT_TESTS OFF CACHE BOOL "" FORCE)
+    set(CPUINFO_BUILD_MOCK_TESTS OFF CACHE BOOL "" FORCE)
+    set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "" FORCE)
     FetchContent_Declare(cpuinfo
         GIT_REPOSITORY https://github.com/pytorch/cpuinfo.git
-        GIT_TAG c61fe919607bbc534d7a5a5707bdd7041e72c5ff)
+        GIT_TAG c61fe919607bbc534d7a5a5707bdd7041e72c5ff
+    )
     FetchContent_MakeAvailable(
         cpuinfo)
 endif()
 
+if (TORCHAO_BUILD_KLEIDIAI)
+    if (NOT TARGET kleidiai)
+        include(FetchContent)
+        # KleidiAI is an open-source library that provides optimized
+        # performance-critical routines, also known as micro-kernels, for artificial
+        # intelligence (AI) workloads tailored for Arm® CPUs.
+        set(KLEIDIAI_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+        set(KLEIDIAI_BUILD_BENCHMARKS OFF CACHE BOOL "" FORCE)
+        FetchContent_Declare(kleidiai
+            GIT_REPOSITORY https://git.gitlab.arm.com/kleidi/kleidiai.git
+            GIT_TAG v1.12.0
+        )
+        FetchContent_MakeAvailable(kleidiai)
+    endif()
+endif()
+
+
 # Build ATen ops
 if(TORCHAO_BUILD_ATEN_OPS)
     find_package(Torch REQUIRED)
@@ -124,6 +140,9 @@ if(TORCHAO_BUILD_ATEN_OPS)
     target_link_torchao_parallel_backend(torchao_ops_aten "${TORCHAO_PARALLEL_BACKEND}")
     if (TORCHAO_BUILD_CPU_AARCH64)
         target_link_libraries(torchao_ops_aten PRIVATE torchao_kernels_aarch64)
+        if (TORCHAO_BUILD_KLEIDIAI)
+            target_link_libraries(torchao_ops_aten PRIVATE kleidiai)
+        endif()
     endif()
     target_link_libraries(torchao_ops_aten PRIVATE cpuinfo)
     target_include_directories(torchao_ops_aten PRIVATE "${TORCH_INCLUDE_DIRS}")
@@ -168,17 +187,16 @@ if(TORCHAO_BUILD_EXECUTORCH_OPS)
 
     list(TRANSFORM _torchao_op_srcs_executorch PREPEND "${CMAKE_CURRENT_SOURCE_DIR}/")
     add_library(torchao_ops_executorch STATIC ${_torchao_op_srcs_executorch})
-    target_link_torchao_parallel_backend(torchao_ops_executorch executorch)
-    target_include_directories(torchao_ops_executorch PRIVATE "${EXECUTORCH_INCLUDE_DIRS}")
+
     target_compile_definitions(torchao_ops_executorch PRIVATE USE_EXECUTORCH=1)
+
+    # This links to ExecuTorch
+    target_link_torchao_parallel_backend(torchao_ops_executorch executorch)
     if (TORCHAO_BUILD_CPU_AARCH64)
         target_link_libraries(torchao_ops_executorch PRIVATE torchao_kernels_aarch64)
+        if (TORCHAO_BUILD_KLEIDIAI)
+            target_link_libraries(torchao_ops_executorch PRIVATE kleidiai)
+        endif()
     endif()
     target_link_libraries(torchao_ops_executorch PRIVATE cpuinfo)
-    install(
-        TARGETS
-            torchao_ops_executorch
-        EXPORT _targets
-        DESTINATION lib
-    )
 endif()

torchao/experimental/kernels/cpu/aarch64/CMakeLists.txt

@@ -12,21 +12,4 @@ if (TORCHAO_BUILD_CPU_AARCH64)
     ${CMAKE_CURRENT_SOURCE_DIR}/quantization/quantize.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/valpacking/interleave.cpp
   )
-  if (TORCHAO_BUILD_KLEIDIAI)
-    include(FetchContent)
-    # KleidiAI is an open-source library that provides optimized
-    # performance-critical routines, also known as micro-kernels, for artificial
-    # intelligence (AI) workloads tailored for Arm® CPUs.
-    FetchContent_Declare(kleidiai
-          GIT_REPOSITORY https://git.gitlab.arm.com/kleidi/kleidiai.git
-          GIT_TAG v1.5.0)
-    FetchContent_MakeAvailable(kleidiai)
-
-    target_link_libraries(torchao_kernels_aarch64 PUBLIC kleidiai)
-  endif()
-
-install(
-  TARGETS torchao_kernels_aarch64
-  DESTINATION lib
-)
 endif()

torchao/experimental/ops/tests/CMakeLists.txt

@@ -29,6 +29,20 @@ endif()
 
 if(TORCHAO_BUILD_KLEIDIAI)
     add_compile_definitions(TORCHAO_ENABLE_KLEIDI=1)
+    # TODO: build tests at top-level so we can use same KleidiAI version
+    if (NOT TARGET kleidiai)
+        include(FetchContent)
+        # KleidiAI is an open-source library that provides optimized
+        # performance-critical routines, also known as micro-kernels, for artificial
+        # intelligence (AI) workloads tailored for Arm® CPUs.
+        set(KLEIDIAI_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+        set(KLEIDIAI_BUILD_BENCHMARKS OFF CACHE BOOL "" FORCE)
+        FetchContent_Declare(kleidiai
+            GIT_REPOSITORY https://git.gitlab.arm.com/kleidi/kleidiai.git
+            GIT_TAG v1.12.0
+        )
+        FetchContent_MakeAvailable(kleidiai)
+    endif()
 endif()
 
 if(TORCHAO_BUILD_ARM_I8MM)
@@ -80,6 +94,13 @@ if (TORCHAO_BUILD_CPU_AARCH64)
     torchao_kernels_aarch64
   )
 endif()
+if (TORCHAO_BUILD_KLEIDIAI)
+  target_link_libraries(
+    test_linear_8bit_act_xbit_weight
+    PRIVATE
+    kleidiai
+  )
+endif()
 target_link_torchao_parallel_backend(test_linear_8bit_act_xbit_weight "${TORCHAO_PARALLEL_BACKEND}")
 
 add_executable(