""" ============================================================================= CS349E · 第 8 讲 · 01_systolic_array.py —— 脉动阵列：TPU 的矩阵单元怎样工作 ============================================================================= 【这个脚本要回答的问题】 TPU v1 的核心是一个 256×256 的乘加阵列（65,536 个 8 位乘加单元），每个周期做 65,536 次乘加。 它怎样把数据送进这么多运算单元，而不被读写存储拖垮？ 演示 1：一个 3×3 的小阵列逐周期跑一遍：数据怎样"斜着"流进去、部分和怎样往下流；输出和 X·W 比较。 演示 2：利用率和批大小：为什么阵列需要"一大批"输入才能忙起来（填充和排空的开销）。 演示 3：权重从哪来：大矩阵切块、每块的权重从 DRAM 读进来——TPU 论文 §7 的 600×600 例子，以及 roofline 拐点从哪来。 演示 4：为什么值得：每次乘加要从缓冲区读几个数、花多少能量（Horowitz 的能量表）。 【运行】python3 01_systolic_array.py （约 2 秒） 【依赖】numpy；本目录的 hw_engines；../lecture01 的 console、hw_specs。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) from hw_engines import SystolicArray, systolic_cycles, tiled_matmul_time # noqa: E402 from console import fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import energy_pj # noqa: E402 # ----------------------------------------------------------------------------- def demo1_wavefront(): print_header("演示 1：3×3 的权重驻留阵列，4 个输入向量（B = 4）——逐周期看哪些 PE 在做乘加") rng = np.random.default_rng(0) X = rng.integers(-3, 4, (4, 3)).astype(float) W = rng.integers(-3, 4, (3, 3)).astype(float) Y, st, frames = SystolicArray(3, 3).run(X, W, trace=True) print("每个小方格是一个 PE（行 i 存 W 的第 i 行）；数字 = 这个周期它在处理第几个输入向量，'.' = 空闲") lines = ["", "", ""] for t, f in enumerate(frames): for i in range(3): cells = [] for j in range(3): b = t - i - j cells.append(str(b) if f[i, j] else ".") lines[i] += " " + "".join(cells) + " " print("周期 " + "".join(f" {t:<3}" for t in range(len(frames)))) for i in range(3): print(f"行 {i} " + lines[i]) print(f"\n输出和 X·W 的最大差：{np.max(np.abs(Y - X @ W)):.1e}；共 {st['cycles']} 个周期" f"（公式 B + R + C − 1 = {systolic_cycles(4, 3, 3)}），有用的乘加 {st['useful_macs']} 次，利用率 {st['utilization']:.0%}") print_observations([ "第 b 个输入向量的第 i 个数在第 b + i 个周期从左边进入第 i 行，然后每个周期向右走一格；部分和每个周期向下走一格。" "同一个输入向量的计算在阵列里形成一条斜着推进的\"波前\"（论文图 4：diagonal wavefront）。", "每个 PE 只和左边、上边的邻居交换数据，没有一个 PE 去读大的存储——这就是\"脉动\"（systolic，像心脏一下一下地泵血）。", "开头几个周期阵列还没填满、最后几个周期在排空，所以 4 个输入时利用率只有 44%。", ]) # ----------------------------------------------------------------------------- def demo2_utilization(): print_header("演示 2：利用率 = B / (B + R + C − 1)——阵列越大，越需要大批") _, st = SystolicArray(32, 32).run(np.ones((100, 32)), np.ones((32, 32))) print(f"（先用逐周期模拟验证公式：32×32 阵列、B = 100：模拟 {st['cycles']} 个周期、利用率 {st['utilization']:.4f}；" f"公式 {100 / systolic_cycles(100, 32, 32):.4f}）") rows = [] for B in (1, 8, 64, 256, 1024, 8192): rows.append([B] + [f"{B / systolic_cycles(B, n, n):.1%}" for n in (32, 128, 256, 512)]) print_table(["批大小 B", "32×32", "128×128", "256×256（TPU v1）", "512×512"], rows) print_observations([ "阵列要先\"填满\"（R + C 个周期）才能每周期做满 R·C 次乘加。只送 1 个输入向量时，256×256 的阵列利用率不到 0.2%。", "这和 decode 的困境是同一件事：每一步只有很少的 token，再大的矩阵引擎也忙不起来（第 1 讲 §4.2）。" "TPU 论文表 4：在 7 ms 的响应时间限制下，TPU 的 MLP0 能用 200 的批，CPU / GPU 只能用 16。", "阵列越大，需要的批越大。TPU v1 到 v5e 都是 128×128 或 256×256；更大不一定更好（演示 3）。", ]) # ----------------------------------------------------------------------------- def demo3_weight_loading(): print_header("演示 3：权重从 DRAM 来——切块、补零、读权重的时间（TPU v1：700 MHz、权重存储 DDR3 34 GB/s）") rows = [] for n in (256, 512): r = tiled_matmul_time(1, 600, 600, n, n) rows.append([f"{n}×{n}", r["tiles"], f"{r['padded_weight_bytes'] / 1024:.0f} KiB", fmt_time(r["load_time"]), "18 µs" if n == 256 else "32 µs"]) print("TPU 论文 §7：LSTM1 有一个 600×600 的权重矩阵。") print_table(["阵列", "块数", "补零后要读的权重", "读权重的时间", "论文的数字"], rows) print("\n一个 4096×4096 的全连接层（int8 权重 16 MiB），不同批大小，256×256 阵列，权重双缓冲（算这块时读下一块）：") rows2 = [] for B in (16, 64, 200, 1024, 1350, 4096): r = tiled_matmul_time(B, 4096, 4096) rows2.append([B, fmt_time(r["time"]), f"{r['utilization']:.1%}", "读权重" if r["load_time"] > r["compute_time"] else "计算"]) print_table(["批大小 B", "时间", "阵列利用率", "瓶颈"], rows2) print_observations([ "600×600 在 256×256 的阵列上要切成 3×3 = 9 块（补零到 768×768），在 512×512 上是 2×2 = 4 块（补零到 1024×1024）——" "大阵列补的零更多，要从 DRAM 读的字节反而多。读权重的时间 17.3 µs vs 30.8 µs，正好是论文说的 18 µs 和 32 µs。" "论文因此说把阵列从 256 扩到 512 对所有应用都没好处（图 11），这是二维的\"内部碎片\"。", "读一块权重（64 KiB）要 1.93 µs，而 B 个输入流过去要 B / 700 MHz：B 小于约 1350 时读权重比计算慢，阵列在等数据。" "这个 1350 就是 TPU roofline 的拐点（每读 1 字节权重要做 1350 次乘加，02 脚本）。", "论文表 3：MLP 和 LSTM 有 44%–62% 的周期在等权重（weight stall），它们的批大小（64–200）远小于 1350。", ]) # ----------------------------------------------------------------------------- def demo4_energy(): print_header("演示 4：为什么值得——每次乘加要读几个数、花多少能量（Horowitz 45 nm 能量表，int8）") mac = energy_pj("int8_mult") + energy_pj("int32_add") sram_byte = energy_pj("sram_32KB_read64") / 8 big_sram_byte = energy_pj("sram_1MB_read64") / 8 dram_byte = energy_pj("dram_read64_low") / 8 print(f"一次 int8 乘 + int32 加 ≈ {mac:.2f} pJ；从 32 KB 的 SRAM 读 1 字节 ≈ {sram_byte:.1f} pJ，1 MB 的 SRAM ≈ {big_sram_byte:.1f} pJ，" f"DRAM ≈ {dram_byte:.0f} pJ（都按 64 位一次读算每字节）") rows = [] R = C = 256 for name, B in (("没有复用：每次乘加读 1 个输入 + 1 个权重（int8）", None), ("256×256 脉动阵列，B = 16", 16), ("256×256 脉动阵列，B = 200", 200), ("256×256 脉动阵列，B = 2048", 2048)): if B is None: ub, wt = 1.0, 1.0 else: ub = (B * R) / (B * R * C) # 每个输入被一行的 C 个 PE 用 wt = (R * C) / (B * R * C) # 每个权重被 B 个输入用 e = mac + ub * big_sram_byte + wt * dram_byte rows.append([name, f"{ub:.4f}", f"{wt:.4f}", f"{e:.2f} pJ", f"{mac / e:.0%}"]) print_table(["做法", "每次乘加读缓冲区（字节）", "每次乘加读 DRAM 权重（字节）", "每次乘加的能量", "其中运算本身"], rows) print("（输入从 24 MiB 的统一缓冲区读，按\"1 MB SRAM\"的能量算；部分和在 PE 之间传、累加器的读写都没算——偏乐观。）") print_observations([ "没有复用时，每次乘加读 2 个字节，能量几乎全花在读数据上（运算只占约 0.2%）：这是第 1 讲 §2 的\"运算几乎免费、搬运昂贵\"。", "脉动阵列里，每个输入被一行的 256 个 PE 用、每个权重被 B 个输入用：读缓冲区的次数降到 1/256，读 DRAM 的次数降到 1/B。" "B = 2048 时运算本身已经占到大头——矩阵引擎的意义就是让每个搬进来的字节被用很多次。", "B 小时 DRAM 读权重仍然主导能量和时间（演示 3）：这就是推理（尤其是 decode）\"受限于内存\"的物理原因。", ]) if __name__ == "__main__": demo1_wavefront() demo2_utilization() demo3_weight_loading() demo4_energy()