Skip to content

Commit 906cd6d

Browse files
larryliu0820facebook-github-bot
larryliu0820
authored and
facebook-github-bot
committed
Provide CMake build for custom ops
Summary: As titled. Adding CMakeLists.txt and an option to register ops in `examples/custom_ops`. Reviewed By: dbort Differential Revision: D48103652 fbshipit-source-id: 2d4c57703f841383ee19e80b2a7e989dbab439b8
1 parent f9a1fd0 commit 906cd6d

File tree

5 files changed

+166
-10
lines changed

5 files changed

+166
-10
lines changed

‎CMakeLists.txt

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,10 @@
4343
cmake_minimum_required(VERSION 3.13)
4444
project(executorch)
4545

46+
# option to register custom ops in `examples/custom_ops`
47+
option(REGISTER_EXAMPLE_CUSTOM_OPS
48+
"Register custom ops defined in examples/custom_ops" OFF)
49+
4650
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
4751
if(NOT CMAKE_CXX_STANDARD)
4852
set(CMAKE_CXX_STANDARD 17)
@@ -268,3 +272,9 @@ add_executable(executor_runner ${_executor_runner__srcs})
268272
target_link_libraries(executor_runner executorch portable_kernels_bindings
269273
gflags)
270274
target_compile_options(executor_runner PUBLIC ${_common_compile_options})
275+
276+
# Generate custom_ops_lib based on REGISTER_EXAMPLE_CUSTOM_OPS
277+
if(REGISTER_EXAMPLE_CUSTOM_OPS)
278+
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/examples/custom_ops)
279+
target_link_libraries(executor_runner custom_ops_lib)
280+
endif()

