[dynamo, nested graph breaks] small fixes to resume function generation (#151056)

Old: ~pack resume function stack + locals into a list: we need to be able to pass frame stack+locals in lists to hand off to nested functions in the future, so we implement this part first.~

We are no longer doing this right now since GraphModule/guard variable naming gets messed up. Going forward, our approach will be to keep the top frame unpacked, but pack the rest of the contents of other frames in a list.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/151056
Approved by: https://github.com/jansel

torch/_dynamo/bytecode_transformation.py

@@ -1607,6 +1607,7 @@ def bytecode_from_template(fn, varname_map=None, noreturn=True, noprefix=True):
         # If we don't reset starts_line, then the generated
         # bytecode's line number will be based on fn's.
         inst.starts_line = None
+        inst.positions = None
         if varname_map and inst.argval in varname_map:
             inst.argval = varname_map[inst.argval]
 

torch/_dynamo/resume_execution.py

@@ -567,78 +567,3 @@ class ContinueExecutionCache:
         return ContinueExecutionCache.lookup(
             meta.code, lineno, new_offset, setup_fn_target_offsets, *args
         )
-
-
-"""
-# partially finished support for with statements
-
-def convert_locals_to_cells(
-        instructions: List[Instruction],
-        code_options: Dict[str, Any]):
-
-    code_options["co_cellvars"] = tuple(
-        var
-        for var in code_options["co_varnames"]
-        if var not in code_options["co_freevars"]
-        and not var.startswith("___stack")
-    )
-    cell_and_free = code_options["co_cellvars"] + code_options["co_freevars"]
-    for inst in instructions:
-        if str(inst.argval).startswith("___stack"):
-            continue
-        elif inst.opname == "LOAD_FAST":
-            inst.opname = "LOAD_DEREF"
-        elif inst.opname == "STORE_FAST":
-            inst.opname = "STORE_DEREF"
-        elif inst.opname == "DELETE_FAST":
-            inst.opname = "DELETE_DEREF"
-        else:
-            continue
-        inst.opcode = dis.opmap[inst.opname]
-        assert inst.argval in cell_and_free, inst.argval
-        inst.arg = cell_and_free.index(inst.argval)
-
-def patch_setup_with(
-    instructions: List[Instruction],
-    code_options: Dict[str, Any]
-):
-    nonlocal need_skip
-    need_skip = True
-    target_index = next(
-        idx for idx, i in enumerate(instructions) if i.offset == offset
-    )
-    assert instructions[target_index].opname == "SETUP_WITH"
-    convert_locals_to_cells(instructions, code_options)
-
-    stack_depth_before = nstack + stack_effect(instructions[target_index].opcode,
-                                               instructions[target_index].arg)
-
-    inside_with = []
-    inside_with_resume_at = None
-    stack_depth = stack_depth_before
-    idx = target_index + 1
-    for idx in range(idx, len(instructions)):
-        inst = instructions[idx]
-        if inst.opname == "BEGIN_FINALLY":
-            inside_with_resume_at = inst
-            break
-        elif inst.target is not None:
-            unimplemented("jump from with not supported")
-        elif inst.opname in ("BEGIN_FINALLY", "WITH_CLEANUP_START", "WITH_CLEANUP_FINISH", "END_FINALLY",
-                             "POP_FINALLY", "POP_EXCEPT",
-                             "POP_BLOCK", "END_ASYNC_FOR"):
-            unimplemented("block ops not supported")
-        inside_with.append(inst)
-        stack_depth += stack_effect(inst.opcode, inst.arg)
-    assert inside_with_resume_at
-
-    instructions = [
-        create_instruction("LOAD_FAST", f"___stack{i}") for i in range(nstack)
-    ] + [
-        create_instruction("SETUP_WITH", target=instructions[target_index].target)
-        ... call the function ...
-        unpack_tuple
-    ] + [
-        create_instruction("JUMP_ABSOLUTE", target=inside_with_resume_at)
-    ]
-"""

torch/_dynamo/symbolic_convert.py

@@ -1427,7 +1427,7 @@ class InstructionTranslatorBase(
                 )
 
         # for continuation functions
-        if name.startswith("___stack"):
+        if name.startswith("__stack"):
             self.symbolic_locals.pop(name)
 
     def LOAD_DEREF(self, inst):