Warmup Guide — torch.compile, FX Graph & Custom Backends

Zero-to-expert primer for Phase 04. Builds the compilation stack from "why graphs at all" through TorchDynamo's bytecode capture, FX IR, guards, backends, Inductor, and Triton — assuming Phase 01's dispatcher knowledge.

Table of Contents

Chapter 1: Eager vs Graph — Why Capture Exists

Eager mode (default PyTorch): each Python line dispatches one kernel immediately. Maximum flexibility — arbitrary control flow, prints, debugger — but the GPU sees one op at a time. Costs: per-op Python/dispatcher overhead (microseconds each, fatal for small ops), and no cross-op optimization — every op writes its result to DRAM and the next op reads it back.

Graph mode: first capture the computation as a data structure, then optimize it as a whole (fuse ops, plan memory, specialize on shapes), then execute the optimized artifact. The entire phase is about the capture-optimize-execute pipeline and where you, the NPU/compiler engineer, insert your own stages.

The capture problem: Python is too dynamic to compile directly. Three historical strategies, all of which you should be able to compare:

  1. Trace (torch.jit.trace, fx.symbolic_trace): run the function with example inputs, record ops. Fast, but control flow is baked in — the trace lies if a branch depends on data.
  2. Script (TorchScript): parse Python source into a static language subset. Honest about control flow, but the subset is painful and the project is effectively frozen.
  3. Bytecode interception (TorchDynamo, Chapter 4): hook CPython frame evaluation, symbolically execute bytecode, and fall back to Python wherever capture fails. This is the approach that finally worked.

Chapter 2: torch.fx — The IR You Will Actually Touch

FX is PyTorch's Python-level intermediate representation: a GraphModule wrapping a Graph of Nodes. Each node has:

  • op: one of placeholder (input), get_attr (parameter fetch), call_function (torch.add, operator.mul...), call_method (x.view), call_module (self.linear), output.
  • target: the thing called (function object, attr name, module path).
  • args / kwargs: inputs, which may be other Nodes — this linking is the dataflow graph.
  • meta: a dict where shape/dtype propagation and your own passes stash information (node.meta["tensor_meta"] after ShapeProp).

How symbolic tracing works: call the module with Proxy objects instead of tensors; every operation on a Proxy records a node and returns a new Proxy. This is also exactly why it fails on data-dependent control flow — if proxy > 0: needs a concrete bool that doesn't exist. (Dynamo solves this differently; FX remains the IR.)

Why FX matters to you: it is the interchange format. Dynamo produces FX graphs; quantization workflows annotate FX graphs; your Lab 02 NPU backend consumes an FX graph. It's Python objects all the way down — printable (graph.print_tabular()), mutable, and testable.

Chapter 3: Graph Passes — Reading, Matching, Rewriting

A pass is a function GraphModule → GraphModule. The three core skills (all in Lab 01):

1. Analysis: walk graph.nodes (topologically ordered), classify, count, propagate shapes (torch.fx.passes.shape_prop.ShapeProp executes the graph with FakeTensors to fill meta).

2. Pattern matching: find subgraphs like conv → batch_norm. The robust way is structural: for each batch_norm node, check node.args[0] is a conv with exactly one user (you!). Single-user checks matter — if the conv output feeds elsewhere, fusing it away is wrong. This single-user discipline is where most homegrown passes have bugs.

