Allow kernel manual registration (#491)

Summary:
X-link: pytorch/pytorch#110086


Exposing a codegen mode for generating a hook for user to register their kernels.

If we pass `--manual-registration` flag to `gen_executorch.py`, we will generate the following files:
1. RegisterKernels.h which declares a `register_all_kernels()` API inside `torch::executor` namespace.
2. RegisterKernelsEverything.cpp which implements `register_all_kernels()` by defining an array of generated kernels.

This way user can depend on the library declared by `executorch_generated_lib` macro (with `manual_registration=True`) and be able to include `RegisterKernels.h`. Then they can manually call `register_all_kernels()` instead of relying on C++ static initialization mechanism which is not available in some embedded systems.

Reviewed By: cccclai

Differential Revision: D49439673

Author: larryliu0820

codegen/templates/RegisterKernels.cpp

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+// ${generated_comment}
+// This implements register_all_kernels() API that is declared in
+// RegisterKernels.h
+#include "RegisterKernels.h"
+#include "${fn_header}" // Generated Function import headers
+
+namespace torch {
+namespace executor {
+
+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");
+    return success_with_kernel_reg;
+  }
+  return Error::Ok;
+}
+
+} // namespace executor
+} // namespace torch

codegen/templates/RegisterKernels.h

@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+// ${generated_comment}
+// Exposing an API for registering all kernels at once.
+#include <executorch/runtime/core/evalue.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/kernel/operator_registry.h>
+#include <executorch/runtime/platform/profiler.h>
+
+namespace torch {
+namespace executor {
+
+Error register_all_kernels();
+
+} // namespace executor
+} // namespace torch

runtime/kernel/test/targets.bzl

@@ -75,6 +75,30 @@ def define_common_targets():
         ],
     )
 
+    executorch_generated_lib(
+        name = "test_manual_registration_lib",
+        deps = [
+            ":executorch_all_ops",
+            "//executorch/kernels/portable:operators",
+        ],
+        functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
+        manual_registration = True,
+        visibility = [
+            "//executorch/...",
+        ],
+    )
+
+    runtime.cxx_test(
+        name = "test_kernel_manual_registration",
+        srcs = [
+            "test_kernel_manual_registration.cpp",
+        ],
+        deps = [
+            "//executorch/runtime/kernel:operator_registry",
+            ":test_manual_registration_lib",
+        ],
+    )
+
     for aten_mode in (True, False):
         aten_suffix = "_aten" if aten_mode else ""
 

runtime/kernel/test/test_kernel_manual_registration.cpp

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Include RegisterKernels.h and call register_all_kernels().
+#include <gtest/gtest.h>
+#include <vector>
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/kernel/kernel_runtime_context.h>
+#include <executorch/runtime/kernel/operator_registry.h>
+#include <executorch/runtime/kernel/test/RegisterKernels.h>
+#include <executorch/runtime/platform/runtime.h>
+#include <executorch/test/utils/DeathTest.h>
+
+using namespace ::testing;
+
+namespace torch {
+namespace executor {
+
+class KernelManualRegistrationTest : public ::testing::Test {
+ public:
+  void SetUp() override {
+    torch::executor::runtime_init();
+  }
+};
+
+TEST_F(KernelManualRegistrationTest, ManualRegister) {
+  Error result = register_all_kernels();
+  // Check that we can find the kernel for foo.
+  EXPECT_EQ(result, Error::Ok);
+  EXPECT_FALSE(hasOpsFn("fpp"));
+  EXPECT_TRUE(hasOpsFn("aten::add.out"));
+}
+
+} // namespace executor
+} // namespace torch

shim/xplat/executorch/codegen/codegen.bzl

@@ -20,6 +20,15 @@ GENERATED_SOURCES = [
     "RegisterCodegenUnboxedKernelsEverything.cpp",
 ]
 
+MANUAL_REGISTRATION_SOURCES = [
+    # buildifier: keep sorted
+    "RegisterKernelsEverything.cpp",
+]
+
+MANUAL_REGISTRATION_HEADERS = [
+    "RegisterKernels.h",
+]
+
 # Fake kernels only return `out` or any other tensor from arguments
 CUSTOM_OPS_DUMMY_KERNEL_SOURCES = ["Register{}Stub.cpp".format(backend) for backend in STATIC_DISPATCH_BACKENDS]
 
@@ -70,10 +79,14 @@ def et_operator_library(
         **kwargs
     )
 
