Skip to content

Plumbing for dtype selection in cmake flow - header generation #11432

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Plumbing for dtype selection in cmake flow - header generation #11432

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 21 additions & 3 deletions examples/selective_build/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,10 @@ option(EXECUTORCH_SELECT_OPS_LIST "Register a list of ops, separated by comma"
option(EXECUTORCH_SELECT_ALL_OPS
"Whether to register all ops defined in portable kernel library." OFF
)

# Option to enable dtype selective build
option(EXECUTORCH_SELECT_DTYPE "Enable dtype selection during build." OFF
)
# ------------------------------- OPTIONS END --------------------------------

#
Expand Down Expand Up @@ -108,16 +112,30 @@ gen_selected_ops(
"${EXECUTORCH_SELECT_OPS_LIST}"
INCLUDE_ALL_OPS
"${EXECUTORCH_SELECT_ALL_OPS}"
DTYPE_SELECT
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: could we use DTYPE_SELECTIVE_BUILD to mirror the buck API here?

"${EXECUTORCH_SELECT_DTYPE}"
)

generate_bindings_for_kernels(
LIB_NAME "select_build_lib" FUNCTIONS_YAML
${EXECUTORCH_ROOT}/kernels/portable/functions.yaml CUSTOM_OPS_YAML
LIB_NAME
"select_build_lib"
FUNCTIONS_YAML
${EXECUTORCH_ROOT}/kernels/portable/functions.yaml
CUSTOM_OPS_YAML
"${_custom_ops_yaml}"
DTYPE_SELECT
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we only need this option in gen_selected_ops at the moment, so let's leave it out of the other APIs for now?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was running into an issue where if the selected_op_variants.h file wasn't include in the sources list in the gen_operators_lib function, cmake was smart enough to delete it. The casing in that function conditionally adds it to the source list if the flag is set.

"${EXECUTORCH_SELECT_DTYPE}"
)

gen_operators_lib(
LIB_NAME "select_build_lib" KERNEL_LIBS ${_kernel_lib} DEPS executorch_core
LIB_NAME
"select_build_lib"
KERNEL_LIBS
${_kernel_lib}
DEPS
executorch_core
DTYPE_SELECT
"${EXECUTORCH_SELECT_DTYPE}"
)

list(TRANSFORM _executor_runner__srcs PREPEND "${EXECUTORCH_ROOT}/")
Expand Down
3 changes: 3 additions & 0 deletions examples/selective_build/test_selective_build.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ test_cmake_select_all_ops() {
rm -rf ${build_dir}
retry cmake -DCMAKE_BUILD_TYPE=Release \
-DEXECUTORCH_SELECT_ALL_OPS=ON \
-DEXECUTORCH_SELECT_DTYPE=ON \
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's add these in once the E2E is working - currently, it won't enable dtype selective build so may be misleading

-DCMAKE_INSTALL_PREFIX=cmake-out \
-DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
-B${build_dir} \
Expand Down Expand Up @@ -123,6 +124,7 @@ test_cmake_select_ops_in_list() {
aten::_native_batch_norm_legit_no_training.out,aten::hardtanh.out,aten::add.out,\
aten::mean.out,aten::view_copy.out,aten::permute_copy.out,aten::addmm.out,\
aten,aten::clone.out" \
-DEXECUTORCH_SELECT_DTYPE=ON \
-DCMAKE_INSTALL_PREFIX=cmake-out \
-DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
-B${build_dir} \
Expand All @@ -146,6 +148,7 @@ test_cmake_select_ops_in_yaml() {
rm -rf ${build_dir}
retry cmake -DCMAKE_BUILD_TYPE=Release \
-DEXECUTORCH_SELECT_OPS_YAML=ON \
-DEXECUTORCH_SELECT_DTYPE=ON \
-DCMAKE_INSTALL_PREFIX=cmake-out \
-DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
-B${build_dir} \
Expand Down
66 changes: 56 additions & 10 deletions tools/cmake/Codegen.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,20 @@
include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)

function(gen_selected_ops)
set(arg_names LIB_NAME OPS_SCHEMA_YAML ROOT_OPS INCLUDE_ALL_OPS)
set(arg_names LIB_NAME OPS_SCHEMA_YAML ROOT_OPS INCLUDE_ALL_OPS DTYPE_SELECT)
cmake_parse_arguments(GEN "" "" "${arg_names}" ${ARGN})

message(STATUS "Generating operator lib:")
message(STATUS " LIB_NAME: ${GEN_LIB_NAME}")
message(STATUS " OPS_SCHEMA_YAML: ${GEN_OPS_SCHEMA_YAML}")
message(STATUS " ROOT_OPS: ${GEN_ROOT_OPS}")
message(STATUS " INCLUDE_ALL_OPS: ${GEN_INCLUDE_ALL_OPS}")
message(STATUS " DTYPE_SELECT: ${GEN_DTYPE_SELECT}")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is GEN_DTYPE_SELECT being defined?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

GEN_DTYPE_SELECT is toggled ON or OFF via the command line switch -DEXECUTORCH_SELECT_DTYPE=ON. The option is defined in examples/selective_build/CMakeLists.txt.


set(_oplist_yaml
${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/selected_operators.yaml
)

file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})

file(GLOB_RECURSE _codegen_tools_srcs "${EXECUTORCH_ROOT}/codegen/tools/*.py")
Expand Down Expand Up @@ -53,6 +55,23 @@ function(gen_selected_ops)
WORKING_DIRECTORY ${EXECUTORCH_ROOT}
)

if(GEN_DTYPE_SELECT)
set(_opvariant_h
${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/selected_op_variants.h
)
set(_gen_opvariant_command "${PYTHON_EXECUTABLE}" -m codegen.tools.gen_selected_op_variants
--yaml-file=${_oplist_yaml}
--output-dir=${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/
)
message("Command - ${_gen_opvariant_command}")
add_custom_command(
COMMENT "Generating selected_op_variants.h for ${GEN_LIB_NAME}"
OUTPUT ${_opvariant_h}
COMMAND ${_gen_opvariant_command}
DEPENDS ${_optlist_yaml} ${_codegen_tools_srcs}
WORKING_DIRECTORY ${EXECUTORCH_ROOT}
)
endif()
endfunction()

# Codegen for registering kernels. Kernels are defined in functions_yaml and
Expand All @@ -62,14 +81,15 @@ endfunction()
# functions_yaml CUSTOM_OPS_YAML custom_ops_yaml )
function(generate_bindings_for_kernels)
set(options ADD_EXCEPTION_BOUNDARY)
set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML)
set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML DTYPE_SELECT)
cmake_parse_arguments(GEN "${options}" "${arg_names}" "" ${ARGN})

message(STATUS "Generating kernel bindings:")
message(STATUS " LIB_NAME: ${GEN_LIB_NAME}")
message(STATUS " FUNCTIONS_YAML: ${GEN_FUNCTIONS_YAML}")
message(STATUS " CUSTOM_OPS_YAML: ${GEN_CUSTOM_OPS_YAML}")
message(STATUS " ADD_EXCEPTION_BOUNDARY: ${GEN_ADD_EXCEPTION_BOUNDARY}")
message(STATUS " DTYPE_SELECT: ${GEN_DTYPE_SELECT}")

# Command to generate selected_operators.yaml from custom_ops.yaml.
file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*")
Expand All @@ -78,6 +98,13 @@ function(generate_bindings_for_kernels)
# By default selective build output is selected_operators.yaml
set(_oplist_yaml ${_out_dir}/selected_operators.yaml)

# If dtype selective build is enable, force header file to be preserved
if(GEN_DTYPE_SELECT)
Copy link
Contributor

@lucylq lucylq Jun 6, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can leave this out for now, as it'll be preserved once we link the header into the portable library.

set(_opvariant_h ${_out_dir}/selected_op_variants.h)
else()
set(_opvariant_h "")
endif()

# Command to codegen C++ wrappers to register custom ops to both PyTorch and
# Executorch runtime.
execute_process(
Expand Down Expand Up @@ -108,6 +135,10 @@ function(generate_bindings_for_kernels)
${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h
)

if(GEN_DTYPE_SELECT)
list(APPEND _gen_command_sources ${_out_dir}/selected_op_variants.h)
endif()

if(GEN_FUNCTIONS_YAML)
list(APPEND _gen_command --functions-yaml-path=${GEN_FUNCTIONS_YAML})
endif()
Expand All @@ -122,8 +153,9 @@ function(generate_bindings_for_kernels)
COMMENT "Generating code for kernel registration"
OUTPUT ${_gen_command_sources}
COMMAND ${_gen_command}
DEPENDS ${_oplist_yaml} ${GEN_CUSTOM_OPS_YAML} ${GEN_FUNCTIONS_YAML}
${_codegen_templates} ${_torchgen_srcs}
DEPENDS ${_oplist_yaml} ${_opvariants_h} ${GEN_CUSTOM_OPS_YAML}
${GEN_FUNCTIONS_YAML} ${_codegen_templates}
${_torchgen_srcs}
WORKING_DIRECTORY ${EXECUTORCH_ROOT}
)
# Make generated file list available in parent scope
Expand Down Expand Up @@ -165,29 +197,43 @@ endfunction()

# Generate a runtime lib for registering operators in Executorch
function(gen_operators_lib)
set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS)
set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS DTYPE_SELECT)
cmake_parse_arguments(GEN "" "" "${multi_arg_names}" ${ARGN})

