Update on "[dynamo, nested graph breaks] prevent excessive recompilations"

williamwen42 · williamwen42 · commit e71853e33cee · 2025-08-11T15:31:30.000-07:00
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames Lucaskabela

[ghstack-poisoned]
diff --git a/test/dynamo/test_nested_graph_breaks.py b/test/dynamo/test_nested_graph_breaks.py
@@ -304,7 +304,6 @@ def f3(x):
         self.assertEqual(cnts.frame_count, 3)
         self.assertEqual(cnts.op_count, 7)
 
-    @unittest.expectedFailure
     @torch._dynamo.config.patch(recompile_limit=1, fail_on_recompile_limit_hit=True)
     def test_no_recompiles(self):
         global f1, f2, f3
diff --git a/torch/_dynamo/resume_execution.py b/torch/_dynamo/resume_execution.py
@@ -311,7 +311,9 @@ def generate(
         stack_ctx_vars: tuple[tuple[int, tuple[Any, ...]], ...],
         argnames_ctx_vars: tuple[tuple[str, tuple[Any, ...]], ...],
         null_idxes: tuple[int, ...],
-        has_nested: bool,
+        # mainly used to ensure distinct code objects per stack trace,
+        # which prevents excessive recompilation of inner frames
+        nested_code_objs: tuple[types.CodeType],
     ) -> types.CodeType:
         assert offset is not None
         assert not (
@@ -332,7 +334,7 @@ def generate(
                 stack_ctx_vars,
                 argnames_ctx_vars,
                 null_idxes,
-                has_nested,
+                nested_code_objs,
             )
 
         is_py311_plus = sys.version_info >= (3, 11)
@@ -466,7 +468,7 @@ def update(
                     )
 
             # Call nested resume function
-            if has_nested:
+            if nested_code_objs:
                 prefix.extend(
                     [
                         # set up __nested_resume_fns[-1] call
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -2506,7 +2506,7 @@ def create_call_resume_at(self, inst, all_stack_locals_metadata):
 
         # build the resume function for each frame
         resume_names = []
-        resume_codes = []
+        resume_codes: list[types.CodeType] = []
         for i, meta in enumerate(all_stack_locals_metadata):
             cur_tx = txes[i]
             if cur_tx is self:
@@ -2594,7 +2594,7 @@ def create_call_resume_at(self, inst, all_stack_locals_metadata):
                 tuple(meta.stack_ctx_args),
                 tuple(meta.locals_ctx_args),
                 tuple(meta.stack_null_idxes),
-                self is not cur_tx,
+                tuple(resume_codes),
             )
             resume_codes.append(new_code)
 
