Update on "thread through specialization to compile_fx"

bobrenjc93 · bobrenjc93 · commit 79825349f554 · 2025-05-01T22:32:28.000-07:00
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
@@ -1498,7 +1498,7 @@ def compile_and_call_fx_graph(self, tx, rv, root):
 
             compiled_fns = []
             with self.restore_global_state():
-                for specialization in backend_specializations
+                for specialization in old_fake_mode.shape_env.backend_specializations
                     compiled_fns.append(self.call_user_compiler(modified_gm, specialization))
 
             from torch.fx._lazy_graph_module import _LazyGraphModule
diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py
@@ -931,6 +931,12 @@ def find_symbol_binding_fx_nodes(
     return r
 
 
+@dataclass
+class BackendSpecialization:
+    symbol: sympy.Symbol
+    hint: int
+    specialization: Callable
+
 # Analogous to ConvertIntSource
 @dataclass(frozen=True)
 class ConvertIntKey:
@@ -3561,6 +3567,8 @@ def _init(
 
         self.trace_asserts = trace_asserts
 
+        self.backend_specializations = []
+
         from torch.fx.experimental.validator import translation_validation_enabled
 
         self._translation_validation_enabled = translation_validation_enabled()
@@ -4044,6 +4052,11 @@ def _produce_dyn_sizes_from_int_tuple(
                 do_not_specialize_zero_one=config.backed_size_oblivious,
                 symbolic_context=symbolic_context,
             )
+            for specialization in symbolic_context.backend_specializations:
+                self.backend_specializations.append(BackendSpecialization(
+                    sym,
+                    *specialization,
+                ))
             if (
                 config.backed_size_oblivious
                 and isinstance(sym, sympy.Symbol)  # could be static
