4.0 KiB
Copy 算子 — 教学视频脚本
时长:15-20 分钟 | 难度:入门
[5min] 概念讲解
开场(30s)
"今天我们从 intro-ops 训练营最简单的算子开始——copy。GPU 上做 copy 和 CPU 上不同——你写不出 memcpy 那种一行代码,因为 GPU 有成百上千个线程同时在工作。怎么让每个线程知道它该搬哪些数据?这就是 grid-stride loop 要解决的问题。"
算子在深度学习中的用途(1min)
"Copy 看起来 trivial,但在深度学习框架中无处不在。Tensor 的 clone()、contiguous()、数据加载中的 host-to-device 拷贝——底层都是类似的 copy kernel。理解 copy 就是理解 GPU 内存模型的第一步。"
算法推导(3.5min)
关键画面: 展示 grid-stride-loop Mermaid 图(docs/diagrams/grid-stride-loop.md)
讲解要点:
- GPU 的线程组织:grid → block → thread
- 每个线程如何计算自己的全局索引:
idx = blockIdx.x * blockDim.x + threadIdx.x - stride = grid 总线程数 =
gridDim.x * blockDim.x - 为什么需要循环?因为线程数可能少于元素数
- 循环条件
i < N保证了越界安全
过渡语: "好,理论就这么多。接下来我们在 IDE 里把它写出来。"
[10min] 代码实操
CUDA Kernel(5min)
打开文件: ops/copy/nvidia/kernel.cuh
边写边讲:
模板参数 T:支持 float、half 等多种数据类型
__global__:告诉 nvcc 这是 GPU 函数
命名空间 oprt::copy::nvidia:遵循项目约定
关键决策点:
- "为什么用
int64_t而不是int?——tensor 可能有超过 2^31 个元素" - "为什么
idx和stride的计算不放在循环条件里?——编译器优化和可读性" - "为什么循环变量叫
i不是idx?——idx是起始位置,i是当前处理位置"
写完后展示完整代码。
TileLang Kernel(5min)
打开文件: ops/copy/tilelang/kernel.py
边写边讲:
@tilelang.jit:JIT 编译装饰器——和 nvcc 不同,这里编译发生在运行时
T.Parallel(N // BLOCK_N):告诉编译器这 N/BLOCK_N 个 tile 可以并行
T.copy:把全局内存的一整块搬进 fragment
对比时刻: "看,TileLang 不需要你写 grid、block、thread——编译器帮你做。这 5 行 Python 生成的 CUDA code 比你手写的可能还高效。"
[3min] 测试验证
跑测试
# NVIDIA
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \
pytest tests/op_tests/test_copy.py -v --backend nvidia
展示: 终端中测试全部绿色的画面。
跑 benchmark
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \
python tests/run_ops.py --op copy --backend nvidia --mode all
解读 benchmark 输出:
- "看这个 Bandwidth——如果你跑不到峰值的 80%+,说明合并访问有问题"
- "copy 是纯内存搬运,它的性能天花板就是 GPU 内存带宽"
[2min] 常见错误演示
错误 1:忘记 grid-stride loop(30s)
// 错误——只处理了前 blockDim.x 个元素
int idx = threadIdx.x;
dst[idx] = src[idx]; // 超过 blockDim.x 的部分没被处理!
"这就是为什么需要 grid-stride loop——如果不写循环,每个线程只处理一个元素,你搬不完整个 tensor。"
错误 2:循环写反条件(30s)
// 错误——死循环
for (int i = idx; i < n; i++) { // 忘记 += stride!
"每次循环 i 都只加 1——第一个线程就把所有活干完了,其他线程白等,而且还死循环。"
错误 3:TileLang 里用 T.Serial(1min)
# 错误——用 T.Serial 跑 copy
for i in T.Serial(N // BLOCK_N):
...
"T.Serial 告诉编译器'这些 tile 必须顺序执行'——你的并行性全丢了。copy 的每个 tile 完全独立,应该用 T.Parallel。"
课后挑战
"把 copy kernel 的带宽从 60% 跑到 90% 以上。提示:检查你的 grid/block 尺寸是否合理。下周 vector_add 见!"