From 5acef861b6d97512322c5d11cc7cf4ecfaa28103 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Tue, 30 Jun 2026 23:38:30 -0700 Subject: [PATCH 1/6] Add named manual kernel registration API Signed-off-by: goutamadwant --- codegen/gen.py | 55 +++++++++- codegen/templates/RegisterKernels.cpp | 6 +- codegen/templates/RegisterKernels.h | 4 +- codegen/test/test_executorch_gen.py | 103 +++++++++++++++++++ shim_et/xplat/executorch/codegen/codegen.bzl | 10 ++ tools/cmake/Codegen.cmake | 50 +++++++-- 6 files changed, 212 insertions(+), 16 deletions(-) diff --git a/codegen/gen.py b/codegen/gen.py index 9b102aa1bf6..1d905e6db9b 100644 --- a/codegen/gen.py +++ b/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" +MANUAL_REGISTRATION_LIB_NAME_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 MANUAL_REGISTRATION_LIB_NAME_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,15 @@ 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", + "--lib-name", + "--lib_name", + help="library name to include in the manual registration API name. " + "Requires --manual-registration.", ) parser.add_argument( "--generate", @@ -977,6 +1021,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 +1062,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 +1073,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: diff --git a/codegen/templates/RegisterKernels.cpp b/codegen/templates/RegisterKernels.cpp index 0d173f52bb3..80a32deb3e4 100644 --- a/codegen/templates/RegisterKernels.cpp +++ b/codegen/templates/RegisterKernels.cpp @@ -7,7 +7,7 @@ */ // ${generated_comment} -// This implements register_all_kernels() API that is declared in +// This implements ${manual_registration_function_name}() API that is declared in // RegisterKernels.h #include "RegisterKernels.h" #include @@ -16,14 +16,14 @@ 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; diff --git a/codegen/templates/RegisterKernels.h b/codegen/templates/RegisterKernels.h index 3c7ecff50b5..4d4fd75dd52 100644 --- a/codegen/templates/RegisterKernels.h +++ b/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 #include #include @@ -16,7 +16,7 @@ namespace torch { namespace executor { -Error register_all_kernels(); +Error ${manual_registration_function_name}(); } // namespace executor } // namespace torch diff --git a/codegen/test/test_executorch_gen.py b/codegen/test/test_executorch_gen.py index cd445acee62..5a4d8ed6298 100644 --- a/codegen/test/test_executorch_gen.py +++ b/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,105 @@ 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_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_requires_cpp_identifier(self) -> None: + with self.assertRaisesRegex(ValueError, "valid C\\+\\+ identifier"): + get_manual_registration_function_name( + manual_registration=True, + manual_registration_lib_name="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: ( diff --git a/shim_et/xplat/executorch/codegen/codegen.bzl b/shim_et/xplat/executorch/codegen/codegen.bzl index 318996784a1..f26d3fc704f 100644 --- a/shim_et/xplat/executorch/codegen/codegen.bzl +++ b/shim_et/xplat/executorch/codegen/codegen.bzl @@ -273,6 +273,7 @@ def _prepare_genrule_and_lib( custom_ops_yaml_path = None, custom_ops_requires_runtime_registration = True, manual_registration = False, + manual_registration_lib_name = None, aten_mode = False, support_exceptions = True): """ @@ -350,6 +351,10 @@ def _prepare_genrule_and_lib( genrule_cmd = genrule_cmd + [ "--manual_registration", ] + if manual_registration_lib_name: + genrule_cmd = genrule_cmd + [ + "--manual-registration-lib-name={}".format(manual_registration_lib_name), + ] if custom_ops_yaml_path: genrule_cmd = genrule_cmd + [ "--custom_ops_yaml_path=" + custom_ops_yaml_path, @@ -828,6 +833,7 @@ def executorch_generated_lib( visibility = [], aten_mode = False, manual_registration = False, + manual_registration_lib_name = None, use_default_aten_ops_lib = True, deps = [], xplat_deps = [], @@ -888,6 +894,9 @@ def executorch_generated_lib( 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. compiler_flags: compiler_flags args to runtime.cxx_library + manual_registration_lib_name: Optional C++ identifier to use when + generating a named manual registration API. If omitted, manual + registration keeps using `register_all_kernels`. dtype_selective_build: In additional to operator selection, dtype selective build further selects the dtypes for each operator. Can be used with model or dict selective build APIs, where dtypes can be specified. @@ -999,6 +1008,7 @@ def executorch_generated_lib( custom_ops_requires_runtime_registration = custom_ops_requires_runtime_registration, aten_mode = aten_mode, manual_registration = manual_registration, + manual_registration_lib_name = manual_registration_lib_name, support_exceptions = support_exceptions, ) diff --git a/tools/cmake/Codegen.cmake b/tools/cmake/Codegen.cmake index 4253fa44dc5..9436540eddc 100644 --- a/tools/cmake/Codegen.cmake +++ b/tools/cmake/Codegen.cmake @@ -154,9 +154,9 @@ endfunction() # custom_ops_yaml. # # Invoked as generate_bindings_for_kernels( LIB_NAME lib_name FUNCTIONS_YAML -# functions_yaml CUSTOM_OPS_YAML custom_ops_yaml ) +# functions_yaml CUSTOM_OPS_YAML custom_ops_yaml [MANUAL_REGISTRATION] ) 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}) @@ -165,6 +165,7 @@ function(generate_bindings_for_kernels) message(STATUS " FUNCTIONS_YAML: ${GEN_FUNCTIONS_YAML}") message(STATUS " CUSTOM_OPS_YAML: ${GEN_CUSTOM_OPS_YAML}") message(STATUS " ADD_EXCEPTION_BOUNDARY: ${GEN_ADD_EXCEPTION_BOUNDARY}") + message(STATUS " MANUAL_REGISTRATION: ${GEN_MANUAL_REGISTRATION}") message(STATUS " DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}") # Command to generate selected_operators.yaml from custom_ops.yaml. @@ -205,11 +206,27 @@ function(generate_bindings_for_kernels) if(GEN_ADD_EXCEPTION_BOUNDARY) set(_gen_command "${_gen_command}" --add-exception-boundary) endif() + if(GEN_MANUAL_REGISTRATION) + list( + APPEND + _gen_command + --manual-registration + --manual-registration-lib-name=${GEN_LIB_NAME} + ) + endif() - set(_gen_command_sources - ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp - ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h - ) + if(GEN_MANUAL_REGISTRATION) + set(_gen_command_sources + ${_out_dir}/RegisterKernelsEverything.cpp + ${_out_dir}/RegisterKernels.h ${_out_dir}/Functions.h + ${_out_dir}/NativeFunctions.h + ) + else() + set(_gen_command_sources + ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp + ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h + ) + endif() if(GEN_FUNCTIONS_YAML) list(APPEND _gen_command --functions-yaml-path=${GEN_FUNCTIONS_YAML}) @@ -268,13 +285,15 @@ endfunction() # Generate a runtime lib for registering operators in Executorch function(gen_operators_lib) + set(options MANUAL_REGISTRATION) set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS DTYPE_SELECTIVE_BUILD) - cmake_parse_arguments(GEN "" "" "${multi_arg_names}" ${ARGN}) + cmake_parse_arguments(GEN "${options}" "" "${multi_arg_names}" ${ARGN}) message(STATUS "Generating operator lib:") message(STATUS " LIB_NAME: ${GEN_LIB_NAME}") message(STATUS " KERNEL_LIBS: ${GEN_KERNEL_LIBS}") message(STATUS " DEPS: ${GEN_DEPS}") + message(STATUS " MANUAL_REGISTRATION: ${GEN_MANUAL_REGISTRATION}") message(STATUS " DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}") set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) @@ -284,9 +303,17 @@ function(gen_operators_lib) add_library(${GEN_LIB_NAME}) - set(_srcs_list ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp - ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h - ) + if(GEN_MANUAL_REGISTRATION) + set(_srcs_list + ${_out_dir}/RegisterKernelsEverything.cpp + ${_out_dir}/RegisterKernels.h ${_out_dir}/Functions.h + ${_out_dir}/NativeFunctions.h + ) + else() + set(_srcs_list ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp + ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h + ) + endif() if(GEN_DTYPE_SELECTIVE_BUILD) list(APPEND _srcs_list ${_opvariant_h}) endif() @@ -373,6 +400,9 @@ function(gen_operators_lib) executorch_target_link_options_shared_lib(${GEN_LIB_NAME}) set(_generated_headers ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h) + if(GEN_MANUAL_REGISTRATION) + list(APPEND _generated_headers ${_out_dir}/RegisterKernels.h) + endif() if(GEN_DTYPE_SELECTIVE_BUILD) list(APPEND _generated_headers ${_opvariant_h}) endif() From 90eede4bf794e9c7461a8dafe9a0746d6fee2ee1 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Tue, 7 Jul 2026 23:22:56 -0700 Subject: [PATCH 2/6] Fix lintrunner formatting Signed-off-by: goutamadwant --- codegen/templates/RegisterKernels.cpp | 4 ++-- codegen/test/test_executorch_gen.py | 4 +--- tools/cmake/Codegen.cmake | 17 ++++++----------- 3 files changed, 9 insertions(+), 16 deletions(-) diff --git a/codegen/templates/RegisterKernels.cpp b/codegen/templates/RegisterKernels.cpp index 80a32deb3e4..5c171d9063c 100644 --- a/codegen/templates/RegisterKernels.cpp +++ b/codegen/templates/RegisterKernels.cpp @@ -7,8 +7,8 @@ */ // ${generated_comment} -// This implements ${manual_registration_function_name}() API that is declared in -// RegisterKernels.h +// This implements ${manual_registration_function_name}() API that is declared +// in RegisterKernels.h #include "RegisterKernels.h" #include #include "${fn_header}" // Generated Function import headers diff --git a/codegen/test/test_executorch_gen.py b/codegen/test/test_executorch_gen.py index 5a4d8ed6298..47b34ba6fd7 100644 --- a/codegen/test/test_executorch_gen.py +++ b/codegen/test/test_executorch_gen.py @@ -394,9 +394,7 @@ def test_register_kernels_cpp_uses_named_function(self) -> None: loc=Location(__file__, 1), valid_tags=set(), ) - backend_indices: dict[ - DispatchKey, dict[OperatorName, BackendMetadata] - ] = { + backend_indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]] = { DispatchKey.CPU: {}, DispatchKey.QuantizedCPU: {}, } diff --git a/tools/cmake/Codegen.cmake b/tools/cmake/Codegen.cmake index 9436540eddc..b98b73cb56f 100644 --- a/tools/cmake/Codegen.cmake +++ b/tools/cmake/Codegen.cmake @@ -207,19 +207,15 @@ function(generate_bindings_for_kernels) set(_gen_command "${_gen_command}" --add-exception-boundary) endif() if(GEN_MANUAL_REGISTRATION) - list( - APPEND - _gen_command - --manual-registration - --manual-registration-lib-name=${GEN_LIB_NAME} + list(APPEND _gen_command --manual-registration + --manual-registration-lib-name=${GEN_LIB_NAME} ) endif() if(GEN_MANUAL_REGISTRATION) set(_gen_command_sources - ${_out_dir}/RegisterKernelsEverything.cpp - ${_out_dir}/RegisterKernels.h ${_out_dir}/Functions.h - ${_out_dir}/NativeFunctions.h + ${_out_dir}/RegisterKernelsEverything.cpp ${_out_dir}/RegisterKernels.h + ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h ) else() set(_gen_command_sources @@ -305,9 +301,8 @@ function(gen_operators_lib) if(GEN_MANUAL_REGISTRATION) set(_srcs_list - ${_out_dir}/RegisterKernelsEverything.cpp - ${_out_dir}/RegisterKernels.h ${_out_dir}/Functions.h - ${_out_dir}/NativeFunctions.h + ${_out_dir}/RegisterKernelsEverything.cpp ${_out_dir}/RegisterKernels.h + ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h ) else() set(_srcs_list ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp From 9ee9ab75296bf227fc897d57f531ca692029ca72 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Thu, 9 Jul 2026 19:50:40 -0700 Subject: [PATCH 3/6] Sanitize manual registration lib names Signed-off-by: goutamadwant --- codegen/gen.py | 27 ++++++++++++++------ codegen/test/test_executorch_gen.py | 15 ++++++++--- shim_et/xplat/executorch/codegen/codegen.bzl | 7 ++--- tools/cmake/Codegen.cmake | 1 + 4 files changed, 36 insertions(+), 14 deletions(-) diff --git a/codegen/gen.py b/codegen/gen.py index 1d905e6db9b..449d23f8f2b 100644 --- a/codegen/gen.py +++ b/codegen/gen.py @@ -81,7 +81,19 @@ DEFAULT_MANUAL_REGISTRATION_FUNCTION_NAME = "register_all_kernels" -MANUAL_REGISTRATION_LIB_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +MANUAL_REGISTRATION_LIB_NAME_SANITIZE_PATTERN = re.compile(r"[^A-Za-z0-9_]+") + + +def sanitize_manual_registration_lib_name(lib_name: str) -> str: + sanitized = MANUAL_REGISTRATION_LIB_NAME_SANITIZE_PATTERN.sub("_", lib_name).strip( + "_" + ) + if not sanitized: + raise ValueError( + "--manual-registration-lib-name must contain at least one " + "alphanumeric or underscore character" + ) + return sanitized def get_manual_registration_function_name( @@ -97,12 +109,10 @@ def get_manual_registration_function_name( "--manual-registration-lib-name requires --manual-registration" ) - if not MANUAL_REGISTRATION_LIB_NAME_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" + sanitized_lib_name = sanitize_manual_registration_lib_name( + manual_registration_lib_name + ) + return f"register_{sanitized_lib_name}_kernels" def _sig_decl_wrapper(sig: CppSignature | ExecutorchCppSignature) -> str: @@ -1002,7 +1012,8 @@ def main() -> None: "--lib-name", "--lib_name", help="library name to include in the manual registration API name. " - "Requires --manual-registration.", + "Characters that are not valid in a C++ identifier are converted to " + "underscores. Requires --manual-registration.", ) parser.add_argument( "--generate", diff --git a/codegen/test/test_executorch_gen.py b/codegen/test/test_executorch_gen.py index 47b34ba6fd7..d86a6b6904a 100644 --- a/codegen/test/test_executorch_gen.py +++ b/codegen/test/test_executorch_gen.py @@ -341,6 +341,15 @@ def test_named_function_name(self) -> None: "register_portable_ops_lib_kernels", ) + def test_named_function_sanitizes_library_name(self) -> None: + self.assertEqual( + get_manual_registration_function_name( + manual_registration=True, + manual_registration_lib_name="portable-ops.lib::debug", + ), + "register_portable_ops_lib_debug_kernels", + ) + def test_named_function_requires_manual_registration(self) -> None: with self.assertRaisesRegex( ValueError, "--manual-registration-lib-name requires" @@ -350,11 +359,11 @@ def test_named_function_requires_manual_registration(self) -> None: manual_registration_lib_name="portable_ops_lib", ) - def test_named_function_requires_cpp_identifier(self) -> None: - with self.assertRaisesRegex(ValueError, "valid C\\+\\+ identifier"): + def test_named_function_requires_name_content(self) -> None: + with self.assertRaisesRegex(ValueError, "alphanumeric or underscore"): get_manual_registration_function_name( manual_registration=True, - manual_registration_lib_name="portable-ops-lib", + manual_registration_lib_name="---", ) diff --git a/shim_et/xplat/executorch/codegen/codegen.bzl b/shim_et/xplat/executorch/codegen/codegen.bzl index f26d3fc704f..129c6a6e0b4 100644 --- a/shim_et/xplat/executorch/codegen/codegen.bzl +++ b/shim_et/xplat/executorch/codegen/codegen.bzl @@ -894,9 +894,10 @@ def executorch_generated_lib( 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. compiler_flags: compiler_flags args to runtime.cxx_library - manual_registration_lib_name: Optional C++ identifier to use when - generating a named manual registration API. If omitted, manual - registration keeps using `register_all_kernels`. + manual_registration_lib_name: Optional library name to include when + generating a named manual registration API. Characters that are not + valid in a C++ identifier are converted to underscores. If omitted, + manual registration keeps using `register_all_kernels`. dtype_selective_build: In additional to operator selection, dtype selective build further selects the dtypes for each operator. Can be used with model or dict selective build APIs, where dtypes can be specified. diff --git a/tools/cmake/Codegen.cmake b/tools/cmake/Codegen.cmake index b98b73cb56f..e16cef97303 100644 --- a/tools/cmake/Codegen.cmake +++ b/tools/cmake/Codegen.cmake @@ -207,6 +207,7 @@ function(generate_bindings_for_kernels) set(_gen_command "${_gen_command}" --add-exception-boundary) endif() if(GEN_MANUAL_REGISTRATION) + # codegen.gen sanitizes GEN_LIB_NAME before embedding it in a C++ symbol. list(APPEND _gen_command --manual-registration --manual-registration-lib-name=${GEN_LIB_NAME} ) From 33b8d67631c8f1997d403f4ea38f7137e616d42e Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Wed, 15 Jul 2026 23:01:52 -0700 Subject: [PATCH 4/6] Remove generic manual registration lib aliases Signed-off-by: goutamadwant --- codegen/gen.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/codegen/gen.py b/codegen/gen.py index 449d23f8f2b..7bdf18a73e2 100644 --- a/codegen/gen.py +++ b/codegen/gen.py @@ -1009,8 +1009,6 @@ def main() -> None: parser.add_argument( "--manual-registration-lib-name", "--manual_registration_lib_name", - "--lib-name", - "--lib_name", help="library name to include in the manual registration API name. " "Characters that are not valid in a C++ identifier are converted to " "underscores. Requires --manual-registration.", From 7c91db794c4855f54e93d10a736c4b20b022d6b8 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Tue, 21 Jul 2026 05:35:36 +0200 Subject: [PATCH 5/6] Harden named manual kernel registration Signed-off-by: goutamadwant --- codegen/gen.py | 24 +++------- codegen/test/test_executorch_gen.py | 14 +++--- examples/portable/custom_ops/CMakeLists.txt | 44 +++++++++++++++++++ .../named_manual_registration_test.cpp | 30 +++++++++++++ .../portable/custom_ops/test_custom_ops.sh | 19 ++++++++ shim_et/xplat/executorch/codegen/codegen.bzl | 3 ++ tools/cmake/Codegen.cmake | 34 +++++++++++--- 7 files changed, 137 insertions(+), 31 deletions(-) create mode 100644 examples/portable/custom_ops/named_manual_registration_test.cpp diff --git a/codegen/gen.py b/codegen/gen.py index 7bdf18a73e2..1b237c068c3 100644 --- a/codegen/gen.py +++ b/codegen/gen.py @@ -81,19 +81,7 @@ DEFAULT_MANUAL_REGISTRATION_FUNCTION_NAME = "register_all_kernels" -MANUAL_REGISTRATION_LIB_NAME_SANITIZE_PATTERN = re.compile(r"[^A-Za-z0-9_]+") - - -def sanitize_manual_registration_lib_name(lib_name: str) -> str: - sanitized = MANUAL_REGISTRATION_LIB_NAME_SANITIZE_PATTERN.sub("_", lib_name).strip( - "_" - ) - if not sanitized: - raise ValueError( - "--manual-registration-lib-name must contain at least one " - "alphanumeric or underscore character" - ) - return sanitized +CPP_IDENTIFIER_PATTERN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") def get_manual_registration_function_name( @@ -109,10 +97,12 @@ def get_manual_registration_function_name( "--manual-registration-lib-name requires --manual-registration" ) - sanitized_lib_name = sanitize_manual_registration_lib_name( - manual_registration_lib_name - ) - return f"register_{sanitized_lib_name}_kernels" + 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: diff --git a/codegen/test/test_executorch_gen.py b/codegen/test/test_executorch_gen.py index d86a6b6904a..15d38c9e8e3 100644 --- a/codegen/test/test_executorch_gen.py +++ b/codegen/test/test_executorch_gen.py @@ -341,14 +341,12 @@ def test_named_function_name(self) -> None: "register_portable_ops_lib_kernels", ) - def test_named_function_sanitizes_library_name(self) -> None: - self.assertEqual( + 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", - ), - "register_portable_ops_lib_debug_kernels", - ) + ) def test_named_function_requires_manual_registration(self) -> None: with self.assertRaisesRegex( @@ -359,11 +357,11 @@ def test_named_function_requires_manual_registration(self) -> None: manual_registration_lib_name="portable_ops_lib", ) - def test_named_function_requires_name_content(self) -> None: - with self.assertRaisesRegex(ValueError, "alphanumeric or underscore"): + 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="---", + manual_registration_lib_name="1_portable_ops_lib", ) diff --git a/examples/portable/custom_ops/CMakeLists.txt b/examples/portable/custom_ops/CMakeLists.txt index 8e679697b47..b9b5d22666b 100644 --- a/examples/portable/custom_ops/CMakeLists.txt +++ b/examples/portable/custom_ops/CMakeLists.txt @@ -56,6 +56,9 @@ option( "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,44 @@ target_link_libraries( 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 + ) + 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() diff --git a/examples/portable/custom_ops/named_manual_registration_test.cpp b/examples/portable/custom_ops/named_manual_registration_test.cpp new file mode 100644 index 00000000000..d79555b025f --- /dev/null +++ b/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 + +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; +} diff --git a/examples/portable/custom_ops/test_custom_ops.sh b/examples/portable/custom_ops/test_custom_ops.sh index 58a7de3a5f2..630aa66ba2c 100644 --- a/examples/portable/custom_ops/test_custom_ops.sh +++ b/examples/portable/custom_ops/test_custom_ops.sh @@ -77,6 +77,24 @@ test_cmake_custom_op_2() { ${build_dir}/custom_ops_executor_runner --model_path="./${model_name}.pte" } +test_cmake_named_manual_registration() { + CMAKE_PREFIX_PATH="$PWD/cmake-out/lib/cmake/ExecuTorch" + + local example_dir=examples/portable/custom_ops + local build_dir=cmake-out/${example_dir} + rm -rf ${build_dir} + retry cmake \ + -DREGISTER_EXAMPLE_CUSTOM_OP=1 \ + -DTEST_NAMED_MANUAL_REGISTRATION=ON \ + -DCMAKE_PREFIX_PATH="$CMAKE_PREFIX_PATH" \ + -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ + -B${build_dir} \ + ${example_dir} + + cmake --build ${build_dir} --target named_manual_registration_test -j9 + ${build_dir}/named_manual_registration_test +} + if [[ -z $PYTHON_EXECUTABLE ]]; then PYTHON_EXECUTABLE=python3 @@ -85,3 +103,4 @@ fi cmake_install_executorch_lib test_cmake_custom_op_1 test_cmake_custom_op_2 +test_cmake_named_manual_registration diff --git a/shim_et/xplat/executorch/codegen/codegen.bzl b/shim_et/xplat/executorch/codegen/codegen.bzl index 129c6a6e0b4..ff4dfbccf34 100644 --- a/shim_et/xplat/executorch/codegen/codegen.bzl +++ b/shim_et/xplat/executorch/codegen/codegen.bzl @@ -303,6 +303,9 @@ def _prepare_genrule_and_lib( }, } """ + if manual_registration_lib_name and not manual_registration: + fail("manual_registration_lib_name requires manual_registration = True") + aten_src_path = runtime.external_dep_location("aten-src-path") genrule_cmd = [ "$(exe //executorch/codegen:gen)", diff --git a/tools/cmake/Codegen.cmake b/tools/cmake/Codegen.cmake index e16cef97303..0fd10469669 100644 --- a/tools/cmake/Codegen.cmake +++ b/tools/cmake/Codegen.cmake @@ -155,6 +155,14 @@ endfunction() # # Invoked as generate_bindings_for_kernels( LIB_NAME lib_name FUNCTIONS_YAML # functions_yaml CUSTOM_OPS_YAML custom_ops_yaml [MANUAL_REGISTRATION] ) +function(_manual_registration_property_name lib_name output_var) + string(SHA256 _lib_name_hash "${lib_name}") + set(${output_var} + "_EXECUTORCH_MANUAL_REGISTRATION_${_lib_name_hash}" + PARENT_SCOPE + ) +endfunction() + function(generate_bindings_for_kernels) set(options ADD_EXCEPTION_BOUNDARY MANUAL_REGISTRATION) set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML DTYPE_SELECTIVE_BUILD) @@ -168,6 +176,19 @@ function(generate_bindings_for_kernels) message(STATUS " MANUAL_REGISTRATION: ${GEN_MANUAL_REGISTRATION}") message(STATUS " DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}") + if(GEN_MANUAL_REGISTRATION AND NOT GEN_LIB_NAME MATCHES + "^[A-Za-z_][A-Za-z0-9_]*$" + ) + message( + FATAL_ERROR + "Manual registration LIB_NAME must be a valid C++ identifier: ${GEN_LIB_NAME}" + ) + endif() + _manual_registration_property_name("${GEN_LIB_NAME}" _registration_property) + set_property( + GLOBAL PROPERTY "${_registration_property}" "${GEN_MANUAL_REGISTRATION}" + ) + # Command to generate selected_operators.yaml from custom_ops.yaml. file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*") @@ -207,7 +228,6 @@ function(generate_bindings_for_kernels) set(_gen_command "${_gen_command}" --add-exception-boundary) endif() if(GEN_MANUAL_REGISTRATION) - # codegen.gen sanitizes GEN_LIB_NAME before embedding it in a C++ symbol. list(APPEND _gen_command --manual-registration --manual-registration-lib-name=${GEN_LIB_NAME} ) @@ -282,15 +302,17 @@ endfunction() # Generate a runtime lib for registering operators in Executorch function(gen_operators_lib) - set(options MANUAL_REGISTRATION) set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS DTYPE_SELECTIVE_BUILD) - cmake_parse_arguments(GEN "${options}" "" "${multi_arg_names}" ${ARGN}) + cmake_parse_arguments(GEN "" "" "${multi_arg_names}" ${ARGN}) + + _manual_registration_property_name("${GEN_LIB_NAME}" _registration_property) + get_property(_manual_registration GLOBAL PROPERTY "${_registration_property}") message(STATUS "Generating operator lib:") message(STATUS " LIB_NAME: ${GEN_LIB_NAME}") message(STATUS " KERNEL_LIBS: ${GEN_KERNEL_LIBS}") message(STATUS " DEPS: ${GEN_DEPS}") - message(STATUS " MANUAL_REGISTRATION: ${GEN_MANUAL_REGISTRATION}") + message(STATUS " MANUAL_REGISTRATION: ${_manual_registration}") message(STATUS " DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}") set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) @@ -300,7 +322,7 @@ function(gen_operators_lib) add_library(${GEN_LIB_NAME}) - if(GEN_MANUAL_REGISTRATION) + if(_manual_registration) set(_srcs_list ${_out_dir}/RegisterKernelsEverything.cpp ${_out_dir}/RegisterKernels.h ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h @@ -396,7 +418,7 @@ function(gen_operators_lib) executorch_target_link_options_shared_lib(${GEN_LIB_NAME}) set(_generated_headers ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h) - if(GEN_MANUAL_REGISTRATION) + if(_manual_registration) list(APPEND _generated_headers ${_out_dir}/RegisterKernels.h) endif() if(GEN_DTYPE_SELECTIVE_BUILD) From 5a2c01ca27c27b4f6c5d7036aa18f62a4e526ad6 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Fri, 24 Jul 2026 11:25:36 +0200 Subject: [PATCH 6/6] Make CMake manual registration explicit --- codegen/gen.py | 4 ++-- examples/portable/custom_ops/CMakeLists.txt | 1 + shim_et/xplat/executorch/codegen/codegen.bzl | 6 ++--- tools/cmake/Codegen.cmake | 25 ++++---------------- 4 files changed, 11 insertions(+), 25 deletions(-) diff --git a/codegen/gen.py b/codegen/gen.py index 1b237c068c3..7014e243ddb 100644 --- a/codegen/gen.py +++ b/codegen/gen.py @@ -1000,8 +1000,8 @@ def main() -> None: "--manual-registration-lib-name", "--manual_registration_lib_name", help="library name to include in the manual registration API name. " - "Characters that are not valid in a C++ identifier are converted to " - "underscores. Requires --manual-registration.", + "The library name must be a valid C++ identifier. Requires " + "--manual-registration.", ) parser.add_argument( "--generate", diff --git a/examples/portable/custom_ops/CMakeLists.txt b/examples/portable/custom_ops/CMakeLists.txt index b9b5d22666b..774f9eefcf9 100644 --- a/examples/portable/custom_ops/CMakeLists.txt +++ b/examples/portable/custom_ops/CMakeLists.txt @@ -148,6 +148,7 @@ if(TEST_NAMED_MANUAL_REGISTRATION) ) gen_operators_lib( LIB_NAME "${lib_name}" KERNEL_LIBS ${lib_name}_kernels DEPS executorch + MANUAL_REGISTRATION ) endfunction() diff --git a/shim_et/xplat/executorch/codegen/codegen.bzl b/shim_et/xplat/executorch/codegen/codegen.bzl index ff4dfbccf34..5f65bcd8c0e 100644 --- a/shim_et/xplat/executorch/codegen/codegen.bzl +++ b/shim_et/xplat/executorch/codegen/codegen.bzl @@ -898,9 +898,9 @@ def executorch_generated_lib( fbcode_deps: Additional fbcode deps, can be used to provide custom operator library. compiler_flags: compiler_flags args to runtime.cxx_library manual_registration_lib_name: Optional library name to include when - generating a named manual registration API. Characters that are not - valid in a C++ identifier are converted to underscores. If omitted, - manual registration keeps using `register_all_kernels`. + generating a named manual registration API. The library name must be + a valid C++ identifier. If omitted, manual registration keeps using + `register_all_kernels`. dtype_selective_build: In additional to operator selection, dtype selective build further selects the dtypes for each operator. Can be used with model or dict selective build APIs, where dtypes can be specified. diff --git a/tools/cmake/Codegen.cmake b/tools/cmake/Codegen.cmake index 0fd10469669..baac09caff4 100644 --- a/tools/cmake/Codegen.cmake +++ b/tools/cmake/Codegen.cmake @@ -155,14 +155,6 @@ endfunction() # # Invoked as generate_bindings_for_kernels( LIB_NAME lib_name FUNCTIONS_YAML # functions_yaml CUSTOM_OPS_YAML custom_ops_yaml [MANUAL_REGISTRATION] ) -function(_manual_registration_property_name lib_name output_var) - string(SHA256 _lib_name_hash "${lib_name}") - set(${output_var} - "_EXECUTORCH_MANUAL_REGISTRATION_${_lib_name_hash}" - PARENT_SCOPE - ) -endfunction() - function(generate_bindings_for_kernels) set(options ADD_EXCEPTION_BOUNDARY MANUAL_REGISTRATION) set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML DTYPE_SELECTIVE_BUILD) @@ -184,11 +176,6 @@ function(generate_bindings_for_kernels) "Manual registration LIB_NAME must be a valid C++ identifier: ${GEN_LIB_NAME}" ) endif() - _manual_registration_property_name("${GEN_LIB_NAME}" _registration_property) - set_property( - GLOBAL PROPERTY "${_registration_property}" "${GEN_MANUAL_REGISTRATION}" - ) - # Command to generate selected_operators.yaml from custom_ops.yaml. file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*") @@ -302,17 +289,15 @@ endfunction() # Generate a runtime lib for registering operators in Executorch function(gen_operators_lib) + set(options MANUAL_REGISTRATION) set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS DTYPE_SELECTIVE_BUILD) - cmake_parse_arguments(GEN "" "" "${multi_arg_names}" ${ARGN}) - - _manual_registration_property_name("${GEN_LIB_NAME}" _registration_property) - get_property(_manual_registration GLOBAL PROPERTY "${_registration_property}") + cmake_parse_arguments(GEN "${options}" "" "${multi_arg_names}" ${ARGN}) message(STATUS "Generating operator lib:") message(STATUS " LIB_NAME: ${GEN_LIB_NAME}") message(STATUS " KERNEL_LIBS: ${GEN_KERNEL_LIBS}") message(STATUS " DEPS: ${GEN_DEPS}") - message(STATUS " MANUAL_REGISTRATION: ${_manual_registration}") + message(STATUS " MANUAL_REGISTRATION: ${GEN_MANUAL_REGISTRATION}") message(STATUS " DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}") set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) @@ -322,7 +307,7 @@ function(gen_operators_lib) add_library(${GEN_LIB_NAME}) - if(_manual_registration) + if(GEN_MANUAL_REGISTRATION) set(_srcs_list ${_out_dir}/RegisterKernelsEverything.cpp ${_out_dir}/RegisterKernels.h ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h @@ -418,7 +403,7 @@ function(gen_operators_lib) executorch_target_link_options_shared_lib(${GEN_LIB_NAME}) set(_generated_headers ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h) - if(_manual_registration) + if(GEN_MANUAL_REGISTRATION) list(APPEND _generated_headers ${_out_dir}/RegisterKernels.h) endif() if(GEN_DTYPE_SELECTIVE_BUILD)