The Hitchhiker's Guide — Phase 04: torch.compile, FX Graph, Custom Backends

"The first time you look at what torch.compile actually does to your model, you will have one of two reactions: either 'this is incredibly clever' or 'this is incredibly frightening.' Both reactions are correct." — anonymous PyTorch contributor

Section 1: torch.compile Architecture — Three Layers

torch.compile is not a single tool. It is a three-layer compiler pipeline:

Python code
     ↓
[Layer 1] TorchDynamo — bytecode rewriting, Python frame capture
     ↓  
FX Graph (Aten/Prims ops only, no Python control flow)
     ↓
[Layer 2] AOTAutograd — ahead-of-time differentiation, lifts grad into the graph
     ↓
Joint forward+backward FX graph
     ↓
[Layer 3] Backend — inductor (default), nvfuser, custom backend
     ↓
Optimized code / Triton kernels

Layer 1: TorchDynamo — Bytecode Rewriting

TorchDynamo intercepts Python execution at the CPython bytecode level. It does not use tracing in the classical sense — it evaluates Python code and builds an FX graph from the operations it observes.

The key mechanism is torch._dynamo.eval_frame.set_eval_frame(), which installs a callback that intercepts CPython's frame.f_back evaluation. When a guarded region encounters Python control flow (an if statement, a for loop with dynamic bounds), Dynamo breaks the graph at that point, compiling everything before the break into one subgraph and resuming Python execution.

Graph breaks are the primary performance enemy of torch.compile. Every graph break means:

1. A Triton/C++ dispatch overhead
2. Loss of fusion opportunities across the break
3. Potential synchronization

Common graph break causes:

# 1. data-dependent control flow
if x.sum() > 0:  # .sum() returns a Python scalar → break
    ...

# 2. non-PyTorch library calls
import numpy as np
y = np.array(x.numpy())  # leaves PyTorch graph → break

# 3. custom __torch_function__ with Python fallback
class MyTensor(torch.Tensor):
    def __add__(self, other):
        return super().__add__(other)  # may break depending on implementation

# 4. print statements
print(x.shape)  # accesses tensor metadata but causes Python escape → break

# 5. in-place ops on tensors that require grad
x += 1  # may cause aliasing issues → break

How to debug graph breaks:

import torch._dynamo
torch._dynamo.config.verbose = True

# Or use explain():
explanation = torch._dynamo.explain(model)(x)
print(explanation.graphs)       # compiled subgraphs
print(explanation.break_reasons) # why each break happened

Guards: for each compiled subgraph, Dynamo generates guard conditions — Python predicates that must be true for the compiled version to be valid. If guards fail (e.g., input shape changed), Dynamo recompiles (within the cache_size_limit).

Example guard for a linear layer:

# Dynamo generates something like:
def guard_fn(input):
    return (
        input.dtype == torch.float32 and
        input.device.type == 'cuda' and
        input.shape[1] == 512  # static shape guard
    )

Dynamic shapes (torch.compile(dynamic=True)) replace static shape guards with symbolic shape guards using torch.fx.experimental.symbolic_shapes.

Section 2: FX — Functional Transformation of PyTorch Graphs

torch.fx is the intermediate representation used throughout torch.compile. It represents a PyTorch program as a directed acyclic graph of nodes.

FX Node Types

Capturing an FX Graph

import torch
import torch.fx as fx

model = torch.nn.Sequential(
    torch.nn.Linear(512, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10)
)

# Symbolic tracing (static, follows Python code)
tracer = fx.symbolic_trace(model)
print(tracer.graph)

# torch.export (newer, handles more cases, produces aten-level graph)
from torch.export import export
exported = export(model, args=(torch.randn(1, 512),))
print(exported.graph)

Writing FX Passes

An FX pass is a function that takes a GraphModule, transforms its graph, and returns the modified GraphModule.

Pattern: Replace all ReLU with LeakyReLU

import torch.fx as fx

def replace_relu_with_leaky(gm: fx.GraphModule) -> fx.GraphModule:
    for node in gm.graph.nodes:
        if node.op == 'call_function' and node.target == torch.nn.functional.relu:
            with gm.graph.inserting_after(node):
                new_node = gm.graph.call_function(
                    torch.nn.functional.leaky_relu,
                    args=node.args,
                    kwargs={'negative_slope': 0.01}
                )
                node.replace_all_uses_with(new_node)
            gm.graph.erase_node(node)
    gm.recompile()
    return gm

Pattern: Fuse Linear + ReLU → LinearReLU

def fuse_linear_relu(gm: fx.GraphModule) -> fx.GraphModule:
    graph = gm.graph
    for node in list(graph.nodes):
        # Match: relu(linear(x))
        if (node.op == 'call_module' and 
            isinstance(gm.get_submodule(node.target), torch.nn.ReLU)):
            linear_node = node.args[0]
            if (linear_node.op == 'call_module' and 
                isinstance(gm.get_submodule(linear_node.target), torch.nn.Linear)):
                # Create fused module
                linear = gm.get_submodule(linear_node.target)
                fused = torch.nn.utils.fusion.fuse_linear_bn_eval(linear, None)
                # ... (register fused module, replace both nodes with one)
    gm.recompile()
    return gm

Pattern: Constant folding