3. Rewriting: create the replacement (e.g., fold BN's $\gamma/\sqrt{\sigma^2+\epsilon}$ scale into conv weight, $\beta - \mu\gamma/\sqrt{\sigma^2+\epsilon}$ into bias), insert with graph.inserting_after(node), reroute consumers via node.replace_all_uses_with(new), graph.erase_node the dead ones, then graph.lint() and gm.recompile().

Conv-BN folding math (Lab 01's centerpiece): BN at inference is an affine transform per channel; an affine transform of a conv output is absorbable into the conv's weights and bias:

$$W' = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} \cdot W, \qquad b' = \frac{\gamma,(b - \mu)}{\sqrt{\sigma^2+\epsilon}} + \beta$$

One kernel instead of two, and — critical for Phase 03 — quantization observers see the folded distribution, which is the one that ships.

Chapter 4: TorchDynamo — Capture Without Tracing's Lies

Mechanism: PEP 523 lets a C extension replace CPython's frame evaluation. When a torch.compile'd function runs, Dynamo intercepts the frame before execution and symbolically executes its bytecode: tensor operations append FX nodes; Python-level operations (list appends, attribute reads) are tracked symbolically; anything un-traceable triggers a graph break (Chapter 5).

The output per frame: (1) an FX graph of the tensor compute, (2) a guard set — runtime predicates under which this graph is valid, (3) rewritten bytecode that calls the compiled graph and resumes Python where capture stopped.

Why bytecode, not source or tracing: bytecode is what actually executes (no parser divergence); symbolic execution sees real Python semantics (closures, globals, builtins); and fallback is natural — any unsupported construct just ends the graph and lets CPython continue. Tracing can't fall back; scripting can't handle full Python. Dynamo's bet: capture most of the program honestly, rather than all of it wrongly.

Chapter 5: Guards and Graph Breaks

Guards make specialization safe. The captured graph assumed things: x.shape == (8, 512), x.dtype == float16, self.training == False, that global flag's value. Each assumption becomes a guard — a cheap predicate checked on every subsequent call. All pass → run the compiled artifact. Any fail → recompile for the new situation (each function holds a cache of (guards → compiled graph) entries; exceeding torch._dynamo.config.cache_size_limit falls back to eager, your classic silent perf bug).

Dynamic shapes: by default Dynamo specializes on exact shapes; on the second shape it sees, it generalizes that dimension to symbolic (s0) with SymInt arithmetic — guards become constraints like s0 > 1 rather than s0 == 8. dynamic=True requests this up front.

Graph breaks happen at: data-dependent control flow on tensor values (if x.sum() > 0), .item()/.tolist() (forces materialization), unsupported builtins/C extensions, prints. The cost is not correctness — it's fragmentation: each fragment compiles separately, killing cross-fragment fusion, with eager Python between. Diagnose with torch._dynamo.explain(fn)(args) or TORCH_LOGS=graph_breaks; the fix is usually rewriting the offending line in tensor terms (torch.where instead of if).

Chapter 6: AOTAutograd and Decompositions

The FX graph Dynamo captures is forward only and in terms of high-level ops. AOTAutograd traces through the autograd machinery ahead of time, producing a joint forward+backward graph, then partitions it into separate forward/backward graphs — choosing what to save vs recompute between them (this is where activation recomputation decisions live). Your backend may therefore receive a backward graph too.

Decompositions: PyTorch has ~2000 operators; no backend wants to implement them all. The _decomp registry rewrites composite ops into a small primitive set (prims / core ATen IR — a few hundred ops). Your Lab 02 backend declares which ops it handles natively and lets decompositions lower the rest. Functionalization also runs here: in-place ops (add_) become pure ops plus explicit copies, giving backends a side-effect-free graph.

Chapter 7: Backends — Where Your Code Plugs In

A backend is just a callable:

def my_backend(gm: torch.fx.GraphModule, example_inputs):
    # inspect/compile gm.graph however you like...
    return compiled_callable        # same signature as gm.forward

torch.compile(model, backend=my_backend). Whatever you return gets called with real tensors; returning gm.forward unchanged is a valid (null) backend — start there.

The real-world pattern your Lab 02 implements is partitioning: walk the graph, mark nodes your NPU supports, group maximal supported subgraphs, compile those to your ISA, and leave the rest to eager. This is exactly how QNN/TensorRT/CoreML integrations work, and the partition boundaries (with their device-transfer costs) are where deployment performance is won or lost — few large islands beat many small ones, every time.

Chapter 8: TorchInductor and Triton

Inductor is the default backend: FX → its own loop-level IR (define-by-run, scheduling fusion decisions) → Triton source for GPU (or C++/OpenMP for CPU) → compiled and cached. You can read everything it generates: TORCH_LOGS=output_code.

Triton is a Python DSL for writing GPU kernels at the block level: you write the program one block of threads executes (tl.load a tile, compute, tl.store), and the compiler handles the within-block parallelism, memory coalescing, and software pipelining that CUDA makes you do by hand. The mental model shift from CUDA: think "tile program with masks", not "one thread's scalar program". Lab 03 has you write the fused kernel by hand that Inductor would generate, then benchmark both.

Chapter 9: Kernel Fusion — The Point of All This

Why does any of this matter? Memory traffic. Unfused y = relu(x + b):

  • kernel 1 reads $x, b$ (2n), writes $t$ (n)
  • kernel 2 reads $t$ (n), writes $y$ (n) → 5n DRAM traffic + 2 launches

Fused: read $x, b$, write $y$ → 3n and 1 launch. For elementwise chains (bias, norm, activation, dropout, residual — i.e., half a transformer layer), fusion is a 1.5–3× end-to-end win because these ops are bandwidth-bound (Phase 07's roofline makes this quantitative: fusion raises arithmetic intensity). Matmuls themselves are usually left to cuBLAS/vendor kernels; the win is fusing the elementwise epilogue into them (Inductor does both).

This is also the NPU story (Phase 06): NPU compilers live and die by how long an op chain they can keep inside on-chip SRAM without a DRAM round-trip.

Lab Walkthrough Guidance

Order: Lab 01 (FX passes) → Lab 02 (backend) → Lab 03 (Triton).

  • Lab 01: print the graph first (print_tabular()); write the analysis pass before the rewrites; for conv-BN folding, verify numerics against the unfused module on random inputs to 1e-6 — then check the single-user edge case (a conv feeding two consumers must not fold).
  • Lab 02: start with the null backend (return gm.forward) to see what Dynamo hands you and how often (add a counter — watch recompiles happen as you vary shapes). Then op classification, then partitioning, then the mock-ISA codegen. Keep a supported_ops set as data, not code — that's how real integrations stay maintainable.
  • Lab 03: get the unfused PyTorch reference and its timing first. Write the Triton kernel with masks for the ragged tail; verify exactness, then benchmark across sizes — small sizes show launch overhead, large show bandwidth; explain both regimes with Chapter 9's math.

Success Criteria

You are ready for Phase 05 when you can, from memory:

  1. Name the six FX node ops and walk a printed graph aloud.
  2. Write conv-BN folding's weight/bias formulas and state the single-user precondition.
  3. Explain how Dynamo captures (bytecode symbolic execution), what a guard is, and the recompile-vs-graph-break distinction.
  4. List four graph-break triggers and the tensor-level rewrite for each.
  5. Describe what AOTAutograd and decompositions transform, and why backends want core ATen IR rather than all ~2000 ops.
  6. Compute the 5n→3n fusion traffic example and generalize it to a longer chain.

Interview Q&A

Q: torch.compile made my model slower. Diagnose. In order: (1) compile-time amortization — measure steady state, not first calls; (2) graph breaks fragmenting the model — torch._dynamo.explain, count fragments; (3) recompile storms from varying shapes/Python values — TORCH_LOGS=recompiles, fix with dynamic=True or by passing values as tensors; (4) cache-limit fallback to eager; (5) tiny workload where launch overhead dominated and CUDA graphs (mode= "reduce-overhead") is the actual fix.

Q: Why did Dynamo succeed where TorchScript failed? TorchScript demanded the whole program fit a static subset — one unsupported construct and you rewrite your codebase. Dynamo captures incrementally with graceful fallback: unsupported code just breaks the graph and runs in Python. Adoption cost went from "rewrite" to "decorate", with a perf gradient (fix breaks as you find them) instead of a correctness cliff.

Q: Your NPU supports conv/matmul/elementwise but not softmax. What does your backend do? Partition: maximal supported subgraphs run on NPU; softmax stays on host. Each boundary costs a device round-trip, so fuse aggressively within islands and consider whether a decomposition of softmax (exp/sum/div — all supported elementwise ops!) lets you keep it on-device — that judgment call is exactly Lab 02's design space.

Q: What does a Meta/Fake kernel do and why does compile need it? It computes output shapes and dtypes without touching data. Dynamo/AOTAutograd trace with FakeTensors — every op must propagate metadata or capture fails; it's also how shape specialization and memory planning happen before any real kernel runs. (You wrote one for warp_softmax in Phase 01 Lab 02.)

References