-def _get_headers(genrule_name, prefix = "", custom_op = None):
+def _get_headers(genrule_name, prefix = "", custom_op = None, manual_registration = False):
+    headers = OPERATOR_HEADERS + (CUSTOM_OPS_NATIVE_FUNCTION_HEADER if custom_op else [])
     return {
         prefix + f: ":{}[{}]".format(genrule_name, f)
-        for f in OPERATOR_HEADERS + (CUSTOM_OPS_NATIVE_FUNCTION_HEADER if custom_op else [])
+        for f in (MANUAL_REGISTRATION_HEADERS if manual_registration else [])
+    }, {
+        prefix + f: ":{}[{}]".format(genrule_name, f)
+        for f in headers
     }
 
 def _prepare_genrule_and_lib(
@@ -82,6 +95,7 @@ def _prepare_genrule_and_lib(
         custom_ops_yaml_path = None,
         custom_ops_aten_kernel_deps = [],
         custom_ops_requires_runtime_registration = True,
+        manual_registration = False,
         aten_mode = False):
     """
     This function returns two dicts `genrules` and `libs`, derived from the arguments being passed
@@ -122,15 +136,18 @@ def _prepare_genrule_and_lib(
         # actually-generated files matches GENERATED_FILES.
     ]
 
+    # Sources for generated kernel registration lib
+    sources = MANUAL_REGISTRATION_SOURCES if manual_registration else GENERATED_SOURCES
+
     # The command will always generate these files.
-    genrule_outs = GENERATED_SOURCES + OPERATOR_HEADERS + (CUSTOM_OPS_NATIVE_FUNCTION_HEADER if custom_ops_yaml_path else [])
+    genrule_outs = sources + OPERATOR_HEADERS + (CUSTOM_OPS_NATIVE_FUNCTION_HEADER if custom_ops_yaml_path else []) + MANUAL_REGISTRATION_HEADERS
 
     genrules = {}
     libs = {}
 
     # if aten_mode is true, we don't need functions_yaml_path
     genrule_name = name + "_combined"
-    headers = _get_headers(genrule_name = genrule_name, custom_op = custom_ops_yaml_path)
+    exported_headers, headers = _get_headers(genrule_name = genrule_name, custom_op = custom_ops_yaml_path, manual_registration = manual_registration)
 
     # need to register ATen ops into Executorch runtime:
     need_reg_aten_ops = aten_mode or functions_yaml_path
@@ -149,6 +166,10 @@ def _prepare_genrule_and_lib(
         ]
     if aten_mode:
         genrule_cmd = genrule_cmd + ["--use_aten_lib"]
+    if manual_registration:
+        genrule_cmd = genrule_cmd + [
+            "--manual_registration",
+        ]
     if custom_ops_yaml_path:
         genrule_cmd = genrule_cmd + [
             "--custom_ops_yaml_path=" + custom_ops_yaml_path,
@@ -160,13 +181,15 @@ def _prepare_genrule_and_lib(
 
     if need_reg_ops:
         libs[name] = {
+            "exported_headers": exported_headers,
             "genrule": genrule_name,
             "headers": headers,
-            "srcs": GENERATED_SOURCES,
+            "srcs": sources,
         }
 
     header_lib = name + "_headers"
     libs[header_lib] = {
+        "exported_headers": exported_headers,
         "headers": headers,
     }
     return genrules, libs
@@ -303,6 +326,7 @@ def executorch_generated_lib(
         custom_ops_requires_runtime_registration = True,
         visibility = [],
         aten_mode = False,
+        manual_registration = False,
         use_default_aten_ops_lib = True,
         deps = [],
         xplat_deps = [],
@@ -350,6 +374,7 @@ def executorch_generated_lib(
         visibility: Visibility of the C++ library targets.
         deps: Additinal deps of the main C++ library. Needs to be in either `//executorch` or `//caffe2` module.
         platforms: platforms args to runtime.cxx_library (only used when in xplat)
+        manual_registration: if true, generate RegisterKernels.cpp and RegisterKernels.h.
         use_default_aten_ops_lib: If `aten_mode` is True AND this flag is True, use `torch_mobile_all_ops` for ATen operator library.
         xplat_deps: Additional xplat deps, can be used to provide custom operator library.
         fbcode_deps: Additional fbcode deps, can be used to provide custom operator library.
@@ -394,6 +419,7 @@ def executorch_generated_lib(
         custom_ops_aten_kernel_deps = custom_ops_aten_kernel_deps,
         custom_ops_requires_runtime_registration = custom_ops_requires_runtime_registration,
         aten_mode = aten_mode,
+        manual_registration = manual_registration,
     )
 
     # genrule for selective build from static operator list
@@ -457,6 +483,7 @@ def executorch_generated_lib(
             # target, and are not meant to be used by targets outside of this
             # directory.
             headers = libs[lib_name]["headers"],
+            exported_headers = libs[lib_name]["exported_headers"],
             exported_preprocessor_flags = ["-DUSE_ATEN_LIB"] if aten_mode else [],
             # link_whole is necessary because the operators register themselves via
             # static initializers that run at program startup.