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# This file contains util functions to generate code for kernel registration for 8# both AOT and runtime. 9 10# Selective build. See codegen/tools/gen_oplist.py for how to use these 11# arguments. 12function(gen_selected_ops) 13 set(arg_names LIB_NAME OPS_SCHEMA_YAML ROOT_OPS INCLUDE_ALL_OPS) 14 cmake_parse_arguments(GEN "" "" "${arg_names}" ${ARGN}) 15 16 message(STATUS "Generating operator lib:") 17 message(STATUS " LIB_NAME: ${GEN_LIB_NAME}") 18 message(STATUS " OPS_SCHEMA_YAML: ${GEN_OPS_SCHEMA_YAML}") 19 message(STATUS " ROOT_OPS: ${GEN_ROOT_OPS}") 20 message(STATUS " INCLUDE_ALL_OPS: ${GEN_INCLUDE_ALL_OPS}") 21 22 set(_oplist_yaml 23 ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/selected_operators.yaml 24 ) 25 file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) 26 27 file(GLOB_RECURSE _codegen_tools_srcs "${EXECUTORCH_ROOT}/codegen/tools/*.py") 28 29 set(_gen_oplist_command "${PYTHON_EXECUTABLE}" -m codegen.tools.gen_oplist 30 --output_path=${_oplist_yaml} 31 ) 32 33 if(GEN_OPS_SCHEMA_YAML) 34 list(APPEND _gen_oplist_command 35 --ops_schema_yaml_path="${GEN_OPS_SCHEMA_YAML}" 36 ) 37 endif() 38 if(GEN_ROOT_OPS) 39 list(APPEND _gen_oplist_command --root_ops="${GEN_ROOT_OPS}") 40 endif() 41 if(GEN_INCLUDE_ALL_OPS) 42 list(APPEND _gen_oplist_command --include_all_operators) 43 endif() 44 45 message("Command - ${_gen_oplist_command}") 46 add_custom_command( 47 COMMENT "Generating selected_operators.yaml for ${GEN_LIB_NAME}" 48 OUTPUT ${_oplist_yaml} 49 COMMAND ${_gen_oplist_command} 50 DEPENDS ${GEN_OPS_SCHEMA_YAML} ${_codegen_tools_srcs} 51 WORKING_DIRECTORY ${EXECUTORCH_ROOT} 52 ) 53 54endfunction() 55 56# Codegen for registering kernels. Kernels are defined in functions_yaml and 57# custom_ops_yaml. 58# 59# Invoked as generate_bindings_for_kernels( LIB_NAME lib_name FUNCTIONS_YAML 60# functions_yaml CUSTOM_OPS_YAML custom_ops_yaml ) 61function(generate_bindings_for_kernels) 62 set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML) 63 cmake_parse_arguments(GEN "" "${arg_names}" "" ${ARGN}) 64 65 message(STATUS "Generating kernel bindings:") 66 message(STATUS " LIB_NAME: ${GEN_LIB_NAME}") 67 message(STATUS " FUNCTIONS_YAML: ${GEN_FUNCTIONS_YAML}") 68 message(STATUS " CUSTOM_OPS_YAML: ${GEN_CUSTOM_OPS_YAML}") 69 70 # Command to generate selected_operators.yaml from custom_ops.yaml. 71 file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*") 72 73 set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) 74 # By default selective build output is selected_operators.yaml 75 set(_oplist_yaml ${_out_dir}/selected_operators.yaml) 76 77 # Command to codegen C++ wrappers to register custom ops to both PyTorch and 78 # Executorch runtime. 79 execute_process( 80 COMMAND 81 "${PYTHON_EXECUTABLE}" -c 82 "from distutils.sysconfig import get_python_lib;print(get_python_lib())" 83 OUTPUT_VARIABLE site-packages-out 84 ERROR_VARIABLE site-packages-out-error 85 RESULT_VARIABLE site-packages-result 86 WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} 87 OUTPUT_STRIP_TRAILING_WHITESPACE 88 ) 89 file(GLOB_RECURSE _torchgen_srcs "${site-packages-out}/torchgen/*.py") 90 set(_gen_command 91 "${PYTHON_EXECUTABLE}" -m torchgen.gen_executorch 92 --source-path=${EXECUTORCH_ROOT}/codegen --install-dir=${_out_dir} 93 --tags-path=${site-packages-out}/torchgen/packaged/ATen/native/tags.yaml 94 --aten-yaml-path=${site-packages-out}/torchgen/packaged/ATen/native/native_functions.yaml 95 --op-selection-yaml-path=${_oplist_yaml} 96 ) 97 98 set(_gen_command_sources 99 ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp 100 ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h 101 ) 102 103 if(GEN_FUNCTIONS_YAML) 104 list(APPEND _gen_command --functions-yaml-path=${GEN_FUNCTIONS_YAML}) 105 endif() 106 if(GEN_CUSTOM_OPS_YAML) 107 list(APPEND _gen_command --custom-ops-yaml-path=${GEN_CUSTOM_OPS_YAML}) 108 list(APPEND _gen_command_sources ${_out_dir}/RegisterCPUCustomOps.cpp 109 ${_out_dir}/RegisterSchema.cpp ${_out_dir}/CustomOpsNativeFunctions.h 110 ) 111 endif() 112 113 add_custom_command( 114 COMMENT "Generating code for kernel registration" 115 OUTPUT ${_gen_command_sources} 116 COMMAND ${_gen_command} 117 DEPENDS ${_oplist_yaml} ${GEN_CUSTOM_OPS_YAML} ${GEN_FUNCTIONS_YAML} 118 ${_codegen_templates} ${_torchgen_srcs} 119 WORKING_DIRECTORY ${EXECUTORCH_ROOT} 120 ) 121 # Make generated file list available in parent scope 122 set(gen_command_sources 123 ${_gen_command_sources} 124 PARENT_SCOPE 125 ) 126endfunction() 127 128# Generate an AOT lib for registering custom ops into PyTorch 129function(gen_custom_ops_aot_lib) 130 cmake_parse_arguments(GEN "" "LIB_NAME" "KERNEL_SOURCES" ${ARGN}) 131 message(STATUS "Generating custom ops aot lib:") 132 message(STATUS " LIB_NAME: ${GEN_LIB_NAME}") 133 foreach(SOURCE IN LISTS GEN_KERNEL_SOURCES) 134 message(STATUS " KERNEL_SOURCE: ${SOURCE}") 135 endforeach() 136 137 set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) 138 add_library( 139 ${GEN_LIB_NAME} SHARED 140 ${_out_dir}/RegisterCPUCustomOps.cpp ${_out_dir}/RegisterSchema.cpp 141 ${_out_dir}/CustomOpsNativeFunctions.h "${GEN_KERNEL_SOURCES}" 142 ) 143 # Find `Torch`. 144 find_package(Torch REQUIRED) 145 # This lib uses ATen lib, so we explicitly enable rtti and exceptions. 146 target_compile_options(${GEN_LIB_NAME} PRIVATE -frtti -fexceptions) 147 target_compile_definitions(${GEN_LIB_NAME} PRIVATE USE_ATEN_LIB=1) 148 include_directories(${TORCH_INCLUDE_DIRS}) 149 target_link_libraries(${GEN_LIB_NAME} PRIVATE torch) 150 151 include(${EXECUTORCH_ROOT}/build/Utils.cmake) 152 153 target_link_options_shared_lib(${GEN_LIB_NAME}) 154 if(TARGET portable_lib) 155 target_link_libraries(${GEN_LIB_NAME} PRIVATE portable_lib) 156 else() 157 target_link_libraries(${GEN_LIB_NAME} PRIVATE executorch_core) 158 endif() 159endfunction() 160 161# Generate a runtime lib for registering operators in Executorch 162function(gen_operators_lib) 163 set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS) 164 cmake_parse_arguments(GEN "" "" "${multi_arg_names}" ${ARGN}) 165 166 message(STATUS "Generating operator lib:") 167 message(STATUS " LIB_NAME: ${GEN_LIB_NAME}") 168 message(STATUS " KERNEL_LIBS: ${GEN_KERNEL_LIBS}") 169 message(STATUS " DEPS: ${GEN_DEPS}") 170 171 set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) 172 173 add_library(${GEN_LIB_NAME}) 174 target_sources( 175 ${GEN_LIB_NAME} 176 PRIVATE ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp 177 ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h 178 ) 179 target_link_libraries(${GEN_LIB_NAME} PRIVATE ${GEN_DEPS}) 180 if(GEN_KERNEL_LIBS) 181 target_link_libraries(${GEN_LIB_NAME} PUBLIC ${GEN_KERNEL_LIBS}) 182 endif() 183 184 target_link_options_shared_lib(${GEN_LIB_NAME}) 185 set(_generated_headers ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h) 186 set_target_properties( 187 ${GEN_LIB_NAME} PROPERTIES PUBLIC_HEADER "${_generated_headers}" 188 ) 189endfunction() 190 191# Merge two kernel yaml files, prioritizing functions from FUNCTIONS_YAML and 192# taking functions from FALLBACK_YAML when no implementation is found. This 193# corresponds to the merge_yaml buck implementation in codegen/tools. 194function(merge_yaml) 195 set(arg_names FUNCTIONS_YAML FALLBACK_YAML OUTPUT_DIR) 196 cmake_parse_arguments(GEN "" "${arg_names}" "" ${ARGN}) 197 message(STATUS "Merging kernel yaml files:") 198 message(STATUS " FUNCTIONS_YAML: ${GEN_FUNCTIONS_YAML}") 199 message(STATUS " FALLBACK_YAML: ${GEN_FALLBACK_YAML}") 200 message(STATUS " OUTPUT_DIR: ${GEN_OUTPUT_DIR}") 201 202 set(_gen_command 203 "${PYTHON_EXECUTABLE}" -m codegen.tools.merge_yaml 204 --functions_yaml_path=${GEN_FUNCTIONS_YAML} 205 --fallback_yaml_path=${GEN_FALLBACK_YAML} --output_dir=${GEN_OUTPUT_DIR} 206 ) 207 208 add_custom_command( 209 COMMENT "Merging kernel yaml files" 210 OUTPUT ${GEN_OUTPUT_DIR}/merged.yaml 211 COMMAND ${_gen_command} 212 DEPENDS ${GEN_FUNCTIONS_YAML} ${GEN_FALLBACK_YAML} 213 WORKING_DIRECTORY ${EXECUTORCH_ROOT} 214 ) 215endfunction() 216