From a4cadab9df17021bc7b7830de39e51013da36473 Mon Sep 17 00:00:00 2001 From: Pruthvi Madugundu Date: Tue, 2 May 2023 16:23:56 -0700 Subject: [PATCH] Revert "[dynamo 3.11] changes to with contexts (#94101)" This reverts commit 1123ab8647a0957f28e31b983d7f2de9277786e8. --- torch/_dynamo/resume_execution.py | 17 ++++++------ torch/_dynamo/symbolic_convert.py | 44 ++++--------------------------- 2 files changed, 14 insertions(+), 47 deletions(-) diff --git a/torch/_dynamo/resume_execution.py b/torch/_dynamo/resume_execution.py index a4d06b81c9f5b..053b819d751f6 100644 --- a/torch/_dynamo/resume_execution.py +++ b/torch/_dynamo/resume_execution.py @@ -97,7 +97,10 @@ def __call__(self, code_options, cleanup): ] else: + # NOTE: copying over for now since more changes are anticipated + with_except_start = create_instruction("WITH_EXCEPT_START") pop_top_after_with_except_start = create_instruction("POP_TOP") + cleanup_complete_jump_target = create_instruction("NOP") def create_load_none(): @@ -107,6 +110,7 @@ def create_load_none(): cleanup[:] = ( [ + create_instruction("POP_BLOCK"), create_load_none(), create_load_none(), create_load_none(), @@ -117,27 +121,24 @@ def create_load_none(): create_instruction( "JUMP_FORWARD", target=cleanup_complete_jump_target ), - create_instruction("PUSH_EXC_INFO"), - create_instruction("WITH_EXCEPT_START"), + with_except_start, create_instruction( "POP_JUMP_FORWARD_IF_TRUE", target=pop_top_after_with_except_start, ), - create_instruction("RERAISE", 2), - create_instruction("COPY", 3), - create_instruction("POP_EXCEPT"), - create_instruction("RERAISE", 1), + create_instruction("RERAISE"), pop_top_after_with_except_start, - create_instruction("POP_EXCEPT"), create_instruction("POP_TOP"), create_instruction("POP_TOP"), + create_instruction("POP_EXCEPT"), + create_instruction("POP_TOP"), cleanup_complete_jump_target, ] + cleanup ) return create_call_function(0, False) + [ - create_instruction("BEFORE_WITH"), + create_instruction("SETUP_WITH", target=with_except_start), create_instruction("POP_TOP"), ] diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py index 162c9fa87b4fb..e585c11ff3713 100644 --- a/torch/_dynamo/symbolic_convert.py +++ b/torch/_dynamo/symbolic_convert.py @@ -102,7 +102,6 @@ def _step_logger(): @dataclasses.dataclass class BlockStackEntry: - id: int target: Instruction stack_index: Optional[int] = None with_context: ContextWrappingVariable = None @@ -879,11 +878,11 @@ def jump(self, inst): def SETUP_LOOP(self, inst): # only exists in python<=3.7 - self.block_stack.append(BlockStackEntry(0, inst.target)) + self.block_stack.append(BlockStackEntry(inst.target)) def SETUP_EXCEPT(self, inst): # only exists in python<=3.7 - self.block_stack.append(BlockStackEntry(0, inst.target)) + self.block_stack.append(BlockStackEntry(inst.target)) def POP_BLOCK(self, inst): self.block_stack.pop() @@ -895,12 +894,10 @@ def SETUP_WITH(self, inst): self.output.guards.update(ctx.guards) if isinstance(self, InstructionTranslator): - self.block_stack.append( - BlockStackEntry(0, inst.target, len(self.stack), ctx) - ) + self.block_stack.append(BlockStackEntry(inst.target, len(self.stack), ctx)) else: # can't restore this while inlining - self.block_stack.append(BlockStackEntry(0, inst.target)) + self.block_stack.append(BlockStackEntry(inst.target)) self.push( WithExitFunctionVariable( ctx, @@ -911,7 +908,7 @@ def SETUP_WITH(self, inst): self.push(ctx.enter(self)) def SETUP_FINALLY(self, inst): - self.block_stack.append(BlockStackEntry(0, inst.target)) + self.block_stack.append(BlockStackEntry(inst.target)) def BEGIN_FINALLY(self, inst): self.push(None) @@ -1572,13 +1569,6 @@ def CALL(self, inst): kwargs = {} self.call_function(fn, args, kwargs) self.kw_names = None - # 3.11 removed POP_BLOCK, so we manually pop the block stack here - if ( - isinstance(fn, WithExitFunctionVariable) - and len(self.block_stack) > 0 - and id(fn) == self.block_stack[-1].id - ): - self.block_stack.pop() def COPY(self, inst): self.push(self.stack[-inst.arg]) @@ -1602,30 +1592,6 @@ def SWAP(self, inst): def CACHE(self, inst): pass - def BEFORE_WITH(self, inst): - ctx = self.pop() - if not isinstance(ctx, ContextWrappingVariable): - unimplemented(f"BEFORE_WITH {ctx}") - self.output.guards.update(ctx.guards) - - exit = WithExitFunctionVariable( - ctx, - inst.target, - **VariableTracker.propagate(ctx), - ) - # 3.11 no longer uses a block stack, but we still keep track of one - # so that we know which contexts are currently active. - if isinstance(self, InstructionTranslator): - self.block_stack.append( - BlockStackEntry(id(exit), inst.target, self.real_stack_len(), ctx) - ) - else: - # can't restore this while inlining - self.block_stack.append(BlockStackEntry(id(exit), inst.target)) - - self.push(exit) - self.push(ctx.enter(self)) - def copy_graphstate(self) -> InstructionTranslatorGraphState: """Create a checkpoint of the current state by copying everything""" return InstructionTranslatorGraphState(