Add build flag to set parallel_backend (#1870)

up

Author: metascroy

torchao/experimental/CMakeLists.txt

@@ -24,6 +24,10 @@ if(NOT TORCHAO_INCLUDE_DIRS)
   set(TORCHAO_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/../..)
 endif()
 
+if(NOT DEFINED TORCHAO_PARALLEL_BACKEND)
+    set(TORCHAO_PARALLEL_BACKEND aten_openmp)
+endif()
+
 include(CMakePrintHelpers)
 
 add_compile_options("-Wall" "-Werror" "-Wno-deprecated")

torchao/experimental/ops/embedding_xbit/CMakeLists.txt

@@ -12,7 +12,7 @@ find_package(Torch REQUIRED)
 add_library(torchao_ops_embedding_xbit_aten OBJECT
     op_embedding_xbit_aten.cpp
 )
-target_link_torchao_parallel_backend(torchao_ops_embedding_xbit_aten "aten_openmp")
+target_link_torchao_parallel_backend(torchao_ops_embedding_xbit_aten "${TORCHAO_PARALLEL_BACKEND}")
 if (TORCHAO_BUILD_CPU_AARCH64)
     target_link_libraries(torchao_ops_embedding_xbit_aten PRIVATE torchao_kernels_aarch64)
 endif()

torchao/experimental/ops/linear_8bit_act_xbit_weight/CMakeLists.txt

@@ -24,7 +24,7 @@ add_library(torchao_ops_linear_8bit_act_xbit_weight_aten OBJECT
     linear_8bit_act_xbit_weight.cpp
     op_linear_8bit_act_xbit_weight_aten.cpp
 )
-target_link_torchao_parallel_backend(torchao_ops_linear_8bit_act_xbit_weight_aten aten_openmp)
+target_link_torchao_parallel_backend(torchao_ops_linear_8bit_act_xbit_weight_aten "${TORCHAO_PARALLEL_BACKEND}")
 
 if(TORCHAO_BUILD_CPU_AARCH64)
     target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE torchao_kernels_aarch64)