Inside torch.compile: From Python to Kernels
torch.compile captures regions of a Python program, turns their tensor operations into optimized executable code, and reuses that code when its assumptions still hold.
The big picture
NO MATCH ↓
TorchDynamo
Traces Python bytecode, captures tensor operations, and records the assumptions needed to reuse the result.
AOTAutograd / graph preparation
Prepares tensor graphs for the backend. When gradients are needed, captures and partitions forward and backward computation. Inference needs no backward graph.
TorchInductor
Lowers and schedules the tensor operations, chooses implementations, and fuses compatible work.
CACHE SPECIALIZATION ↓
MATCH ↓
With default graph-break behavior, unsupported Python can run eagerly between compiled regions. fullgraph=True raises on a graph break. A guard miss can trigger tracing and compilation; compiler caches may avoid rebuilding everything.
Read the left path for a call that needs a new specialization. Read the right path for a call that can reuse one. Expand a compiler stage for its role.
Three boundaries matter: Dynamo captures the program, Inductor chooses how to execute its tensor operations, and guards decide when the compiled result can be reused. AOTAutograd sits between capture and the backend to prepare graphs, including backward computation when needed. This diagram shows the default Inductor route; other backends can implement a different compilation path.
Why this machinery exists
Eager execution handles tensor operations as Python reaches them. Capturing a larger region gives a compiler room to combine compatible operations and reduce intermediate memory traffic and execution overhead. The preparation costs time, so reuse matters.
A small function makes that trade-off visible. Start with a multiply, an add, and a ReLU:
import torch
def transform(x):
return torch.relu(x * 2.0 + 1.0)
For x = [-1.0, -0.25, 0.5], the result is [0.0, 0.5, 2.0]. The math is straightforward. The interesting part is how much work the runtime does to carry it out.
Why compile three operations?
In ordinary eager execution, the multiply produces a tensor, the add consumes it and produces another, and ReLU produces the output. On a GPU, separate pointwise operations typically involve separate kernel launches and intermediate reads and writes. The arithmetic itself is cheap. Moving the values and submitting the work can take longer.
Imagine a simple implementation that materializes both intermediate tensors in device memory. Each stage reads an array and writes an array: six array-sized transfers in total. A fused implementation could read the input once, keep the intermediate values in registers, and write the output once. That gives us two array-sized transfers in this simplified model. Actual memory traffic depends on caching and the generated implementation; this is a way to reason about the opportunity, not a measured threefold speedup.
A compiler can see the relationship among the operations and choose a combined implementation. That takes preparation. If we execute the function only once, the compilation cost may outweigh any savings. Repeated execution is where reuse becomes important.
Run the actual Python
The Python example was checked with PyTorch 2.14.0, Python 3.12.8, and an Apple Silicon CPU. CUDA execution has not been verified locally.
Download the notebook, open Colab, and choose File → Upload notebook. Run its cells in order. It uses CUDA when available and CPU otherwise; a GPU is optional. The notebook prints the installed PyTorch version and device so the output has a concrete context. Local CPU compilation also requires a working C++ compiler.
x = torch.randn(1024, device="cuda" if torch.cuda.is_available() else "cpu")
compiled = torch.compile(transform, backend="inductor", fullgraph=True)
expected = transform(x)
actual = compiled(x)
torch.testing.assert_close(actual, expected)
The call to torch.compile creates a wrapper. With this ordinary just-in-time path, the first invocation supplies the inputs needed for capture and compilation. Calling compiled(x) again can reuse the resulting specialization. Existing persistent caches can reduce how much compiler work a new process has to do, so “first call” and “completely cold compilation” are different measurements.
fullgraph=True asks capture to succeed as one graph or raise an error. It does not request one GPU kernel. We use it here so an unexpected graph break cannot quietly change the example into several captured regions. Dynamo core concepts
See what Dynamo captured
A small custom backend lets us print the graph:
captures = []
def inspect_backend(gm, example_inputs):
captures.append(gm.code)
print(gm.code)
return gm.forward
traced = torch.compile(
transform,
backend=inspect_backend,
fullgraph=True,
dynamic=False,
)
traced(torch.randn(8))
Look for a multiplication, an addition, and a ReLU connected by intermediate values. Variable names and formatting can differ between releases. The graph records the operations to perform, rather than storing the numerical answer for this particular random input.
This backend returns the graph’s Python callable. It is an inspection tool, not an optimizing backend. The earlier compiled function uses Inductor; traced lets us observe Dynamo without confusing graph capture with kernel generation. Custom backends
Change the values, then change the shape
Run these calls after defining traced:
for size in [8, 8, 16, 8]:
probe = torch.randn(size)
torch.testing.assert_close(traced(probe), transform(probe))
print(size, len(captures))
With dynamic=False in a fresh session, this example is intended to capture two specializations: one for shape (8,) and another for (16,). The second call has new random values but the same shape. The last call returns to the earlier shape and can reuse its specialization.
That distinction is central to this function: its tensor contents vary at runtime, while static shape assumptions influence reuse. Changing dtype, stride, or relevant Python state can also affect guards. Dynamic shape behavior is a separate policy; we deliberately disable it here to make the experiment easy to inspect.
The notebook reports the actual capture count and checks the result. If the behavior changes in a later release, the assertion exposes that difference rather than presenting a fixed transcript as a live run.
Measure the work you mean to measure
The downloadable example reports the first compiled invocation separately from warmed execution. It checks numerical agreement, warms up the compiled function, then uses torch.utils.benchmark.Timer to compare medians. GPU work must finish before a wall-clock measurement is meaningful; the example synchronizes around the first-call measurement, and the benchmark utility handles accelerator synchronization for its measurements. Benchmark utility
For this experiment, ask whether repeated calls save enough time to recover startup cost. If a compiled call saves s seconds and compilation adds c seconds, a rough break-even point is c / s calls, provided s is positive and inputs continue to reuse the same compiled work.
A tiny tensor may show little benefit because overhead dominates. A larger tensor gives memory traffic more weight. Try both before drawing a conclusion. The walkthrough has no fixed speedup claim: the notebook measures your runtime, and the architecture explains which costs to investigate.
The next useful step is to inspect Inductor’s generated output for this same function and compare it with the fused-loop sketch above. That connects a plausible optimization to what the compiler actually emitted.