""" ============================================================================= CS349E · 第 10 讲 · 01_naive_and_coalescing.py —— 第一个矩阵乘法内核：朴素版与合并访存 ============================================================================= 【这个脚本要回答的问题】 矩阵乘法的"标准答案"（cuBLAS）在 A6000 上能跑到 23 TFLOPS，朴素写法只有 0.3 TFLOPS，差 75 倍。 差距从哪来？第一步只改了一行"线程 → 输出元素"的映射，就快了 6 倍——为什么？ 演示 1：目标：4096³ 的 FP32 矩阵乘法在 A6000 上的 roofline（它本应受限于计算）。 演示 2：朴素内核 vs 合并访存内核：在模拟器上真的算一遍，数每个 warp 碰到几个 cache line。 演示 3：把计数放大到 4096³，代入性能模型，和 Boehm 博客的实测对照。 演示 4：边界：矩阵不是 32 的倍数时为什么必须有边界检查。 【运行】python3 01_naive_and_coalescing.py （约 2 秒） 【依赖】numpy；本目录的 kernel_lab（它导入 ../prereq2 的 simt_lab）；../lecture01 的 console。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from kernel_lab import A6000, KERNELS, PERF_A6000, PERF_H100, launch, per_fma, predict, run_kernel, scale_counts # noqa: E402 from simt_lab import cdiv # noqa: E402 from console import fmt_bytes, fmt_num, fmt_time, print_header, print_observations, print_table # noqa: E402 N_BIG = 4096 N_SIM = 128 def demo1_target(): print_header("演示 1：目标——4096³ 的 FP32 矩阵乘法应该受限于计算") n = N_BIG flops = 2 * n ** 3 min_bytes = 3 * n * n * 4 rows = [] for perf in (PERF_A6000, PERF_H100): peak = perf.peak_flops() ridge = peak / perf.dram_bw rows.append([perf.name, f"{peak / 1e12:.1f} TFLOPS", f"{perf.dram_bw / 1e9:.0f} GB/s", f"{ridge:.0f}", fmt_time(flops / peak), fmt_time(min_bytes / perf.dram_bw)]) print(f"FLOPs = 2·n³ = {fmt_num(flops)}；最少要搬的字节 = 3·n²·4 = {fmt_bytes(min_bytes)}；" f"算术强度上限 = {flops / min_bytes:.0f} FLOP/字节") print_table(["GPU", "FP32 峰值（CUDA core）", "显存带宽", "临界强度 FLOP/字节", "按峰值算的时间", "只搬一遍数据的时间"], rows) cub = 23249.6e9 print(f"\nBoehm 博客里 cuBLAS 的实测：{cub / 1e12:.1f} TFLOPS = {fmt_time(flops / cub)}，" f"是 A6000 FP32 峰值的 {cub / PERF_A6000.peak_flops():.0%}。") print_observations([ f"强度上限 {flops / min_bytes:.0f} FLOP/字节远高于临界强度（A6000 约 50），所以这个矩阵乘法**应该**受限于计算：" "理想的内核把每个数从显存只读一次，然后在片上反复使用。", "下面每一步优化做的都是同一件事：让数据在离计算单元更近的地方（L1 → 共享内存 → 寄存器）被多用几次。", "cuBLAS 也只跑到峰值的六成左右——FP32 的 CUDA core 峰值很难接近（时钟在满载时会下降、还有地址计算等开销）；" "后面所有\"百分之几\"都是相对 cuBLAS 说的。", ]) def demo2_coalescing(): print_header(f"演示 2：朴素 vs 合并访存（模拟器上真的算 {N_SIM}³，并数事件）") rng = np.random.default_rng(0) A = rng.standard_normal((N_SIM, N_SIM)).astype(np.float32) B = rng.standard_normal((N_SIM, N_SIM)).astype(np.float32) ref = A.astype(np.float64) @ B.astype(np.float64) rows = [] stats = {} for spec in KERNELS[:2]: C, st = run_kernel(spec, A, B) stats[spec.name] = st pf = per_fma(st) warps_iters = st.fma / 32 rows.append([spec.name, f"{np.abs(C - ref).max():.1e}", f"{st.gld_instr / warps_iters:.2f}", f"{st.gld_lines / warps_iters:.2f}", f"{st.gld_sectors / warps_iters:.2f}", f"{st.gld_sectors * 32 / st.fma:.1f}", f"{pf['warp_instr'] * 1000:.0f}"]) print_table(["内核", "最大误差", "每个 warp 每步：load 指令", "碰到的 128B line", "碰到的 32B sector", "每次乘加经过 L1 的字节（sector × 32）", "每 1000 次乘加的 warp 指令"], rows) st1, st2 = stats[KERNELS[0].name], stats[KERNELS[1].name] print(f"\n朴素版一个 warp 每步：读 A 碰 32 个 line（32 个线程读 32 个不同的行）+ 读 B 碰 1 个 line（32 个线程同一个地址）；" f"合并版：读 A 1 个 line（同一个地址，广播）+ 读 B 1 个 line（连续 32 个 float = 128 字节）。") print(f"line 数之比：{st1.gld_lines / st2.gld_lines:.1f} 倍。") print_observations([ "两个内核做的乘加、读的数**完全一样**，只是\"哪个线程算哪个元素\"不同：朴素版 threadIdx.x 连续的 32 个线程负责 32 个不同的行，" "读 A 时地址相差一整行（K·4 字节）；合并版让它们负责同一行的连续 32 列。", "朴素版读 A 时每个线程只要 4 字节，却要动用一个 32 字节的 sector，只有 1/8 有用；" "每次乘加经过 L1 的字节从 33 降到 5（读 B 的那条在两个版本里都很便宜：朴素版是 32 个线程读同一个地址，合并版是连续 128 字节）。", f"更要紧的是 L1 要处理的 cache line 数：朴素版每步 {st1.gld_lines / (st1.fma / 32):.0f} 个、合并版 {st2.gld_lines / (st2.fma / 32):.0f} 个" "——前置讲 P2 §5 的合并访存在矩阵乘法里的样子。", ]) return stats def demo3_model(stats): print_header("演示 3：放大到 4096³，代入性能模型，和实测对照") rows = [] for spec in KERNELS[:2]: st = stats[spec.name] big = scale_counts(st, N_BIG) p = predict(spec, big, N_BIG, PERF_A6000, A6000) rows.append([spec.name, fmt_time(p["t_fma"]), fmt_time(p["t_issue"]), fmt_time(p["t_l1"]), fmt_time(p["t_dram"]), p["bound"], f"{p['gflops']:.0f}", f"{spec.boehm:.0f}", f"{spec.boehm / p['gflops']:.2f}"]) print_table(["内核", "FP32 运算", "指令发射", "L1（每周期 1 个 line）", "显存", "瓶颈", "模型 GFLOP/s", "Boehm 实测 GFLOP/s", "实测 / 模型"], rows) print_observations([ "朴素版：模型说瓶颈是 L1——每 32 次乘加要处理 33 个 cache line，L1 每周期只能处理 1 个；" "按这个算出约 290 GFLOP/s，和实测的 309 很接近（Boehm 用 profiler 看到的显存吞吐只有 15 GB/s，" "说明慢的不是显存，而是 L1 里一个个 line 地排队）。", "合并版：L1 的压力降了约 16 倍，模型的上限跳到约 4.8 TFLOPS，可实测只有 2.0 TFLOPS（模型的四成）。" "我们的模型只数吞吐，没算延迟、没数地址计算指令；这一步之后的瓶颈已经不在 L1，而在\"每次乘加都要发两条 load\"本身" "（下一个脚本把数据先搬进共享内存）。", "模型给的是上限：实测接近模型，说明模型抓到了瓶颈；实测远低于模型，说明还有模型没看见的东西——两种情况都有用。", ]) def _naive_no_check(blk, A, B, C, M, N, K): """朴素内核去掉边界检查：线程号超出矩阵的线程也去读写。""" x = blk.bx * blk.bdx + blk.tx y = blk.by * blk.bdy + blk.ty acc = np.zeros(blk.n_threads, dtype=np.float32) for i in range(K): acc += blk.load(A, x * K + i) * blk.load(B, i * N + y) blk.store(C, x * N + y, acc) def demo4_bounds(): print_header("演示 4：边界检查——M = 100、N = 70、K = 50（都不是 32 的倍数）") rng = np.random.default_rng(1) M, N, K = 100, 70, 50 A = rng.standard_normal((M, K)).astype(np.float32) B = rng.standard_normal((K, N)).astype(np.float32) rows = [] for spec in KERNELS[:2]: C, st = run_kernel(spec, A, B) rows.append([spec.name + "（有边界检查）", f"{np.abs(C - A @ B).max():.1e}", f"{st.simt_efficiency():.1%}", "—"]) C = np.zeros((M, N), dtype=np.float32) try: launch(_naive_no_check, (cdiv(M, 32), cdiv(N, 32)), (32, 32), args=(A, B, C, M, N, K)) msg = "没有报错" except IndexError as e: msg = str(e).split("：")[0] rows.append(["朴素版（去掉边界检查）", "—", "—", msg]) print_table(["内核", "最大误差", "SIMT 效率（活跃线程比例）", "模拟器报告"], rows) print_observations([ f"网格按 ⌈100/32⌉ × ⌈70/32⌉ = {cdiv(M, 32)} × {cdiv(N, 32)} 个块启动，总线程数比元素多；多出来的线程必须什么都不做。", "去掉边界检查后，第一个越界的线程就会读到 A 之外的地址。模拟器直接报错；真实 GPU 上可能悄悄读到别的数组的数据、" "算出错误结果，也可能非法访问后崩溃——这是 KernelBench 里最常见的一类错误（第 14 讲）。", "SIMT 效率低于 100%：边缘那些块里有一部分线程被 mask 掉、在空转。矩阵越小、块越大，这部分浪费越明显" "（cuBLAS 为不同形状准备了不同的分块，第 13 讲的自动调优就是在选这个）。", "后面的 k3–k6 和 Boehm 的原版一样假设形状是块的整数倍；生产代码要么补边界检查，要么把输入补齐（padding）。", ]) if __name__ == "__main__": demo1_target() st = demo2_coalescing() demo3_model(st) demo4_bounds()