Added support for --lib-name in manual kernel registration

Author: Tanish2101

codegen/gen.py

@@ -323,6 +323,7 @@ def gen_unboxing(
     use_aten_lib: bool,
     kernel_index: ETKernelIndex,
     manual_registration: bool,
+    lib_name: Optional[str] = None,
     add_exception_boundary: bool = False,
 ) -> None:
     # Iterable type for write_sharded is a Tuple of (native_function, (kernel_key, metadata))
@@ -339,7 +340,9 @@ def key_func(
 
     header = ["Functions.h" if use_aten_lib else "NativeFunctions.h"]
     filename = (
-        "RegisterKernels.cpp"
+        f"register_{lib_name}_kernels.cpp"
+        if manual_registration and lib_name
+        else "RegisterKernels.cpp"
         if manual_registration
         else "RegisterCodegenUnboxedKernels.cpp"
     )
@@ -356,9 +359,10 @@ def key_func(
             "fn_header": (
                 header if unbox_kernel_entry == items[0] else []
             ),  # Only write header once
+            "lib_name": lib_name or "all",
         },
         num_shards=1,
-        sharded_keys={"unboxed_kernels", "fn_header"},
+        sharded_keys={"unboxed_kernels", "fn_header", "lib_name"},
     )
 
 
@@ -953,6 +957,12 @@ def main() -> None:
         help="a boolean flag to indicate whether we want to manually call"
         "register_kernels() or rely on static init. ",
     )
+    parser.add_argument(
+        "--lib-name",
+        type=str,
+        default=None,
+        help="Optional library name used to customize the generated register_<lib>_kernels() function and file names.",
+    )
     parser.add_argument(
         "--generate",
         type=str,
@@ -1015,6 +1025,7 @@ def main() -> None:
             kernel_index=kernel_index,
             manual_registration=options.manual_registration,
             add_exception_boundary=options.add_exception_boundary,
+            lib_name=options.lib_name,
         )
         if custom_ops_native_functions:
             gen_custom_ops(

codegen/templates/RegisterKernels.cpp

@@ -15,14 +15,14 @@
 namespace torch {
 namespace executor {
 
-Error register_all_kernels() {
+Error register_${lib_name}_kernels() {
   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 ${lib_name} kernels");
     return success_with_kernel_reg;
   }
   return Error::Ok;

codegen/templates/RegisterKernels.h

@@ -16,7 +16,7 @@
 namespace torch {
 namespace executor {
 
-Error register_all_kernels();
+Error register_${lib_name}_kernels();
 
 } // namespace executor
 } // namespace torch

docs/source/kernel-library-selective-build.md

@@ -90,3 +90,37 @@ cmake -D… -DSELECT_OPS_YAML=ON
 ```
 
 To select from either an operator name list or a schema yaml from kernel library.
+
+## Manual Kernel Registration with '--lib-name'
+ExecuTorch now supports generating library-specific kernel registration APIs using the '--lib-name' option along with '--manual-registration' during codegen. This allows applications to avoid using static initialization or linker flags like '-force_load' when linking in kernel libraries.
+
+## Motivation
+In environments like Xcode, using static libraries requires developers to manually specify '-force_load' flags to ensure kernel registration code is executed. This is inconvenient and error-prone.
+
+By passing a library name to the codegen script, developers can generate explicit registration functions and headers, which they can call directly in their application.
+
+## How to Use
+Run the codegen script like this:
+
+```
+python -m codegen.gen \
+  --functions-yaml-path=path/to/functions.yaml \
+  --manual-registration \
+  --lib-name=custom
+```
+This will generate:
+
+'register_custom_kernels.cpp' defines 'register_custom_kernels()' with only the kernels selected and 'register_custom_kernels.h' declares the function for inclusion in your application
+
+Then in your application, call:
+
+```
+#include "register_custom_kernels.h"
+
+register_custom_kernels(); // Registers only the "custom" kernels
+```
+
+This avoids relying on static initialization and enables you to register only the kernels you want.
+
+### Compatibility
+If '--lib-name' is not passed, the default behavior remains unchanged, the codegen script will generate a general 'RegisterKernels.cpp' and 'register_all_kernels()' function.

tools/cmake/Codegen.cmake

@@ -142,6 +142,10 @@ function(generate_bindings_for_kernels)
     set(_gen_command "${_gen_command}" --add-exception-boundary)
   endif()
 
+  if(GEN_LIB_NAME)
+    list(APPEND _gen_command --lib-name=${GEN_LIB_NAME})
+  endif()
+
   set(_gen_command_sources
       ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp
       ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h