Compute Graphs and Compilation

The Graph Was Always There

Autograd already builds a compute graph — to run backward.

Eager: walk it one kernel at a time (a launch + a memory round trip each). Capture it, and a compiler can rewrite across nodes. That is the whole idea.

Two Capture Philosophies

torch.compile: capture Python bytecode, graph-break to Python when stuck. jax.jit: trace to a jaxpr — no breaks, but recompile on new shapes.

Breaks vs. Retraces

def f(x):
    x = x * 2
    if x.sum() > 0:        # Data-dependent: forces a graph break
        return x + 1
    return x - 1

explanation = torch._dynamo.explain(f)(torch.randn(8, device=d2l.try_gpu()))
print(f'graph breaks: {explanation.graph_break_count}')

@jax.jit
def f(x):
    print('this prints once, at trace time')  # Not in the compiled graph
    return jnp.sin(x) + 1

x = jnp.arange(4.0)
print(jax.make_jaxpr(lambda x: jnp.sin(x) + 1)(x))  # The captured graph
_ = f(x); _ = f(x)   # Second call: no print — the trace is cached

{ lambda ; a:f32[4]. let b:f32[4] = sin a; c:f32[4] = add b 1.0:f32[] in (c,) }
this prints once, at trace time

The print vanishes: tracing sees tensor ops only. Purity is the price of having no graph breaks.

Fusion Cures the Bandwidth Chain

The bleeding chain from §13.1, cured in one line:

x = jax.random.normal(jax.random.PRNGKey(0), (4000, 4000))

def gelu_ish(x):
    return 0.5 * x * (1 + jnp.tanh(0.8 * (x + 0.04 * x**3)))

compiled = jax.jit(gelu_ish)
# First call compiles; then verify the rewrite: same answer, then faster
assert jnp.allclose(gelu_ish(x), compiled(x), atol=1e-6)

print(d2l.Benchmark(lambda: gelu_ish(x), desc='eager'))
print(d2l.Benchmark(lambda: compiled(x), desc='compiled'))

eager: 1.92 ms/call
compiled: 0.18 ms/call

Same answer first (the allclose), then faster: close to an order of magnitude, entirely on the bytes side — one memory round trip instead of one per op. FlashAttention (§10.5) is this idea by hand.

Compile the Training Step

def loss_fn(params, X):
    h = X
    for W, b in params[:-1]:
        h = jax.nn.gelu(h @ W + b)
    W, b = params[-1]
    return (h @ W + b).sum()

def train_step(params, X, lr=0.01):  # Loss, gradients, AND the update
    loss, grads = jax.value_and_grad(loss_fn)(params, X)
    return loss, jax.tree.map(lambda p, g: p - lr * g, params, grads)

key = jax.random.PRNGKey(0)
shapes = [(1024, 1024), (1024, 1024), (1024, 1024)]
params = [(jax.random.normal(k, s) * 0.03, jnp.zeros(s[1]))
          for k, s in zip(jax.random.split(key, 3), shapes)]
X = jax.random.normal(key, (512, 1024))

t0 = time.perf_counter()
compiled = jax.jit(train_step).lower(params, X).compile()  # AOT: compile now
print(f'ahead-of-time compile: {time.perf_counter() - t0:.1f} s')

print(d2l.Benchmark(lambda: train_step(params, X), desc='eager'))
print(d2l.Benchmark(lambda: compiled(params, X), desc='compiled'))

ahead-of-time compile: 2.2 s
eager: 19.29 ms/call
compiled: 0.18 ms/call

Fixed cost on call one; repaid over every step after. torch.compile(net) captures forward+backward (the optimizer stays eager); jax.jit swallows the whole step, update included. JAX’s AOT lower().compile() also lets you inspect the plan.

The Overhead Regime: Capture & Replay

Deep stack of thin layers — launches dominate arithmetic:

# XLA already amortizes launches: the whole jitted graph is one dispatch.
deep_params = [jax.random.normal(k, (256, 256)) * 0.06
               for k in jax.random.split(jax.random.PRNGKey(1), 60)]
x = jax.random.normal(jax.random.PRNGKey(2), (64, 256))

def deep_fwd(params, x):
    for W in params:
        x = jnp.tanh(x @ W)
    return x

compiled = jax.jit(deep_fwd)
compiled(deep_params, x).block_until_ready()

print(d2l.Benchmark(lambda: deep_fwd(deep_params, x), desc='eager'))
print(d2l.Benchmark(lambda: compiled(deep_params, x), desc='jit'))

eager: 9.42 ms/call
jit: 0.43 ms/call

reduce-overhead records the launches into a CUDA graph and replays them as one. XLA amortizes launches by construction — no such knob needed.

When Not to Compile

short runs — compile cost never repaid
graph breaks in the hot loop — little captured
shape churn (JAX) — recompiles every step
already compute-bound — nothing to fuse

Diagnose first (§13.1). Compile for bandwidth and overhead, not reflexively.