Add named manual kernel registration API

codegen/gen.py

@@ -2,6 +2,7 @@
 
 import argparse
 import os
+import re
 from collections import defaultdict
 from dataclasses import dataclass
 from pathlib import Path
@@ -79,6 +80,31 @@
     from torchgen.selective_build.selector import SelectiveBuilder
 
 
+DEFAULT_MANUAL_REGISTRATION_FUNCTION_NAME = "register_all_kernels"
+CPP_IDENTIFIER_PATTERN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
+
+
+def get_manual_registration_function_name(
+    *,
+    manual_registration: bool,
+    manual_registration_lib_name: str | None,
+) -> str:
+    if not manual_registration_lib_name:
+        return DEFAULT_MANUAL_REGISTRATION_FUNCTION_NAME
+
+    if not manual_registration:
+        raise ValueError(
+            "--manual-registration-lib-name requires --manual-registration"
+        )
+
+    if not CPP_IDENTIFIER_PATTERN.fullmatch(manual_registration_lib_name):
+        raise ValueError(
+            "--manual-registration-lib-name must be a valid C++ identifier"
+        )
+
+    return f"register_{manual_registration_lib_name}_kernels"
+
+
 def _sig_decl_wrapper(sig: CppSignature | ExecutorchCppSignature) -> str:
     """
     A wrapper function to basically get `sig.decl(include_context=True)`.
@@ -329,6 +355,9 @@ def gen_unboxing(
     use_aten_lib: bool,
     kernel_index: ETKernelIndex,
     manual_registration: bool,
+    manual_registration_function_name: str = (
+        DEFAULT_MANUAL_REGISTRATION_FUNCTION_NAME
+    ),
     add_exception_boundary: bool = False,
 ) -> None:
     # Iterable type for write_sharded is a Tuple of (native_function, (kernel_key, metadata))
@@ -364,6 +393,9 @@ def key_func(
             ),  # Only write header once
         },
         num_shards=1,
+        base_env={
+            "manual_registration_function_name": manual_registration_function_name,
+        },
         sharded_keys={"unboxed_kernels", "fn_header"},
     )
 
@@ -484,6 +516,9 @@ def gen_headers(
     kernel_index: ETKernelIndex,
     cpu_fm: FileManager,
     use_aten_lib: bool,
+    manual_registration_function_name: str = (
+        DEFAULT_MANUAL_REGISTRATION_FUNCTION_NAME
+    ),
 ) -> None:
     """Generate headers.
 
@@ -534,6 +569,7 @@ def gen_headers(
         "RegisterKernels.h",
         lambda: {
             "generated_comment": "@" + "generated by torchgen/gen_executorch.py",
+            "manual_registration_function_name": manual_registration_function_name,
         },
     )
     headers = {
@@ -958,7 +994,14 @@ def main() -> None:
         "--manual-registration",
         action="store_true",
         help="a boolean flag to indicate whether we want to manually call"
-        "register_kernels() or rely on static init. ",
+        "register_all_kernels() or rely on static init. ",
+    )
+    parser.add_argument(
+        "--manual-registration-lib-name",
+        "--manual_registration_lib_name",
+        help="library name to include in the manual registration API name. "
+        "The library name must be a valid C++ identifier. Requires "
+        "--manual-registration.",
     )
     parser.add_argument(
         "--generate",
@@ -977,6 +1020,13 @@ def main() -> None:
     )
     options = parser.parse_args()
     assert options.tags_path, "tags.yaml is required by codegen yaml parsing."
+    try:
+        manual_registration_function_name = get_manual_registration_function_name(
+            manual_registration=options.manual_registration,
+            manual_registration_lib_name=options.manual_registration_lib_name,
+        )
+    except ValueError as exc:
+        parser.error(str(exc))
 
     selector = get_custom_build_selector(
         options.op_registration_whitelist,
@@ -1011,6 +1061,7 @@ def main() -> None:
             kernel_index=kernel_index,
             cpu_fm=cpu_fm,
             use_aten_lib=options.use_aten_lib,
+            manual_registration_function_name=manual_registration_function_name,
         )
 
     if "sources" in options.generate:
@@ -1021,6 +1072,7 @@ def main() -> None:
             use_aten_lib=options.use_aten_lib,
             kernel_index=kernel_index,
             manual_registration=options.manual_registration,
+            manual_registration_function_name=manual_registration_function_name,
             add_exception_boundary=options.add_exception_boundary,
         )
         if custom_ops_native_functions:

codegen/templates/RegisterKernels.cpp

@@ -7,23 +7,23 @@
  */
 
 // ${generated_comment}
