[experimental][kleidi] compile-time preprocessor switch for kleidi tests

pytorch · Oct 11, 2024 · f6e22fb · f6e22fb
1 parent d4bb3ed
commit f6e22fb
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/torchao/experimental/kernels/cpu/aarch64/CMakeLists.txt b/torchao/experimental/kernels/cpu/aarch64/CMakeLists.txt
@@ -20,8 +20,6 @@ FetchContent_MakeAvailable(kleidiai)
 CMAKE_DEPENDENT_OPTION(BUILD_KLEIDI "Download, build, and link against Arm KleidiAI library"
   OFF "CMAKE_SYSTEM_PROCESSOR STREQUAL \"arm64\"" ON)
 
-add_compile_definitions("TORCHAO_ENABLE_KLEIDI=$<BOOL:${BUILD_KLEIDI}>")
-
 if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
   add_library(
     torchao_kernels_aarch64
@@ -31,6 +29,9 @@ if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
     ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/valpacking/interleave.cpp
   )
   if (BUILD_KLEIDI)
+    # Temporarily exposing this to the parent scope until we wire
+    # this up properly from the top level
+    set(TORCHAO_ENABLE_KLEIDI ON PARENT_SCOPE)
     message(STATUS "Building with Kleidi")
     target_link_libraries(torchao_kernels_aarch64 PUBLIC kleidiai)
   endif()

diff --git a/torchao/experimental/kernels/cpu/aarch64/tests/CMakeLists.txt b/torchao/experimental/kernels/cpu/aarch64/tests/CMakeLists.txt
@@ -32,8 +32,15 @@ add_library(
 if(NOT TORCHAO_INCLUDE_DIRS)
   set(TORCHAO_INCLUDE_DIRS ${TORCHAO_LIBRARIES})
 endif()
+
 add_subdirectory(${TORCHAO_LIBRARIES}/torchao/experimental/kernels/cpu/aarch64 ${CMAKE_CURRENT_BINARY_DIR}/torchao_kernels_aarch64)
 
+# The TORCHAO_ENABLE_KLEIDI cmake variable should be set by `torchao_kernels_aarch64"
+# This is a temporary work around.
+if(TORCHAO_ENABLE_KLEIDI)
+  add_compile_definitions(TORCHAO_ENABLE_KLEIDI)
+endif()
+
 enable_testing()
 
 add_executable(test_quantization test_quantization.cpp)

diff --git a/torchao/experimental/kernels/cpu/aarch64/tests/test_linear.cpp b/torchao/experimental/kernels/cpu/aarch64/tests/test_linear.cpp
@@ -14,8 +14,10 @@
 #include <torchao/experimental/kernels/cpu/aarch64/linear/linear.h>
 #include <torchao/experimental/kernels/cpu/aarch64/tests/test_utils.h>
 
+#ifdef TORCHAO_ENABLE_KLEIDI
 #include <torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h>
 #include <torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p8x8_1x8x32_neon_dotprod.h>
+#endif
 
 float kTol = 0.0001;
 
@@ -354,9 +356,7 @@ TEST(
   }
 }
 
-// #ifdef TORCHAO_ENABLE_KLEIDI
-// TODO: Wire up the the compile defination for TORCHAO_ENABLE_KLEIDI
-
+#ifdef TORCHAO_ENABLE_KLEIDI
 template <bool has_bias, bool has_clamp>
 void test_kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod(
     int m,
@@ -585,5 +585,5 @@ TEST(
       true /*has_clamp*/>(
       /*m=*/11, /*k=*/128, /*n=*/182, /*group_size=*/128);
 }
-// #endif // defined(TORCHAO_ENABLE_KLEIDI)
+#endif // TORCHAO_ENABLE_KLEIDI
 #endif // defined(__aarch64__) || defined(__ARM_NEON)