""" ============================================================================= CS349E · 第 10 讲 · 02_shared_memory_tiling.py —— 共享内存分块：把数据搬到片上再用 32 遍 ============================================================================= 【这个脚本要回答的问题】 合并访存之后，每次乘加仍要从全局内存（经过 L1）读两个数。把 A、B 的小块先搬进共享内存， 块里的线程就能反复使用它们。全局访问少了 32 倍——为什么实测只快了 1.5 倍？ 演示 1：kernel 3 的事件计数：全局访问降了多少，共享内存访问是多少，模型说瓶颈在哪。 演示 2：块大小 BS = 8、16、32：全局流量 ∝ 1/BS，可共享内存的访问次数不变——"一个线程只算一个输出"的天花板。 演示 3：忘了 __syncthreads() 会怎样（模拟器的竞争检测）。 演示 4：占用率：1024 个线程的块为什么每个 SM 只放得下一个。 【运行】python3 02_shared_memory_tiling.py （约 3 秒） 【依赖】numpy；本目录的 kernel_lab；../lecture01 的 console。 """ import os import sys from dataclasses import replace 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, per_fma, predict, run_kernel, scale_counts # noqa: E402 from simt_lab import occupancy # noqa: E402 from console import fmt_time, print_header, print_observations, print_table # noqa: E402 N_BIG = 4096 N_SIM = 128 K2, K3 = KERNELS[1], KERNELS[2] def _mats(n, seed=0): rng = np.random.default_rng(seed) A = rng.standard_normal((n, n)).astype(np.float32) B = rng.standard_normal((n, n)).astype(np.float32) return A, B, A.astype(np.float64) @ B.astype(np.float64) def k3_spec(bs): """块大小为 bs 的 kernel 3：bs×bs 个线程、两块 bs×bs 的共享内存。""" return replace(K3, name=f"k3（BS = {bs}）", block=(bs * bs,), tile=(bs, bs), smem=2 * bs * bs * 4, multiple=bs) def demo1_counts(): print_header(f"演示 1：kernel 3 的事件计数（模拟器上 {N_SIM}³），和 kernel 2 对比") A, B, ref = _mats(N_SIM) rows, res = [], {} for spec in (K2, K3): C, st = run_kernel(spec, A, B) pf = per_fma(st) big = scale_counts(st, N_BIG) p = predict(spec, big, N_BIG, PERF_A6000, A6000) res[spec.name] = (st, p) rows.append([spec.name, f"{np.abs(C - ref).max():.1e}", f"{pf['gld_instr'] * 1000:.1f}", f"{pf['gld_lines'] * 1000:.2f}", f"{pf['lds_instr'] * 1000:.1f}", f"{pf['sts_instr'] * 1000:.2f}", f"{pf['lds_wavefronts'] * 1000:.1f}", f"{st.barriers // st.blocks}"]) print_table(["内核", "最大误差", "全局 load 指令 / 1000 乘加", "L1 line / 1000 乘加", "LDS 指令 / 1000 乘加", "STS 指令 / 1000 乘加", "共享内存轮数 / 1000 乘加", "每块 sync 次数"], rows) print() rows = [] for spec in (K2, K3): st, p = res[spec.name] rows.append([spec.name, fmt_time(p["t_issue"]), fmt_time(p["t_l1"]), fmt_time(p["t_smem"]), p["bound"], f"{p['gflops']:.0f}", f"{spec.boehm:.0f}", f"{spec.boehm / p['gflops']:.2f}"]) print_table(["内核（4096³，A6000）", "指令发射", "L1", "共享内存", "模型瓶颈", "模型 GFLOP/s", "实测", "实测 / 模型"], rows) st2, st3 = res[K2.name][0], res[K3.name][0] print_observations([ f"全局 load 指令从每 1000 次乘加 {per_fma(st2)['gld_instr'] * 1000:.1f} 条降到 {per_fma(st3)['gld_instr'] * 1000:.2f} 条（32 倍）：" "块里 1024 个线程每人搬 A、B 各一个数，搬进来的 32×32 个数每个被 32 个线程用到。", "代价是每次乘加要从共享内存读 2 个数（As 那个是一个 warp 读同一个地址 = 广播，Bs 那个是连续 32 个字，都没有 bank 冲突），" "共享内存每周期只能送 128 字节（一轮），于是瓶颈从 L1 挪到了共享内存：模型的上限约 4.7 TFLOPS，比 kernel 2 的模型上限还略低。", "实测从 2.0 涨到 3.0 TFLOPS（模型的约六成）。Boehm 的 profiler 显示最多的停顿是 \"Stall MIO Throttle\"——" "共享内存指令的队列满了，和模型说的瓶颈一致；它还说占用率（66%）不是主要问题（演示 4）。", ]) def demo2_block_size(): print_header(f"演示 2：块大小 BS = 8、16、32（{N_SIM}³ 上计数，4096³ 上估计）") A, B, ref = _mats(N_SIM) rows = [] for bs in (8, 16, 32): spec = k3_spec(bs) C, st = run_kernel(spec, A, B, BS=bs) pf = per_fma(st) p = predict(spec, scale_counts(st, N_BIG), N_BIG, PERF_A6000, A6000) rows.append([f"BS = {bs}（{bs * bs} 线程）", f"{np.abs(C - ref).max():.1e}", f"{pf['gld_lines'] * 1000:.2f}", f"{pf['lds_wavefronts'] * 1000:.1f}", f"{pf['barriers'] * 1000:.2f}", fmt_time(p["t_l1"]), fmt_time(p["t_smem"]), fmt_time(p["t_dram"]), p["bound"], f"{p['gflops']:.0f}"]) print_table(["块大小", "最大误差", "L1 line / 1000 乘加", "共享内存轮数 / 1000 乘加", "sync / 1000 乘加", "L1", "共享内存", "显存", "模型瓶颈", "模型 GFLOP/s"], rows) print_observations([ "块越大，从全局内存搬的越少（L1 的 line 数 ∝ 1/BS，显存流量也随块变大而下降）——这是第 9 讲\"分块决定外层流量\"的 GPU 版本。", "但每次乘加从共享内存读 2 个数这件事和 BS 无关：共享内存那一列三种块大小一模一样。" "只要\"一个线程只算 C 的一个元素\"，每次乘加就要从共享内存取两个操作数，上限就被共享内存的 128 字节/周期卡住。", "要突破它，必须让**寄存器**也复用数据：一个线程算多个输出，从共享内存读进寄存器的一个数被用好几次（下一个脚本）。", ]) def demo3_races(): print_header("演示 3：去掉一次 __syncthreads() 会怎样") A, B, ref = _mats(64, seed=3) rows = [] for drop, what in ((None, "两次都有"), ("first", "去掉第一次（装完块之后）"), ("second", "去掉第二次（用完块之后）")): spec = k3_spec(32) C, st = run_kernel(spec, A, B, BS=32, drop_sync=drop) first = st.races[0].split("：", 1)[1] if st.races else "—" rows.append([what, f"{np.abs(C - ref).max():.1e}", st.race_count, first[:70]]) print_table(["写法", "最大误差（模拟器里）", "检测到的竞争次数", "第一条报告"], rows) print_observations([ "模拟器让所有线程严格同步地一条一条执行，所以**算出来的数居然是对的**——真实 GPU 上 warp 之间不同步，" "去掉 sync 的程序可能这次对、下次错，取决于调度。这正是这类 bug 难查的原因：测试通过不代表没有竞争。", "去掉第一次：有的 warp 还没把数据写进共享内存，别的 warp 就开始读（写后读，RAW）。", "去掉第二次：有的 warp 还在读这一轮的块，跑得快的 warp 已经开始写下一轮的数据（读后写，WAR）。", "所以模拟器不看结果，而是记录\"谁在哪个同步阶段读写过哪个位置\"来报告竞争（前置讲 P2 §6.2）；" "真实 GPU 上对应的工具是 compute-sanitizer 的 racecheck。", ]) def demo4_occupancy(): print_header("演示 4：占用率——每个 SM 能同时放几个块（A6000：每 SM 最多 48 个 warp、16 个块、64K 个寄存器、100 KB 共享内存）") rows = [] for name, threads, regs, smem in (("k3，BS = 32", 1024, 40, 8192), ("k3，BS = 16", 256, 40, 2048), ("k3，BS = 32，每线程 64 个寄存器", 1024, 64, 8192), ("k3，BS = 32，每线程 72 个寄存器", 1024, 72, 8192)): o = occupancy(threads, regs, smem, A6000) rows.append([name, threads, regs, smem, o["blocks_per_sm"], o["active_warps"], f"{o['occupancy']:.0%}", o["limiter"], " / ".join(f"{k} {v}" for k, v in o["limits"].items())]) print_table(["配置", "每块线程", "每线程寄存器（假设）", "每块共享内存（字节）", "每 SM 块数", "活跃 warp", "占用率", "卡在", "各项限制下的块数"], rows) print_observations([ "1024 个线程 = 32 个 warp，A6000 每个 SM 最多 48 个 warp，所以只放得下 1 个块：占用率 32/48 = 67%，和 Boehm 的 profiler 一致。", "BS = 16 的块小，一个 SM 能放 6 个，占用率 100%，可演示 2 的模型说它并不更快：瓶颈是共享内存的吞吐，不是延迟；" "占用率只在\"在途的请求不够多、延迟藏不住\"时才是问题（前置讲 P2 §7.2）。", "1024 线程的块，每线程 64 个寄存器时 64 × 1024 = 65536，正好占满寄存器堆；到 72 个就一个块也放不下、内核根本启动不了" "（真实 CUDA 会报 too many resources requested for launch）。下一个脚本里每个线程要存 64 个累加器，所以块要变小（256 个线程）。", ]) if __name__ == "__main__": demo1_counts() demo2_block_size() demo3_races() demo4_occupancy()