saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

torchgen: move dispatch_helpers out of RegisterDispatchDefinitions.ini (#144363)

Browse source 
The dispatch_helpers should be generated once, not once per kernel namespace.

Differential Revision: [D67925497](https://our.internmc.facebook.com/intern/diff/D67925497/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/144363
Approved by: https://github.com/bdhirsh
This commit is contained in:
Scott Wolchok 2025-01-09 15:00:20 -08:00 committed by PyTorch MergeBot
parent 7a93a58b3c
commit 603e1c0b02
4 changed files with 10 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
aten/src/ATen/templates/RegisterDispatchDefinitions.ini
View file

@ -5,8 +5,6 @@ ${ns_prologue}
// at namespace already. // at namespace already.
namespace { namespace {
${dispatch_helpers}
${dispatch_anonymous_definitions} ${dispatch_anonymous_definitions}
${static_init_dispatch_registrations} ${static_init_dispatch_registrations}

6
aten/src/ATen/templates/RegisterDispatchKey.cpp
View file

@ -51,5 +51,11 @@ $external_backend_headers
$dispatch_headers $dispatch_headers
$ops_headers $ops_headers
namespace at {
namespace {
$dispatch_helpers
} // namespace
} // namespace at
// See template file RegisterDispatchDefinitions.ini // See template file RegisterDispatchDefinitions.ini
$dispatch_definitions $dispatch_definitions

7
torchgen/gen.py
View file

@ -1642,9 +1642,6 @@ TORCH_LIBRARY_IMPL({namespace}, {dispatch_key}, m) {{
lambda: { lambda: {
"ns_prologue": ns_helper.prologue, "ns_prologue": ns_helper.prologue,
"ns_epilogue": ns_helper.epilogue, "ns_epilogue": ns_helper.epilogue,
"dispatch_helpers": dest.gen_registration_helpers(backend_idx)
if gen_dispatch_helpers
else [],
"dispatch_anonymous_definitions": anonymous_definitions[ "dispatch_anonymous_definitions": anonymous_definitions[
kernel_namespace kernel_namespace
], ],
@ -2331,7 +2328,9 @@ def gen_source_files(
backend_index, per_operator_headers, rocm backend_index, per_operator_headers, rocm
), ),
"ops_headers": operator_headers(), "ops_headers": operator_headers(),
"dispatch_helpers": "", "dispatch_helpers": dest.gen_registration_helpers(backend_index)
if gen_dispatch_helpers
else [],
"dispatch_definitions": dispatch_definitions, "dispatch_definitions": dispatch_definitions,
}, },
) )

2
torchgen/gen_backend_stubs.py
View file

@ -499,6 +499,7 @@ TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {
"dispatch_headers": dest.gen_registration_headers( "dispatch_headers": dest.gen_registration_headers(
backend_index, per_operator_headers=per_operator_headers, rocm=False backend_index, per_operator_headers=per_operator_headers, rocm=False
), ),
"dispatch_helpers": dest.gen_registration_helpers(backend_index),
"dispatch_definitions": fm.substitute_with_template( "dispatch_definitions": fm.substitute_with_template(
"RegisterDispatchDefinitions.ini", "RegisterDispatchDefinitions.ini",
lambda: { lambda: {
@ -506,7 +507,6 @@ TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {
"ns_epilogue": ns_helper.epilogue, "ns_epilogue": ns_helper.epilogue,
"static_init_dispatch_registrations": static_init_dispatch_registrations, "static_init_dispatch_registrations": static_init_dispatch_registrations,
"deferred_dispatch_registrations": deferred_dispatch_registrations, "deferred_dispatch_registrations": deferred_dispatch_registrations,
"dispatch_helpers": dest.gen_registration_helpers(backend_index),
"dispatch_namespace": dispatch_key.lower(), "dispatch_namespace": dispatch_key.lower(),
"dispatch_namespaced_definitions": "", "dispatch_namespaced_definitions": "",
"dispatch_anonymous_definitions": list( "dispatch_anonymous_definitions": list(