-// This implements register_all_kernels() API that is declared in
-// RegisterKernels.h
+// This implements ${manual_registration_function_name}() API that is declared
+// in RegisterKernels.h
 #include "RegisterKernels.h"
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include "${fn_header}" // Generated Function import headers
 
 namespace torch {
 namespace executor {
 
-Error register_all_kernels() {
+Error ${manual_registration_function_name}() {
   Kernel kernels_to_register[] = {
       ${unboxed_kernels} // Generated kernels
   };
   Error success_with_kernel_reg =
       ::executorch::runtime::register_kernels({kernels_to_register});
   if (success_with_kernel_reg != Error::Ok) {
-    ET_LOG(Error, "Failed register all kernels");
+    ET_LOG(Error, "Failed to register kernels");
     return success_with_kernel_reg;
   }
   return Error::Ok;

codegen/templates/RegisterKernels.h

@@ -7,7 +7,7 @@
  */
 
 // ${generated_comment}
-// Exposing an API for registering all kernels at once.
+// Exposing an API for registering generated 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>
@@ -16,7 +16,7 @@
 namespace torch {
 namespace executor {
 
-Error register_all_kernels();
+Error ${manual_registration_function_name}();
 
 } // namespace executor
 } // namespace torch

codegen/test/test_executorch_gen.py

@@ -14,6 +14,9 @@
 from executorch.codegen.gen import (
     ComputeCodegenUnboxedKernels,
     gen_functions_declarations,
+    gen_headers,
+    gen_unboxing,
+    get_manual_registration_function_name,
     parse_yaml_files,
     translate_native_yaml,
 )
@@ -29,6 +32,7 @@
     OperatorName,
 )
 from torchgen.selective_build.selector import SelectiveBuilder
+from torchgen.utils import FileManager
 
 
 TEST_YAML = """
@@ -318,6 +322,110 @@ def tearDown(self) -> None:
             pass
 
 
+class TestManualRegistrationFunctionName(unittest.TestCase):
+    def test_default_function_name(self) -> None:
+        self.assertEqual(
+            get_manual_registration_function_name(
+                manual_registration=True,
+                manual_registration_lib_name=None,
+            ),
+            "register_all_kernels",
+        )
+
+    def test_named_function_name(self) -> None:
+        self.assertEqual(
+            get_manual_registration_function_name(
+                manual_registration=True,
+                manual_registration_lib_name="portable_ops_lib",
+            ),
+            "register_portable_ops_lib_kernels",
+        )
+
+    def test_named_function_rejects_invalid_library_name(self) -> None:
+        with self.assertRaisesRegex(ValueError, "valid C\\+\\+ identifier"):
+            get_manual_registration_function_name(
+                manual_registration=True,
+                manual_registration_lib_name="portable-ops.lib::debug",
+            )
+
+    def test_named_function_requires_manual_registration(self) -> None:
+        with self.assertRaisesRegex(
+            ValueError, "--manual-registration-lib-name requires"
+        ):
+            get_manual_registration_function_name(
+                manual_registration=False,
+                manual_registration_lib_name="portable_ops_lib",
+            )
+
+    def test_named_function_rejects_leading_digit(self) -> None:
+        with self.assertRaisesRegex(ValueError, "valid C\\+\\+ identifier"):
+            get_manual_registration_function_name(
+                manual_registration=True,
+                manual_registration_lib_name="1_portable_ops_lib",
+            )
+
+
+class TestManualRegistrationTemplates(unittest.TestCase):
+    def setUp(self) -> None:
+        self.template_dir = os.path.join(
+            os.path.dirname(os.path.dirname(__file__)), "templates"
+        )
+        self.function_name = "register_portable_ops_lib_kernels"
+
+    def test_register_kernels_header_uses_named_function(self) -> None:
+        with tempfile.TemporaryDirectory() as tempdir:
+            gen_headers(
+                native_functions=[],
+                gen_custom_ops_header=False,
+                custom_ops_native_functions=[],
+                selector=SelectiveBuilder.get_nop_selector(),
+                kernel_index=ETKernelIndex(index={}),  # type: ignore[arg-type]
+                cpu_fm=FileManager(tempdir, self.template_dir, False),
+                use_aten_lib=False,
+                manual_registration_function_name=self.function_name,
+            )
+
+            with open(os.path.join(tempdir, "RegisterKernels.h")) as f:
+                header = f.read()
+
+        self.assertIn(f"Error {self.function_name}();", header)
+        self.assertNotIn("Error register_all_kernels();", header)
+
+    def test_register_kernels_cpp_uses_named_function(self) -> None:
+        with tempfile.TemporaryDirectory() as tempdir:
+            native_function, backend_index = NativeFunction.from_yaml(
+                {
+                    "func": "custom_1::op_1() -> bool",
+                    "dispatch": {"CPU": "kernel_1"},
+                },
+                loc=Location(__file__, 1),
+                valid_tags=set(),
+            )
+            backend_indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]] = {
+                DispatchKey.CPU: {},
+                DispatchKey.QuantizedCPU: {},
+            }
+            BackendIndex.grow_index(backend_indices, backend_index)
+
+            gen_unboxing(
+                native_functions=[native_function],
+                cpu_fm=FileManager(tempdir, self.template_dir, False),
+                selector=SelectiveBuilder.from_yaml_dict(
+                    {"include_all_operators": True}
+                ),
+                use_aten_lib=False,
+                kernel_index=ETKernelIndex.from_backend_indices(backend_indices),
+                manual_registration=True,
+                manual_registration_function_name=self.function_name,
+            )
+
+            with open(os.path.join(tempdir, "RegisterKernelsEverything.cpp")) as f:
+                source = f.read()
+
+        self.assertIn(f"Error {self.function_name}() {{", source)
+        self.assertNotIn("Error register_all_kernels() {", source)
+
+
 class TestGenFunctionsDeclarations(unittest.TestCase):
     def setUp(self) -> None:
         (

examples/portable/custom_ops/CMakeLists.txt

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -56,6 +56,9 @@
   "Register whether custom op 1 (my_ops::mul3) or custom op 2 (my_ops::mul4) \
   or no custom op at all." OFF
 )
+option(TEST_NAMED_MANUAL_REGISTRATION
+       "Build two custom-op libraries with named manual registration." OFF
+)
 # ------------------------------- OPTIONS END --------------------------------
 
 #
@@ -129,3 +132,45 @@
 target_compile_options(
   custom_ops_executor_runner PUBLIC ${_common_compile_options}
 )
+
+if(TEST_NAMED_MANUAL_REGISTRATION)
+  function(add_manual_registration_lib lib_name op_name kernel_source)
+    gen_selected_ops(LIB_NAME "${lib_name}" ROOT_OPS "${op_name}")
+    generate_bindings_for_kernels(
+      LIB_NAME "${lib_name}" CUSTOM_OPS_YAML
+      ${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml MANUAL_REGISTRATION
+    )
+
+    add_library(${lib_name}_kernels ${kernel_source})
+    target_link_libraries(${lib_name}_kernels PRIVATE executorch)
+    target_compile_options(
+      ${lib_name}_kernels PUBLIC ${_common_compile_options}
+    )
+    gen_operators_lib(
+      LIB_NAME "${lib_name}" KERNEL_LIBS ${lib_name}_kernels DEPS executorch
+      MANUAL_REGISTRATION
+    )
+  endfunction()
+
+  add_manual_registration_lib(
+    manual_ops_1_lib "my_ops::mul3.out"
+    ${CMAKE_CURRENT_LIST_DIR}/custom_ops_1_out.cpp
+  )
+  add_manual_registration_lib(
+    manual_ops_2_lib "my_ops::mul4.out"
+    ${CMAKE_CURRENT_LIST_DIR}/custom_ops_2_out.cpp
+  )
+
+  add_executable(
+    named_manual_registration_test named_manual_registration_test.cpp
+  )
+  add_dependencies(
+    named_manual_registration_test manual_ops_1_lib manual_ops_2_lib
+  )
+  target_include_directories(
+    named_manual_registration_test PRIVATE ${CMAKE_CURRENT_BINARY_DIR}
+  )
+  target_link_libraries(
+    named_manual_registration_test manual_ops_1_lib manual_ops_2_lib executorch
+  )
+endif()

examples/portable/custom_ops/named_manual_registration_test.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.
+ */
+
+#include "manual_ops_1_lib/RegisterKernels.h"
+#include "manual_ops_2_lib/RegisterKernels.h"
+
+#include <executorch/runtime/kernel/operator_registry.h>
+
+int main() {
+  if (torch::executor::register_manual_ops_1_lib_kernels() !=
+      executorch::runtime::Error::Ok) {
+    return 1;
+  }
+  if (torch::executor::register_manual_ops_2_lib_kernels() !=
+      executorch::runtime::Error::Ok) {
+    return 2;
+  }
+  if (!executorch::runtime::registry_has_op_function("my_ops::mul3.out")) {
+    return 3;
+  }
+  if (!executorch::runtime::registry_has_op_function("my_ops::mul4.out")) {
+    return 4;
+  }
+  return 0;
+}