From f6e12bdb4545a37e49556bc10a3994f6ea9670e1 Mon Sep 17 00:00:00 2001 From: William Wen Date: Mon, 28 Jul 2025 11:42:23 -0700 Subject: [PATCH] [dynamo, nested graph breaks] use CALL_FUNCTION_EX when calling resume function [ghstack-poisoned] --- torch/_dynamo/symbolic_convert.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py index b9215d77dee1..99bc871aabb9 100644 --- a/torch/_dynamo/symbolic_convert.py +++ b/torch/_dynamo/symbolic_convert.py @@ -75,7 +75,6 @@ create_dup_top, create_instruction, create_jump_absolute, - create_reverse, create_swap, get_code_keys, Instruction, @@ -2515,7 +2514,6 @@ def create_call_resume_at(self, inst, push, all_stack_locals_metadata): stack_len = len(self.stack) - len( all_stack_locals_metadata[-1].stack_null_idxes ) - nargs = stack_len + len(argnames) new_code: types.CodeType = ContinueExecutionCache.lookup( self.f_code, @@ -2605,16 +2603,24 @@ def create_call_resume_at(self, inst, push, all_stack_locals_metadata): [ create_instruction("POP_TOP"), create_instruction("BUILD_LIST", arg=len(argnames)), + *create_swap(3), + # live_locals, frames[-1][0], frames + create_instruction("BUILD_LIST", arg=1), + *create_swap(2), + # live_locals, [frames], frames[-1][0] + create_instruction("LIST_EXTEND", arg=1), + *create_swap(2), create_instruction("LIST_EXTEND", arg=1), - # UNPACK_SEQUENCE reverses elements - create_instruction("UNPACK_SEQUENCE", arg=nargs), - *create_reverse(nargs), ] ) - # frames, *(stack + live locals) + # [frames, *(stack + live locals)] - cg.extend_output(create_call_function(nargs + 1, False)) - cg.append_output(create_instruction("RETURN_VALUE")) + cg.extend_output( + [ + create_instruction("CALL_FUNCTION_EX", arg=0), + create_instruction("RETURN_VALUE"), + ] + ) return cg.get_instructions() def should_compile_partial_graph(self):