{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-1 fw=all -->\n",
    "\n",
    "# Compute Graphs and Compilation\n",
    "\n",
    "that section left a program bleeding. An unfused chain of\n",
    "elementwise operations — the kind every activation function and\n",
    "normalization layer contains — ran each operation as its own GPU kernel,\n",
    "and each kernel paid a full round trip to memory: read the input, compute\n",
    "one cheap thing, write the output, only for the next kernel to read it\n",
    "straight back. The diagnosis was bandwidth-bound, overpaying in bytes by a\n",
    "factor of however many operations the chain contains. This section applies\n",
    "the fix, in one line, and then explains what the line did.\n",
    "\n",
    "The fix is *compilation*: instead of executing the network one operation\n",
    "at a time as Python reaches it (*eager* execution), capture the whole\n",
    "computation as a graph and hand it to a compiler, which can see across\n",
    "operations and rewrite them — fusing that elementwise chain into a single\n",
    "kernel that reads once, computes everything, and writes once. Both of our\n",
    "frameworks do this; they do it in interestingly different ways, and the\n",
    "differences are worth understanding because each has a characteristic\n",
    "failure mode you will meet in practice.\n",
    "\n",
    "*Prerequisites: the three regimes and the* `d2l.Benchmark` *timer of*\n",
    "that section*; the kernel-launch and memory-round-trip costs\n",
    "of* that section*. This section retires the old imperative-\n",
    "versus-symbolic \"hybridize\" framing in favor of the modern one: eager by\n",
    "default, with a tracing compiler you switch on.*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "compilation-compute-graphs-and-compilation"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import numpy as np\n",
    "import time"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-2 fw=all -->\n",
    "\n",
    "## The Graph Was Always There\n",
    "\n",
    "A neural network *is* a computation graph — you have been building them\n",
    "since that section. When you write `y = relu(x @ W + b)`,\n",
    "autograd records a directed graph of operations precisely so it can walk\n",
    "it backward for gradients (the figure). The graph is not\n",
    "new; what is new is *who gets to see it*. (Two footnotes keep the picture\n",
    "honest: what autograd records is a *tape* for the backward pass, not a\n",
    "compiler's whole-program IR; and not every node costs a kernel — a view\n",
    "launches nothing, and a library matmul is already many operations fused\n",
    "behind one launch.)\n",
    "\n",
    "![The compute graph of a two-layer network. Autograd already builds this\n",
    "to run the backward pass; a compiler that sees the whole graph can rewrite\n",
    "it before running the forward pass.](../img/mdl-perf-compute-graph.svg)\n",
    "\n",
    "Eager execution walks this graph one node at a time, as Python reaches\n",
    "each line. Every node becomes a kernel launch (5–15 µs of CPU-to-GPU\n",
    "latency, that section) and, for a memory-bound op, a round trip\n",
    "to HBM. Python sees one operation, dispatches it, and moves on; it never\n",
    "learns that the `add` feeding the `relu` could have been done *inside*\n",
    "the same kernel. A compiler that captures the whole graph before running\n",
    "it can — that is the entire source of its advantage.\n",
    "\n",
    "A word of history, at footnote altitude: the frameworks once forced a\n",
    "choice between *imperative* execution (flexible, debuggable, slow) and\n",
    "*symbolic* graph construction (fast, rigid, awkward). That war is over.\n",
    "Everyone converged on the same answer — **eager by default, with a\n",
    "tracing compiler you turn on** — and the rest of this section is about the\n",
    "two leading implementations of that answer.\n",
    "\n",
    "## Capture: Two Philosophies\n",
    "\n",
    "The frameworks capture the graph in genuinely different ways, and the\n",
    "difference determines how each one breaks.\n",
    "\n",
    "**PyTorch — `torch.compile`, capture from Python bytecode.** `torch.compile`\n",
    "wraps a module or function and, on first call, its front-end (TorchDynamo)\n",
    "*inspects the Python bytecode* as it runs, extracting the tensor\n",
    "operations into a graph while leaving the surrounding Python alone\n",
    "[@Ansel.Yang.He.ea.2024]. When it meets Python it cannot trace into\n",
    "— a data-dependent branch, a `print`, an unsupported library call — it\n",
    "does not fail; it inserts a **graph break**, runs that piece in ordinary\n",
    "Python, and resumes capturing after it (the figure,\n",
    "top). The captured pieces are compiled; the breaks are the price. It also\n",
    "plants *guards*: cheap runtime checks that the assumptions the graph was\n",
    "compiled under (tensor shapes, dtypes) still hold, re-compiling if they do\n",
    "not.\n",
    "\n",
    "![Two capture pipelines. `torch.compile` extracts a graph from Python\n",
    "bytecode and falls back to Python on a graph break; `jax.jit` traces the\n",
    "function into a jaxpr and recompiles when input shapes or dtypes\n",
    "change.](../img/mdl-perf-compile-pipelines.svg)\n",
    "\n",
    "Let's see a graph break happen. A branch on a tensor's *value* forces one,\n",
    "because the trace cannot know which way it goes until runtime:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-3 fw=all -->\n",
    "\n",
    "**JAX — `jax.jit`, capture by tracing.** `jax.jit` takes the opposite\n",
    "stance: it *traces* the function by calling it once with abstract\n",
    "placeholder values that record every operation they touch, producing a\n",
    "typed intermediate representation called a **jaxpr**, which XLA lowers to\n",
    "optimized device code (the figure, bottom). Tracing\n",
    "sees tensor operations only — anything that is not a traced array\n",
    "operation, like a Python `print`, runs *once at trace time* and then\n",
    "vanishes from the compiled function:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "compilation-capture-two-philosophies-2"
   },
   "source": [
    "@jax.jit\n",
    "def f(x):\n",
    "    print('this prints once, at trace time')  # Not in the compiled graph\n",
    "    return jnp.sin(x) + 1\n",
    "\n",
    "x = jnp.arange(4.0)\n",
    "print(jax.make_jaxpr(lambda x: jnp.sin(x) + 1)(x))  # The captured graph\n",
    "_ = f(x); _ = f(x)   # Second call: no print — the trace is cached"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[4]\u001b[39m. \u001b[34;1mlet\u001b[39;22m b\u001b[35m:f32[4]\u001b[39m = sin a; c\u001b[35m:f32[4]\u001b[39m = add b 1.0:f32[] \u001b[34;1min \u001b[39;22m(c,) }\nthis prints once, at trace time\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-4 fw=all -->\n",
    "\n",
    "The purity requirement is the flip side of having no graph breaks: because\n",
    "the whole function must be traceable, JAX cannot fall back to Python\n",
    "mid-graph — but in exchange it never silently splits your function into\n",
    "compiled fragments. Its characteristic footgun is elsewhere: the trace is\n",
    "specialized to the *shapes and dtypes* it saw, so calling the compiled\n",
    "function with a new shape triggers a full **retrace and recompile**. In a\n",
    "training loop with variable-length batches this can mean recompiling every\n",
    "step. The escape hatches are to keep shapes static (the\n",
    "`drop_remainder=True` data-loading discipline of that section\n",
    "was exactly this, paying off now), to mark genuinely-constant arguments\n",
    "with `static_argnums`, and to express data-dependent control flow with\n",
    "`lax.cond`/`lax.scan` so it lives *inside* the graph rather than breaking\n",
    "it. The contrast in one line: **`torch.compile` bends around Python and\n",
    "you watch for breaks; `jax.jit` demands purity and you watch for\n",
    "recompiles.**\n",
    "\n",
    "## What the Compiler Does: Fusion\n",
    "\n",
    "Now the payoff. Recall the bleeding elementwise chain from\n",
    "that section — a handful of cheap operations, each its\n",
    "own kernel, each a full memory round trip. Compilation fuses them into\n",
    "one kernel that reads the input once, does all the arithmetic in\n",
    "registers, and writes once:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "compilation-what-the-compiler-does-fusion"
   },
   "source": [
    "x = jax.random.normal(jax.random.PRNGKey(0), (4000, 4000))\n",
    "\n",
    "def gelu_ish(x):\n",
    "    return 0.5 * x * (1 + jnp.tanh(0.8 * (x + 0.04 * x**3)))\n",
    "\n",
    "compiled = jax.jit(gelu_ish)\n",
    "# First call compiles; then verify the rewrite: same answer, then faster\n",
    "assert jnp.allclose(gelu_ish(x), compiled(x), atol=1e-6)\n",
    "\n",
    "print(d2l.Benchmark(lambda: gelu_ish(x), desc='eager'))\n",
    "print(d2l.Benchmark(lambda: compiled(x), desc='compiled'))"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "eager: 1.92 ms/call\ncompiled: 0.18 ms/call\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-5 fw=all -->\n",
    "\n",
    "The `assert` is not decoration: a compiler that rewrites your arithmetic\n",
    "must first return the same answer — here to within $10^{-6}$, the slack\n",
    "that reordered floating-point arithmetic legitimately needs — and only\n",
    "then is \"faster\" worth having. It is: the fused version is dramatically\n",
    "faster — close to an order of magnitude\n",
    "on this chain — and the reason is entirely on the bytes side of the\n",
    "roofline: the chain performs the same arithmetic either way, but eager\n",
    "execution makes one memory round trip *per operation* while the fused\n",
    "kernel makes a single round trip for the whole chain. Fuse eight\n",
    "elementwise ops and you cut roughly eight memory traversals to one. This\n",
    "is the general rule — **fusion trades kernel launches and memory traffic\n",
    "for free arithmetic in registers** — and it is why compilation helps most\n",
    "exactly where that section said you are bandwidth- or\n",
    "overhead-bound, and barely at all where you are already compute-bound (a\n",
    "single large matmul is already one well-tuned kernel; there is nothing to\n",
    "fuse).\n",
    "\n",
    "When the compiler cannot fuse enough — when a memory access pattern needs\n",
    "restructuring, not just merging — people write the kernel by hand. The\n",
    "hand-written FlashAttention kernel of that section is\n",
    "exactly this: the same \"keep intermediates on-chip, never round-trip the\n",
    "big matrix\" idea, executed by an expert for a pattern the general\n",
    "compiler cannot discover. Writing such kernels is its own craft — Triton\n",
    "[@Tillet.Kung.Cox.2019] (an important backend of `torch.compile`'s\n",
    "Inductor, which also draws on template and library kernels) and Pallas\n",
    "let you author them in Python-like syntax — and it\n",
    "is deliberately out of scope for this book (that section\n",
    "drew that fence). The point here is that the compiler gets you most of the\n",
    "fusion win automatically, for free, on the code you already wrote.\n",
    "\n",
    "## Compiling the Training Step, Measured\n",
    "\n",
    "Fusing one elementwise chain is a demonstration; the real use is\n",
    "compiling the training step. It is worth being precise about what each\n",
    "framework captures, because they differ. In PyTorch,\n",
    "`torch.compile(net)` captures the model's forward *and* the backward that\n",
    "flows through it, but the optimizer's `opt.step()` below stays eager —\n",
    "compiling the optimizer too is possible, just not what the one-liner\n",
    "gives you. In JAX, nothing stops the jitted function from *being* the\n",
    "whole step: loss, gradients, and the parameter update, one compiled\n",
    "program. Either way, the lesson to internalize is the *shape* of the\n",
    "cost: a fixed price on the first call, repaid over every step after:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "compilation-whole-step-compilation-measured"
   },
   "source": [
    "def loss_fn(params, X):\n",
    "    h = X\n",
    "    for W, b in params[:-1]:\n",
    "        h = jax.nn.gelu(h @ W + b)\n",
    "    W, b = params[-1]\n",
    "    return (h @ W + b).sum()\n",
    "\n",
    "def train_step(params, X, lr=0.01):  # Loss, gradients, AND the update\n",
    "    loss, grads = jax.value_and_grad(loss_fn)(params, X)\n",
    "    return loss, jax.tree.map(lambda p, g: p - lr * g, params, grads)\n",
    "\n",
    "key = jax.random.PRNGKey(0)\n",
    "shapes = [(1024, 1024), (1024, 1024), (1024, 1024)]\n",
    "params = [(jax.random.normal(k, s) * 0.03, jnp.zeros(s[1]))\n",
    "          for k, s in zip(jax.random.split(key, 3), shapes)]\n",
    "X = jax.random.normal(key, (512, 1024))\n",
    "\n",
    "t0 = time.perf_counter()\n",
    "compiled = jax.jit(train_step).lower(params, X).compile()  # AOT: compile now\n",
    "print(f'ahead-of-time compile: {time.perf_counter() - t0:.1f} s')\n",
    "\n",
    "print(d2l.Benchmark(lambda: train_step(params, X), desc='eager'))\n",
    "print(d2l.Benchmark(lambda: compiled(params, X), desc='compiled'))"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "ahead-of-time compile: 2.2 s\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "eager: 19.29 ms/call\ncompiled: 0.18 ms/call\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-6 fw=all -->\n",
    "\n",
    "The compiled step is faster in steady state — a fused training step makes\n",
    "far fewer trips to memory and issues far fewer launches — and the first\n",
    "call carries a visible one-time cost while the compiler works: a fraction\n",
    "of a second to a couple of seconds on this toy step, depending on the\n",
    "framework and the state of its compile cache, and about two seconds for\n",
    "the real Transformer of that section. On a short\n",
    "experiment that price may never be repaid; on a real training run of\n",
    "thousands of steps it is amortized to nothing. The JAX tab also shows off its\n",
    "ahead-of-time staging: `lower(...).compile()` runs the compiler *now*, as\n",
    "an explicit step rather than lazily on first call, returning a compiled\n",
    "object you can then introspect — `compiled.memory_analysis()` reports the\n",
    "memory the compiler *planned* before a single byte is allocated, a theme\n",
    "we develop in that section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-7 fw=jax -->\n",
    "\n",
    "Look at the JAX ratio once more: the un-jitted step is about two orders\n",
    "of magnitude slower than the compiled one, because every operation in the\n",
    "loss, the gradient, *and* the update dispatches separately. The practical\n",
    "rule for JAX users is blunt — never run a training step un-jitted except\n",
    "to debug. Eager JAX is a development surface; essentially all of its\n",
    "performance lives behind `jax.jit`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-8 fw=all -->\n",
    "\n",
    "## The Overhead Regime: Capture and Replay\n",
    "\n",
    "Fusion attacks the bandwidth regime. The third regime —\n",
    "*overhead-bound*, where the GPU sits idle waiting for Python to launch the\n",
    "next tiny kernel — needs a different weapon. A model that is a deep stack\n",
    "of thin layers issues a great many small kernels, and if each takes\n",
    "longer to *launch* than to *run*, the device starves no matter how little\n",
    "work it truly has (the figure). Compiling for fusion\n",
    "helps some, but the decisive fix is to stop involving Python at all:\n",
    "record the entire sequence of launches once, then *replay* it with a\n",
    "single command. This is what CUDA graphs do, and `torch.compile` exposes\n",
    "them through one mode flag:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "compilation-the-overhead-regime-capture-and-replay"
   },
   "source": [
    "# XLA already amortizes launches: the whole jitted graph is one dispatch.\n",
    "deep_params = [jax.random.normal(k, (256, 256)) * 0.06\n",
    "               for k in jax.random.split(jax.random.PRNGKey(1), 60)]\n",
    "x = jax.random.normal(jax.random.PRNGKey(2), (64, 256))\n",
    "\n",
    "def deep_fwd(params, x):\n",
    "    for W in params:\n",
    "        x = jnp.tanh(x @ W)\n",
    "    return x\n",
    "\n",
    "compiled = jax.jit(deep_fwd)\n",
    "compiled(deep_params, x).block_until_ready()\n",
    "\n",
    "print(d2l.Benchmark(lambda: deep_fwd(deep_params, x), desc='eager'))\n",
    "print(d2l.Benchmark(lambda: compiled(deep_params, x), desc='jit'))"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "eager: 9.42 ms/call\njit: 0.43 ms/call\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=compilation-md-9 fw=all -->\n",
    "\n",
    "The `reduce-overhead` mode captures the model's kernel launches into a\n",
    "CUDA graph and replays the whole thing per call, collapsing a\n",
    "hundred-odd launch latencies into one. The catch is that replay is\n",
    "*rigid*: a replayed graph is a fixed sequence of kernels on fixed memory\n",
    "addresses, so the input shape must not change between calls (change it\n",
    "and PyTorch re-captures) — and it is why we time the forward pass under\n",
    "`torch.no_grad()`: autograd's saved-for-backward activations are fresh\n",
    "allocations on every call, and their changing addresses would force a\n",
    "re-capture each time. The JAX tab needs no separate\n",
    "mechanism and says so: a jitted function is *already* a single dispatched\n",
    "executable, so XLA amortizes launch overhead by construction — the\n",
    "absence of a \"reduce-overhead\" knob in JAX is not a missing feature but a\n",
    "consequence of compile-by-tracing.\n",
    "\n",
    "## When Compilation Hurts\n",
    "\n",
    "Compilation is not free and not always worth it. The checklist:\n",
    "\n",
    "* **Short runs.** If you train for a few dozen steps, the seconds of\n",
    "  compile time may exceed everything you save. Compilation pays over\n",
    "  thousands of steps, not tens.\n",
    "* **Graph breaks in hot loops.** A `torch.compile`d function riddled with\n",
    "  data-dependent branches compiles many small fragments and falls back to\n",
    "  Python between them, keeping little of the benefit. `torch._dynamo.explain`\n",
    "  (used above) counts the breaks; the fix is to remove value-dependent\n",
    "  control flow from the hot path.\n",
    "* **Shape churn (JAX).** A `jax.jit`ted function called with a new shape\n",
    "  every step recompiles every step. Pad to fixed shapes, or accept that\n",
    "  this particular function should not be jitted.\n",
    "* **Already compute-bound.** If the profiler shows one big matmul\n",
    "  dominating, there is nothing to fuse and no overhead to hide.\n",
    "  Compiling is usually harmless, but measure rather than assume —\n",
    "  compilation can also regress time or memory — and do not expect a win.\n",
    "\n",
    "Diagnosis still comes first. Compilation is the fix for the bandwidth and\n",
    "overhead regimes; reach for it when that section's method\n",
    "points there, not reflexively. One last note on portability: the same\n",
    "capture machinery that compiles a graph can also *export* it — PyTorch's\n",
    "`torch.export` emits a portable captured graph, and JAX lowers jitted\n",
    "functions to the StableHLO interchange format; either path lets a model\n",
    "run outside Python entirely, in a serving runtime or on a phone.\n",
    "That is the modern successor to the old model-serialization story, and it\n",
    "belongs to deployment rather than training; we note it and move on.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* A neural network is already a compute graph (autograd builds it for the\n",
    "  backward pass). Eager execution walks it one kernel at a time; a\n",
    "  compiler that captures the whole graph can rewrite it — chiefly by\n",
    "  *fusing* operations so intermediates stay on-chip.\n",
    "* `torch.compile` captures from Python bytecode and inserts *graph breaks*\n",
    "  where it meets untraceable Python; `jax.jit` traces to a shape-\n",
    "  specialized *jaxpr* with no breaks but *recompiles on new shapes*. Watch\n",
    "  breaks in one, recompiles in the other.\n",
    "* Fusion is the fix for the bandwidth regime — close to an order of\n",
    "  magnitude on our unfused elementwise chain — because it trades memory\n",
    "  round trips for free register arithmetic. It barely helps\n",
    "  already-compute-bound kernels.\n",
    "* `torch.compile(mode=\"reduce-overhead\")` uses CUDA graphs to collapse\n",
    "  many small launches into one replay, curing the overhead regime; XLA\n",
    "  amortizes launches by construction and needs no such knob. Both require\n",
    "  static shapes.\n",
    "* Compilation costs seconds on the first call and repays over thousands of\n",
    "  steps. Reach for it when the method diagnoses bandwidth or overhead, not\n",
    "  reflexively.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Introduce a data-dependent `if` into a `torch.compile`d function, find\n",
    "   the break with `torch._dynamo.explain`, then rewrite the control flow\n",
    "   with `torch.where` and confirm the break count drops to zero. What\n",
    "   happened to the steady-state time?\n",
    "1. Force a `jax.jit` retrace by calling a jitted function with three\n",
    "   different input lengths in a loop, and time it. Fix it by padding\n",
    "   every input to a common length (mask the padding out of the result)\n",
    "   and confirm the recompiles stop. Then explain why `static_argnums` is\n",
    "   *not* a fix here: a static argument becomes part of the compilation\n",
    "   cache key, so every new value triggers a fresh trace-and-compile —\n",
    "   verify this by putting a counter in the function and watching it tick.\n",
    "   When *is* `static_argnums` the right tool?\n",
    "1. Sweep the depth of the thin-layer stack from 10 to 200 at fixed width\n",
    "   256 and plot eager time and `reduce-overhead` time against depth. At\n",
    "   what depth does capture-and-replay start to win, and why does the\n",
    "   crossover exist?\n",
    "1. Compile the matmul sweep of that section. For which\n",
    "   sizes does the compiled time equal the eager time, and why? (Hint:\n",
    "   what is a single large matmul already?)\n",
    "1. Time the *first* call and the tenth call of a compiled mid-sized\n",
    "   training step. How many steps of steady-state savings are needed to\n",
    "   repay the first-call compile cost? Relate the answer to the \"when\n",
    "   compilation hurts\" checklist."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
