use backend specializations in compile_and_call_fx_graph

bobrenjc93 · bobrenjc93 · commit 9914ca319d29 · 2025-05-01T07:14:49.000-07:00
ghstack-source-id: f0e905e Pull Request resolved: #152601
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
@@ -1496,8 +1496,11 @@ def compile_and_call_fx_graph(self, tx, rv, root):
                 # a lot of fake_tensor ownership assumptions and runs afoul of detect_fake_mode
                 self.tracing_context.fake_mode = backend_fake_mode
 
+            compiled_fns = []
             with self.restore_global_state():
-                compiled_fn = self.call_user_compiler(gm)
+                for specialization in backend_specializations
+                    modified_gm = specialize(gm, specialization)
+                    compiled_fns.append(self.call_user_compiler(modified_gm))
 
             from torch.fx._lazy_graph_module import _LazyGraphModule
 
