""" ============================================================================= CS349E · 第 9 讲 · 01_loop_nests.py —— 一切从循环嵌套开始：分块、循环顺序与"驻留" ============================================================================= 【这个脚本要回答的问题】 卷积、全连接、Transformer 的线性层，都是同一种循环嵌套。加速器之间的区别，可以全部写成对这个循环嵌套的变换。 演示 1：几种层写成矩阵乘法 C[M,N] += A[M,K]·B[K,N] 以后的形状、乘加次数，以及"每个数最多能被用几次"。 演示 2：循环顺序决定谁"驻留"：两级存储、一个缓冲区放每个张量的一个块，换一下循环顺序，三个张量各被装进来几次。 演示 3：分块大小决定 DRAM 流量：同一个循环顺序，块越大复用越多，但要放得进片上缓冲区。 演示 4：解析公式（驻留规则）和逐次迭代的暴力计数逐个比较。 【运行】python3 01_loop_nests.py （约 1 秒） 【依赖】numpy；本目录的 loopnest；../lecture01 的 console（通过 loopnest 不导入，这里自己加路径）。 """ 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 loopnest import DIMS, TENSORS, Mapping, access_counts, brute_force_fills, versions # noqa: E402 from console import fmt_bytes, fmt_num, print_header, print_observations, print_table # noqa: E402 # ----------------------------------------------------------------------------- def demo1_layers(): print_header("演示 1：几种层都是矩阵乘法 C[M,N] += A[M,K]·B[K,N]") layers = [ ("AlexNet CONV3（批 16，im2col）", 16 * 13 * 13, 384, 256 * 3 * 3, "M = 批 × 输出 13×13，N = 输出通道 384，K = 输入通道 256 × 3×3"), ("全连接层，批 1", 1, 4096, 4096, "M = 批大小"), ("全连接层，批 16", 16, 4096, 4096, ""), ("Llama-3.1-8B 的 W_up，prefill 2048 个 token", 2048, 14336, 4096, "M = token 数"), ("Llama-3.1-8B 的 W_up，decode 批 16", 16, 14336, 4096, "每个序列 1 个 token"), ] rows = [] for name, M, N, K, note in layers: macs = M * N * K ideal = M * K + K * N + M * N rows.append([name, M, N, K, fmt_num(macs), N, M, f"{macs / ideal:.1f}", note]) print_table(["层", "M", "N", "K", "乘加次数", "A 的每个数最多用几次", "B（权重）最多用几次", "乘加 / 最少要搬的数", "说明"], rows) print_observations([ "卷积通过 im2col（把每个输出位置要看的那一小块输入展开成一行）变成矩阵乘法：Interstellar 的 7 层循环" "（批、输出通道、输入通道、输出高宽、卷积核高宽）可以归到 M、N、K 三组。", "\"乘加 / 最少要搬的数\"是理想情况下每搬一个数做几次乘加（只搬一次每个输入、写一次每个输出）。" "卷积和 prefill 有上百到上千次的复用；批 1 的全连接层只有约 1 次——权重每个数只被用一次，怎么排都省不掉从 DRAM 读权重。", "这就是第 1 讲的 roofline 在循环层面的样子：复用的上限由层的形状决定，能不能达到这个上限由分块和循环顺序决定（演示 2–3）。", ]) # ----------------------------------------------------------------------------- def demo2_order(): print_header("演示 2：循环顺序决定谁\"驻留\"（M = N = K = 8，切成 4 的块：外层循环每维 2 次，缓冲区每个张量放一个 4×4 的块）") rows = [] for order in (("M", "N", "K"), ("M", "K", "N"), ("K", "N", "M")): loops = [(d, 2) for d in order] cells = [" → ".join(order)] for t in TENSORS: cells.append(versions(loops, t)) stat = [t for t in TENSORS if versions(loops, t) == 4] cells.append({"C": "输出驻留（C 一直在算同一块）", "A": "A 驻留", "B": "B（权重）驻留"}[stat[0]] if stat else "-") rows.append(cells) print_table(["外层循环顺序（外 → 内）", "A 的块装几次", "B 的块装几次", "C 的块装几次", "最内层循环和谁无关"], rows) print_observations([ "一共 2×2×2 = 8 次外层迭代。最内层的循环和哪个张量无关，那个张量的块在这一圈里不变，不用重新装——它\"驻留\"在缓冲区里，" "只装 4 次（= 和它有关的两个循环的上界之积），其余两个张量各装 8 次。", "M→N→K：K 在最内层，C[M,N] 和 K 无关 → 输出驻留（output stationary），部分和一直在缓冲区里累加，只在算完时写回一次。" "K→N→M：M 在最内层，B[K,N] 和 M 无关 → 权重驻留（weight stationary），TPU 的脉动阵列就是这样（第 8 讲）。", "所以\"weight / output / input stationary\"这些名字，本质上只是循环顺序的不同。", ]) # ----------------------------------------------------------------------------- def demo3_blocking(): print_header("演示 3：分块大小决定 DRAM 流量（M = 256、N = 512、K = 512，循环顺序 M → N → K，没有 PE 阵列）") M, N, K = 256, 512, 512 compulsory = M * K + K * N + M * N rows = [] for tm, tn, tk in ((8, 8, 8), (32, 32, 32), (64, 64, 64), (64, 128, 128), (128, 128, 64), (256, 512, 64)): m = Mapping({"M": 1, "N": 1, "K": 1}, {"M": tm, "N": tn, "K": tk}, {"M": M // tm, "N": N // tn, "K": K // tk}, (), (), DIMS, ("M", "N", "K")) c = access_counts(m) dram = sum(c[t]["dram_read"] + c[t]["dram_write"] for t in TENSORS) buf = (tm * tk + tk * tn + tm * tn) * 2 rows.append([f"{tm}×{tn}×{tk}", fmt_bytes(buf, binary=True), fmt_num(dram), f"{dram / compulsory:.1f}×", f"{c['MACs'] / dram:.0f}"]) print(f"最少要搬 {fmt_num(compulsory)} 个数（每个输入读一次、每个输出写一次）") print_table(["GB 里的块 M×N×K", "三个块一共占（16 位）", "DRAM 读写的数", "是最少的几倍", "每搬一个数做几次乘加"], rows) print_observations([ "块越大，同一个数在片上被用的次数越多，DRAM 流量越接近\"每个数只搬一次\"的下限（8×8×8 的块要多搬 32 倍）；" "代价是片上缓冲区要放得下——最后一行做到了下限，但要 352 KiB，放不进 Eyeriss 那样 128 KB 的缓冲区。", "第 8 讲 H100 的数据块（64×64 到 256×256）、TPU 的 VMEM、NKI 的 SBUF 分块，都是在做这件事：在容量限制下把复用做大。", "同样的块大小，换一个循环顺序结果会不同（演示 2）。分块 + 顺序 + 每一级的容量，三件事要一起选——这就是映射（mapping）的搜索空间。", ]) # ----------------------------------------------------------------------------- def demo4_verify(): print_header("演示 4：驻留规则（解析公式）vs 逐次迭代的暴力计数") rng = np.random.default_rng(42) rows = [] agree = 0 for _ in range(8): dram = {d: int(rng.choice([1, 2, 3])) for d in DIMS} gb = {d: int(rng.choice([1, 2, 4])) for d in DIMS} og = tuple(rng.permutation(DIMS)) od = tuple(rng.permutation(DIMS)) m = Mapping({d: 1 for d in DIMS}, gb, dram, (), (), og, od) bf = brute_force_fills(m) loops_gb = [(d, m.dram[d]) for d in od] loops_rf = loops_gb + [(d, m.gb[d]) for d in og] an = {t: (versions(loops_rf, t), versions(loops_gb, t)) for t in TENSORS} ok = all(an[t] == bf[t] for t in TENSORS) agree += ok rows.append([f"DRAM {''.join(od)} {[dram[d] for d in od]}", f"GB {''.join(og)} {[gb[d] for d in og]}", " / ".join(f"{bf[t][0]}" for t in TENSORS), " / ".join(f"{an[t][0]}" for t in TENSORS), "一致" if ok else "不一致"]) print_table(["DRAM 层循环（外→内）与上界", "GB 层循环与上界", "暴力计数：RF 块版本 A/B/C", "公式", ""], rows) print(f"\n8 个随机映射里 {agree} 个一致（loopnest 的自检另外跑了 20 个，含空间展开）。") print_observations([ "暴力计数真的把所有外层循环跑一遍，记下每个张量的块下标什么时候变；公式只看\"最内层的相关循环在哪\"。两者完全一致。", "有了这条规则，访问次数不用模拟就能算出来——Interstellar 和 Timeloop 这类工具能在几秒钟里评估几十万种映射，靠的就是它（03、04 脚本）。", ]) if __name__ == "__main__": demo1_layers() demo2_order() demo3_blocking() demo4_verify()