CLAUDE.md

1# torch/_dynamo
2 
3TorchDynamo is a Python-level JIT compiler that captures PyTorch programs into
4FX graphs by symbolically executing Python bytecode. It hooks into CPython's
5PEP 523 frame evaluation API to intercept execution, traces operations into an
6FX graph, compiles the graph with a backend (e.g. Inductor), and generates new
7bytecode that calls the compiled code.
8 
9## Architecture Overview
10 
11The compilation pipeline, in execution order:
12 
131. **`eval_frame.py`** — Runtime entry point. `torch.compile()` wraps a
14   function in an `OptimizedModule`. At runtime, the C extension
15   (`torch._C._dynamo.eval_frame`) intercepts Python frames via PEP 523.
162. **`convert_frame.py`** — `ConvertFrameAssert.__call__` checks caches,
17   handles recompilation limits, calls `_compile()` → `trace_frame()`.
183. **`symbolic_convert.py`** — The heart of Dynamo. `InstructionTranslator`
19   symbolically executes bytecode instruction-by-instruction. Maintains a
20   symbolic `stack` (list of `VariableTracker`s) and `symbolic_locals` (dict of
21   name → `VariableTracker`). Opcodes are dispatched via a `dispatch_table`
22   built by `BytecodeDispatchTableMeta`.
234. **`output_graph.py`** — `OutputGraph` owns the FX graph being built (via
24   `SubgraphTracer`), the `SideEffects` tracker, guards, shape environment,
25   and graph args. `compile_subgraph()` finalizes the graph, calls the backend,
26   and generates output bytecode.
275. **`codegen.py`** — `PyCodegen` emits output bytecode: loads graph inputs
28   (via `Source.reconstruct()`), calls the compiled graph, unpacks outputs, and
29   replays side effects.
306. **`resume_execution.py`** — Generates continuation functions for execution
31   after graph breaks.
32 
33## Key Abstractions
34 
35### VariableTracker (`variables/`)
36 
37Every Python value encountered during tracing is wrapped in a `VariableTracker`
38subclass. Key interface: `as_python_constant()`, `as_proxy()`,
39`call_function()`, `call_method()`, `getattro_impl()`, `reconstruct()`.
40 
41Key fields: `source` (where the value came from, for guards) and
42`mutation_type` (whether/how mutations are tracked).
43 
44**Factory**: `VariableTracker.build(tx, value, source=...)` dispatches to
45`VariableBuilder` (sourced values needing guards) or `SourcelessBuilder`
46(values created during tracing).
47 
48Key subclass families in `variables/`: `TensorVariable` / `SymNodeVariable`
49(tensor.py), `ConstantVariable` (constant.py), `ListVariable` /
50`TupleVariable` (lists.py), `ConstDictVariable` (dicts.py), `SetVariable` (sets.py),
51`UserFunctionVariable` (functions.py), `BuiltinVariable` (builtin.py),
52`NNModuleVariable` (nn_module.py), `UserDefinedObjectVariable`
53(user_defined.py), `TorchHigherOrderOperatorVariable` (higher_order_ops.py),
54`LazyVariableTracker` (lazy.py). `VariableBuilder` and `SourcelessBuilder` are
55in builder.py.
56 
57### Source (`source.py`)
58 
59Tracks value provenance — how to access a value at runtime. Used for guard
60generation (`source.make_guard(GuardBuilder.XXX)`) and bytecode reconstruction
61(`source.reconstruct(codegen)`). Root sources: `LocalSource`, `GlobalSource`.
62Chained sources: `AttrSource`, `GetItemSource`, `NNModuleSource`, etc.
63 
64### Guards (`guards.py`)
65 
66Runtime conditions that must hold for cached compiled code to be reused.
67Install via `install_guard(source.make_guard(GuardBuilder.TYPE_MATCH))`.
68Common types: `TYPE_MATCH`, `ID_MATCH`, `EQUALS_MATCH`, `TENSOR_MATCH`,
69`SEQUENCE_LENGTH`. At finalization, `CheckFunctionManager` builds a tree of
70C++ `GuardManager` objects for fast runtime checking.
71 
72### Side Effects (`side_effects.py`)
73 
74Tracks mutations during tracing (attribute stores, list mutations, cell
75variable updates, tensor hooks) and replays them as bytecode after graph
76execution. The `MutationType` system (`variables/base.py`) controls what
77mutations are allowed: `ValueMutationNew/Existing`,
78`AttributeMutationNew/Existing`, or `None` (immutable). The `scope` field
79prevents cross-scope mutations inside higher-order operators.
80 
81### Other key files
82 
83- `trace_rules.py` — inline/skip/graph-break decisions per function
84- `exc.py` — exception hierarchy: `Unsupported` (graph break), `RestartAnalysis`
85  (restart tracing), `ObservedException` (user exceptions during tracing),
86  `BackendCompilerFailed`
87- `config.py` — configuration flags, supports `config.patch()` context
88  manager/decorator
89- `bytecode_transformation.py` / `bytecode_analysis.py` — low-level bytecode
90  manipulation, liveness analysis
91- `pgo.py` — profile-guided optimization for dynamic shapes
92- `polyfills/` — traceable replacements for stdlib functions
93- `repro/` — reproduction/minification tools
94 
95## Graph Breaks
96 
97Call `unimplemented()` (from `exc.py`) to trigger a graph break:
98 
99```python
100from torch._dynamo.exc import unimplemented
101from torch._dynamo import graph_break_hints
102 
103unimplemented(
104    gb_type="short_category_name",
105    context=f"dynamic details: {value}",
106    explanation="Human-readable explanation of why this breaks the graph.",
107    hints=[*graph_break_hints.SUPPORTABLE],
108)
109```
110 
111- `gb_type`: Context-free category (no dynamic strings).
112- `context`: Developer-facing details (can be dynamic).
113- `explanation`: User-facing explanation (can be dynamic).
114- `hints`: Use constants from `graph_break_hints.py`: `SUPPORTABLE`,
115  `FUNDAMENTAL`, `DIFFICULT`, `DYNAMO_BUG`, `USER_ERROR`,
116  `CAUSED_BY_EARLIER_GRAPH_BREAK`.
117 
118The `break_graph_if_unsupported` decorator on instruction handlers catches
119`Unsupported`, logs the graph break, updates the `SpeculationLog`, and restarts
120analysis. On the second pass, the partial graph is compiled at the break point
121and a resume function handles the rest.
122 
123## Testing
124 
125Tests live in `test/dynamo/`. Use `torch._dynamo.test_case.TestCase` as base
126class — it calls `torch._dynamo.reset()` in setUp/tearDown and patches config
127for strict error checking.
128 
129```bash
130python test/dynamo/test_misc.py                       # whole file
131python test/dynamo/test_misc.py MiscTests.test_foo    # single test
132python test/dynamo/test_misc.py -k test_foo           # pattern match
133```
134 
135### Common patterns
136 
137The default backend to `torch.compile()` is `backend="eager"`.
138 
139**CompileCounter** — count compilations and graph ops:
140```python
141cnt = torch._dynamo.testing.CompileCounter()
142 
143@torch.compile(backend=cnt)
144def fn(x):
145    return x + 1
146 
147fn(torch.randn(10))
148self.assertEqual(cnt.frame_count, 1)
149self.assertEqual(cnt.op_count, 1)
150```
151 
152**fullgraph=True** — assert no graph breaks:
153```python
154torch.compile(fn, backend="eager", fullgraph=True)(x)
155```
156 
157**EagerAndRecordGraphs** — inspect captured FX graphs:
158```python
159backend = torch._dynamo.testing.EagerAndRecordGraphs()
160torch.compile(fn, backend=backend)(x)
161graph = backend.graphs[0]
162```
163 
164**normalize_gm + assertExpectedInline** — snapshot test graph output:
165```python
166from torch._dynamo.testing import normalize_gm
167self.assertExpectedInline(
168    normalize_gm(backend.graphs[0].print_readable(False)),
169    """\
170expected output here
171""",
172)
173```
174 
175Call `torch._dynamo.reset()` within a test when testing multiple compilation
176scenarios in a single test method. The base class handles setUp/tearDown reset
177automatically.
178 
179## Debugging
180 
181### TORCH_LOGS
182 
183```bash
184TORCH_LOGS="graph_breaks" python script.py       # see graph breaks
185TORCH_LOGS="guards,recompiles" python script.py  # see guards and recompilation reasons
186TORCH_LOGS="graph_code" python script.py         # see captured FX graph code
187TORCH_LOGS="+dynamo" python script.py            # full debug logging
188TORCH_LOGS="bytecode" python script.py           # see bytecode transformations
189```
190 
191### Reproducing crashes
192 
193When writing a minimal repro for a Dynamo crash:
194 
1951. **Capture `TORCH_LOGS="+dynamo"` for the known-failing case.** Look for
196   the frame ID (e.g. `[3/0_1]`), INLINING/FAILED/COMPILING/Restart events,
197   speculation behavior, and graph break reasons.
1982. **Capture the same logs for your repro attempt** and diff against the
199   failing case. The divergence point tells you what condition you're missing.
2003. **Match each condition from the crash traceback:**
201   - Does the code need a speculation checkpoint? Add tensor ops before the
202     failing code so `compile_subgraph` creates one.
203   - Does the function need to be inlined on retry (not skipped)? The graph
204     break type matters — `step_unsupported` keeps inlining on retry while
205     `unimplemented` may skip the function.
206   - Does the crash happen in a resume function? Add a `graph_break()` earlier
207     so PEP 523 compiles the resume as a fresh frame.
208 
209### Structured tracing (for production)
210 
211```bash
212TORCH_TRACE=/path/to/dir python script.py  # explicit trace directory
213```
214 
215Analyze with `tlparse`.
216 
217### Compile-time profiling
218 
219`TORCH_COMPILE_DYNAMO_PROFILER=1` prints per-function cumtime/tottime
220(cProfile-style) showing where Dynamo spends time during tracing. Set to a
221file path instead to save a profile loadable by `snakeviz`.
222 
223### comptime.breakpoint() (`comptime.py`)
224 
225Drops into pdb during **compilation** to inspect Dynamo state. Call
226`comptime.breakpoint()` in user code; in the pdb session use `ctx`
227(`ComptimeContext`) to call `print_locals()`, `print_bt()`, `print_graph()`,
228or `get_local("x").as_fake()`.
229 
230### Bytecode Debugger (`bytecode_debugger.py`)
231 
232pdb-like debugger for stepping through Dynamo-generated bytecode. Useful for
233debugging segfaults (no Python traceback) and codegen errors.
234 
235```python
236with torch._dynamo.bytecode_debugger.debug():
237    my_compiled_fn(x)
238```
239 
240**Programmatic breakpoints** (no graph break): call
241`torch._dynamo.bytecode_debugger.breakpoint()` in user code, or
242`codegen.extend_output(create_breakpoint())` in codegen. Auto-activates
243without an explicit `debug()` wrapper.
244 
245**Segfault debugging**: `v` (verbose) then `c` (continue) — every instruction
246is printed with `flush=True` before execution, so the last line before a crash
247is the culprit. On exceptions, the debugger stops at the faulting instruction
248automatically.
249 
250## C++ Runtime (`torch/csrc/dynamo/`)
251 
252The C/C++ layer implements the PEP 523 frame evaluation hook, the cache, and
253the guard evaluation tree. Performance-critical runtime on every Python frame.
254 
255### Frame Evaluation
256 
257**`eval_frame.c`** — Installs a custom frame evaluation function via
258`_PyInterpreterState_SetEvalFrameFunc`. A thread-local callback controls
259behavior: `None` (disabled), `Py_False` (run-only / cache lookup), or a
260callable (full Dynamo).
261 
262**`eval_frame_cpp.cpp`** — `dynamo__custom_eval_frame` is called for every
263frame: gets `ExtraState` from the code object, builds a `FrameLocalsMapping`
264(O(1) access to locals without dict materialization), evaluates guards via
265`run_root_guard_manager()` across all `CacheEntry`s (LRU ordered). On cache
266hit, executes compiled code via a shadow frame (`dynamo_eval_custom_code_impl`
267copies `localsplus` into a new frame with the compiled code object). On miss,
268calls the Python callback to trigger compilation.
269 
270### Cache
271 
272**`extra_state.cpp/.h`** — `ExtraState` is attached per code object via
273`_PyCode_SetExtra`. Contains a `cache_entry_list` (LRU linked list),
274`frame_state` (dynamic shapes detection), and `FrameExecStrategy`.
275 
276**`cache_entry.cpp/.h`** — Each `CacheEntry` stores a `RootGuardManager*`
277(raw C++ pointer for fast guard eval), the compiled code object, and the
278backend.
279 
280### Guard Evaluation Tree (`guards.cpp`)
281 
282Guards are organized as a C++ tree (~7800 lines) mirroring the data access
283pattern. `RootGuardManager` is the root, receiving a `FrameLocalsMapping`.
284Each `GuardManager` node has leaf guards and child accessors.
285 
286**LeafGuard** subclasses: `TYPE_MATCH` (Py_TYPE pointer comparison),
287`ID_MATCH`, `EQUALS_MATCH`, `TENSOR_MATCH` (dtype/device/shape/strides/dispatch
288keys in C++), `DICT_VERSION`, `GLOBAL_STATE` (grad mode, autocast, etc.).
289 
290**GuardAccessor** subclasses define tree edges: `FrameLocalsGuardAccessor`
291(O(1) index), `GetAttrGuardAccessor`, `DictGetItemGuardAccessor`,
292`GlobalsGuardAccessor`, etc.
293 
294Key optimizations: fail-fast accessor reordering, dict version tag matching to
295skip subtrees, `FrameLocalsMapping` avoids dict construction,
296`check_nopybind()` avoids pybind11 overhead.
297 
298### Other C++ files
299 
300- `framelocals_mapping.cpp` — O(1) frame locals/cells/freevars access
301- `cpython_defs.c` — copied CPython internals for frame manipulation
302- `init.cpp` — `torch._C._dynamo` module and pybind11 bindings
303- `debug_macros.h` — `DEBUG_TRACE`, `NULL_CHECK`, `INSPECT(...)` (drops into
304  pdb from C)
305- `compiled_autograd.cpp/.h` — compiled autograd engine
306