forked from ccf-ai-infra/Intro-ops
246 lines
6.8 KiB
Plaintext
246 lines
6.8 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"# 01 — Copy 算子\n",
|
||
"\n",
|
||
"## 学习目标\n",
|
||
"\n",
|
||
"1. 理解 **grid-stride loop**：为什么它可以处理任意大小的 tensor\n",
|
||
"2. 完成 NVIDIA CUDA kernel 的 TODO\n",
|
||
"3. 完成 TileLang kernel 的 TODO\n",
|
||
"4. 验证正确性 + 查看 benchmark 结果\n",
|
||
"\n",
|
||
"Copy 是最简单的算子——把数据从 `src` 搬到 `dst`，不涉及任何计算。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 概念导入：Grid-Stride Loop\n",
|
||
"\n",
|
||
"GPU 有成千上万个线程，但 tensor 大小可能更大。Grid-stride loop 让每个线程处理多个元素：\n",
|
||
"\n",
|
||
"```python\n",
|
||
"# 伪代码\n",
|
||
"for i in range(thread_idx, N, grid_total_threads):\n",
|
||
" dst[i] = src[i]\n",
|
||
"```\n",
|
||
"\n",
|
||
"- `grid_total_threads = gridDim.x * blockDim.x`（所有线程总数）\n",
|
||
"- 步长 = grid_total_threads，每个线程“跨步”处理\n",
|
||
"- 循环条件 `i < N` 保证不越界\n",
|
||
"\n",
|
||
"详细图示见 [docs/diagrams/grid-stride-loop.md](../docs/diagrams/grid-stride-loop.md)。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## PyTorch 参考实现\n",
|
||
"\n",
|
||
"先看看用 PyTorch 怎么做 copy——训练营的目标就是实现和它一样的功能："
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import torch\n",
|
||
"\n",
|
||
"# PyTorch 版 copy\n",
|
||
"def pytorch_copy(src: torch.Tensor) -> torch.Tensor:\n",
|
||
" return src.clone()\n",
|
||
"\n",
|
||
"# 测试\n",
|
||
"src = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
|
||
"expected = pytorch_copy(src)\n",
|
||
"print(f\"输入: {src[:5]}\")\n",
|
||
"print(f\"输出: {expected[:5]}\")\n",
|
||
"print(f\"一致: {torch.allclose(src, expected)}\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## NVIDIA CUDA Kernel\n",
|
||
"\n",
|
||
"打开 `ops/copy/nvidia/kernel.cuh`，你会看到这个骨架。请填写 `TODO` 部分："
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 以下代码来自 ops/copy/nvidia/kernel.cuh，仅供参考——请直接编辑该文件\n",
|
||
"\n",
|
||
"KERNEL_SKELETON = \"\"\"\n",
|
||
"template <typename T>\n",
|
||
"__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {\n",
|
||
" // TODO: implement a grid-stride loop copy kernel.\n",
|
||
" //\n",
|
||
" // Suggested steps:\n",
|
||
" // 1. Compute the global thread index.\n",
|
||
" // int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n",
|
||
" // 2. Compute the grid-wide stride.\n",
|
||
" // int64_t stride = gridDim.x * blockDim.x;\n",
|
||
" // 3. Loop over i = idx; i < n; i += stride.\n",
|
||
" // 4. Copy src[i] to dst[i].\n",
|
||
"}\n",
|
||
"\"\"\"\n",
|
||
"print(KERNEL_SKELETON)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 你的任务\n",
|
||
"\n",
|
||
"编辑 `ops/copy/nvidia/kernel.cuh`，完成 TODO 后，重新构建并运行下面的验证："
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import subprocess, sys\n",
|
||
"\n",
|
||
"result = subprocess.run([\n",
|
||
" sys.executable, \"-m\", \"pytest\",\n",
|
||
" \"tests/op_tests/test_copy.py\", \"-v\",\n",
|
||
" \"--backend\", \"nvidia\"\n",
|
||
"], capture_output=True, text=True)\n",
|
||
"\n",
|
||
"print(result.stdout)\n",
|
||
"if result.returncode == 0:\n",
|
||
" print(\"✓ NVIDIA copy kernel 测试全部通过！\")\n",
|
||
"else:\n",
|
||
" print(result.stderr[-500:] if result.stderr else \"\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## TileLang Kernel\n",
|
||
"\n",
|
||
"打开 `ops/copy/tilelang/kernel.py`，填写 TODO："
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"TILELANG_SKELETON = \"\"\"\n",
|
||
"import tilelang\n",
|
||
"import tilelang.language as T\n",
|
||
"\n",
|
||
"@tilelang.jit\n",
|
||
"def copy_kernel(src, BLOCK_N: int, dtype):\n",
|
||
" N = T.const(\"N\")\n",
|
||
" src: T.Tensor((N,), dtype)\n",
|
||
" out = T.empty((N,), dtype)\n",
|
||
"\n",
|
||
" # TODO: implement a tile-wise copy kernel.\n",
|
||
" #\n",
|
||
" # Suggested steps:\n",
|
||
" # 1. Launch one TileLang kernel over the N // BLOCK_N tiles.\n",
|
||
" # for i in T.Parallel(N // BLOCK_N):\n",
|
||
" # 2. Use T.copy to move one tile from src to out.\n",
|
||
" # tile = T.copy(src[i * BLOCK_N : (i + 1) * BLOCK_N])\n",
|
||
" # out[i * BLOCK_N : (i + 1) * BLOCK_N] = tile\n",
|
||
"\n",
|
||
" return out\n",
|
||
"\"\"\"\n",
|
||
"print(TILELANG_SKELETON)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 验证 TileLang 版本\n",
|
||
"result = subprocess.run([\n",
|
||
" sys.executable, \"-m\", \"pytest\",\n",
|
||
" \"tests/op_tests/test_copy.py\", \"-v\",\n",
|
||
" \"--backend\", \"tilelang\"\n",
|
||
"], capture_output=True, text=True)\n",
|
||
"\n",
|
||
"print(result.stdout)\n",
|
||
"if result.returncode == 0:\n",
|
||
" print(\"✓ TileLang copy kernel 测试全部通过！\")\n",
|
||
"else:\n",
|
||
" print(result.stderr[-500:] if result.stderr else \"\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 对比：你的实现 vs PyTorch"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from operator_runtime import copy\n",
|
||
"\n",
|
||
"src = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
|
||
"my_out = copy(src, backend=\"nvidia\")\n",
|
||
"torch_out = src.clone()\n",
|
||
"\n",
|
||
"print(f\"输入: {src[:5]}\")\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",
|
||
"- [ ] 理解 grid-stride loop 原理\n",
|
||
"- [ ] 理解 global memory 合并访问（coalesced access）\n",
|
||
"- [ ] NVIDIA `kernel.cuh` TODO 完成\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
|
||
}
|