Skip to content

Commit 7857dfd

Browse files
Erik-LundellErik-Lundell
authored
Share selective prim-ops helper (pytorch#21732)
As pointed out in pytorch#21714, the helper is useful for all builds. Therefore, move the helper from ArmRunnerUtils to Codegen.cmake. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Erik Lundell <erik.lundell@arm.com>
1 parent df3d715 commit 7857dfd

2 files changed

Lines changed: 77 additions & 57 deletions

File tree

‎backends/arm/cmake/ArmRunnerUtils.cmake‎

Lines changed: 1 addition & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -46,62 +46,6 @@ function(verify_targets_exist)
4646
endforeach()
4747
endfunction()
4848

49-
# Private function to isolate logic for building prim ops library.
50-
function(_arm_runner_create_prim_ops_lib)
51-
# Parse arguments.
52-
set(options INCLUDE_ALL_OPS)
53-
set(one_value_args LIB_NAME SELECTED_OPS_YAML)
54-
set(multi_value_args DEPS)
55-
cmake_parse_arguments(
56-
ARG "${options}" "${one_value_args}" "${multi_value_args}" ${ARGN}
57-
)
58-
59-
set(_out_dir "${CMAKE_CURRENT_BINARY_DIR}/${ARG_LIB_NAME}")
60-
set(_sources ${EXECUTORCH_ROOT}/kernels/prim_ops/register_prim_ops.cpp)
61-
62-
# Generate selected operator list if needed.
63-
if(NOT ARG_INCLUDE_ALL_OPS)
64-
set(_selected_prim_ops_header "${_out_dir}/selected_prim_ops.h")
65-
set(_gen_selected_prim_ops_script
66-
"${EXECUTORCH_ROOT}/codegen/tools/gen_selected_prim_ops.py"
67-
)
68-
set(_gen_selected_prim_ops_command
69-
"${PYTHON_EXECUTABLE}" -m codegen.tools.gen_selected_prim_ops
70-
--op-selection-yaml-path=${ARG_SELECTED_OPS_YAML}
71-
--output-dir=${_out_dir}
72-
)
73-
add_custom_command(
74-
COMMENT "Generating selected_prim_ops.h for ${ARG_LIB_NAME}"
75-
OUTPUT ${_selected_prim_ops_header}
76-
COMMAND ${_gen_selected_prim_ops_command}
77-
DEPENDS ${ARG_SELECTED_OPS_YAML} ${_gen_selected_prim_ops_script}
78-
WORKING_DIRECTORY ${EXECUTORCH_ROOT}
79-
)
80-
list(APPEND _sources ${_selected_prim_ops_header})
81-
endif()
82-
83-
# Add prim ops registration library.
84-
add_library(${ARG_LIB_NAME} ${_sources})
85-
target_include_directories(
86-
${ARG_LIB_NAME} PRIVATE ${EXECUTORCH_ROOT}/kernels/prim_ops ${_out_dir}
87-
)
88-
if(NOT ARG_INCLUDE_ALL_OPS)
89-
target_compile_definitions(
90-
${ARG_LIB_NAME} PRIVATE ET_PRIM_OPS_SELECTIVE_BUILD
91-
EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD
92-
)
93-
endif()
94-
target_link_libraries(${ARG_LIB_NAME} PRIVATE ${ARG_DEPS})
95-
executorch_target_link_options_shared_lib(${ARG_LIB_NAME})
96-
97-
# Add prim ops kernel library.
98-
add_library(
99-
${ARG_LIB_NAME}_impl ${EXECUTORCH_ROOT}/kernels/prim_ops/et_copy_index.cpp
100-
${EXECUTORCH_ROOT}/kernels/prim_ops/et_view.cpp
101-
)
102-
target_link_libraries(${ARG_LIB_NAME}_impl PRIVATE ${ARG_DEPS})
103-
endfunction()
104-
10549
#[[
10650
Create a selected operator registration library for one operator family.
10751
@@ -221,7 +165,7 @@ function(arm_runner_create_selected_ops_lib)
221165
if(_arm_runner_include_all_ops)
222166
list(APPEND _arm_runner_prim_ops_args INCLUDE_ALL_OPS)
223167
endif()
224-
_arm_runner_create_prim_ops_lib(${_arm_runner_prim_ops_args})
168+
gen_selected_prim_ops_lib(${_arm_runner_prim_ops_args})
225169
return()
226170
endif()
227171

‎tools/cmake/Codegen.cmake‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# Copyright (c) Meta Platforms, Inc. and affiliates.
22
# All rights reserved.
3+
# Copyright 2026 Arm Limited and/or its affiliates.
34
#
45
# This source code is licensed under the BSD-style license found in the
56
# LICENSE file in the root directory of this source tree.
@@ -150,6 +151,81 @@ function(gen_selected_max_kernel_num)
150151
)
151152
endfunction()
152153

