Intro-ops/notebooks/02_vector_add_kernel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 02 — Vector Add 算子\n",
    "\n",
    "## 学习目标\n",
    "\n",
    "1. 理解逐元素并行：每个线程独立处理一对 `(a[i], b[i])`\n",
    "2. 理解 tile-level 并行（`T.Parallel` vs `T.Serial` 的区别）\n",
    "3. 完成 NVIDIA + TileLang 两个 kernel 的 TODO\n",
    "\n",
    "Vector add 和 copy 结构相似，但引入了**双输入**和**逐元素计算**。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概念导入：逐元素并行\n",
    "\n",
    "```\n",
    "a = [a0, a1, a2, ..., aN-1]\n",
    "b = [b0, b1, b2, ..., bN-1]\n",
    "         ↓ 逐元素相加\n",
    "c = [a0+b0, a1+b1, a2+b2, ..., aN-1+bN-1]\n",
    "```\n",
    "\n",
    "每个 `c[i] = a[i] + b[i]` 完全独立——不需要线程间通信。\n",
    "Grid-stride loop 同样适用。\n",
    "\n",
    "详细图示见 [docs/diagrams/thread-grid-layout.md](../docs/diagrams/thread-grid-layout.md)。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch 参考实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "a = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
    "b = torch.randn_like(a)\n",
    "expected = a + b\n",
    "\n",
    "print(f\"a[:5]:   {a[:5]}\")\n",
    "print(f\"b[:5]:   {b[:5]}\")\n",
    "print(f\"a+b[:5]: {expected[:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NVIDIA CUDA Kernel\n",
    "\n",
    "打开 `ops/vector_add/nvidia/kernel.cuh`，你需要完成两部分：\n",
    "\n",
    "1. **`add_values<T>` 辅助函数**：返回 `a + b`（泛型版本和 half 特化版）\n",
    "2. **`vector_add_contiguous_kernel`**：grid-stride loop 遍历，`out[i] = add_values(a[i], b[i])`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "KERNEL_SKELETON = \"\"\"\n",
    "template <typename T>\n",
    "__device__ T add_values(T a, T b) {\n",
    "    // TODO: return the elementwise sum for generic types.\n",
    "}\n",
    "\n",
    "template <>\n",
    "__device__ inline half add_values<half>(half a, half b) {\n",
    "    // TODO: return the half-precision elementwise sum.\n",
    "    // 提示：half 的加法用 __hadd(a, b)\n",
    "}\n",
    "\n",
    "template <typename T>\n",
    "__global__ void vector_add_contiguous_kernel(\n",
    "    T *out, const T *a, const T *b, int64_t n) {\n",
    "    // TODO: grid-stride loop, out[i] = add_values(a[i], b[i])\n",
    "}\n",
    "\"\"\"\n",
    "print(KERNEL_SKELETON)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证 NVIDIA 版本\n",
    "import subprocess, sys\n",
    "result = subprocess.run([\n",
    "    sys.executable, \"-m\", \"pytest\",\n",
    "    \"tests/op_tests/test_vector_add.py\", \"-v\",\n",
    "    \"--backend\", \"nvidia\"\n",
    "], capture_output=True, text=True)\n",
    "print(result.stdout)\n",
    "if result.returncode == 0:\n",
    "    print(\"✓ NVIDIA vector_add kernel 测试全部通过！\")\n",
    "else:\n",
    "    print(result.stderr[-500:] if result.stderr else \"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TileLang Kernel\n",
    "\n",
    "打开 `ops/vector_add/tilelang/kernel.py`，注意这里使用了 `T.Parallel`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TILELANG_SKELETON = \"\"\"\n",
    "@tilelang.jit\n",
    "def vector_add_kernel(a, b, BLOCK_N: int, dtype):\n",
    "    N = T.const(\"N\")\n",
    "    a: T.Tensor((N,), dtype)\n",
    "    b: T.Tensor((N,), dtype)\n",
    "    out = T.empty((N,), dtype)\n",
    "\n",
    "    # TODO: implement a tile-wise vector add kernel.\n",
    "    #\n",
    "    # Suggested steps:\n",
    "    # 1. T.Parallel(N // BLOCK_N) to iterate over tiles\n",
    "    # 2. Compute base = tile_idx * BLOCK_N\n",
    "    # 3. Inner T.Parallel(BLOCK_N): out[base+i] = a[base+i] + b[base+i]\n",
    "\n",
    "    return out\n",
    "\"\"\"\n",
    "print(TILELANG_SKELETON)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `T.Parallel` vs `T.Serial`\n",
    "\n",
    "| | `T.Parallel` | `T.Serial` |\n",
    "|--|-------------|------------|\n",
    "| 含义 | 循环迭代可以并行 | 循环迭代必须顺序执行 |\n",
    "| 何时用 | 迭代之间无数据依赖 | 迭代之间有数据依赖（如累加器） |\n",
    "| vector_add | ✓ 每个 `(a[i], b[i])` 独立 | N/A |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证 TileLang 版本\n",
    "result = subprocess.run([\n",
    "    sys.executable, \"-m\", \"pytest\",\n",
    "    \"tests/op_tests/test_vector_add.py\", \"-v\",\n",
    "    \"--backend\", \"tilelang\"\n",
    "], capture_output=True, text=True)\n",
    "print(result.stdout)\n",
    "if result.returncode == 0:\n",
    "    print(\"✓ TileLang vector_add kernel 测试全部通过！\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 对比验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator_runtime import vector_add\n",
    "\n",
    "a = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
    "b = torch.randn_like(a)\n",
    "\n",
    "my_out = vector_add(a, b, backend=\"nvidia\")\n",
    "torch_out = a + b\n",
    "\n",
    "print(f\"我的 kernel: {my_out[:5]}\")\n",
    "print(f\"PyTorch:     {torch_out[:5]}\")\n",
    "print(f\"一致: {torch.allclose(my_out, torch_out)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 检查清单\n",
    "\n",
    "- [ ] 理解逐元素并行的线程网格布局\n",
    "- [ ] 理解 `T.Parallel` vs `T.Serial` 的区别\n",
    "- [ ] NVIDIA `kernel.cuh` TODO 完成（add_values + kernel 函数）\n",
    "- [ ] TileLang `kernel.py` TODO 完成\n",
    "- [ ] 两种后端测试全部通过\n",
    "- [ ] benchmark 跑通"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}