{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 01 · Tensors and the bytes they move\n",
    "\n",
    "Inference performance is mostly an accounting problem: how many bytes cross the memory bus, and how\n",
    "fast that bus is. Before any of that, you need to say exactly how large a tensor is, and know when an\n",
    "operation copies one.\n",
    "\n",
    "Work top to bottom. Every `# PREDICT:` line is a place to commit a number before running the cell.\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. How big is a tensor, really?\n",
    "\n",
    "`shape` gives the logical size. `element_size()` and `nbytes` give what the hardware actually moves.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.zeros(1024, 1024, dtype=torch.float32)\n",
    "\n",
    "# PREDICT: how many MB? ______\n",
    "print('elements  ', x.numel())\n",
    "print('bytes/elem', x.element_size())\n",
    "print('MB        ', x.nbytes / 1e6)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same tensor in lower precision. This is the whole quantization argument in one cell: the shape is\n",
    "unchanged, the traffic is not.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dt in [torch.float32, torch.bfloat16, torch.float16, torch.int8]:\n",
    "    t = torch.zeros(1024, 1024, dtype=dt)\n",
    "    print(f'{str(dt):<18} {t.nbytes / 1e6:6.2f} MB')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "A 7-billion-parameter model in bfloat16 — how many GB of weights? Compute it, then hold onto the\n",
    "number. In notebook 03 it becomes the reason token generation is slow.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = 7e9\n",
    "# PREDICT: GB in bf16 ______\n",
    "print(params * 2 / 1e9, 'GB')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Views versus copies\n",
    "\n",
    "Some operations return a new *view* over the same storage and move nothing. Others allocate and copy.\n",
    "Telling them apart is the difference between a free reshape and a hidden full pass over memory.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.arange(12).reshape(3, 4)\n",
    "b = a.t()               # transpose: a view\n",
    "d = a.t().contiguous()  # forced into a new layout: a copy\n",
    "\n",
    "print('a storage', a.untyped_storage().data_ptr())\n",
    "print('b storage', b.untyped_storage().data_ptr(), '<- same as a, so b is a view')\n",
    "print('d storage', d.untyped_storage().data_ptr(), '<- different, so d is a copy')\n",
    "print()\n",
    "print('a stride', a.stride(), 'contiguous', a.is_contiguous())\n",
    "print('b stride', b.stride(), 'contiguous', b.is_contiguous())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Strides are the mechanism.** A transpose does not move data; it swaps the step sizes used to walk\n",
    "the same buffer. That is why it is free — and also why the *next* kernel may be slower, because it now\n",
    "walks memory in a pattern the hardware likes less. You will measure that cost later.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Classify each of these before running it: view or copy?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(64, 128)\n",
    "base = x.untyped_storage().data_ptr()\n",
    "\n",
    "ops = {\n",
    "    'x.view(128, 64)':     lambda: x.view(128, 64),\n",
    "    'x.reshape(128, 64)':  lambda: x.reshape(128, 64),\n",
    "    'x.t()':               lambda: x.t(),\n",
    "    'x[:32]':              lambda: x[:32],\n",
    "    'x[:, :32]':           lambda: x[:, :32],\n",
    "    'x + 0':               lambda: x + 0,\n",
    "    'x.to(torch.float16)': lambda: x.to(torch.float16),\n",
    "}\n",
    "\n",
    "for name, fn in ops.items():\n",
    "    out = fn()\n",
    "    shared = out.untyped_storage().data_ptr() == base\n",
    "    print(f'{name:<22}', 'view (no copy)' if shared else 'COPY (allocates)')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What this buys you\n",
    "\n",
    "You can now answer, for any tensor, *how many bytes does touching this cost*. That is the numerator of\n",
    "arithmetic intensity — the ratio that decides whether a kernel is worth optimising for maths or for\n",
    "memory. Notebook 02 gets you the denominator: honest timing.\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
}