Define torchao op library by srcs instead of object libraries

torchao/experimental/CMakeLists.txt

@@ -36,9 +36,13 @@ endif()
 add_compile_options("-Wall" "-Werror" "-Wno-deprecated" "-Wno-shorten-64-to-32")
 
 include(CMakePrintHelpers)
+include(${CMAKE_CURRENT_SOURCE_DIR}/Utils.cmake)
+
 message("TORCHAO_INCLUDE_DIRS: ${TORCHAO_INCLUDE_DIRS}")
 include_directories(${TORCHAO_INCLUDE_DIRS})
 
+
+# Build cpu/aarch64 kernels
 if(TORCHAO_BUILD_CPU_AARCH64)
     message(STATUS "Building with cpu/aarch64")
     add_compile_definitions(TORCHAO_BUILD_CPU_AARCH64)
@@ -77,29 +81,50 @@ if(TORCHAO_BUILD_CPU_AARCH64)
         add_compile_definitions(TORCHAO_ENABLE_ARM_I8MM)
     endif()
 
-    # 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)
     endif()
+
+    # Defines torchao_kernels_aarch64
+    add_subdirectory(kernels/cpu/aarch64)
 endif()
 
-# Add quantized operation dir
-add_subdirectory(ops/linear_8bit_act_xbit_weight)
-add_subdirectory(ops/embedding_xbit)
-
-# ATen ops lib
-if (TORCHAO_BUILD_ATEN_OPS)
-    add_library(torchao_ops_aten SHARED)
-    target_link_libraries(
-        torchao_ops_aten PRIVATE
-        torchao_ops_linear_8bit_act_xbit_weight_aten
-        torchao_ops_embedding_xbit_aten
+
+
+if (NOT TARGET cpuinfo)
+    # For some reason cpuinfo package has unused functions/variables
+    # TODO (T215533422): fix upstream
+    add_compile_options(-Wno-unused-function -Wno-unused-variable)
+    include(FetchContent)
+    FetchContent_Declare(cpuinfo
+        GIT_REPOSITORY https://github.com/pytorch/cpuinfo.git
+        GIT_TAG c61fe919607bbc534d7a5a5707bdd7041e72c5ff)
+    FetchContent_MakeAvailable(
+        cpuinfo)
+endif()
+
+# Build ATen ops
+if(TORCHAO_BUILD_ATEN_OPS)
+    find_package(Torch REQUIRED)
+    set(_torchao_op_srcs_aten)
+    list(APPEND _torchao_op_srcs_aten
+        ops/embedding_xbit/op_embedding_xbit_aten.cpp
+        ops/linear_8bit_act_xbit_weight/linear_8bit_act_xbit_weight.cpp
+        ops/linear_8bit_act_xbit_weight/op_linear_8bit_act_xbit_weight_aten.cpp
     )
+    list(TRANSFORM _torchao_op_srcs_aten PREPEND "${CMAKE_CURRENT_SOURCE_DIR}/")
+    add_library(torchao_ops_aten SHARED ${_torchao_op_srcs_aten})
+    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)
+    endif()
+    target_link_libraries(torchao_ops_aten PRIVATE cpuinfo)
+    target_include_directories(torchao_ops_aten PRIVATE "${TORCH_INCLUDE_DIRS}")
+    target_link_libraries(torchao_ops_aten PRIVATE "${TORCH_LIBRARIES}")
+    target_compile_definitions(torchao_ops_aten PRIVATE USE_ATEN=1)
 
-    # Add MPS support if enabled
+     # Add MPS support if enabled
     if (TORCHAO_BUILD_MPS_OPS)
         message(STATUS "Building with MPS support")
         add_subdirectory(ops/mps)
@@ -114,18 +139,38 @@ if (TORCHAO_BUILD_ATEN_OPS)
     )
 endif()
 
-# Build executorch lib if enabled
+
+# Build ExecuTorch ops
 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
-        torchao_ops_embedding_xbit_executorch
+    # 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
+    if(NOT DEFINED EXECUTORCH_INCLUDE_DIRS AND NOT DEFINED EXECUTORCH_LIBRARIES)
+        message(WARNING "EXECUTORCH_INCLUDE_DIRS and EXECUTORCH_LIBRARIES are not defined. Looking for ExecuTorch.")
+        find_package(ExecuTorch HINTS ${CMAKE_PREFIX_PATH}/executorch/share/cmake)
+    endif()
+    set(_torchao_op_srcs_executorch)
+    list(APPEND _torchao_op_srcs_executorch
+        ops/embedding_xbit/op_embedding_xbit_executorch.cpp
+        ops/linear_8bit_act_xbit_weight/linear_8bit_act_xbit_weight.cpp
+        ops/linear_8bit_act_xbit_weight/op_linear_8bit_act_xbit_weight_executorch.cpp
     )
+    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)
+    target_link_libraries(torchao_ops_executorch PRIVATE "${EXECUTORCH_LIBRARIES}")
+    if (TORCHAO_BUILD_CPU_AARCH64)
+        target_link_libraries(torchao_ops_executorch PRIVATE torchao_kernels_aarch64)
+    endif()
+    target_link_libraries(torchao_ops_executorch PRIVATE cpuinfo)
     install(
         TARGETS
             torchao_ops_executorch
-            torchao_ops_linear_8bit_act_xbit_weight_executorch
-            torchao_ops_embedding_xbit_executorch
         EXPORT _targets
         DESTINATION lib
     )