Add examples for manual kernel registration

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

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})

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");

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 "")

examples/executor_runner/executor_runner.cpp

@@ -30,7 +30,9 @@
 #include <executorch/runtime/platform/profiler.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/util/util.h>
-
+#if defined(EXECUTORCH_MANUAL_KERNEL_REG)
+#include <executorch/examples/selective_build/RegisterKernels.h>
+#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:";

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",
         ],

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

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 = [],

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

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 "")

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 "")