From 603e1c0b02a546c06ed95a08eda982d7e208ec98 Mon Sep 17 00:00:00 2001
From: Scott Wolchok <swolchok@fb.com>
Date: Thu, 9 Jan 2025 15:00:20 -0800
Subject: [PATCH] torchgen: move dispatch_helpers out of
 RegisterDispatchDefinitions.ini (#144363)

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
---
 aten/src/ATen/templates/RegisterDispatchDefinitions.ini | 2 --
 aten/src/ATen/templates/RegisterDispatchKey.cpp         | 6 ++++++
 torchgen/gen.py                                         | 7 +++----
 torchgen/gen_backend_stubs.py                           | 2 +-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/aten/src/ATen/templates/RegisterDispatchDefinitions.ini b/aten/src/ATen/templates/RegisterDispatchDefinitions.ini
index 3bf7f9b1bb3..97c921de18f 100644
--- a/aten/src/ATen/templates/RegisterDispatchDefinitions.ini
+++ b/aten/src/ATen/templates/RegisterDispatchDefinitions.ini
@@ -5,8 +5,6 @@ ${ns_prologue}
 // at namespace already.
 namespace {
 
-${dispatch_helpers}
-
 ${dispatch_anonymous_definitions}
 
 ${static_init_dispatch_registrations}
diff --git a/aten/src/ATen/templates/RegisterDispatchKey.cpp b/aten/src/ATen/templates/RegisterDispatchKey.cpp
index c25b513061e..f55b3714293 100644
--- a/aten/src/ATen/templates/RegisterDispatchKey.cpp
+++ b/aten/src/ATen/templates/RegisterDispatchKey.cpp
@@ -51,5 +51,11 @@ $external_backend_headers
 $dispatch_headers
 $ops_headers
 
+namespace at {
+namespace {
+$dispatch_helpers
+} // namespace
+} // namespace at
+
 // See template file RegisterDispatchDefinitions.ini
 $dispatch_definitions
diff --git a/torchgen/gen.py b/torchgen/gen.py
index d08c2ff2cb8..e7ee60511de 100644
--- a/torchgen/gen.py
+++ b/torchgen/gen.py
@@ -1642,9 +1642,6 @@ TORCH_LIBRARY_IMPL({namespace}, {dispatch_key}, m) {{
                 lambda: {
                     "ns_prologue": ns_helper.prologue,
                     "ns_epilogue": ns_helper.epilogue,
-                    "dispatch_helpers": dest.gen_registration_helpers(backend_idx)
-                    if gen_dispatch_helpers
-                    else [],
                     "dispatch_anonymous_definitions": anonymous_definitions[
                         kernel_namespace
                     ],
@@ -2331,7 +2328,9 @@ def gen_source_files(
                     backend_index, per_operator_headers, rocm
                 ),
                 "ops_headers": operator_headers(),
-                "dispatch_helpers": "",
+                "dispatch_helpers": dest.gen_registration_helpers(backend_index)
+                if gen_dispatch_helpers
+                else [],
                 "dispatch_definitions": dispatch_definitions,
             },
         )
diff --git a/torchgen/gen_backend_stubs.py b/torchgen/gen_backend_stubs.py
index 299a41081da..bc9e2959ab7 100644
--- a/torchgen/gen_backend_stubs.py
+++ b/torchgen/gen_backend_stubs.py
@@ -499,6 +499,7 @@ TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {
             "dispatch_headers": dest.gen_registration_headers(
                 backend_index, per_operator_headers=per_operator_headers, rocm=False
             ),
+            "dispatch_helpers": dest.gen_registration_helpers(backend_index),
             "dispatch_definitions": fm.substitute_with_template(
                 "RegisterDispatchDefinitions.ini",
                 lambda: {
@@ -506,7 +507,6 @@ TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {
                     "ns_epilogue": ns_helper.epilogue,
                     "static_init_dispatch_registrations": static_init_dispatch_registrations,
                     "deferred_dispatch_registrations": deferred_dispatch_registrations,
-                    "dispatch_helpers": dest.gen_registration_helpers(backend_index),
                     "dispatch_namespace": dispatch_key.lower(),
                     "dispatch_namespaced_definitions": "",
                     "dispatch_anonymous_definitions": list(