message(STATUS "Generating operator lib:")
message(STATUS " LIB_NAME: ${GEN_LIB_NAME}")
message(STATUS " KERNEL_LIBS: ${GEN_KERNEL_LIBS}")
message(STATUS " DEPS: ${GEN_DEPS}")
message(STATUS " DTYPE_SELECT: ${GEN_DTYPE_SELECT}")

set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})

add_library(${GEN_LIB_NAME})
target_sources(
${GEN_LIB_NAME}
PRIVATE ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp
${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h
)
if(GEN_DTYPE_SELECT)
target_sources(
${GEN_LIB_NAME}
PRIVATE ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp
${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h
${_out_dir}/selected_op_variants.h
)
else()
target_sources(
${GEN_LIB_NAME}
PRIVATE ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp
${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h
)
endif()

target_link_libraries(${GEN_LIB_NAME} PRIVATE ${GEN_DEPS})
if(GEN_KERNEL_LIBS)
target_link_libraries(${GEN_LIB_NAME} PUBLIC ${GEN_KERNEL_LIBS})
endif()

target_link_options_shared_lib(${GEN_LIB_NAME})
set(_generated_headers ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h)
if(GEN_DTYPE_SELECT)
list(APPEND _generated_headers ${_out_dir}/selected_op_variants.h)
endif()
set_target_properties(
${GEN_LIB_NAME} PROPERTIES PUBLIC_HEADER "${_generated_headers}"
)
Expand Down
Loading