{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 02 · Timing a GPU without lying to yourself\n",
    "\n",
    "CUDA kernels launch asynchronously. `time.time()` around a GPU call usually measures how long it took\n",
    "to *queue* the work rather than to do it, and the number it prints is confidently wrong.\n",
    "\n",
    "This notebook builds the benchmark harness you will reuse for the rest of the course.\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": "markdown",
   "metadata": {},
   "source": [
    "## 1. The wrong way, demonstrated\n",
    "\n",
    "Run this. The first number is nonsense.\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": [
    "x = torch.randn(4096, 4096, device=dev, dtype=torch.float16)\n",
    "\n",
    "# PREDICT: how long does one 4096x4096 matmul take, in ms? ______\n",
    "\n",
    "t0 = time.perf_counter()\n",
    "y = x @ x\n",
    "naive = (time.perf_counter() - t0) * 1e3\n",
    "\n",
    "if dev == 'cuda':\n",
    "    torch.cuda.synchronize()\n",
    "honest = (time.perf_counter() - t0) * 1e3\n",
    "\n",
    "print(f'no synchronize: {naive:8.3f} ms   <- timed the launch, not the work')\n",
    "print(f'synchronized:   {honest:8.3f} ms')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. The harness\n",
    "\n",
    "Three things separate a real measurement from a plausible one:\n",
    "\n",
    "1. **Warm up.** The first call pays for autotuning, allocation and lazy initialisation.\n",
    "2. **Synchronize**, or use CUDA events, so you time the work rather than the queue.\n",
    "3. **Repeat and take the median.** A single sample on a shared machine is noise.\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",
    "\n",
    "print(f'{bench(lambda: x @ x):.3f} ms')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Turn time into a rate you can sanity-check\n",
    "\n",
    "A matmul of two `n x n` matrices is `2 * n**3` floating point operations. Divide by time and you get\n",
    "achieved FLOP/s, which you can compare against what the card claims on paper.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def matmul_tflops(n, dtype=torch.float16):\n",
    "    a = torch.randn(n, n, device=dev, dtype=dtype)\n",
    "    ms = bench(lambda: a @ a)\n",
    "    return ms, (2 * n ** 3) / (ms * 1e-3) / 1e12\n",
    "\n",
    "print(f'{\"n\":>6} {\"ms\":>10} {\"TFLOP/s\":>10}')\n",
    "for n in [512, 1024, 2048, 4096]:\n",
    "    ms, tf = matmul_tflops(n)\n",
    "    print(f'{n:>6} {ms:>10.3f} {tf:>10.1f}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Read the shape of that table, not the absolute numbers.** Small matrices get poor TFLOP/s because\n",
    "the GPU spends its time on launch overhead and never fills up. As `n` grows the same hardware looks\n",
    "dramatically better. This is the first appearance of a rule that governs everything later: *a kernel\n",
    "has no single speed.* It has a speed at a size.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Extend the sweep until TFLOP/s flattens, and write down the `n` where it plateaus. That is roughly\n",
    "where your card stops being launch- and occupancy-limited and starts being genuinely compute-limited.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PREDICT: the n where TFLOP/s flattens ______\n",
    "# your sweep here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Keep the harness\n",
    "\n",
    "You will want `bench` in every later notebook. It is short on purpose — retype it rather than importing\n",
    "it. Typing it a few times is how the three rules stop being a checklist and become a habit.\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
}