{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 04 · Fusion: paying for traffic you did not need\n",
    "\n",
    "`relu(x * 2 + 1) * 0.5` looks like one line. Eagerly, it is several kernels, and every intermediate\n",
    "makes a full round trip to HBM and back for no reason.\n",
    "\n",
    "Fusion is the fix — and it is why FlashAttention was a breakthrough without changing a single\n",
    "mathematical step. This notebook makes the cost visible, then removes it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, time\n",
    "\n",
    "dev = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print('torch', torch.__version__, '| device', dev)\n",
    "if dev == 'cuda':\n",
    "    p = torch.cuda.get_device_properties(0)\n",
    "    print(p.name, '|', round(p.total_memory / 1e9, 1), 'GB')\n",
    "else:\n",
    "    print('No GPU. In Colab: Runtime > Change runtime type > GPU.')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bench(fn, warmup=10, iters=50):\n",
    "    \"\"\"Median wall time of fn() in ms, measured on the GPU's own clock.\"\"\"\n",
    "    for _ in range(warmup):\n",
    "        fn()\n",
    "    if dev != 'cuda':\n",
    "        ts = []\n",
    "        for _ in range(iters):\n",
    "            t0 = time.perf_counter(); fn(); ts.append((time.perf_counter() - t0) * 1e3)\n",
    "        ts.sort(); return ts[len(ts) // 2]\n",
    "    torch.cuda.synchronize()\n",
    "    ts = []\n",
    "    for _ in range(iters):\n",
    "        s = torch.cuda.Event(enable_timing=True)\n",
    "        e = torch.cuda.Event(enable_timing=True)\n",
    "        s.record(); fn(); e.record()\n",
    "        torch.cuda.synchronize()\n",
    "        ts.append(s.elapsed_time(e))\n",
    "    ts.sort(); return ts[len(ts) // 2]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Count the traffic before measuring it\n",
    "\n",
    "A chain of elementwise ops moves roughly two reads and a write *per step* when unfused. Fused, the\n",
    "intermediates never leave the chip: you read the input once and write the answer once.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Predict first.** Commit your number and one sentence of reasoning *before* running the cell.\n",
    "> The gap between prediction and measurement is the lesson; skip the prediction and the lesson is\n",
    "> invisible.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 1 << 24  # ~16.7M elements\n",
    "x = torch.randn(n, device=dev, dtype=torch.float16)\n",
    "\n",
    "def chain(t):\n",
    "    return torch.relu(t * 2.0 + 1.0) * 0.5\n",
    "\n",
    "# PREDICT: speedup from fusing this chain ______ x\n",
    "eager_ms = bench(lambda: chain(x))\n",
    "print(f'eager {eager_ms:.3f} ms   ({x.nbytes / 1e6:.1f} MB of input)')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Let the compiler fuse it\n",
    "\n",
    "`torch.compile` captures the graph and generates one fused kernel. The first call pays for compilation,\n",
    "which is exactly why the harness warms up before timing.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compiled = torch.compile(chain)\n",
    "_ = compiled(x)  # compile now, so we do not time the compiler\n",
    "\n",
    "comp_ms = bench(lambda: compiled(x))\n",
    "print(f'eager    {eager_ms:8.3f} ms')\n",
    "print(f'compiled {comp_ms:8.3f} ms   ({eager_ms / comp_ms:.2f}x)')\n",
    "print()\n",
    "print(f'fused traffic floor: {2 * x.nbytes / 1e6:.1f} MB (one read, one write)')\n",
    "print(f'compiled achieved:   {2 * x.nbytes / (comp_ms * 1e-3) / 1e9:.0f} GB/s')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Explain the number in bytes, not adjectives.** If the speedup is close to the ratio of traffic you\n",
    "removed, your model of the kernel is right. If it is far off, something else is the limit — often\n",
    "launch overhead at small sizes, which is the third regime.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Where fusion stops helping\n",
    "\n",
    "Shrink the problem and the picture inverts: the work becomes too small to matter and you are timing the\n",
    "cost of launching kernels at all.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'{\"elements\":>12} {\"eager ms\":>10} {\"compiled\":>10} {\"speedup\":>9}')\n",
    "for shift in [12, 16, 20, 24]:\n",
    "    m = 1 << shift\n",
    "    t = torch.randn(m, device=dev, dtype=torch.float16)\n",
    "    c = torch.compile(chain)\n",
    "    _ = c(t)\n",
    "    e_ms = bench(lambda: chain(t))\n",
    "    c_ms = bench(lambda: c(t))\n",
    "    print(f'{m:>12} {e_ms:>10.4f} {c_ms:>10.4f} {e_ms / c_ms:>9.2f}x')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Write one paragraph — in bytes and launches, not adjectives — explaining the shape of that table. Why\n",
    "is the speedup small at 4096 elements, largest in the middle, and then flat?\n",
    "\n",
    "Then predict what changes if you switch to float32, and commit the prediction before you run it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PREDICT: in float32 the speedup will ______ because ______\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Where this goes next\n",
    "\n",
    "You have now met all three regimes by measuring them: compute-bound (large matmul), memory-bound\n",
    "(elementwise chains), and overhead-bound (tiny tensors). Module 2 stops inferring the regime from wall\n",
    "time and reads it straight off a profile.\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}