Torch Compile

`torch.compile` is the just-in-time graph compiler that PyTorch 2.0 wraps around an ordinary eager model to make it run as fused, optimized kernels: you write `model = torch.compile(model)` and change nothing else. The first call traces the model and compiles it; later calls run the compiled code. It keeps the debuggability of eager mode while recovering most of the speed that used to require hand-written kernels or a separate graph framework, and unlike the older TorchScript it almost never asks you to rewrite your model to make tracing succeed.

The frontend, TorchDynamo, captures a graph by hooking CPython itself. Using the frame-evaluation API (PEP 523), Dynamo intercepts the bytecode of your function and symbolically traces it into an FX graph of tensor operations. The important design choice is what happens when it hits something it cannot trace — data-dependent control flow, a `print`, a `.item()`, an unsupported library call. Instead of failing, Dynamo inserts a graph break: it compiles the graph up to that point, runs the offending line in normal eager Python, then resumes capturing a fresh graph afterward. This is why `torch.compile` is safe to drop onto arbitrary code; the worst case is simply less of the model gets fused.

AOTAutograd captures the backward pass ahead of time, so training compiles too. A plain graph capture only sees the forward computation, but most of the cost of training is in the backward pass. AOTAutograd traces forward and backward together into a joint graph, which lets the compiler fuse and schedule gradients as aggressively as activations and decide which intermediate tensors to save versus recompute. Beneath it, PrimTorch decomposes PyTorch's roughly two thousand operators into a small, stable set of about two hundred and fifty primitives, so a backend only has to implement the primitives rather than the entire sprawling API surface.

TorchInductor, the default backend, lowers that graph to real fused kernels. Inductor performs operator fusion, memory planning, and buffer reuse, then generates code: Triton for the GPU and vectorized C++/OpenMP for the CPU. This is the concrete link between the two most important compilation tools in PyTorch — the graph that Dynamo captured is ultimately emitted as Triton kernels, so `torch.compile` is, at the bottom, a Triton kernel generator wrapped in a Python-capture frontend. The fusion is where the speedup comes from: pointwise chains, normalizations, and activation functions collapse into single passes that keep data in registers instead of streaming it back and forth through HBM.

Guards and recompilation are the mental model you actually tune against. When Dynamo compiles a graph it installs *guards* on the properties it assumed — tensor shapes, dtypes, Python attribute values. If a later call violates a guard (a new sequence length, say), that specialization is invalid and Dynamo compiles a new one. Left unmanaged this causes recompilation storms, so `torch.compile` supports dynamic shapes to compile one shape-agnostic graph instead of one per size. The other lever is graph breaks: because each break caps how much can be fused, keeping `.item()`, host-side prints, and untraceable calls out of the hot path is the main way to make a compiled model faster, and `fullgraph=True` turns any graph break into a hard error so you can find them.

Read `torch.compile` through a *capture-then-fuse* lens rather than a *magic-speed-flag* lens: Dynamo decides how much of your Python it can turn into a graph, AOTAutograd extends that to the backward pass, and Inductor turns the result into the same Triton kernels you could have written by hand — so the two things you actually control are how many graph breaks you leave in the hot path and how often guards force a recompile.