‎examples/custom_ops/CMakeLists.txt

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# Example CMakeLists.txt for registering custom ops into Executorch. In this
8+
# example we have custom ops `my_ops::mul3.out` implemented in C++ in
9+
# `examples/custom_ops/custom_ops_1.cpp`. We also have it registered into EXIR
10+
# in `examples/custom_ops/custom_ops_1.py`. This CMakeLists.txt runs a script to
11+
# generate wrapper code based on the operator-kernel binding defined in
12+
# `examples/custom_ops/custom_ops.yaml`. Then creates a library that contains
13+
# both binding wrapper and the implementation source file. This library can be
14+
# linked into Executorch binary (`executor_runner` in this example) and it is
15+
# ready to run models containing that custom op.
16+
cmake_minimum_required(VERSION 3.13)
17+
18+
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
19+
if(NOT CMAKE_CXX_STANDARD)
20+
set(CMAKE_CXX_STANDARD 17)
21+
endif()
22+
23+
if(NOT PYTHON_EXECUTABLE)
24+
set(PYTHON_EXECUTABLE python3)
25+
endif()
26+
# Source root directory for executorch.
27+
if(NOT EXECUTORCH_ROOT)
28+
set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
29+
endif()
30+
# Source root directory for pytorch.
31+
if(NOT TORCH_ROOT)
32+
set(TORCH_ROOT ${EXECUTORCH_ROOT}/third-party/pytorch)
33+
endif()
34+
# Command to generate selected_operators.yaml from custom_ops.yaml.
35+
set(_oplist_yaml ${CMAKE_CURRENT_BINARY_DIR}/selected_operators.yaml)
36+
file(GLOB_RECURSE _codegen_tools_srcs "${EXECUTORCH_ROOT}/codegen/tools/*.py")
37+
file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*")
38+
file(GLOB_RECURSE _torchgen_srcs "${TORCH_ROOT}/torchgen/*.py")
39+
40+
set(_gen_oplist_command
41+
"${PYTHON_EXECUTABLE}" -m codegen.tools.gen_oplist
42+
--output_path=${_oplist_yaml}
43+
--ops_schema_yaml_path=${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml)
44+
45+
# Command to codegen C++ wrappers to register custom ops to both PyTorch and
46+
# Executorch runtime.
47+
set(_gen_command
48+
"${PYTHON_EXECUTABLE}" -m torchgen.gen_executorch
49+
--source-path=${EXECUTORCH_ROOT}/codegen
50+
--install-dir=${CMAKE_CURRENT_BINARY_DIR}
51+
--tags-path=${TORCH_ROOT}/aten/src/ATen/native/tags.yaml
52+
--aten-yaml-path=${TORCH_ROOT}/aten/src/ATen/native/native_functions.yaml
53+
--op-selection-yaml-path=${_oplist_yaml}
54+
--custom-ops-yaml-path=${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml)
55+
56+
set(_gen_command_sources
57+
${CMAKE_CURRENT_BINARY_DIR}/RegisterCodegenUnboxedKernelsEverything.cpp
58+
${CMAKE_CURRENT_BINARY_DIR}/RegisterCPUCustomOps.cpp
59+
${CMAKE_CURRENT_BINARY_DIR}/RegisterSchema.cpp
60+
${CMAKE_CURRENT_BINARY_DIR}/Functions.h
61+
${CMAKE_CURRENT_BINARY_DIR}/NativeFunctions.h
62+
${CMAKE_CURRENT_BINARY_DIR}/CustomOpsNativeFunctions.h)
63+
message(STATUS "Generating selected operator list ${_gen_oplist_command}")
64+
65+
add_custom_command(
66+
COMMENT "Generating selected_operators.yaml for custom ops"
67+
OUTPUT ${_oplist_yaml}
68+
COMMAND ${_gen_oplist_command}
69+
DEPENDS ${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml ${_codegen_tools_srcs}
70+
WORKING_DIRECTORY ${EXECUTORCH_ROOT})
71+
72+
add_custom_command(
73+
COMMENT "Generating code for custom operator registration"
74+
OUTPUT ${_gen_command_sources}
75+
COMMAND ${_gen_command}
76+
DEPENDS ${_oplist_yaml} ${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml
77+
${_codegen_templates} ${_torchgen_srcs}
78+
WORKING_DIRECTORY ${EXECUTORCH_ROOT})
79+
# Prepare for C++ libraries.
80+
81+
# 1. TODO: C++ library to register custom ops into PyTorch.
82+
# ~~~
83+
# add_library(custom_ops_aot_lib SHARED
84+
# ${OUTPUT_DIRECTORY}/RegisterCPUCustomOps.cpp
85+
# ${OUTPUT_DIRECTORY}/RegisterSchema.cpp
86+
# ${OUTPUT_DIRECTORY}/CustomOpsNativeFunctions.h)
87+
# ~~~
88+
89+
# Find `Torch`.
90+
# ~~~
91+
# find_package(Torch REQUIRED)
92+
# target_link_libraries(custom_ops_aot_lib PUBLIC Torch)
93+
# ~~~
94+
95+
# 1. C++ library to register custom ops into Executorch runtime.
96+
97+
add_library(custom_ops_lib)
98+
target_sources(
99+
custom_ops_lib
100+
PRIVATE
101+
${CMAKE_CURRENT_BINARY_DIR}/RegisterCodegenUnboxedKernelsEverything.cpp
102+
${CMAKE_CURRENT_BINARY_DIR}/Functions.h
103+
${CMAKE_CURRENT_BINARY_DIR}/NativeFunctions.h
104+
${CMAKE_CURRENT_BINARY_DIR}/CustomOpsNativeFunctions.h
105+
${CMAKE_CURRENT_LIST_DIR}/custom_ops_1.cpp)
106+
107+
target_link_libraries(custom_ops_lib PRIVATE executorch)
108+
109+
# Ensure that the load-time constructor functions run. By default, the linker
110+
# would remove them since there are no other references to them.
111+
if((CMAKE_CXX_COMPILER_ID MATCHES "AppleClang")
112+
OR (APPLE AND CMAKE_CXX_COMPILER_ID MATCHES "Clang"))
113+
target_link_options(custom_ops_lib INTERFACE
114+
"-Wl,-force_load,$<TARGET_FILE:custom_ops_lib>")
115+
elseif(CMAKE_CXX_COMPILER_ID MATCHES "Clang|GNU")
116+
target_link_options(
117+
custom_ops_lib INTERFACE
118+
"-Wl,--whole-archive,$<TARGET_FILE:custom_ops_lib>,--no-whole-archive")
119+
endif()

‎examples/custom_ops/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,3 +29,5 @@ After the model is exported by EXIR, we need C++ implementations of these custom
2929
kernel_name: custom::mul3_out_impl # sub-namespace native:: is auto-added
3030
```
3131
For how to write these YAML entries, please refer to [`kernels/portable/README.md`](https://github.com/pytorch/executorch/blob/main/kernels/portable/README.md).
32+
33+
Currently we provide 2 build systems that links `my_ops::mul3.out` kernel (written in `custom_ops_1.cpp`) to Executor runtime: buck2 and CMake. Both instructions are listed in `examples/custom_ops/test_custom_ops.sh`.

‎examples/custom_ops/test_custom_ops.sh

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,43 @@
55
# This source code is licensed under the BSD-style license found in the
66
# LICENSE file in the root directory of this source tree.
77

8-
# Test the end-to-end flow of using custom operator in a PyTorch model and use EXIR to capture and export a model file. Then use `executor_runner` demo C++ binary to run the model.
8+
# Test the end-to-end flow of using custom operator in a PyTorch model and use
9+
# EXIR to capture and export a model file. Then use `executor_runner` demo C++
10+
# binary to run the model.
911

10-
test_custom_op_1() {
11-
echo 'Exporting custom_ops_1.pte'
12-
python3 -m examples.custom_ops.custom_ops_1
12+
set -e
13+
14+
test_buck2_custom_op_1() {
15+
local model_name='custom_ops_1'
16+
echo "Exporting ${model_name}.pte"
17+
python3 -m "examples.custom_ops.${model_name}"
1318
# should save file custom_ops_1.pte
1419

1520
echo 'Running executor_runner'
16-
buck2 run //fbcode/executorch/examples/executor_runner:executor_runner -- --model_path=./custom_ops_1.pte
21+
buck2 run //fbcode/executorch/examples/executor_runner:executor_runner \
22+
--config=executorch.include_custom_ops=1 -- --model_path="./${model_name}.pte"
1723
# should give correct result
1824

19-
echo 'Removing custom_ops_1.pte'
20-
rm ./custom_ops_1.pte
25+
echo "Removing ${model_name}.pte"
26+
rm "./${model_name}.pte"
27+
}
28+
29+
test_cmake_custom_op_1() {
30+
local model_name='custom_ops_1'
31+
echo "Exporting ${model_name}.pte"
32+
python3 -m "examples.custom_ops.${model_name}"
33+
# should save file custom_ops_1.pte
34+
(rm -rf cmake-out \
35+
&& mkdir cmake-out \
36+
&& cd cmake-out \
37+
&& cmake -DBUCK2=buck2 -DBUILD_EXAMPLE_CUSTOM_OPS=ON ..)
38+
39+
echo 'Building executor_runner'
40+
cmake --build cmake-out -j9
41+
42+
echo 'Running executor_runner'
43+
cmake-out/executor_runner --model_path="./${model_name}.pte"
2144
}
2245

23-
test_custom_op_1
46+
test_buck2_custom_op_1
47+
test_cmake_custom_op_1

‎examples/executor_runner/targets.bzl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ def define_common_targets():
77
TARGETS and BUCK files that call this function.
88
"""
99

10+
include_custom_ops = native.read_config("executorch", "include_custom_ops", "0") == "1"
11+
1012
# Test driver for models, uses all portable kernels.
1113
runtime.cxx_binary(
1214
name = "executor_runner",
@@ -17,8 +19,7 @@ def define_common_targets():
1719
"//executorch/extension/data_loader:file_data_loader",
1820
"//executorch/util:util",
1921
"//executorch/kernels/portable:generated_lib_all_ops",
20-
"//executorch/examples/custom_ops:generated_lib",
21-
],
22+
] + (["//executorch/examples/custom_ops:generated_lib"] if include_custom_ops else []),
2223
external_deps = [
2324
"gflags",
2425
],

0 commit comments

Comments
 (0)