def constant_fold(gm: fx.GraphModule) -> fx.GraphModule:
    """Evaluate ops with constant inputs at compile time."""
    graph = gm.graph
    for node in list(graph.nodes):
        if node.op == 'call_function':
            # Check if all inputs are constants (graph inputs are not)
            if all(isinstance(arg, (int, float, bool)) or 
                   (isinstance(arg, fx.Node) and arg.op == 'get_attr')
                   for arg in node.args):
                try:
                    # Execute the op with concrete values
                    const_value = node.target(*[
                        gm.get_parameter(a.target) if isinstance(a, fx.Node) else a
                        for a in node.args
                    ])
                    # Replace node with constant
                    with graph.inserting_before(node):
                        const_node = graph.get_attr(f'_const_{node.name}')
                    gm.register_buffer(f'_const_{node.name}', const_value)
                    node.replace_all_uses_with(const_node)
                    graph.erase_node(node)
                except Exception:
                    pass  # can't fold — skip
    gm.recompile()
    return gm

Section 3: Registering a Custom torch.compile Backend

A custom backend receives an FX GraphModule and returns a callable (typically a Python function wrapping compiled code).

from torch._dynamo.backends.common import aot_autograd
import torch._dynamo as dynamo

def my_custom_backend(gm: torch.fx.GraphModule, example_inputs):
    """
    gm: FX GraphModule (after AOTAutograd, contains aten-level ops)
    example_inputs: list of example input tensors (for shape/dtype info)
    
    Must return: callable with same signature as gm.forward
    """
    print(f"Custom backend received graph with {len(list(gm.graph.nodes))} nodes")
    
    # Option 1: Apply transformations and return gm itself
    gm = my_optimization_pass(gm)
    return gm
    
    # Option 2: JIT compile the graph (e.g., via TorchScript)
    scripted = torch.jit.script(gm)
    return scripted
    
    # Option 3: Lower to custom ISA and return simulation
    npu_program = lower_to_npu(gm, example_inputs)
    return npu_program.simulate  # callable that simulates execution

# Register the backend
@dynamo.register_backend
def my_backend(gm, example_inputs):
    return my_custom_backend(gm, example_inputs)

# Use it:
model = MyModel()
compiled = torch.compile(model, backend="my_backend")
output = compiled(x)

Using AOT Autograd with a custom backend:

from torch._dynamo.backends.common import aot_autograd

def my_inference_backend(gm, example_inputs):
    # Receives joint graph (fwd+bwd combined)
    # Lower to hardware
    return lower_to_hardware(gm)

my_backend_with_aot = aot_autograd(fw_compiler=my_inference_backend)
compiled = torch.compile(model, backend=my_backend_with_aot)

Section 4: Triton Kernels — When and How

Triton is a Python-based GPU kernel language that compiles to PTX. It is the backend for TorchInductor's generated kernels.

Key concepts:

• Programs: Triton launches many programs (like CUDA blocks); each program handles a tile of data
• tl.program_id(axis): which tile this program handles
• tl.load/tl.store: vectorized memory ops with optional masking
• tl.dot: matrix multiply of tiles (uses tensor cores)
• tl.constexpr: compile-time constant (like CUDA constexpr template param)

Softmax kernel pattern:

import triton
import triton.language as tl

@triton.jit
def softmax_kernel(
    input_ptr, output_ptr,
    n_rows, n_cols,
    BLOCK_SIZE: tl.constexpr  # must be power of 2, ≥ n_cols
):
    row_idx = tl.program_id(0)
    row_start = input_ptr + row_idx * n_cols
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    
    # Load row (with masking for out-of-bounds)
    row = tl.load(row_start + cols, mask=mask, other=float('-inf'))
    
    # Numerically stable softmax
    row_max = tl.max(row, axis=0)
    row -= row_max  # broadcast
    row_exp = tl.exp(row)
    row_sum = tl.sum(row_exp, axis=0)
    softmax_output = row_exp / row_sum
    
    # Store result
    output_start = output_ptr + row_idx * n_cols
    tl.store(output_start + cols, softmax_output, mask=mask)

Autotuning:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 128}),
        triton.Config({'BLOCK_SIZE': 256}),
        triton.Config({'BLOCK_SIZE': 512}),
        triton.Config({'BLOCK_SIZE': 1024}),
    ],
    key=['n_cols']  # re-tune when n_cols changes
)
@triton.jit
def softmax_kernel(..., BLOCK_SIZE: tl.constexpr):
    ...

Integrating Triton kernel as a custom torch.op:

class TritonSoftmax(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        n_rows, n_cols = x.shape
        output = torch.empty_like(x)
        BLOCK_SIZE = triton.next_power_of_2(n_cols)
        softmax_kernel[(n_rows,)](
            x, output, n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE
        )
        ctx.save_for_backward(output)
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        output, = ctx.saved_tensors
        # Softmax backward: grad_in = output * (grad_out - (grad_out * output).sum(dim=-1, keepdim=True))
        ...

triton_softmax = TritonSoftmax.apply

Section 5: FX Graph Optimization Passes — Reference List

These are the passes TorchInductor applies (and you can apply in a custom backend):

Section 6: Interview Pitfalls

Section 7: Resources

1. PyTorch torch.compile tutorial — https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html
2. TorchDynamo overview — Jason Ansel et al., "TorchDynamo: Towards Principled Python JIT Compilation" (Meta AI, 2022)
3. TorchInductor — blog post on the default backend architecture: https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes
4. Triton paper — Philippe Tillet et al., "Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations" (MLSys 2019)
5. FX documentation — https://pytorch.org/docs/stable/fx.html
6. torch.export documentation — https://pytorch.org/docs/stable/export.html
7. AOTAutograd internals — https://github.com/pytorch/pytorch/blob/main/torch/_functorch/aot_autograd.py (read the docstring)