up

Author: metascroy

setup.py

@@ -69,11 +69,11 @@ def use_debug_mode():
 
 import torch
 from torch.utils.cpp_extension import (
+    CUDA_HOME,
+    IS_WINDOWS,
     BuildExtension,
     CppExtension,
-    CUDA_HOME,
     CUDAExtension,
-    IS_WINDOWS,
 )
 
 

torchao/experimental/CMakeLists.txt

@@ -37,18 +37,15 @@ if(TORCHAO_BUILD_CPU_AARCH64)
     message(STATUS "Building with cpu/aarch64")
     add_compile_definitions(TORCHAO_BUILD_CPU_AARCH64)
 
+    # Defines torchao_kernels_aarch64
+    add_subdirectory(kernels/cpu/aarch64)
+
     if(TORCHAO_BUILD_KLEIDIAI)
         message(STATUS "Building with Arm KleidiAI library")
-        add_compile_definitions(TORCHAO_ENABLE_KLEIDI=1)
+        add_compile_definitions(TORCHAO_ENABLE_KLEIDI)
     endif()
 endif()
 
-
-if (TORCHAO_BUILD_CPU_AARCH64)
-    # Defines torchao_kernels_aarch64
-    add_subdirectory(kernels/cpu/aarch64)
-endif()
-
 add_subdirectory(ops/linear_8bit_act_xbit_weight)
 add_subdirectory(ops/embedding_xbit)
 
@@ -78,7 +75,6 @@ if(TORCHAO_BUILD_EXECUTORCH_OPS)
     install(
         TARGETS
             torchao_ops_executorch
-            # torchao_kernels_aarch64
             torchao_ops_linear_8bit_act_xbit_weight_executorch
             torchao_ops_embedding_xbit_executorch
         EXPORT _targets

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

@@ -22,9 +22,6 @@ if (TORCHAO_BUILD_CPU_AARCH64)
           GIT_TAG v1.2.0)
     FetchContent_MakeAvailable(kleidiai)
 
-    # Temporarily exposing this to the parent scope until we wire
-    # this up properly from the top level
-    set(TORCHAO_BUILD_KLEIDI ON PARENT_SCOPE)
     target_link_libraries(torchao_kernels_aarch64 PUBLIC kleidiai)
   endif()
 

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

@@ -40,8 +40,7 @@ endif()
 
 add_subdirectory(${TORCHAO_LIBRARIES}/torchao/experimental/kernels/cpu/aarch64 ${CMAKE_CURRENT_BINARY_DIR}/torchao_kernels_aarch64)
 
-# The TORCHAO_BUILD_KLEIDI cmake variable should be set by `torchao_kernels_aarch64"
-if(TORCHAO_BUILD_KLEIDI)
+if(TORCHAO_BUILD_KLEIDIAI)
   add_compile_definitions(TORCHAO_ENABLE_KLEIDI)
 endif()