Skip to content

Commit 9a768ab

Browse files
metascroyfacebook-github-bot
authored andcommitted
Fix up CMakeLists and reorganize some code locations
Summary: This diff cleans up CMakeFiles and reorganizes some of the code directories. Archetecture-specific kernel code remains in experimental/kernels/cpu/aarch64. This is not changed. Code related to ops (parallel.h, memory.h, "high-level" op interface, examples/torch_custom_op, etc) is moved to experimental/ops. The example code that builds the custom ops for ATen/ExecuTorch (experimental/kernels/cpu/linear/examples/torch_custom_op) is moved out of examples and into a more descriptive directory: experimental/ops/linear/linear_a8wxdq_op. These are the kernels that will be used in torchchat initially. The "high-level" op interface is in ops/linear/channelwise_8bit_activation_groupwise_lowbit_weight.h. A library for it is defined in the CMakeFiles.txt. It uses no templates and is not specific to the universal kernels and can be used with KleidiAI kernels. The ops in experimental/ops/linear/linear_a8wxdq_op use the high-level library and the universal kernels in experimental/kernels/cpu/aarch64 to define custom ops. Finally, there is a CMakeLists at the top-level directory experimental. The script build_torchao_ops.sh builds the torchao_ops for different platforms: ``` sh build_torchao_ops.sh ATEN ``` For ExecuTorch, you must define the environment variables EXECUTORCH_INCLUDE_DIRS and EXECUTORCH_LIBRARIES manually before calling: ``` EXECUTORCH_INCLUDE_DIRS=$HOME EXECUTORCH_LIBRARIES=$HOME/executorch/cmake-out/lib/libexecutorch_no_prim_ops.a sh build_torchao_ops.sh EXECUTORCH ``` Reviewed By: digantdesai Differential Revision: D62711903
1 parent d267622 commit 9a768ab

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

44 files changed

+521
-431
lines changed

torchao/experimental/CMakeLists.txt

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
project(torchao)
8+
9+
cmake_minimum_required(VERSION 3.19)
10+
11+
set(CMAKE_CXX_STANDARD 17)
12+
13+
if (NOT CMAKE_BUILD_TYPE)
14+
set(CMAKE_BUILD_TYPE Release)
15+
endif()
16+
17+
18+
# Source root directory for torchao/experimental
19+
if(NOT TORCHAO_ROOT)
20+
set(TORCHAO_ROOT ${CMAKE_CURRENT_SOURCE_DIR})
21+
endif()
22+
23+
if(NOT TORCHAO_INCLUDE_DIRS)
24+
set(TORCHAO_INCLUDE_DIRS ${TORCHAO_ROOT}/../..)
25+
endif()
26+
27+
if (NOT TORCHAO_PARALLEL_BACKEND)
28+
if (TORCHAO_OP_TARGET STREQUAL "ATEN")
29+
set(TORCHAO_PARALLEL_BACKEND "ATEN_OPENMP")
30+
elseif(TORCHAO_OP_TARGET STREQUAL "EXECUTORCH")
31+
set(TORCHAO_PARALLEL_BACKEND "PTHREADPOOL")
32+
else()
33+
message(TORCHAO_PARALLEL_BACKEND "TORCHAO_PARALLEL_BACKEND is not set. Please set it directly or set TORCHAO_OP_TARGET to get a default.")
34+
endif()
35+
endif()
36+
37+
include(CMakePrintHelpers)
38+
39+
add_compile_options("-Wall" "-Werror")
40+
41+
include(CMakePrintHelpers)
42+
message("TORCHAO_INCLUDE_DIRS: ${TORCHAO_INCLUDE_DIRS}")
43+
include_directories(${TORCHAO_INCLUDE_DIRS})
44+
45+
if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
46+
# Defines target torchao_kernels_aarch64
47+
add_subdirectory(${TORCHAO_ROOT}/kernels/cpu/aarch64)
48+
add_subdirectory(${TORCHAO_ROOT}/ops/linear)
49+
add_subdirectory(${TORCHAO_ROOT}/ops/linear/linear_a8wxdq_op)
50+
endif()
Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,14 @@
55
# This source code is licensed under the license found in the
66
# LICENSE file in the root directory of this source tree.
77

8-
SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)
9-
export TORCHAO_INCLUDE_DIRS=${SCRIPT_DIR}/../../../../../../..
10-
118
export CMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')"
129
echo "CMAKE_PREFIX_PATH: ${CMAKE_PREFIX_PATH}"
1310
export CMAKE_OUT=/tmp/cmake-out/torchao
14-
cmake -DTORCHAO_INCLUDE_DIRS=${TORCHAO_INCLUDE_DIRS} \
15-
-DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} \
16-
-DPLATFORM="ATEN" \
17-
-S ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/linear/examples/torch_custom_op \
11+
cmake -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} \
12+
-DCMAKE_INSTALL_PREFIX=${CMAKE_OUT} \
13+
-DTORCHAO_OP_TARGET="$1" \
14+
-DEXECUTORCH_LIBRARIES=${EXECUTORCH_LIBRARIES} \
15+
-DEXECUTORCH_INCLUDE_DIRS=${EXECUTORCH_INCLUDE_DIRS} \
16+
-S . \
1817
-B ${CMAKE_OUT}
19-
cmake --build ${CMAKE_OUT}
18+
cmake --build ${CMAKE_OUT} --target install --config Release

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@
44
# This source code is licensed under the license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
add_library(
8-
kernel_aarch64
9-
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/find_min_and_max.cpp
10-
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/compute_sum.cpp
11-
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/quantization/quantize.cpp
12-
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/valpacking/interleave.cpp
13-
)
7+
if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
8+
add_library(
9+
torchao_kernels_aarch64
10+
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/find_min_and_max.cpp
11+
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/compute_sum.cpp
12+
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/quantization/quantize.cpp
13+
${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/valpacking/interleave.cpp
14+
)
15+
endif()

torchao/experimental/kernels/cpu/linear/benchmarks/CMakeLists.txt

Lines changed: 0 additions & 57 deletions
This file was deleted.

torchao/experimental/kernels/cpu/linear/examples/CMakeLists.txt

Lines changed: 0 additions & 38 deletions
This file was deleted.

torchao/experimental/kernels/cpu/linear/examples/torch_custom_op/CMakeLists.txt

Lines changed: 0 additions & 58 deletions
This file was deleted.

torchao/experimental/kernels/cpu/linear/examples/torch_custom_op/run_custom_op.py

Lines changed: 0 additions & 69 deletions
This file was deleted.

torchao/experimental/kernels/cpu/linear/tests/CMakeLists.txt

Lines changed: 0 additions & 41 deletions
This file was deleted.

torchao/experimental/kernels/cpu/linear/tests/build_and_run_tests.sh

Lines changed: 0 additions & 12 deletions
This file was deleted.

0 commit comments

Comments
 (0)