{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 03 · Arithmetic intensity, and the line that decides everything\n",
    "\n",
    "Every kernel is limited by one of three things: the arithmetic it must do, the bytes it must move, or\n",
    "the overhead of being launched at all. **Arithmetic intensity** — FLOPs performed per byte moved —\n",
    "tells you which, and it is the number to reach for before optimising anything.\n",
    "\n",
    "Low intensity means memory-bound: compute units idle while data crawls in from HBM. High intensity\n",
    "means compute-bound: the data is already close and the arithmetic is the cost.\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. Two kernels at opposite ends\n",
    "\n",
    "Elementwise addition reads two values and writes one, to do a single add. A matmul reads `2n**2` values\n",
    "and does `2n**3` operations. These are not slightly different — they differ by a factor of `n`.\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 = 4096\n",
    "a = torch.randn(n, n, device=dev, dtype=torch.float16)\n",
    "b = torch.randn(n, n, device=dev, dtype=torch.float16)\n",
    "\n",
    "# PREDICT: matmul is ______ x slower than add, on identical tensors\n",
    "add_ms = bench(lambda: a + b)\n",
    "mm_ms = bench(lambda: a @ b)\n",
    "print(f'add    {add_ms:8.3f} ms')\n",
    "print(f'matmul {mm_ms:8.3f} ms   ({mm_ms / add_ms:.0f}x slower)')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Put numbers on why\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "elem = n * n\n",
    "bpe = a.element_size()\n",
    "\n",
    "add_flops, add_bytes = elem, 3 * elem * bpe          # read a, read b, write out\n",
    "mm_flops, mm_bytes = 2 * n ** 3, 3 * elem * bpe      # same traffic, vastly more maths\n",
    "\n",
    "print(f'add    intensity {add_flops / add_bytes:10.2f} FLOP/byte')\n",
    "print(f'matmul intensity {mm_flops / mm_bytes:10.2f} FLOP/byte')\n",
    "print()\n",
    "print(f'add    achieved  {add_bytes / (add_ms * 1e-3) / 1e9:8.0f} GB/s')\n",
    "print(f'matmul achieved  {mm_flops / (mm_ms * 1e-3) / 1e12:8.1f} TFLOP/s')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Notice what each kernel is good at.** The add reaches a bandwidth figure close to what the card can\n",
    "physically sustain: it runs as fast as memory allows and no faster, and extra compute units would do\n",
    "nothing for it. The matmul reaches a meaningful fraction of peak TFLOP/s, because it reuses every byte\n",
    "it loads roughly `n` times.\n",
    "\n",
    "**This is the entire argument for why decode is the hard part of LLM inference.** Generating one token\n",
    "reads every weight in the model to do a comparatively tiny amount of arithmetic. Intensity sits near\n",
    "the floor, so decode runs at the speed of memory — and a faster card with the same bandwidth changes\n",
    "nothing.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. The roofline, drawn from your own measurements\n",
    "\n",
    "Sweep a matmul across sizes and watch intensity climb until the kernel stops being memory-bound.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'{\"n\":>6} {\"FLOP/byte\":>12} {\"TFLOP/s\":>10} {\"GB/s\":>10}')\n",
    "for m in [128, 256, 512, 1024, 2048, 4096]:\n",
    "    t = torch.randn(m, m, device=dev, dtype=torch.float16)\n",
    "    ms = bench(lambda: t @ t)\n",
    "    flops = 2 * m ** 3\n",
    "    byts = 3 * m * m * t.element_size()\n",
    "    print(f'{m:>6} {flops / byts:>12.1f} {flops / (ms * 1e-3) / 1e12:>10.1f} {byts / (ms * 1e-3) / 1e9:>10.0f}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Find your card's **ridge point**. Look up its peak memory bandwidth and peak FP16 throughput, then\n",
    "compute the arithmetic intensity at which it stops being memory-bound. Any kernel below that line is a\n",
    "memory problem, no matter how you write it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peak_bw_gbs = 0   # fill in from your card's spec sheet\n",
    "peak_tflops = 0   # fill in\n",
    "\n",
    "if peak_bw_gbs and peak_tflops:\n",
    "    ridge = (peak_tflops * 1e12) / (peak_bw_gbs * 1e9)\n",
    "    print(f'ridge point: {ridge:.1f} FLOP/byte')\n",
    "    print('below this you are memory-bound, and only reducing traffic helps')\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
}