# mypy: allow-untyped-defs """ This module provides utilities for generating Python bytecode in PyTorch's Dynamo system. It includes functionality for: - Constructing bytecode sequences for Python operations - Managing stack operations and variable tracking - Handling graph outputs and their conversions - Supporting different Python versions (3.11+, 3.12+, 3.13+) - Converting high-level operations to low-level bytecode instructions - Managing constant loading and attribute access - Supporting function creation and closure handling """ import collections import dataclasses import re import sys import types from collections import Counter from typing import Optional, Union import torch.nn from torch.utils._ordered_set import OrderedSet from . import graph_break_hints, utils from .bytecode_transformation import ( add_push_null, add_push_null_call_function_ex, create_call_function, create_call_method, create_dup_top, create_instruction, create_load_const, create_load_method, create_rot_n, Instruction, ) from .exc import IncorrectUsage, unimplemented_v2 from .source import AttrSource, ChainedSource, DictGetItemSource, Source from .utils import is_safe_constant, rot_n_helper from .variables.base import ValueMutationExisting, VariableTracker from .variables.functions import ( ContextlibContextManagerLocalGeneratorObjectVariable, LocalGeneratorObjectVariable, ) from .variables.nn_module import NNModuleVariable from .variables.tensor import ( NumpyNdarrayVariable, SymNodeVariable, TensorVariable, UnspecializedPythonVariable, ) from .variables.torch_function import TensorWithTFOverrideVariable @dataclasses.dataclass class GraphOutputEntry: index: int variable: VariableTracker class PyCodegen: """ Helper class uses for constructing Python bytecode """ def __init__( self, tx=None, root: Optional[torch.nn.Module] = None, graph_output_var: Optional[str] = None, tempvars=None, overridden_sources=None, ) -> None: self.root = root self.top_of_stack: Optional[Union[VariableTracker, Source]] = None self.uses: Counter[VariableTracker] = collections.Counter() self.graph_outputs: dict[int, GraphOutputEntry] = {} self._output: list[Instruction] = [] # This determines which VariableTracker/Source should be stored as # locals, and maps the VariableTracker/Source to the local variable # name. Note that it could map to None initially, in which case we'll # overwrite it to map to real temporary names via `add_cache`. self.tempvars = tempvars or {} self.tx = tx self.graph_output_var = graph_output_var self.code_options = self.tx.output.code_options self.cell_and_freevars = self.tx.cell_and_freevars self.new_var = self.tx.output.new_var self.value_from_source: bool = True # This serves as a way for codegen to use a different source; we need # this because sometimes we can't easily modify the original source # without affecting other components, e.g., guards. self.overridden_sources: dict[Source, Source] = overridden_sources or {} def restore_stack(self, stack_values, *, value_from_source=True): prev = self.value_from_source self.value_from_source &= value_from_source try: self.foreach(stack_values) finally: self.value_from_source = prev def graph_output_vars(self): return [x.variable for x in self.graph_outputs.values()] def call_reconstruct(self, value): res = value.reconstruct(self) assert res is None, f"reconstruct!=None {value}" def add_push_null(self, gen_fn, call_function_ex=False): """ `gen_fn` generates instructions via PyCodegen methods that push a single callable to the stack. `add_push_null` pushes a NULL to the stack before or after the instructions generated by `gen_fn`, depending on Python version. Will attempt to use the NULL push bit for instructions with such bits (LOAD_GLOBAL 3.11+, LOAD_ATTR 3.12+, LOAD_SUPER_ATTR). """ old_len = len(self._output) if sys.version_info < (3, 13): # gen_fn may DUP_TOP instead if TOS is not cleared. # Will cause problems since NULL will be pushed right # before the generated instructions in <= 3.12 self.clear_tos() gen_fn() # inplace modify self._output added_insts = self._output[old_len:] del self._output[old_len:] if call_function_ex: self._output.extend(add_push_null_call_function_ex(added_insts)) else: self._output.extend(add_push_null(added_insts)) if sys.version_info >= (3, 13): # NULL will be at top of stack self.clear_tos() def __call__(self, value, allow_cache=True): """ Generate code such that top-of-stack (TOS) is set to value. `allow_cache` controls the behavior in the following manner. `value` can either be a VariableTracker or a Source. If `value` is a `Source`, `allow_cache` must be True (invariant asserted below). If the source was reconstructed earlier, we will reuse the generated code by loading from top of stack or tempvars. If `value` is a `VariableTracker`, we have the following cases: 1) `allow_cache=True` a) If the value.source is not None, we will emit the code based on `value.source` to handle aliasing. b) If value.source is None (example reconstructing a local list returned by the compiled function), we will reconstruct the variable tracker (w/o any source) to emit bytecode that generates a new python object. In both cases of value.source being None or not, if the value was reconstructed earlier, we will reuse the generated code by loading from top of stack or tempvars. 2) `allow_cache=False` - This is a special case (allow_cache defaults to True). a) If the value.source is not None, we reconstruct the variable tracker and emit a new python object. You might wonder what about aliasing? The place where we use this config also has the followup code where the original python object is assigned to this new python value to handle aliasing (check side_effects.py and search for allow_cache=False). b) If value.source is None, this is not allowed. TODO - assert this. Notable effects: 1. `self.top_of_stack` will be set to `value`, if we don't codegen `value` based on source. 2. `self.uses[value]` will increment, if we don't codegen `value` based on source or cache/top-of-stack reuse; in other words, if we codegen as if `value` is modelling some brand new python value. """ if isinstance(value, Source): # If the source needs to be overridden, use the new one. source = self.overridden_sources.get(value, value) assert allow_cache is True, "allow_cache must be True for Source" if self.top_of_stack is value: self._output.append(create_dup_top()) return if self.tempvars.get(source) is not None: self._output.append(self.create_load(self.tempvars[source])) self.top_of_stack = source return try: self.call_reconstruct(source) except NotImplementedError: unimplemented_v2( gb_type="Reconstruction failure: source.reconstruct not implemented", context=str(source), explanation=f"Dynamo has no bytecode reconstruction implemented for {type(source)} variable {source}.", hints=[*graph_break_hints.DYNAMO_BUG], ) self._output.append(create_dup_top()) self.add_cache(source) self.top_of_stack = source return assert isinstance(value, VariableTracker) output = self._output graph_outputs = self.graph_outputs if allow_cache: if self.top_of_stack is value: output.append(create_dup_top()) return if self.tempvars.get(value) is not None: output.append(self.create_load(self.tempvars[value])) self.top_of_stack = value return if value.is_realized() and isinstance( value, ContextlibContextManagerLocalGeneratorObjectVariable ): raise IncorrectUsage( "NYI: Returning a @contextmanager object from a torch.compile function" ) # Dynamo normally prefers codegen from source to account for aliasing. if ( value.source is not None and allow_cache and not ( value.is_realized() and isinstance(value, LocalGeneratorObjectVariable) ) ): # There's a corner case for export: for instance, if the computation # graph is just identity on an input tensor, Dynamo would just emit # a `LOAD_FAST` from the input source, rather than generating an # identity FX graph. # # However, export wants to maximize graph capture; in the case # above, export _wants to_ obtain an identity FX graph (despite it # appears unnecessarily expensive for `torch.compile`), so we have # the following option to override Dynamo's preference for codegen # from source. Morever, this option applies recursively, for cases # like input tensor being returned in a new dictionary. # # And why the `ValueMutationExisting` check? Not sure, so leaving it # to keep the old behavior, as when `value_from_source` was # introduced. TODO sort out the invariants among side effect, # codegen and export. if ( isinstance(value.mutation_type, ValueMutationExisting) or self.value_from_source ): return self(value.source) if value.is_python_constant() and is_safe_constant(value.as_python_constant()): output.append(self.create_load_const(value.as_python_constant())) elif isinstance(value, TensorWithTFOverrideVariable): graph_outputs_key = self.add_graph_output(value) self.add_push_null( lambda: self.load_import_from(utils.__name__, "to_subclass") ) self.load_graph_output(graph_outputs[graph_outputs_key].index) output.append( self.create_load_global( value.global_mangled_class_name(self.tx), add=True ) ) output.extend(create_call_function(2, False)) elif ( isinstance(value, SymNodeVariable) and value.python_type() == float and not self.tx.export ): # This is a little unusual; force the output convention to be a # Tensor here. Don't do this for export because this is # apparently load bearing for export tests (but I am a bit # doubtful it actually works in the real world) # NB: It works to add_graph_output on a computed expression # as_tensor here, because we memoize as_tensor calls on # SymNodeVariable! graph_outputs_key = self.add_graph_output( value.as_tensor(self.tx, torch.float64) ) def gen_fn(): self.load_graph_output(graph_outputs[graph_outputs_key].index) output.append(self.create_load_attr("item")) self.add_push_null(gen_fn) output.extend(create_call_function(0, False)) elif isinstance( value, ( TensorVariable, SymNodeVariable, UnspecializedPythonVariable, NumpyNdarrayVariable, ), ): graph_outputs_key = self.add_graph_output(value) if isinstance(value, NumpyNdarrayVariable): self.add_push_null( lambda: self.load_import_from(utils.__name__, "to_numpy_helper") ) self.load_graph_output(graph_outputs[graph_outputs_key].index) output.extend(create_call_function(1, False)) elif isinstance(value, UnspecializedPythonVariable) and value.need_unwrap: def gen_fn(): self.load_graph_output(graph_outputs[graph_outputs_key].index) output.append(self.create_load_attr("item")) self.add_push_null(gen_fn) output.extend(create_call_function(0, False)) else: self.load_graph_output(graph_outputs[graph_outputs_key].index) elif isinstance(value, NNModuleVariable): parts = value.module_key.split(".") if parts[0] in self.code_options["co_varnames"]: output.append(self.create_load(parts[0])) parts = parts[1:] else: assert self.root is not None output.append(self.create_load_const_unchecked(self.root)) for part in parts: output.append(self.create_load_attr(part)) else: self.uses[value] += 1 try: self.call_reconstruct(value) except NotImplementedError: unimplemented_v2( gb_type="Reconstruction failure", context=str(value), explanation=f"Dynamo has no bytecode reconstruction implemented for sourceless variable {value}.", hints=[ "If Dynamo attempting to trace a return statement and your code is attempting to return a variable " "that Dynamo cannot reconstruct, then remove it from the return statement.", *graph_break_hints.CAUSED_BY_EARLIER_GRAPH_BREAK, "Report an issue to PyTorch if you need reconstrtuction support. Note that objects that don't have" "reconstruction rules may be fundamentally unreconstructable.", ], ) if allow_cache and value in self.tempvars: self._output.append(create_dup_top()) self.add_cache(value) self.top_of_stack = value def add_graph_output(self, value): graph_outputs_key = id(value.as_proxy()) if graph_outputs_key not in self.graph_outputs: self.graph_outputs[graph_outputs_key] = GraphOutputEntry( len(self.graph_outputs), value ) return graph_outputs_key def load_graph_output(self, index): output = self._output output.append(self.create_load(self.graph_output_var)) output.append(self.create_load_const(index)) output.append(self.create_binary_subscr()) def add_cache(self, value): var = self.new_var() self.tempvars[value] = var self._output.append(self.create_store(var)) def foreach(self, items): for i in items: self(i) def create_binary_subscr(self) -> Instruction: return create_instruction("BINARY_SUBSCR") def setup_globally_cached(self, name, value): """Store value in a new global""" name = re.sub(r"[^a-zA-Z0-9_]+", "_", name) f_globals = self.tx.f_globals if name in f_globals: assert id(f_globals[name]) == id(value) else: f_globals[name] = value return [self.create_load_global(name, add=True)] def clear_tos(self): self.top_of_stack = None def append_output(self, inst): assert isinstance(inst, Instruction) self._output.append(inst) self.clear_tos() def extend_output(self, insts): assert all(isinstance(x, Instruction) for x in insts) self._output.extend(insts) self.clear_tos() def get_instructions(self) -> list[Instruction]: return self._output def create_load(self, name) -> Instruction: assert name in self.code_options["co_varnames"], f"{name} missing" return create_instruction("LOAD_FAST", argval=name) def create_load_closure(self, name) -> Instruction: assert name in self.cell_and_freevars() inst_name = "LOAD_FAST" if sys.version_info >= (3, 13) else "LOAD_CLOSURE" return create_instruction(inst_name, argval=name) def create_load_deref(self, name) -> Instruction: assert name in self.cell_and_freevars() return create_instruction("LOAD_DEREF", argval=name) def create_store(self, name) -> Instruction: assert name in self.code_options["co_varnames"], f"{name} missing" return create_instruction("STORE_FAST", argval=name) def create_store_deref(self, name) -> Instruction: assert name in self.cell_and_freevars() return create_instruction("STORE_DEREF", argval=name) def create_load_global(self, name, add=False) -> Instruction: if add: self.tx.output.update_co_names(name) assert name in self.code_options["co_names"], f"{name} not in co_names" return create_instruction("LOAD_GLOBAL", argval=name) def create_load_const(self, value) -> Instruction: return create_load_const(value) def create_load_const_unchecked(self, value) -> Instruction: return create_load_const(value, checked=False) def load_method(self, name): self.tx.output.update_co_names(name) self.append_output(create_load_method(name)) def call_method(self, nargs): self.extend_output(create_call_method(nargs)) def create_load_attr(self, name) -> Instruction: if name not in self.code_options["co_names"]: self.code_options["co_names"] += (name,) return create_instruction("LOAD_ATTR", argval=name) def load_attr(self, name): self.append_output(self.create_load_attr(name)) def create_load_attrs(self, names): return [self.create_load_attr(name) for name in names.split(".")] def create_store_attr(self, name) -> Instruction: if name not in self.code_options["co_names"]: self.code_options["co_names"] += (name,) return create_instruction("STORE_ATTR", argval=name) def store_attr(self, name): self.append_output(self.create_store_attr(name)) def load_function_name(self, fn_name, push_null, num_on_stack=0): """Load the global fn_name on the stack num_on_stack down""" output = [] if push_null and sys.version_info >= (3, 11): output.extend(add_push_null(self.create_load_global(fn_name, add=True))) if num_on_stack > 0: output.extend( [ *self.rot_n(num_on_stack + 2), *self.rot_n(num_on_stack + 2), ] ) else: output.extend( [ self.create_load_global(fn_name, add=True), *self.rot_n(num_on_stack + 1), ] ) return output def rot_n(self, n): try: return create_rot_n(n) except AttributeError: # desired rotate bytecode doesn't exist, generate equivalent bytecode return [ create_instruction("BUILD_TUPLE", arg=n), self.create_load_const_unchecked(rot_n_helper(n)), *create_rot_n(2), create_instruction("CALL_FUNCTION_EX", arg=0), create_instruction("UNPACK_SEQUENCE", arg=n), ] def pop_null(self): # POP_TOP doesn't work for null, so we pop nulls by pushing in a # nop function, calling it (which consumes the null), and popping the result. assert sys.version_info >= (3, 11) return [ self.create_load_const_unchecked(lambda: None), # 3.13 swapped NULL and callable *( (create_instruction("SWAP", arg=2),) if sys.version_info >= (3, 13) else () ), *create_call_function(0, False), create_instruction("POP_TOP"), ] def pop_top(self): self.append_output(create_instruction("POP_TOP")) def call_function(self, nargs: int, push_null: bool): self.extend_output(create_call_function(nargs, push_null=push_null)) def dup_top(self): self.append_output(create_dup_top()) def store(self, varname): self.append_output(self.create_store(varname)) def load_deref(self, varname): self.append_output(self.create_load_deref(varname)) def make_function_with_closure( self, fn_name: str, code: types.CodeType, push_null: bool, num_on_stack=0 ): freevars = code.co_freevars assert freevars output = self._output def gen_fn(): # Emitting `LOAD_FAST/LOAD_CLOSURE` with names in `co_freevars` # requires that in the generated bytecode, these cells would keep # their original local names, which we ensure via # `CellVariable.local_name`. for var in freevars: assert var in self.cell_and_freevars() output.append(self.create_load_closure(var)) output.append(create_instruction("BUILD_TUPLE", arg=len(freevars))) output.append(self.create_load_const(code)) if sys.version_info < (3, 11): output.append(self.create_load_const(fn_name)) if sys.version_info >= (3, 13): output.extend( [ create_instruction("MAKE_FUNCTION"), create_instruction("SET_FUNCTION_ATTRIBUTE", arg=0x08), ] ) else: output.append(create_instruction("MAKE_FUNCTION", arg=0x08)) if push_null and sys.version_info >= (3, 11): self.add_push_null(gen_fn) output.extend(self.rot_n(num_on_stack + 2)) output.extend(self.rot_n(num_on_stack + 2)) else: gen_fn() output.extend(self.rot_n(num_on_stack + 1)) self.clear_tos() def create_load_python_module(self, mod) -> Instruction: """ Generate a LOAD_GLOBAL instruction to fetch a given python module. """ output = self.tx.output global_scope = output.global_scope name = re.sub(r"^.*[.]", "", mod.__name__) if global_scope.get(name, None) is mod: return self.create_load_global(name, add=True) prefix = f"___module_{name}" global_name = self.tx.output.install_global_by_id(prefix, mod) return self.create_load_global(global_name, add=True) def mark_source_temp(self, source: Source) -> None: """ Mark a source as a temp variable, so that it can be reused. """ if source not in self.tempvars: self.tempvars[source] = None def make_call_generated_code(self, fn_name: str) -> None: """Call the generated code function stored in fn_name""" self.extend_output(self.load_function_name(fn_name, True)) graphargs = self.tx.output.graphargs seen_sources: OrderedSet[Source] = OrderedSet() def collect_temp_source(source): if source in seen_sources: # This source is used atleast twice, so it can be reused self.mark_source_temp(source) # Dont trace source further. This prevents us from marking too # many nodes as temp sources. return seen_sources.add(source) if isinstance(source, ChainedSource): collect_temp_source(source.base) if isinstance(source, DictGetItemSource) and isinstance( source.index, Source ): collect_temp_source(source.index) # Collect all the sources that are used more than once, so that we can # generate tmp variables in the generated pre-graph bytecode. This # essentially implements CSE. for arg in graphargs: if arg.source is not None: collect_temp_source(arg.source) for arg in graphargs: if arg.pass_arg_as_tensor: self.add_push_null( lambda: self.extend_output( [ self.create_load_python_module(torch), self.create_load_attr("_as_tensor_fullprec"), ] ) ) self.call_reconstruct(arg) self.extend_output(create_call_function(1, False)) else: self.call_reconstruct(arg) self.extend_output(create_call_function(len(graphargs), False)) def load_import_from(self, module_name, object_name) -> None: source = AttrSource(self.tx.import_source(module_name), object_name) # Note: This approach is somewhat aggressive because typically, a source is marked # as a tempvar only when it is used more than once. In this case, we're marking it # as a tempvar without performing that analysis. However, this is a simple solution, # and in many cases, load imports are reused multiple times. self.mark_source_temp(source) self(source) def create_call_function_kw(self, nargs, kw_names, push_null) -> list[Instruction]: if sys.version_info >= (3, 13): output = create_call_function(nargs, push_null) assert output[-1].opname == "CALL" output.insert(-1, self.create_load_const(kw_names)) output[-1] = create_instruction("CALL_KW", arg=nargs) return output elif sys.version_info >= (3, 11): output = create_call_function(nargs, push_null) if sys.version_info >= (3, 12): idx = -1 expected_inst = "CALL" else: idx = -2 expected_inst = "PRECALL" assert output[idx].opname == expected_inst kw_names_inst = create_instruction("KW_NAMES", argval=kw_names) output.insert(idx, kw_names_inst) return output return [ self.create_load_const(kw_names), create_instruction("CALL_FUNCTION_KW", arg=nargs), ] def create_delete(self, value) -> Instruction: return create_instruction("DELETE_FAST", argval=value)