Update cmake files

torchao/experimental/CMakeLists.txt

@@ -3,6 +3,7 @@
 #
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
+
 cmake_minimum_required(VERSION 3.19)
 
 project(torchao)
@@ -13,14 +14,13 @@ if (NOT CMAKE_BUILD_TYPE)
     set(CMAKE_BUILD_TYPE Release)
 endif()
 
-option(TORCHAO_OP_EXECUTORCH_BUILD "Building torchao ops for ExecuTorch." OFF)
+option(TORCHAO_BUILD_EXECUTORCH_OPS "Building torchao ops for ExecuTorch." OFF)
 
 
 if(NOT TORCHAO_INCLUDE_DIRS)
   set(TORCHAO_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/../..)
 endif()
 
-
 include(CMakePrintHelpers)
 
 add_compile_options("-Wall" "-Werror" "-Wno-deprecated")
@@ -41,14 +41,15 @@ if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
     )
     install(
         TARGETS torchao_ops_aten
+        EXPORT _targets
         DESTINATION lib
     )
-    if(TORCHAO_OP_EXECUTORCH_BUILD)
+    if(TORCHAO_BUILD_EXECUTORCH_OPS)
         add_library(torchao_ops_executorch STATIC)
         target_link_libraries(torchao_ops_executorch PRIVATE torchao_ops_linear_8bit_act_xbit_weight_executorch)
-
         install(
-            TARGETS torchao_ops_executorch
+            TARGETS torchao_ops_executorch torchao_ops_linear_8bit_act_xbit_weight_executorch torchao_kernels_aarch64
+            EXPORT _targets
             DESTINATION lib
         )
     endif()

torchao/experimental/build_torchao_ops.sh

@@ -13,14 +13,14 @@ TARGET="${1}"
 export CMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
 echo "CMAKE_PREFIX_PATH: ${CMAKE_PREFIX_PATH}"
 if [[ $TARGET == "executorch" ]]; then
-    TORCHAO_OP_EXECUTORCH_BUILD=ON
+    TORCHAO_BUILD_EXECUTORCH_OPS=ON
 else
-    TORCHAO_OP_EXECUTORCH_BUILD=OFF
+    TORCHAO_BUILD_EXECUTORCH_OPS=OFF
 fi
-export CMAKE_OUT=cmake-out/torchao
+export CMAKE_OUT=cmake-out
 cmake -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} \
     -DCMAKE_INSTALL_PREFIX=${CMAKE_OUT} \
-    -DTORCHAO_OP_EXECUTORCH_BUILD="${TORCHAO_OP_EXECUTORCH_BUILD}" \
+    -DTORCHAO_BUILD_EXECUTORCH_OPS="${TORCHAO_BUILD_EXECUTORCH_OPS}" \
     -S . \
     -B ${CMAKE_OUT}
-cmake --build  ${CMAKE_OUT} --target install --config Release
+cmake --build  ${CMAKE_OUT} -j 16 --target install --config Release

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

@@ -4,21 +4,8 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-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 35e156d62d1d7e4d27a39f56ed7770a665628b31) # same as xnnpack for now, TODO - revisit this
-
-FetchContent_MakeAvailable(kleidiai)
-
-# Disabled by default. Force enable if we are on a suitable system.
-# TODO: Introduce ISA specific flags for i8mm.
-CMAKE_DEPENDENT_OPTION(BUILD_KLEIDI "Download, build, and link against Arm KleidiAI library"
-  OFF "CMAKE_SYSTEM_PROCESSOR STREQUAL \"arm64\"" ON)
+option(TORCHAO_BUILD_KLEIDIAI "Download, build, and link against Arm KleidiAI library" OFF)
 
 if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
   add_library(
@@ -28,7 +15,16 @@ if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
     ${CMAKE_CURRENT_SOURCE_DIR}/quantization/quantize.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/valpacking/interleave.cpp
   )
-  if (BUILD_KLEIDI)
+  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 35e156d62d1d7e4d27a39f56ed7770a665628b31) # same as xnnpack for now, TODO - revisit this
+    FetchContent_MakeAvailable(kleidiai)
+
     # Temporarily exposing this to the parent scope until we wire
     # this up properly from the top level
     set(TORCHAO_ENABLE_KLEIDI ON PARENT_SCOPE)

torchao/experimental/ops/linear_8bit_act_xbit_weight/CMakeLists.txt

@@ -19,20 +19,18 @@ target_include_directories(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE
 target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE "${TORCH_LIBRARIES}")
 target_compile_definitions(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE USE_ATEN=1)
 
-if(TORCHAO_OP_EXECUTORCH_BUILD)
-    find_package(ExecuTorch REQUIRED HINTS ${CMAKE_PREFIX_PATH}/executorch/share/cmake)
+if(TORCHAO_BUILD_EXECUTORCH_OPS)
+    # ExecuTorch package is not required, but EXECUTORCH_INCLUDE_DIRS and EXECUTORCH_LIBRARIES must
+    # be defined and EXECUTORCH_LIBRARIES must include the following libraries installed by ExecuTorch:
+    # libexecutorch.a
+    # libextension_threadpool.a
+    # libcpuinfo.a
+    # libpthreadpool.a
+    find_package(ExecuTorch HINTS ${CMAKE_PREFIX_PATH}/executorch/share/cmake)
+    file(GLOB _SRCS "${CMAKE_CURRENT_SOURCE_DIR}/op_linear_8bit_act_xbit_weight_executorch/*.cpp")
     add_library(torchao_ops_linear_8bit_act_xbit_weight_executorch OBJECT
         linear_8bit_act_xbit_weight.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w2s.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w2sz.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w3s.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w3sz.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w4s.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w4sz.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w5s.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w5sz.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w6s.cpp
-        op_linear_8bit_act_xbit_weight_executorch/w6sz.cpp
+        ${_SRCS}
     )
     target_link_torchao_parallel_backend(torchao_ops_linear_8bit_act_xbit_weight_executorch executorch)
     target_include_directories(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE "${EXECUTORCH_INCLUDE_DIRS}")