{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5966b4c5",
   "metadata": {},
   "source": [
    "# Demo 4 — When does modern ML become fast?\n",
    "\n",
    "**Effective Numerical Programming**\n",
    "\n",
    "In this demo, a small neural network is used as a *numerical workload* rather than\n",
    "primarily as a machine-learning exercise.\n",
    "\n",
    "We will compare:\n",
    "\n",
    "1. one-sample-at-a-time execution,\n",
    "2. batched tensor execution,\n",
    "3. CPU and GPU execution, when a CUDA GPU is available,\n",
    "4. eager PyTorch and `torch.compile()`.\n",
    "\n",
    "The main question is:\n",
    "\n",
    "> **Where does the speed come from: the neural network, vectorization/batching,\n",
    "> the accelerator, or compilation?**\n",
    "\n",
    "The network learns the two-dimensional function\n",
    "\n",
    "$$\n",
    "f(x,y)=e^{-0.1(x^2+y^2)}\n",
    "\\left[\\sin(3x)\\cos(2y)+0.3\\sin(xy)\\right].\n",
    "$$\n",
    "\n",
    "This is deliberately a regression problem rather than image classification:\n",
    "the interesting part of the exercise is how the numerical calculation is expressed\n",
    "and executed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "117546ac",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "Run the next cell and inspect the PyTorch version and available device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f7be147",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "print(\"PyTorch:\", torch.__version__)\n",
    "print(\"CUDA available:\", torch.cuda.is_available())\n",
    "if torch.cuda.is_available():\n",
    "    print(\"GPU:\", torch.cuda.get_device_name(0))\n",
    "\n",
    "torch.manual_seed(1234)\n",
    "np.random.seed(1234)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d093b7f7",
   "metadata": {},
   "source": [
    "## 2. Generate a synthetic scientific data set\n",
    "\n",
    "**Task 1.** Complete `target_function(X)`.\n",
    "\n",
    "`X` has shape `(N, 2)`. Return a tensor of shape `(N, 1)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78875818",
   "metadata": {},
   "outputs": [],
   "source": [
    "def target_function(X):\n",
    "    # X[:, 0] is x and X[:, 1] is y\n",
    "    # TODO: implement\n",
    "    pass\n",
    "\n",
    "\n",
    "N_train = 20_000\n",
    "X_train = 8.0 * torch.rand(N_train, 2) - 4.0\n",
    "y_train = target_function(X_train)\n",
    "\n",
    "print(X_train.shape, y_train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8caba77",
   "metadata": {},
   "source": [
    "## 3. Define a small multilayer perceptron\n",
    "\n",
    "**Task 2.** Construct a network\n",
    "\n",
    "`2 → 64 → 64 → 1`\n",
    "\n",
    "using `nn.Linear` and `nn.Tanh`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "878f898a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    # TODO\n",
    ")\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88f3ec6a",
   "metadata": {},
   "source": [
    "## 4. Train the network\n",
    "\n",
    "The training loop is given because training itself is not the main subject of this demo.\n",
    "\n",
    "**Task 3.** Run the cell and observe the loss. Why do we train in batches rather than\n",
    "processing one training sample at a time?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab4ffc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=2e-3)\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "batch_size = 512\n",
    "epochs = 20\n",
    "\n",
    "model.train()\n",
    "for epoch in range(epochs):\n",
    "    perm = torch.randperm(N_train)\n",
    "    running_loss = 0.0\n",
    "\n",
    "    for start in range(0, N_train, batch_size):\n",
    "        idx = perm[start:start + batch_size]\n",
    "        xb = X_train[idx]\n",
    "        yb = y_train[idx]\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        pred = model(xb)\n",
    "        loss = loss_fn(pred, yb)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item() * len(xb)\n",
    "\n",
    "    if epoch % 5 == 0 or epoch == epochs - 1:\n",
    "        print(f\"epoch {epoch:2d}: MSE = {running_loss/N_train:.6e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a12e0e82",
   "metadata": {},
   "source": [
    "## 5. Check that the model learned something\n",
    "\n",
    "**Task 4.** Plot the true and predicted function along the line `y = 0.7`.\n",
    "\n",
    "A neural network that is fast but wrong is not useful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17767ec2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "\n",
    "x = torch.linspace(-4, 4, 500)\n",
    "X_line = torch.column_stack((x, torch.full_like(x, 0.7)))\n",
    "\n",
    "# TODO:\n",
    "# y_true = ...\n",
    "# y_pred = ...\n",
    "\n",
    "# plt.plot(...)\n",
    "# plt.legend()\n",
    "# plt.xlabel(\"x\")\n",
    "# plt.ylabel(\"f(x, 0.7)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "703d6ad9",
   "metadata": {},
   "source": [
    "## 6. Scalar-style execution versus batching\n",
    "\n",
    "We now evaluate the *same trained model* in two ways.\n",
    "\n",
    "**Task 5.**\n",
    "Implement\n",
    "\n",
    "- `predict_loop(model, X)`: process one row at a time in a Python loop;\n",
    "- `predict_batch(model, X)`: process all rows with one model call.\n",
    "\n",
    "Both functions must return the same numerical result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f0144d",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def predict_loop(model, X):\n",
    "    # TODO\n",
    "    pass\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def predict_batch(model, X):\n",
    "    # TODO\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65026e3e",
   "metadata": {},
   "source": [
    "**Task 6.** Benchmark both methods for `N = 10, 100, 1000, 10000`.\n",
    "\n",
    "Before timing, make one untimed call to each function.\n",
    "\n",
    "Questions:\n",
    "\n",
    "1. Which implementation is faster?\n",
    "2. Does batching change the mathematical algorithm?\n",
    "3. Where does the speedup come from?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c1268b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def median_time(fn, repeat=5):\n",
    "    times = []\n",
    "    for _ in range(repeat):\n",
    "        t0 = time.perf_counter()\n",
    "        fn()\n",
    "        times.append(time.perf_counter() - t0)\n",
    "    return np.median(times)\n",
    "\n",
    "\n",
    "Ns = [10, 100, 1000, 10_000]\n",
    "\n",
    "# TODO: benchmark and print a small table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32345b2a",
   "metadata": {},
   "source": [
    "## 7. CPU versus GPU\n",
    "\n",
    "This section is optional if no CUDA GPU is available.\n",
    "\n",
    "GPU operations are asynchronous, so ordinary wall-clock timing can be misleading.\n",
    "We synchronize the GPU before reading the clock.\n",
    "\n",
    "**Task 7.** Complete `time_forward()` and compare CPU and GPU for increasing `N`.\n",
    "\n",
    "Include the cost of copying the input to the GPU in one set of timings, and exclude it\n",
    "in another set.\n",
    "\n",
    "Questions:\n",
    "\n",
    "1. Is the GPU always faster?\n",
    "2. Approximately where is the crossover on your machine?\n",
    "3. Why does including host-to-device transfer change the result?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bb3dd85",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sync_if_cuda(device):\n",
    "    if device.type == \"cuda\":\n",
    "        torch.cuda.synchronize(device)\n",
    "\n",
    "\n",
    "def time_forward(model, X, device, repeat=10):\n",
    "    # TODO\n",
    "    pass\n",
    "\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    # TODO: make CPU and GPU copies of the model and benchmark\n",
    "    pass\n",
    "else:\n",
    "    print(\"No CUDA GPU: skip this section.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9ceefa",
   "metadata": {},
   "source": [
    "## 8. `torch.compile()`: compilation has a cost\n",
    "\n",
    "`torch.compile()` attempts to turn PyTorch operations into optimized compiled kernels.\n",
    "\n",
    "**Task 8.**\n",
    "\n",
    "1. Compile the trained model with `torch.compile(model)`.\n",
    "2. Time the **first call**.\n",
    "3. Then time repeated calls after compilation.\n",
    "4. Compare with ordinary eager execution.\n",
    "\n",
    "Questions:\n",
    "\n",
    "- Why should the first compiled call *not* be used as the steady-state benchmark?\n",
    "- For what kind of workload is compilation worth paying for?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71af5a31",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_big = 8.0 * torch.rand(200_000, 2) - 4.0\n",
    "\n",
    "# TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23768c5c",
   "metadata": {},
   "source": [
    "## 9. Final interpretation\n",
    "\n",
    "Write short answers to the following.\n",
    "\n",
    "1. What was the most important optimization in this demo?\n",
    "2. Which optimizations change only how the same arithmetic is executed?\n",
    "3. Why can a CPU beat a GPU for a small workload?\n",
    "4. Why is correct GPU timing more subtle than CPU timing?\n",
    "5. Why can `torch.compile()` make a program slower if the compiled function is called\n",
    "   only once or a few times?\n",
    "6. What lesson from this exercise applies to numerical computing outside machine learning?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}