diff --git a/CMakeLists.txt b/CMakeLists.txt index eb4c196668a..e5f0361a330 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -278,6 +278,30 @@ if(EXECUTORCH_BUILD_PTHREADPOOL) ) endif() +if(EXECUTORCH_BUILD_KERNELS_TORCHAO) + set(TORCHAO_BUILD_ATEN_OPS OFF) + set(TORCHAO_BUILD_EXECUTORCH_OPS ON) + set(TORCHAO_BUILD_CPU_AARCH64 ON) + set(TORCHAO_ENABLE_ARM_NEON_DOT ON) + + list(APPEND TORCHAO_INCLUDE_DIRS + ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/pthreadpool/include + ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/cpuinfo/include + ${EXECUTORCH_ROOT}/third-party/ao + ) + + set(EXECUTORCH_INCLUDE_DIRS ${TORCHAO_INCLUDE_DIRS}) + + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/ao/torchao/experimental) + executorch_target_link_options_shared_lib(torchao_ops_executorch) + list(APPEND _executorch_kernels torchao_ops_executorch) +endif() + +if(EXECUTORCH_BUILD_TESTS) + set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON) + include(CTest) +endif() + # TODO(dbort): Fix these warnings and remove this flag. set(_common_compile_options -Wno-deprecated-declarations -fPIC) diff --git a/scripts/build_apple_frameworks.sh b/scripts/build_apple_frameworks.sh index 36dc032862c..7e85e2b4b88 100755 --- a/scripts/build_apple_frameworks.sh +++ b/scripts/build_apple_frameworks.sh @@ -125,6 +125,11 @@ libquantized_kernels.a,\ libquantized_ops_lib.a,\ :" +FRAMEWORK_KERNELS_TORCHAO="kernels_torchao:\ +libtorchao_ops_executorch.a,\ +libtorchao_kernels_aarch64.a,\ +:" + usage() { echo "Usage: $0 [OPTIONS]" echo "Build frameworks for Apple platforms." @@ -137,6 +142,7 @@ usage() { echo " --mps Only build the Metal Performance Shaders backend." echo " --optimized Only build the Optimized kernels." echo " --quantized Only build the Quantized kernels." + echo " --torchao Only build the TorchAO kernels." echo " --xnnpack Only build the XNNPACK backend." echo exit 0 @@ -154,6 +160,7 @@ set_cmake_options_override() { "-DEXECUTORCH_BUILD_MPS=OFF" "-DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=OFF" "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED=OFF" + "-DEXECUTORCH_BUILD_KERNELS_TORCHAO=OFF" "-DEXECUTORCH_BUILD_XNNPACK=OFF" ) fi @@ -184,6 +191,7 @@ for arg in "$@"; do --mps) set_cmake_options_override "EXECUTORCH_BUILD_MPS" ;; --optimized) set_cmake_options_override "EXECUTORCH_BUILD_KERNELS_OPTIMIZED" ;; --quantized) set_cmake_options_override "EXECUTORCH_BUILD_KERNELS_QUANTIZED" ;; + --torchao) set_cmake_options_override "EXECUTORCH_BUILD_KERNELS_TORCHAO" ;; --xnnpack) set_cmake_options_override "EXECUTORCH_BUILD_XNNPACK" ;; *) echo -e "\033[31m[error] unknown option: ${arg}\033[0m" @@ -311,6 +319,7 @@ for mode in "${MODES[@]}"; do append_framework_flag "EXECUTORCH_BUILD_KERNELS_LLM" "$FRAMEWORK_KERNELS_LLM" "$mode" append_framework_flag "EXECUTORCH_BUILD_KERNELS_OPTIMIZED" "$FRAMEWORK_KERNELS_OPTIMIZED" "$mode" append_framework_flag "EXECUTORCH_BUILD_KERNELS_QUANTIZED" "$FRAMEWORK_KERNELS_QUANTIZED" "$mode" + append_framework_flag "EXECUTORCH_BUILD_KERNELS_TORCHAO" "$FRAMEWORK_KERNELS_TORCHAO" "$mode" cd "${OUTPUT_DIR}" "$SOURCE_ROOT_DIR"/scripts/create_frameworks.sh "${FRAMEWORK_FLAGS[@]}" diff --git a/tools/cmake/executorch-config.cmake b/tools/cmake/executorch-config.cmake index d87d3693ad8..6c27e8ba616 100644 --- a/tools/cmake/executorch-config.cmake +++ b/tools/cmake/executorch-config.cmake @@ -84,6 +84,8 @@ set(optional_lib_list quantized_kernels quantized_ops_lib quantized_ops_aot_lib + torchao_ops_executorch + torchao_kernels_aarch64 ) foreach(lib ${optional_lib_list}) diff --git a/tools/cmake/preset/apple_common.cmake b/tools/cmake/preset/apple_common.cmake index d58cc44a751..27212a166ed 100644 --- a/tools/cmake/preset/apple_common.cmake +++ b/tools/cmake/preset/apple_common.cmake @@ -29,3 +29,4 @@ set_overridable_option(EXECUTORCH_BUILD_EXTENSION_TENSOR ON) set_overridable_option(EXECUTORCH_BUILD_KERNELS_LLM ON) set_overridable_option(EXECUTORCH_BUILD_KERNELS_OPTIMIZED ON) set_overridable_option(EXECUTORCH_BUILD_KERNELS_QUANTIZED ON) +set_overridable_option(EXECUTORCH_BUILD_KERNELS_TORCHAO ON)