Final changes done based on review comments

Author: Tanish2101

codegen/gen.py

@@ -340,9 +340,7 @@ def key_func(
 
     header = ["Functions.h" if use_aten_lib else "NativeFunctions.h"]
     filename = (
-        f"register_{lib_name}_kernels.cpp"
-        if manual_registration and lib_name
-        else "RegisterKernels.cpp"
+        "RegisterKernels.cpp"
         if manual_registration
         else "RegisterCodegenUnboxedKernels.cpp"
     )
@@ -359,10 +357,11 @@ def key_func(
             "fn_header": (
                 header if unbox_kernel_entry == items[0] else []
             ),  # Only write header once
-            "lib_name": lib_name or "all",
+            "lib_name": lib_name or "", 
+            "use_lib_name_in_register": bool(lib_name),
         },
         num_shards=1,
-        sharded_keys={"unboxed_kernels", "fn_header", "lib_name"},
+        sharded_keys={"unboxed_kernels", "fn_header", "lib_name", "use_lib_name_in_register"},
     )
 
 
@@ -1015,7 +1014,8 @@ def main() -> None:
             cpu_fm=cpu_fm,
             use_aten_lib=options.use_aten_lib,
         )
-
+    if options.lib_name and not options.manual_registration:
+        raise ValueError("--lib-name can only be used with --manual-registration")
     if "sources" in options.generate:
         gen_unboxing(
             native_functions=native_functions,

codegen/templates/RegisterKernels.cpp

@@ -15,16 +15,30 @@
 namespace torch {
 namespace executor {
 
-Error register_${lib_name}_kernels() {
+#if USE_LIB_NAME_IN_REGISTER
+Error register_kernels_${lib_name}() {
+#else
+Error register_all_kernels() {
+#endif
+
   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 to register ${lib_name} kernels");
+    #if USE_LIB_NAME_IN_REGISTER
+      ET_LOG(Error, "Failed to register %zu kernels for %s (from %s)",
+            sizeof(kernels_to_register) / sizeof(Kernel),
+            "${lib_name}",
+            __FILE__);
+    #else
+      ET_LOG(Error, "Failed to register %zu kernels (from %s)",
+            sizeof(kernels_to_register) / sizeof(Kernel),
+            __FILE__);
+    #endif
     return success_with_kernel_reg;
-  }
+}
   return Error::Ok;
 }
 

codegen/templates/RegisterKernels.h

@@ -16,7 +16,11 @@
 namespace torch {
 namespace executor {
 
-Error register_${lib_name}_kernels();
+#if USE_LIB_NAME_IN_REGISTER
+Error register_kernels_${lib_name}();
+#else
+Error register_all_kernels();
+#endif
 
 } // namespace executor
 } // namespace torch

codegen/test/test_executorch_gen.py

@@ -9,6 +9,9 @@
 import os
 import tempfile
 import unittest
+import shutil
+import subprocess
+import sys
 
 import yaml
 from executorch.codegen.gen import (
@@ -693,3 +696,31 @@ def test_codegen_unboxed_default_kernel_key_selected(self) -> None:
         )
 
         self.assertEqual(expected_str, result)
+
+class TestGenMainArgumentChecks(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.mkdtemp()
+        self.dummy_yaml = os.path.join(self.temp_dir, "dummy.yaml")
+        with open(self.dummy_yaml, "w") as f:
+            f.write("- tag: dummy\n")
+
+    def tearDown(self):
+        shutil.rmtree(self.temp_dir)
+
+    def test_lib_name_without_manual_registration_raises(self):
+        script_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../gen.py"))
+        result = subprocess.run(
+            [
+                sys.executable,
+                script_path,
+                "--lib-name", "foo",
+                "--tags-path", self.dummy_yaml,
+                "--aten-yaml-path", self.dummy_yaml,
+                "--functions-yaml-path", self.dummy_yaml,
+            ],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        self.assertNotEqual(result.returncode, 0)
+        self.assertIn("--lib-name can only be used with --manual-registration", result.stderr)

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

@@ -45,7 +45,7 @@ gen_selected_ops(
   ROOT_OPS              # comma separated operator names to be selected
   INCLUDE_ALL_OPS       # boolean flag to include all operators
   OPS_FROM_MODEL        # path to a pte file of model to select operators from
-  DTYPE_SELECTIVE_BUILD # boolean flag to enable dtye selection
+  DTYPE_SELECTIVE_BUILD # boolean flag to enable dtype selection
 )
 ```
 
@@ -112,14 +112,14 @@ python -m codegen.gen \
 ```
 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
+`register_kernels_custom.cpp` defines `register_kernels_custom()` with only the kernels selected and `register_kernels_custom.h` declares the function for inclusion in your application
 
 Then in your application, call:
 
 ```
-#include "register_custom_kernels.h"
+#include "register_kernels_custom.h"
 
-register_custom_kernels(); // Registers only the "custom" kernels
+register_kernels_custom(); // Registers only the "custom" kernels
 ```
 
 This avoids relying on static initialization and enables you to register only the kernels you want.

tools/cmake/Codegen.cmake

@@ -92,7 +92,7 @@ endfunction()
 # Invoked as generate_bindings_for_kernels( LIB_NAME lib_name FUNCTIONS_YAML
 # functions_yaml CUSTOM_OPS_YAML custom_ops_yaml )
 function(generate_bindings_for_kernels)
-  set(options ADD_EXCEPTION_BOUNDARY)
+  set(options ADD_EXCEPTION_BOUNDARY MANUAL_REGISTRATION)
   set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML DTYPE_SELECTIVE_BUILD)
   cmake_parse_arguments(GEN "${options}" "${arg_names}" "" ${ARGN})
 
@@ -102,11 +102,17 @@ function(generate_bindings_for_kernels)
   message(STATUS "  CUSTOM_OPS_YAML: ${GEN_CUSTOM_OPS_YAML}")
   message(STATUS "  ADD_EXCEPTION_BOUNDARY: ${GEN_ADD_EXCEPTION_BOUNDARY}")
   message(STATUS "  DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}")
+  message(STATUS "  MANUAL_REGISTRATION: ${GEN_MANUAL_REGISTRATION}")
 
   # Command to generate selected_operators.yaml from custom_ops.yaml.
   file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*")
 
-  set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})
+  if(GEN_LIB_NAME)
+    set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})
+  else()
+    set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/codegen_output)
+  endif()
+  
   # By default selective build output is selected_operators.yaml
   set(_oplist_yaml ${_out_dir}/selected_operators.yaml)
 
@@ -142,7 +148,7 @@ function(generate_bindings_for_kernels)
     set(_gen_command "${_gen_command}" --add-exception-boundary)
   endif()
 
-  if(GEN_LIB_NAME)
+  if(GEN_LIB_NAME AND GEN_MANUAL_REGISTRATION)
     list(APPEND _gen_command --lib-name=${GEN_LIB_NAME})
   endif()