Day 18 - PyCodegen: How Dynamo Writes New Bytecode Back to CPython

PyCodegen writes "call the compiled result" as new instructions: load __compiled_fn, stage inputs by Source, CALL, unpack outputs, replay side effects, RETURN; bytecode_transformation.py then assembles the instruction list into a legal code object handed back to CPython

Day 12 said what the eval hook returns to CPython is "rewritten bytecode"; the five posts since were all analysis: translate, wrap, record premises, record mutations, harvest the graph. This post is synthesis: how that new bytecode actually gets emitted, instruction by instruction.

The new bytecode's task list

The original function's code object gets swapped for an equivalent rewrite that does exactly six things:

LOAD_GLOBAL __compiled_fn_1 (Day 17 already planted it in globals)
stage each graph input onto the stack, following its Source
CALL, collect the output tuple
unpack, put each output where it belongs
replay side effects (Day 16's ledger)
RETURN_VALUE

CPython runs this as-is, never knowing it is executing a compiler's output.

PyCodegen: a code generator that takes shortcuts

The core interface is tiny: hand PyCodegen a VariableTracker and it emits the shortest instruction sequence that "gets this value onto the stack". The shortcuts rank by priority:

case	emits
value has a Source	`source.reconstruct()`: load from the original location, `LOAD_FAST x`, or `LOAD_GLOBAL cfg` + `LOAD_ATTR scale`
value is a graph output	fetch from the stashed output tuple: `LOAD_FAST graph_out_0` plus an index
plain constant	`LOAD_CONST`
container born during tracing	reconstruction code: `BUILD_LIST`, `BUILD_MAP`

The first row is the key saving: a value with a Source is already reachable in the frame, so why make the graph output an extra copy. The Source chain makes its third appearance here: Day 15 used it to generate guards, Day 17 to name inputs, today to generate load code. One chain, three outputs.

There is also a tempvars cache: a value needed twice gets a STORE_FAST into a temporary on first emit, then LOAD_FAST afterwards, never rebuilt.

Read the before and after

def f(x, n):
    return x * n + 1

TORCH_LOGS="bytecode" python demo.py

Both get printed (excerpted and tidied; instruction names vary a little across Python versions):

ORIGINAL BYTECODE f
  LOAD_FAST    x
  LOAD_FAST    n
  BINARY_OP    *
  LOAD_CONST   1
  BINARY_OP    +
  RETURN_VALUE

MODIFIED BYTECODE f
  LOAD_GLOBAL  __compiled_fn_1
  LOAD_FAST    x
  CALL         1
  STORE_FAST   graph_out_0
  LOAD_FAST    graph_out_0
  LOAD_CONST   0
  BINARY_SUBSCR
  RETURN_VALUE

Reading points: n is never passed to __compiled_fn_1, it was baked into a constant (Day 14), so the graph's only input is x; the graph returns a tuple and [0] fetches the lone return value; with side effects present, replay code would sit right before RETURN_VALUE.

The toolbox underneath: bytecode_transformation.py

Emitting instructions is easy; assembling them back into a legal code object is hard, and all the hard parts live in this file:

Jump virtualization: a jump in raw bytecode is written as a numeric position, "jump to byte 84". That encoding is brittle: insert or delete any instruction in the middle and everything after it shifts, invalidating every hardcoded number at once. So Instruction stores its jump target as a reference to another Instruction object, a bookmark rather than a page number: however the pages move, the bookmark stays on the right one. Insert and delete freely mid-flight; references get converted back to real offsets only at assembly.
EXTENDED_ARG: an instruction's argument field is one byte, holding 0 through 255. A bigger argument (a jump farther than 255, say) needs an EXTENDED_ARG instruction padded in front to supply the extra bits. The trouble: padding one in makes the bytecode longer, every offset shifts, and a jump that was exactly 250 may get pushed past 255, so it needs padding too, which shifts everything again. The only way out is to rescan repeatedly until a full pass adds nothing (a fixed point).
Stack size recomputation, the 3.11+ exception table, the linetable: the line number mapping must be rebuilt so tracebacks still point at your source.
Cross-version differences: 3.11 wants PUSH_NULL before CALL, LOAD_GLOBAL flags change every release; all of it hides behind helpers like create_call_function, so the PyCodegen layer above is written once and runs on every version.

Numeric jump offsets all break the moment an instruction is inserted; references to Instruction objects keep pointing at the target no matter what

The final exit is transform_code_object: it eats the original code object plus the new instruction list and produces a legal new code object. That is the thing Day 12's eval hook hands back to CPython.

Sweeping up

After emission comes one round of bytecode_analysis: liveness finds STORE_FASTs nobody reads and pulls them (remove_dead_code); jumps to the very next instruction get pulled too (remove_pointless_jumps). The generators upstream get to emit carelessly because the janitor cleans up, which is far cheaper than making every emission path careful on its own.

Key takeaways

The final sweep is not a checker. bytecode_analysis never verifies that the generated code is correct; whatever the generators emit must already be correct. What it runs is liveness analysis: work backwards from the results to find STORE_FASTs whose values nobody ever reads and jumps that only hop to the very next instruction, then pull them. In other words it sweeps waste, not errors. The division of labor is cheap because the complexity is paid once: there are many emission paths, and if each had to be careful never to emit waste, the cost of that care would multiply by the number of paths; centralize it in one janitor and every path gets to emit carelessly. And behind the Source chain's three appearances (guards on Day 15, input naming on Day 17, load code here) sits one common thread: the most crucial thing Dynamo records about a value is not what the value is, but the access path for how to fetch it at runtime. With the path in hand, the value can be verified, addressed, and retrieved again.

At this point the "translates cleanly all the way" route is fully open: intercept, translate, wrap, record premises, record mutations, harvest, write code. But Day 13 already said translation can raise its hand at any moment. The next post (Day 19) lays out the full graph break machinery: how the two halves around the break get stitched, how resume functions are generated with today's exact toolbox, why SpeculationLog needs two passes, and how fullgraph=True and explain() help you hunt breaks down.