154+
# Generate a prim ops registration library. By default the library includes only
155+
# operators from SELECTED_OPS_YAML. INCLUDE_ALL_OPS disables prim-op selective
156+
# build and registers every prim op.
157+
function(gen_selected_prim_ops_lib)
158+
# Parse arguments.
159+
set(options INCLUDE_ALL_OPS)
160+
set(one_value_args LIB_NAME SELECTED_OPS_YAML)
161+
set(multi_value_args DEPS)
162+
cmake_parse_arguments(
163+
GEN "${options}" "${one_value_args}" "${multi_value_args}" ${ARGN}
164+
)
165+
166+
message(STATUS "Generating selected prim ops lib:")
167+
message(STATUS " LIB_NAME: ${GEN_LIB_NAME}")
168+
message(STATUS " SELECTED_OPS_YAML: ${GEN_SELECTED_OPS_YAML}")
169+
message(STATUS " INCLUDE_ALL_OPS: ${GEN_INCLUDE_ALL_OPS}")
170+
message(STATUS " DEPS: ${GEN_DEPS}")
171+
172+
if(NOT GEN_LIB_NAME)
173+
message(FATAL_ERROR "gen_selected_prim_ops_lib: LIB_NAME is required")
174+
endif()
175+
if(NOT GEN_INCLUDE_ALL_OPS AND NOT GEN_SELECTED_OPS_YAML)
176+
message(
177+
FATAL_ERROR
178+
"gen_selected_prim_ops_lib: SELECTED_OPS_YAML is required unless INCLUDE_ALL_OPS is set"
179+
)
180+
endif()
181+
182+
set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})
183+
file(MAKE_DIRECTORY ${_out_dir})
184+
185+
# Generate selected operator list if needed.
186+
set(_sources ${EXECUTORCH_ROOT}/kernels/prim_ops/register_prim_ops.cpp)
187+
if(NOT GEN_INCLUDE_ALL_OPS)
188+
set(_selected_prim_ops_header ${_out_dir}/selected_prim_ops.h)
189+
set(_gen_selected_prim_ops_script
190+
${EXECUTORCH_ROOT}/codegen/tools/gen_selected_prim_ops.py
191+
)
192+
set(_gen_selected_prim_ops_command
193+
"${PYTHON_EXECUTABLE}" -m codegen.tools.gen_selected_prim_ops
194+
--op-selection-yaml-path=${GEN_SELECTED_OPS_YAML}
195+
--output-dir=${_out_dir}
196+
)
197+
add_custom_command(
198+
COMMENT "Generating selected_prim_ops.h for ${GEN_LIB_NAME}"
199+
OUTPUT ${_selected_prim_ops_header}
200+
COMMAND ${_gen_selected_prim_ops_command}
201+
DEPENDS ${GEN_SELECTED_OPS_YAML} ${_gen_selected_prim_ops_script}
202+
WORKING_DIRECTORY ${EXECUTORCH_ROOT}
203+
)
204+
list(APPEND _sources ${_selected_prim_ops_header})
205+
endif()
206+
207+
# Add prim ops registration library.
208+
add_library(${GEN_LIB_NAME} ${_sources})
209+
target_include_directories(
210+
${GEN_LIB_NAME} PRIVATE ${EXECUTORCH_ROOT}/kernels/prim_ops ${_out_dir}
211+
)
212+
if(NOT GEN_INCLUDE_ALL_OPS)
213+
target_compile_definitions(
214+
${GEN_LIB_NAME} PRIVATE ET_PRIM_OPS_SELECTIVE_BUILD
215+
EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD
216+
)
217+
endif()
218+
target_link_libraries(${GEN_LIB_NAME} PRIVATE ${GEN_DEPS})
219+
executorch_target_link_options_shared_lib(${GEN_LIB_NAME})
220+
221+
# Add prim ops kernel library.
222+
add_library(
223+
${GEN_LIB_NAME}_impl ${EXECUTORCH_ROOT}/kernels/prim_ops/et_copy_index.cpp
224+
${EXECUTORCH_ROOT}/kernels/prim_ops/et_view.cpp
225+
)
226+
target_link_libraries(${GEN_LIB_NAME}_impl PRIVATE ${GEN_DEPS})
227+
endfunction()
228+
153229
# Codegen for registering kernels. Kernels are defined in functions_yaml and
154230
# custom_ops_yaml.
155231
#

0 commit comments

Comments
 (0)