Skip to content

Commit 7f451c1

Browse files
larryliu0820facebook-github-bot
authored andcommitted
[executorch][kernel reg] Allow kernel manual registration (pytorch#110086)
Summary: X-link: pytorch/executorch#491 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. Test Plan: Rely on the unit test: ``` buck2 test fbcode//executorch/runtime/kernel/test:test_kernel_manual_registration ``` Reviewed By: cccclai Differential Revision: D49439673
1 parent d0f82cd commit 7f451c1

File tree

2 files changed

+43
-2
lines changed

2 files changed

+43
-2
lines changed

test/edge/templates/RegisterKernels.h

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
// ${generated_comment}
10+
// Exposing an API for registering all kernels at once.
11+
#include <executorch/runtime/core/evalue.h>
12+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
13+
#include <executorch/runtime/kernel/operator_registry.h>
14+
#include <executorch/runtime/platform/profiler.h>
15+
16+
namespace torch {
17+
namespace executor {
18+
19+
Error register_all_kernels();
20+
21+
} // namespace executor
22+
} // namespace torch

torchgen/gen_executorch.py

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,7 @@ def gen_unboxing(
253253
selector: SelectiveBuilder,
254254
use_aten_lib: bool,
255255
kernel_index: ETKernelIndex,
256+
manual_registration: bool,
256257
) -> None:
257258
# Iterable type for write_sharded is a Tuple of (native_function, (kernel_key, metadata))
258259
def key_func(
@@ -267,9 +268,13 @@ def key_func(
267268
]
268269

269270
header = ["Functions.h" if use_aten_lib else "NativeFunctions.h"]
270-
271+
filename = (
272+
"RegisterKernels.cpp"
273+
if manual_registration
274+
else "RegisterCodegenUnboxedKernels.cpp"
275+
)
271276
cpu_fm.write_sharded(
272-
"RegisterCodegenUnboxedKernels.cpp",
277+
filename,
273278
items,
274279
key_fn=key_func,
275280
env_callable=lambda unbox_kernel_entry: {
@@ -448,6 +453,12 @@ def gen_headers(
448453
),
449454
},
450455
)
456+
cpu_fm.write(
457+
"RegisterKernels.h",
458+
lambda: {
459+
"generated_comment": "@" + "generated by torchgen/gen_executorch.py",
460+
},
461+
)
451462
headers = {
452463
"headers": [
453464
"#include <executorch/runtime/core/exec_aten/exec_aten.h> // at::Tensor etc.",
@@ -863,6 +874,13 @@ def main() -> None:
863874
help="a boolean flag to indicate whether we use ATen kernels or not, in the future this flag will be per "
864875
"operator",
865876
)
877+
parser.add_argument(
878+
"--manual_registration",
879+
"--manual-registration",
880+
action="store_true",
881+
help="a boolean flag to indicate whether we want to maually call"
882+
"register_kernels() or rely on static init. ",
883+
)
866884
parser.add_argument(
867885
"--generate",
868886
type=str,
@@ -916,6 +934,7 @@ def main() -> None:
916934
selector=selector,
917935
use_aten_lib=options.use_aten_lib,
918936
kernel_index=kernel_index,
937+
manual_registration=options.manual_registration,
919938
)
920939
if custom_ops_native_functions:
921940
gen_custom_ops(

0 commit comments

Comments
 (0)