diff --git a/CMakeLists.txt b/CMakeLists.txt index 8bb5da588e4..f40e8a65134 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -103,6 +103,9 @@ if(BUILD_SELECTIVE_BUILD_TEST) # Option to register ops from yaml file option(SELECT_OPS_YAML "Register all the ops from a given yaml file" OFF) + + # Option to manually register kernels + option(MANUAL_REGISTRATION "Manually register kernels" OFF) endif() # Build xnn_executor_runner which depends on XNNPACK diff --git a/build/Codegen.cmake b/build/Codegen.cmake index 53b7e2fdd00..dd6d39f501d 100644 --- a/build/Codegen.cmake +++ b/build/Codegen.cmake @@ -36,9 +36,27 @@ function(gen_selected_ops ops_schema_yaml root_ops include_all_ops) endfunction() +# Return generated source files +function(get_generated_sources list_name manual_registration) + set(_temp_list + ${CMAKE_CURRENT_BINARY_DIR}/Functions.h + ${CMAKE_CURRENT_BINARY_DIR}/NativeFunctions.h) + # If manually register ops, change source files + if(manual_registration) + list(APPEND _temp_list + ${CMAKE_CURRENT_BINARY_DIR}/RegisterKernelsEverything.cpp) + list(APPEND _temp_list ${CMAKE_CURRENT_BINARY_DIR}/RegisterKernels.h) + else() + list(APPEND _temp_list + ${CMAKE_CURRENT_BINARY_DIR}/RegisterCodegenUnboxedKernelsEverything.cpp) + endif() + + set(${list_name} ${_temp_list} PARENT_SCOPE) +endfunction() + # Codegen for registering kernels. Kernels are defined in functions_yaml and # custom_ops_yaml -function(generate_bindings_for_kernels functions_yaml custom_ops_yaml) +function(generate_bindings_for_kernels functions_yaml custom_ops_yaml manual) # Command to generate selected_operators.yaml from custom_ops.yaml. file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*") file(GLOB_RECURSE _torchgen_srcs "${TORCH_ROOT}/torchgen/*.py") @@ -55,10 +73,12 @@ function(generate_bindings_for_kernels functions_yaml custom_ops_yaml) --aten-yaml-path=${TORCH_ROOT}/aten/src/ATen/native/native_functions.yaml --op-selection-yaml-path=${_oplist_yaml}) - set(_gen_command_sources - ${CMAKE_CURRENT_BINARY_DIR}/RegisterCodegenUnboxedKernelsEverything.cpp - ${CMAKE_CURRENT_BINARY_DIR}/Functions.h - ${CMAKE_CURRENT_BINARY_DIR}/NativeFunctions.h) + get_generated_sources("_gen_command_sources" "${manual}") + + # If manually register ops, append --manual_registration and change out files + if(manual) + list(APPEND _gen_command --manual_registration) + endif() if(functions_yaml) list(APPEND _gen_command --functions-yaml-path=${functions_yaml}) @@ -105,14 +125,10 @@ function(gen_custom_ops_aot_lib lib_name kernel_sources) endfunction() # Generate a runtime lib for registering operators in Executorch -function(gen_operators_lib lib_name kernel_lib deps) +function(gen_operators_lib lib_name kernel_lib deps manual_registration) add_library(${lib_name}) - target_sources( - ${lib_name} - PRIVATE - ${CMAKE_CURRENT_BINARY_DIR}/RegisterCodegenUnboxedKernelsEverything.cpp - ${CMAKE_CURRENT_BINARY_DIR}/Functions.h - ${CMAKE_CURRENT_BINARY_DIR}/NativeFunctions.h) + get_generated_sources("_sources" "${manual_registration}") + target_sources(${lib_name} PRIVATE ${_sources}) target_link_libraries(${lib_name} PRIVATE ${deps}) if(kernel_lib) target_link_libraries(${lib_name} INTERFACE ${kernel_lib}) diff --git a/codegen/templates/RegisterKernels.cpp b/codegen/templates/RegisterKernels.cpp index 2313a30a307..8f0ff794f89 100644 --- a/codegen/templates/RegisterKernels.cpp +++ b/codegen/templates/RegisterKernels.cpp @@ -15,10 +15,11 @@ namespace torch { namespace executor { +static const Kernel kernels_to_register[] = { + ${unboxed_kernels} // Generated kernels +}; + Error register_all_kernels() { - Kernel kernels_to_register[] = { - ${unboxed_kernels} // Generated kernels - }; Error success_with_kernel_reg = register_kernels(kernels_to_register); if (success_with_kernel_reg != Error::Ok) { ET_LOG(Error, "Failed register all kernels"); diff --git a/examples/custom_ops/CMakeLists.txt b/examples/custom_ops/CMakeLists.txt index 6c65e92892d..f8606292618 100644 --- a/examples/custom_ops/CMakeLists.txt +++ b/examples/custom_ops/CMakeLists.txt @@ -43,7 +43,7 @@ elseif(REGISTER_EXAMPLE_CUSTOM_OP EQUAL 2) gen_selected_ops("" "my_ops::mul4.out" "") endif() # Expect gen_selected_ops output file to be selected_operators.yaml -generate_bindings_for_kernels("" ${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml) +generate_bindings_for_kernels("" ${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml "") message("Generated files ${gen_command_sources}") # Prepare for C++ libraries. @@ -71,4 +71,4 @@ add_library(custom_kernels ${kernel_sources}) target_link_libraries(custom_kernels PRIVATE executorch) target_compile_options(custom_kernels PUBLIC ${_common_compile_options}) -gen_operators_lib("custom_ops_lib" custom_kernels executorch) +gen_operators_lib("custom_ops_lib" custom_kernels executorch "") diff --git a/examples/executor_runner/executor_runner.cpp b/examples/executor_runner/executor_runner.cpp index 9286585646d..3e00e99f787 100644 --- a/examples/executor_runner/executor_runner.cpp +++ b/examples/executor_runner/executor_runner.cpp @@ -30,7 +30,9 @@ #include #include #include - +#if defined(EXECUTORCH_MANUAL_KERNEL_REG) +#include +#endif static uint8_t method_allocator_pool[4 * 1024U * 1024U]; // 4 MB DEFINE_string( @@ -47,7 +49,9 @@ using torch::executor::util::FileDataLoader; int main(int argc, char** argv) { runtime_init(); - +#if defined(EXECUTORCH_MANUAL_KERNEL_REG) + register_all_kernels(); +#endif gflags::ParseCommandLineFlags(&argc, &argv, true); if (argc != 1) { std::string msg = "Extra commandline args:"; diff --git a/examples/executor_runner/targets.bzl b/examples/executor_runner/targets.bzl index 7903507337d..0fc0fdb127a 100644 --- a/examples/executor_runner/targets.bzl +++ b/examples/executor_runner/targets.bzl @@ -9,15 +9,19 @@ def define_common_targets(): # Wraps a commandline executable that can be linked against any desired # kernel or backend implementations. Contains a main() function. + select_ops = native.read_config("executorch", "select_ops", None) + runtime.cxx_library( name = "executor_runner_lib", srcs = ["executor_runner.cpp"], deps = [ "//executorch/runtime/executor:program", + "//executorch/runtime/kernel:operator_registry", "//executorch/extension/data_loader:file_data_loader", "//executorch/extension/evalue_util:print_evalue", "//executorch/util:util", - ], + ] + (["//executorch/examples/selective_build:select_add_out_manual_lib"] if select_ops == "add_manual" else []), + preprocessor_flags = ["-DEXECUTORCH_MANUAL_KERNEL_REG"] if select_ops == "add_manual" else [], external_deps = [ "gflags", ], diff --git a/examples/selective_build/CMakeLists.txt b/examples/selective_build/CMakeLists.txt index c1ba4f262cc..2f9ce0a5dbd 100644 --- a/examples/selective_build/CMakeLists.txt +++ b/examples/selective_build/CMakeLists.txt @@ -49,14 +49,28 @@ elseif(SELECT_OPS_YAML) list(APPEND _kernel_lib custom_kernels) endif() generate_bindings_for_kernels(${EXECUTORCH_ROOT}/kernels/portable/functions.yaml - "${_custom_ops_yaml}") -gen_operators_lib("select_build_lib" ${_kernel_lib} executorch) + "${_custom_ops_yaml}" ${MANUAL_REGISTRATION}) +gen_operators_lib( + "select_build_lib" ${_kernel_lib} executorch ${MANUAL_REGISTRATION}) set(_updated__srcs) foreach(_src ${_executor_runner__srcs}) list(APPEND _updated__srcs "${EXECUTORCH_ROOT}/${_src}") endforeach() +if(${MANUAL_REGISTRATION}) + # Create include directory with the header RegisterKernels.h + set(_include_dir "${CMAKE_BINARY_DIR}/examples/selective_build/include") + add_custom_command(TARGET select_build_lib POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy + ${CMAKE_CURRENT_BINARY_DIR}/RegisterKernels.h + ${_include_dir}/executorch/examples/selective_build/RegisterKernels.h) + + target_include_directories( + select_build_lib INTERFACE ${_include_dir}) + # Set preprocessor flags for manual registration + add_definitions(-DEXECUTORCH_MANUAL_KERNEL_REG) +endif() # # selective_build_test: test binary to allow different operator libraries to # link to diff --git a/examples/selective_build/targets.bzl b/examples/selective_build/targets.bzl index 73d3ff81795..82cb910e52c 100644 --- a/examples/selective_build/targets.bzl +++ b/examples/selective_build/targets.bzl @@ -60,6 +60,25 @@ def define_common_targets(): # Select all ops from a given model # TODO(larryliu0820): Add this + # Select add.out and manually register + et_operator_library( + name = "select_add_out", + ops = [ + "aten::add.out", + ], + ) + + executorch_generated_lib( + name = "select_add_out_manual_lib", + functions_yaml_target = "//executorch/kernels/portable:functions.yaml", + manual_registration = True, + deps = [ + "//executorch/kernels/portable:operators", + ":select_add_out", + ], + visibility = ["PUBLIC"], + ) + # ~~~ Test binary for selective build ~~~ select_ops = native.read_config("executorch", "select_ops", None) lib = [] @@ -69,6 +88,8 @@ def define_common_targets(): lib.append(":select_ops_in_list_lib") elif select_ops == "yaml": lib.append(":select_ops_from_yaml_lib") + elif select_ops == "add_manual": + lib.append(":select_add_out_manual_lib") runtime.cxx_binary( name = "selective_build_test", srcs = [], diff --git a/examples/selective_build/test_selective_build.sh b/examples/selective_build/test_selective_build.sh index 4a818dc9872..764cb38d348 100644 --- a/examples/selective_build/test_selective_build.sh +++ b/examples/selective_build/test_selective_build.sh @@ -53,6 +53,18 @@ test_buck2_select_ops_from_yaml() { rm "./custom_ops_1.pte" } +test_buck2_select_add_manual() { + echo "Exporting add" + ${PYTHON_EXECUTABLE} -m examples.export.export_example --model_name="add" + + echo "Running selective build test" + $BUCK run //examples/selective_build:selective_build_test \ + --config=executorch.select_ops=add_manual -- --model_path=./add.pte + + echo "Removing add.pte" + rm "./add.pte" +} + test_cmake_select_all_ops() { echo "Exporting MobilenetV3" ${PYTHON_EXECUTABLE} -m examples.export.export_example --model_name="mv3" @@ -127,6 +139,30 @@ test_cmake_select_ops_in_yaml() { rm "./custom_ops_1.pte" } +test_cmake_select_add_manual() { + echo "Exporting add" + ${PYTHON_EXECUTABLE} -m examples.export.export_example --model_name="add" + # -DCMAKE_BUILD_TYPE=Release \ + + (rm -rf cmake-out \ + && mkdir cmake-out \ + && cd cmake-out \ + && retry cmake -DBUCK2="$BUCK" \ + -DBUILD_SELECTIVE_BUILD_TEST=ON \ + -DSELECT_OPS_LIST="aten::add.out" \ + -DMANUAL_REGISTRATION=ON \ + -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" ..) + + echo "Build selective build test" + cmake --build cmake-out -j9 + + echo 'Running selective build test' + cmake-out/examples/selective_build/selective_build_test --model_path="./add.pte" + + echo "Removing add.pte" + rm "./add.pte" +} + if [[ -z $BUCK ]]; then BUCK=buck2 @@ -142,9 +178,11 @@ then test_cmake_select_all_ops test_cmake_select_ops_in_list test_cmake_select_ops_in_yaml + test_cmake_select_add_manual elif [[ $1 == "buck2" ]]; then test_buck2_select_all_ops test_buck2_select_ops_in_list test_buck2_select_ops_from_yaml + test_buck2_select_add_manual fi diff --git a/kernels/portable/CMakeLists.txt b/kernels/portable/CMakeLists.txt index a0f1712b4e3..4860da9df58 100644 --- a/kernels/portable/CMakeLists.txt +++ b/kernels/portable/CMakeLists.txt @@ -41,7 +41,7 @@ list(FILTER _portable_kernels__srcs EXCLUDE REGEX "codegen") # Executorch (for runtime). Here select all ops in functions.yaml gen_selected_ops("${CMAKE_CURRENT_LIST_DIR}/functions.yaml" "" "") # Expect gen_selected_ops output file to be selected_operators.yaml -generate_bindings_for_kernels(${CMAKE_CURRENT_SOURCE_DIR}/functions.yaml "") +generate_bindings_for_kernels(${CMAKE_CURRENT_SOURCE_DIR}/functions.yaml "" "") message("Generated files ${gen_command_sources}") # @@ -57,4 +57,4 @@ target_compile_options(portable_kernels PUBLIC ${_common_compile_options}) # # portable_ops_lib: Register portable_ops_lib ops kernels into Executorch # runtime -gen_operators_lib("portable_ops_lib" portable_kernels executorch) +gen_operators_lib("portable_ops_lib" portable_kernels executorch "") diff --git a/kernels/quantized/CMakeLists.txt b/kernels/quantized/CMakeLists.txt index cfd041b8818..ffd15254328 100644 --- a/kernels/quantized/CMakeLists.txt +++ b/kernels/quantized/CMakeLists.txt @@ -39,7 +39,7 @@ file(GLOB_RECURSE _quantized_kernels__srcs # Executorch (for runtime). Here select all ops in quantized.yaml gen_selected_ops("${CMAKE_CURRENT_LIST_DIR}/quantized.yaml" "" "") # Expect gen_selected_ops output file to be selected_operators.yaml -generate_bindings_for_kernels("" ${CMAKE_CURRENT_SOURCE_DIR}/quantized.yaml) +generate_bindings_for_kernels("" ${CMAKE_CURRENT_SOURCE_DIR}/quantized.yaml "") message("Generated files ${gen_command_sources}") # Build a AOT library to register quantized ops into PyTorch. @@ -56,4 +56,4 @@ target_compile_options(quantized_kernels PUBLIC ${_common_compile_options}) # Build a library for _quantized_kernels_srcs # # quantized_ops_lib: Register quantized ops kernels into Executorch runtime -gen_operators_lib("quantized_ops_lib" quantized_kernels executorch) +gen_operators_lib("quantized_ops_lib" quantized_kernels executorch "")