| import copy |
| import dataclasses |
| import sys |
| import types |
| from typing import Any, cast, Dict, List, Optional, Tuple |
| |
| from .bytecode_transformation import ( |
| create_call_function, |
| create_call_method, |
| create_dup_top, |
| create_instruction, |
| create_jump_absolute, |
| Instruction, |
| InstructionExnTabEntry, |
| transform_code_object, |
| unique_id, |
| ) |
| from .utils import ExactWeakKeyDictionary |
| |
| # taken from code.h in cpython |
| CO_OPTIMIZED = 0x0001 |
| CO_NEWLOCALS = 0x0002 |
| CO_VARARGS = 0x0004 |
| CO_VARKEYWORDS = 0x0008 |
| CO_NESTED = 0x0010 |
| CO_GENERATOR = 0x0020 |
| CO_NOFREE = 0x0040 |
| CO_COROUTINE = 0x0080 |
| CO_ITERABLE_COROUTINE = 0x0100 |
| CO_ASYNC_GENERATOR = 0x0200 |
| |
| |
| @dataclasses.dataclass(frozen=True) |
| class ReenterWith: |
| stack_index: int |
| target_values: Optional[Tuple[Any, ...]] = None |
| |
| # If we do not want to destroy the stack, we can do the same thing as a |
| # `SETUP_WITH` block, only that we store the context manager in a local_symbol |
| def try_except(self, code_options, cleanup: List[Instruction]): |
| """ |
| Codegen based off of: |
| load args |
| enter context |
| try: |
| (rest) |
| finally: |
| exit context |
| """ |
| load_args = [] |
| if self.target_values: |
| load_args = [ |
| create_instruction("LOAD_CONST", argval=val) |
| for val in self.target_values |
| ] |
| ctx_name = unique_id(f"___context_manager_{self.stack_index}") |
| if ctx_name not in code_options["co_varnames"]: |
| code_options["co_varnames"] += (ctx_name,) |
| for name in ["__enter__", "__exit__"]: |
| if name not in code_options["co_names"]: |
| code_options["co_names"] += (name,) |
| |
| except_jump_target = create_instruction( |
| "NOP" if sys.version_info < (3, 11) else "PUSH_EXC_INFO" |
| ) |
| cleanup_complete_jump_target = create_instruction("NOP") |
| |
| setup_finally = [ |
| *load_args, |
| *create_call_function(len(load_args), True), |
| create_instruction("STORE_FAST", argval=ctx_name), |
| create_instruction("LOAD_FAST", argval=ctx_name), |
| create_instruction("LOAD_METHOD", argval="__enter__"), |
| *create_call_method(0), |
| create_instruction("POP_TOP"), |
| ] |
| |
| if sys.version_info < (3, 11): |
| setup_finally.append( |
| create_instruction("SETUP_FINALLY", target=except_jump_target) |
| ) |
| else: |
| exn_tab_begin = create_instruction("NOP") |
| exn_tab_end = create_instruction("NOP") |
| exn_tab_begin.exn_tab_entry = InstructionExnTabEntry( |
| exn_tab_begin, |
| exn_tab_end, |
| except_jump_target, |
| self.stack_index + 1, |
| False, |
| ) |
| setup_finally.append(exn_tab_begin) |
| |
| def create_reset(): |
| return [ |
| create_instruction("LOAD_FAST", argval=ctx_name), |
| create_instruction("LOAD_METHOD", argval="__exit__"), |
| create_instruction("LOAD_CONST", argval=None), |
| create_dup_top(), |
| create_dup_top(), |
| *create_call_method(3), |
| create_instruction("POP_TOP"), |
| ] |
| |
| if sys.version_info < (3, 9): |
| epilogue = [ |
| create_instruction("POP_BLOCK"), |
| create_instruction("BEGIN_FINALLY"), |
| except_jump_target, |
| *create_reset(), |
| create_instruction("END_FINALLY"), |
| ] |
| elif sys.version_info < (3, 11): |
| epilogue = [ |
| create_instruction("POP_BLOCK"), |
| *create_reset(), |
| create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target), |
| except_jump_target, |
| *create_reset(), |
| create_instruction("RERAISE"), |
| cleanup_complete_jump_target, |
| ] |
| else: |
| finally_exn_tab_end = create_instruction("RERAISE", arg=0) |
| finally_exn_tab_target = create_instruction("COPY", arg=3) |
| except_jump_target.exn_tab_entry = InstructionExnTabEntry( |
| except_jump_target, |
| finally_exn_tab_end, |
| finally_exn_tab_target, |
| self.stack_index + 2, |
| True, |
| ) |
| epilogue = [ |
| exn_tab_end, |
| *create_reset(), |
| create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target), |
| except_jump_target, # PUSH_EXC_INFO |
| *create_reset(), |
| finally_exn_tab_end, # RERAISE 0 |
| finally_exn_tab_target, # COPY 3 |
| create_instruction("POP_EXCEPT"), |
| create_instruction("RERAISE", arg=1), |
| cleanup_complete_jump_target, |
| ] |
| |
| cleanup[:] = epilogue + cleanup |
| return setup_finally |
| |
| def __call__(self, code_options, cleanup): |
| """ |
| Codegen based off of: |
| with ctx(args): |
| (rest) |
| """ |
| load_args = [] |
| if self.target_values: |
| load_args = [ |
| create_instruction("LOAD_CONST", argval=val) |
| for val in self.target_values |
| ] |
| if sys.version_info < (3, 9): |
| with_cleanup_start = create_instruction("WITH_CLEANUP_START") |
| begin_finally = create_instruction("BEGIN_FINALLY") |
| cleanup[:] = [ |
| create_instruction("POP_BLOCK"), |
| begin_finally, |
| with_cleanup_start, |
| create_instruction("WITH_CLEANUP_FINISH"), |
| create_instruction("END_FINALLY"), |
| ] + cleanup |
| |
| return [ |
| *load_args, |
| create_instruction("CALL_FUNCTION", arg=len(load_args)), |
| create_instruction("SETUP_WITH", target=with_cleanup_start), |
| create_instruction("POP_TOP"), |
| ], None |
| elif sys.version_info < (3, 11): |
| 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") |
| |
| cleanup[:] = [ |
| create_instruction("POP_BLOCK"), |
| create_instruction("LOAD_CONST", argval=None), |
| create_instruction("DUP_TOP"), |
| create_instruction("DUP_TOP"), |
| create_instruction("CALL_FUNCTION", arg=3), |
| create_instruction("POP_TOP"), |
| create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target), |
| with_except_start, |
| create_instruction( |
| "POP_JUMP_IF_TRUE", target=pop_top_after_with_except_start |
| ), |
| create_instruction("RERAISE"), |
| pop_top_after_with_except_start, |
| create_instruction("POP_TOP"), |
| create_instruction("POP_TOP"), |
| create_instruction("POP_EXCEPT"), |
| create_instruction("POP_TOP"), |
| cleanup_complete_jump_target, |
| ] + cleanup |
| |
| return [ |
| *load_args, |
| create_instruction("CALL_FUNCTION", arg=len(load_args)), |
| create_instruction("SETUP_WITH", target=with_except_start), |
| create_instruction("POP_TOP"), |
| ], None |
| else: |
| pop_top_after_with_except_start = create_instruction("POP_TOP") |
| cleanup_complete_jump_target = create_instruction("NOP") |
| |
| def create_load_none(): |
| return create_instruction("LOAD_CONST", argval=None) |
| |
| exn_tab_1_begin = create_instruction("POP_TOP") |
| exn_tab_1_end = create_instruction("NOP") |
| exn_tab_1_target = create_instruction("PUSH_EXC_INFO") |
| exn_tab_2_end = create_instruction("RERAISE", arg=2) |
| exn_tab_2_target = create_instruction("COPY", arg=3) |
| |
| exn_tab_1_begin.exn_tab_entry = InstructionExnTabEntry( |
| exn_tab_1_begin, |
| exn_tab_1_end, |
| exn_tab_1_target, |
| self.stack_index + 1, |
| True, |
| ) |
| exn_tab_1_target.exn_tab_entry = InstructionExnTabEntry( |
| exn_tab_1_target, |
| exn_tab_2_end, |
| exn_tab_2_target, |
| self.stack_index + 3, |
| True, |
| ) |
| pop_top_after_with_except_start.exn_tab_entry = InstructionExnTabEntry( |
| pop_top_after_with_except_start, |
| pop_top_after_with_except_start, |
| exn_tab_2_target, |
| self.stack_index + 3, |
| True, |
| ) |
| |
| cleanup[:] = [ |
| exn_tab_1_end, |
| create_load_none(), |
| create_load_none(), |
| create_load_none(), |
| *create_call_function(2, False), |
| create_instruction("POP_TOP"), |
| create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target), |
| exn_tab_1_target, # PUSH_EXC_INFO |
| create_instruction("WITH_EXCEPT_START"), |
| create_instruction( |
| "POP_JUMP_FORWARD_IF_TRUE", |
| target=pop_top_after_with_except_start, |
| ), |
| exn_tab_2_end, # RERAISE 2 |
| exn_tab_2_target, # COPY 3 |
| create_instruction("POP_EXCEPT"), |
| create_instruction("RERAISE", arg=1), |
| pop_top_after_with_except_start, |
| create_instruction("POP_EXCEPT"), |
| create_instruction("POP_TOP"), |
| create_instruction("POP_TOP"), |
| cleanup_complete_jump_target, |
| ] + cleanup |
| |
| return [ |
| *load_args, |
| *create_call_function(len(load_args), True), |
| create_instruction("BEFORE_WITH"), |
| exn_tab_1_begin, # POP_TOP |
| ], exn_tab_1_target |
| |
| |
| @dataclasses.dataclass |
| class ResumeFunctionMetadata: |
| code: types.CodeType |
| instructions: List[Instruction] = dataclasses.field(default_factory=list) |
| # Python 3.11+ fields |
| # NOTE: Python 3.11 removed blocks, but for our purposes, a "block" consists |
| # of instructions of all exception table entries that have the same target. |
| |
| # map from PUSH_EXC_INFO's in the prefix to original block target offset |
| prefix_block_target_offset_remap: List[int] = dataclasses.field( |
| default_factory=list |
| ) |
| # map from new block target offsets to original block target offsets |
| block_target_offset_remap: Optional[Dict[int, int]] = None |
| |
| |
| def _filter_iter(l1, l2, cond): |
| """ |
| Two-pointer conditional filter. |
| e.g. _filter_iter(insts, sorted_offsets, lambda i, o: i.offset == o) |
| returns the instructions with offsets in sorted_offsets |
| """ |
| it = iter(l2) |
| res = [] |
| try: |
| cur = next(it) |
| for val in l1: |
| if cond(val, cur): |
| res.append(val) |
| cur = next(it) |
| except StopIteration: |
| pass |
| return res |
| |
| |
| class ContinueExecutionCache: |
| cache = ExactWeakKeyDictionary() |
| generated_code_metadata = ExactWeakKeyDictionary() |
| |
| @classmethod |
| def lookup(cls, code, lineno, *key): |
| if code not in cls.cache: |
| cls.cache[code] = dict() |
| key = tuple(key) |
| if key not in cls.cache[code]: |
| cls.cache[code][key] = cls.generate(code, lineno, *key) |
| return cls.cache[code][key] |
| |
| @classmethod |
| def generate( |
| cls, |
| code, |
| lineno, |
| offset: int, |
| setup_fn_target_offsets: Tuple[int], # only used in Python 3.11+ |
| nstack: int, |
| argnames: Tuple[str], |
| setup_fns: Tuple[ReenterWith], |
| null_idxes: Tuple[int], |
| ) -> types.CodeType: |
| assert offset is not None |
| assert not ( |
| code.co_flags |
| & (CO_GENERATOR | CO_COROUTINE | CO_ITERABLE_COROUTINE | CO_ASYNC_GENERATOR) |
| ) |
| assert code.co_flags & CO_OPTIMIZED |
| if code in ContinueExecutionCache.generated_code_metadata: |
| return cls.generate_based_on_original_code_object( |
| code, |
| lineno, |
| offset, |
| setup_fn_target_offsets, |
| nstack, |
| argnames, |
| setup_fns, |
| null_idxes, |
| ) |
| |
| is_py311_plus = sys.version_info >= (3, 11) |
| meta = ResumeFunctionMetadata(code) |
| |
| def update(instructions: List[Instruction], code_options: Dict[str, Any]): |
| meta.instructions = copy.deepcopy(instructions) |
| |
| args = [f"___stack{i}" for i in range(nstack)] |
| args.extend(v for v in argnames if v not in args) |
| freevars = tuple(code_options["co_cellvars"] or []) + tuple( |
| code_options["co_freevars"] or [] |
| ) |
| code_options[ |
| "co_name" |
| ] = f"torch_dynamo_resume_in_{code_options['co_name']}_at_{lineno}" |
| if is_py311_plus: |
| qualified_path = code_options["co_qualname"].rsplit(".", maxsplit=1) |
| if len(qualified_path) == 1: |
| code_options["co_qualname"] = code_options["co_name"] |
| else: |
| assert len(qualified_path) == 2 |
| module_name, co_name = qualified_path |
| code_options[ |
| "co_qualname" |
| ] = f"{module_name}.torch_dynamo_resume_in_{co_name}_at_{lineno}" |
| code_options["co_firstlineno"] = lineno |
| code_options["co_cellvars"] = tuple() |
| code_options["co_freevars"] = freevars |
| code_options["co_argcount"] = len(args) |
| code_options["co_posonlyargcount"] = 0 |
| code_options["co_kwonlyargcount"] = 0 |
| code_options["co_varnames"] = tuple( |
| args + [v for v in code_options["co_varnames"] if v not in args] |
| ) |
| code_options["co_flags"] = code_options["co_flags"] & ~( |
| CO_VARARGS | CO_VARKEYWORDS |
| ) |
| target = next(i for i in instructions if i.offset == offset) |
| |
| prefix = [] |
| if is_py311_plus: |
| if freevars: |
| prefix.append( |
| create_instruction("COPY_FREE_VARS", arg=len(freevars)) |
| ) |
| prefix.append(create_instruction("RESUME", arg=0)) |
| |
| cleanup: List[Instruction] = [] |
| hooks = {fn.stack_index: fn for fn in setup_fns} |
| hook_target_offsets = { |
| fn.stack_index: setup_fn_target_offsets[i] |
| for i, fn in enumerate(setup_fns) |
| } |
| offset_to_inst = {inst.offset: inst for inst in instructions} |
| # map old hook targets to new targets generated by the hook |
| old_hook_target_remap = {} |
| null_idxes_i = 0 |
| for i in range(nstack): |
| while ( |
| null_idxes_i < len(null_idxes) |
| and null_idxes[null_idxes_i] == i + null_idxes_i |
| ): |
| prefix.append(create_instruction("PUSH_NULL")) |
| null_idxes_i += 1 |
| prefix.append(create_instruction("LOAD_FAST", argval=f"___stack{i}")) |
| if i in hooks: |
| hook = hooks.pop(i) |
| hook_insts, exn_target = hook(code_options, cleanup) |
| prefix.extend(hook_insts) |
| if is_py311_plus: |
| hook_target_offset = hook_target_offsets.pop(i) |
| old_hook_target = offset_to_inst[hook_target_offset] |
| meta.prefix_block_target_offset_remap.append(hook_target_offset) |
| old_hook_target_remap[old_hook_target] = exn_target |
| if is_py311_plus: |
| # reverse the mapping since targets of later/nested contexts are inserted |
| # into the mapping later, but show up earlier in the prefix. |
| meta.prefix_block_target_offset_remap = list( |
| reversed(meta.prefix_block_target_offset_remap) |
| ) |
| |
| assert not hooks |
| |
| prefix.append(create_jump_absolute(target)) |
| |
| # because the line number table monotonically increases from co_firstlineno |
| # remove starts_line for any instructions before the graph break instruction |
| # this will ensure the instructions after the break have the correct line numbers |
| for inst in instructions: |
| if inst.offset == target.offset: |
| break |
| inst.starts_line = None |
| if sys.version_info >= (3, 11): |
| inst.positions = None |
| |
| if cleanup: |
| prefix.extend(cleanup) |
| prefix.extend(cls.unreachable_codes(code_options)) |
| |
| # remap original instructions' exception table entries |
| if old_hook_target_remap: |
| assert is_py311_plus |
| for inst in instructions: |
| if ( |
| inst.exn_tab_entry |
| and inst.exn_tab_entry.target in old_hook_target_remap |
| ): |
| inst.exn_tab_entry.target = old_hook_target_remap[ |
| inst.exn_tab_entry.target |
| ] |
| |
| # TODO(jansel): add dead code elimination here |
| instructions[:] = prefix + instructions |
| |
| new_code = transform_code_object(code, update) |
| ContinueExecutionCache.generated_code_metadata[new_code] = meta |
| return new_code |
| |
| @staticmethod |
| def unreachable_codes(code_options) -> List[Instruction]: |
| """Codegen a `raise None` to make analysis work for unreachable code""" |
| return [ |
| create_instruction("LOAD_CONST", argval=None), |
| create_instruction("RAISE_VARARGS", arg=1), |
| ] |
| |
| @classmethod |
| def generate_based_on_original_code_object( |
| cls, code, lineno, offset: int, setup_fn_target_offsets: Tuple[int, ...], *args |
| ): |
| """ |
| This handles the case of generating a resume into code generated |
| to resume something else. We want to always generate starting |
| from the original code object so that if control flow paths |
| converge we only generated 1 resume function (rather than 2^n |
| resume functions). |
| """ |
| |
| meta: ResumeFunctionMetadata = ContinueExecutionCache.generated_code_metadata[ |
| code |
| ] |
| new_offset = None |
| |
| def find_new_offset( |
| instructions: List[Instruction], code_options: Dict[str, Any] |
| ): |
| nonlocal new_offset |
| (target,) = (i for i in instructions if i.offset == offset) |
| # match the functions starting at the last instruction as we have added a prefix |
| (new_target,) = ( |
| i2 |
| for i1, i2 in zip(reversed(instructions), reversed(meta.instructions)) |
| if i1 is target |
| ) |
| assert target.opcode == new_target.opcode |
| new_offset = new_target.offset |
| |
| transform_code_object(code, find_new_offset) |
| |
| if sys.version_info >= (3, 11): |
| # setup_fn_target_offsets currently contains the target offset of |
| # each setup_fn, based on `code`. When we codegen the resume function |
| # based on the original code object, `meta.code`, the offsets in |
| # setup_fn_target_offsets must be based on `meta.code` instead. |
| if not meta.block_target_offset_remap: |
| block_target_offset_remap = meta.block_target_offset_remap = {} |
| |
| def remap_block_offsets( |
| instructions: List[Instruction], code_options: Dict[str, Any] |
| ): |
| # NOTE: each prefix block generates exactly one PUSH_EXC_INFO, |
| # so we can tell which block a prefix PUSH_EXC_INFO belongs to, |
| # by counting. Then we can use meta.prefix_block-target_offset_remap |
| # to determine where in the original code the PUSH_EXC_INFO offset |
| # replaced. |
| prefix_blocks: List[Instruction] = [] |
| for inst in instructions: |
| if len(prefix_blocks) == len( |
| meta.prefix_block_target_offset_remap |
| ): |
| break |
| if inst.opname == "PUSH_EXC_INFO": |
| prefix_blocks.append(inst) |
| |
| # offsets into prefix |
| for inst, o in zip( |
| prefix_blocks, meta.prefix_block_target_offset_remap |
| ): |
| block_target_offset_remap[cast(int, inst.offset)] = o |
| |
| # old bytecode targets are after the prefix PUSH_EXC_INFO's |
| old_start_offset = ( |
| cast(int, prefix_blocks[-1].offset) if prefix_blocks else -1 |
| ) |
| # offsets into old bytecode |
| old_inst_offsets = sorted( |
| n for n in setup_fn_target_offsets if n > old_start_offset |
| ) |
| targets = _filter_iter( |
| instructions, old_inst_offsets, lambda inst, o: inst.offset == o |
| ) |
| new_targets = _filter_iter( |
| zip(reversed(instructions), reversed(meta.instructions)), |
| targets, |
| lambda v1, v2: v1[0] is v2, |
| ) |
| for new, old in zip(new_targets, targets): |
| block_target_offset_remap[old.offset] = new[1].offset |
| |
| transform_code_object(code, remap_block_offsets) |
| |
| # if offset is not in setup_fn_target_offsets, it is an error |
| setup_fn_target_offsets = tuple( |
| meta.block_target_offset_remap[n] for n in setup_fn_target_offsets |
| ) |
| 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) |
| ] |
| """ |
