TPU v1 完全拆解（二）：脉动阵列、weight-stationary 数据流与 Roofline 性能归因

数据移动比计算更贵#

脉动阵列不是 TPU 的发明，它 1978 年就由卡内基梅隆大学的 H. T. Kung 和 C. E. Leiserson 提出（Kung 在 1982 年的《Why Systolic Architectures?》中做了系统总结）。要理解它为什么出现，先要理解当时的困境，而这个困境到现在依然是 AI 芯片设计的核心约束。

传统处理器（冯·诺依曼结构）执行一条乘加指令的流程是：取指令 → 从寄存器/内存取操作数 → 计算 → 写回结果。对矩阵乘来说，这个流程有个致命问题：数据搬运的时间和能量远大于计算本身。TPU 论文里引用了一个关键事实：读一次大型 SRAM 消耗的能量，比做一次算术运算还要多（论文引用 [Kun80][Ram91][Ovt15b] 三个出处支持这个论断）。具体量级可以参考后来的经典数据（Horowitz 2014 年发表在 ISSCC 的功耗表）：一次 8 位整型乘法约 0.2 pJ，而一次 32 KB SRAM 读取约 5 pJ，是乘法的 25 倍；DRAM 读取更是 640 pJ，是乘法的 3200 倍。

Kung 的洞察是：如果能减少”把数据搬来搬去”的次数，即使计算单元简单一点，整体能耗和性能也会大幅改善。他注意到很多信号处理应用（FIR 滤波、矩阵乘、卷积）有一个共同特点——数据复用率极高：同一个输入要参与很多次不同的乘加。如果能把数据”留在”计算单元附近反复使用，而不是每次用完就存回内存、下次再取，就可以把数据搬运量缩小一个数量级。

“心脏泵血”式的解决方案#

Kung 给出的方案是：把大量简单的乘加单元（PE，Processing Element）排列成规则网格，让数据像血液一样从阵列边缘流入，流经所有 PE 后从另一侧流出，中间不停顿、不写回内存。这就是”systolic”（脉动的）这个名字的由来——数据在阵列中一波一波地推进，如同心脏泵血。

脉动阵列相对传统结构有三个根本区别：

1. 数据在 PE 之间直接传递，不走内存。每个 PE 只和相邻的 PE 通信（连线短、功耗低、布线简单），计算结果通过相邻 PE 接力传递，只有最终结果才写回。
2. 每个 PE 极其简单。一个 PE 只需要一个乘法器、一个加法器/累加器、几个寄存器，没有指令、没有控制逻辑。整个阵列用同一个控制信号锁步（lockstep）运行，所有 PE 在同一拍做同一件事。SIMD 控制方式下，控制逻辑由整个阵列共享，平均到每个 PE 上的开销几乎可以忽略——对比 MIMD 架构（每个 PE 自带一套指令解码和调度逻辑，开销是数千晶体管量级），省下的面积全部可以换成计算单元。
3. 数据复用发生在 PE 内部。一个数被读进阵列后，在流动过程中被沿途的所有 PE 反复使用，不需要重复访问寄存器堆或内存。

用论文里的说法：TPU 的 MXU 使用脉动执行（systolic execution）来省电，手段是减少对 Unified Buffer 的读写次数——UB 是 24 MiB 的大 SRAM，读写它比做算术贵得多。脉动阵列让一个激活值被读入一次后，在阵列里被 256 个 PE 轮流使用，把”读一次数据做一次运算”变成”读一次数据做几百次运算”。

顺带说一个历史背景：脉动阵列在 20 世纪 80–90 年代曾被 CMU 的 Warp、iWarp 等项目做成真实芯片，但最终没有商业成功——原因不是硬件不行，而是编程太难：通用 CPU 只需要程序员写串行代码，而脉动阵列要求程序员/编译器把算法显式映射成数据流动的时序，当时没有编译器能做好这件事。Kung 本人也在论文里承认这一点。TPU 之所以能让脉动阵列”复活”，恰恰是因为它把软件问题绕开了：推理模型被固定成”权重静态、数据流式”这一种模式，编译器只需要处理矩阵乘一种映射，而矩阵乘到脉动阵列的映射是数学上最规整的。领域专用设计的另一个侧面是：它把”难的部分”从硬件挪到了软件栈，而软件栈只为一类负载服务。

为什么 TPU 选 weight-stationary#

TPU 选 WS 的理由非常直白，结合它的目标负载（推理）可以逐条列出：

1. 推理的权重是只读的。推理不更新权重，权重从 DDR3 读进阵列后永远不需要写回，天然适合”驻留”。
2. 权重复用次数 = 批大小。一层全连接层，权重矩阵是 $N \times N$（8 位 = $N \times N$ 字节），一批 $B$ 个样本的矩阵乘是 $B \times N \times N$ 次乘加。权重驻留时，每个权重字节被复用 $B$ 次，于是”每读取 1 字节权重能换来 $B$ 次运算”。论文 Table 1 给出的操作强度（ops/weight byte）恰好等于各模型的批大小：MLP0 是 200、MLP1 是 168、LSTM0 是 64、LSTM1 是 96——和批大小完全一致，这不是巧合，这正是 WS 数据流的定义性质。
3. 激活值沿行流动，天然被复用 256 次。一个激活值从左侧进入阵列后，会在这一行经过 256 个 PE（对应 256 个不同的权重），相当于一次 UB 读取支撑 256 次乘加。UB 读写的次数被压缩了约 256 倍——这正是论文说的”用脉动执行减少对 UB 的读写来省电”。

对比之下，如果选 OS 模式，部分和驻留在 PE 里，激活值和权重都要流动，对权重带宽的要求就上去了；如果选 IS 模式，权重流动会破坏”权重只读一次”的最大优势。在”静态权重 + 大批量推理”这个特定场景里，WS 是数学上最优的选择。

阶段一：权重预加载（Weight Preload）#

MXU 内部能放 64 KiB 的权重块（tile，恰好是 256×256 个 8 位权重）。执行矩阵乘之前，权重必须先搬进阵列——从 Weight FIFO 逐拍移入，256 字节/拍，正好 256 拍移完一个 tile（256×256 字节 ÷ 256 字节/拍）。这 256 拍的加载时间就是论文里说的”把 tile 移进来需要 256 拍”（the 256 cycles it takes to shift a tile in）。

为了掩盖这 256 拍，MXU 里放了两份权重 tile 做双缓冲：一个 tile 在计算时，另一个 tile 在后台从 FIFO 加载。计算一个 tile 需要 B 拍（见阶段三），只要 $B > 256$，双缓冲就能完全隐藏权重加载时间。论文选 4096 组累加器的推导也与此相关：达到峰值需要的操作强度约 1350，凑整到 2048，再翻倍到 4096，就是为了让编译器在跑峰值时还有空间做双缓冲。

阶段二：波前传播（Wavefront Propagation）#

权重就位后，激活数据开始从左侧流入。但注意：流进阵列的激活不是立刻全部生效的。左上角的 PE 第一拍就拿到了第一个激活和第一个权重，但第 256 列的 PE 要等激活沿着行传播 255 拍才能拿到数据。于是每个”计算前沿”在阵列里呈对角线形状推进——这就是”对角波前”。一个 256 元素乘加操作的波前从左上角到右下角需要约 512 拍（256 拍横移 + 256 拍纵移）的流水填充（pipeline fill）时间，之后阵列进入稳态，每拍吞吐 256 个乘加。

论文强调”控制与数据被流水化，制造出 256 个输入同时读出的假象”——意思是：从软件/编译器视角看，MXU 每拍读入 256 个激活、瞬间产生 256 个部分和，内部的波前传播延迟被完全隐藏了。正确性层面，软件根本感知不到脉动结构的存在；但性能层面，软件要操心阵列的延迟（比如层与层之间的依赖、权重加载的节奏）。

这里值得停下来和 GPU 的执行模型做一次对照，因为两者的本质差异正好解释了后续所有性能数据。GPU 的 SIMT 模型里，一个矩阵乘被拆成几十万个线程，每个线程独立地从寄存器/显存取数、计算、写回——并行度来自大量独立线程，控制分散在每个线程里；而脉动阵列的并行度来自数据在空间上的流动：65536 个 PE 不需要知道自己在算矩阵的哪一块，它们只是接收邻居的数据、做乘加、传给下一个邻居。GPU 靠”很多简单的执行单元 + 复杂调度”赢下通用性，脉动阵列靠”规则数据流 + 几乎为零的调度”赢下面积和能耗。代价是：脉动阵列的延迟不是可调度的，而是几何的——波前要横穿整个阵列（约 512 拍），这条延迟路径无法像 GPU 那样用更多线程来隐藏，只能靠编译器预先排好数据到达的节奏。

阶段三：稳态流水（Steady State）#

每拍，MXU 从左侧读入 256 个激活值，产生 256 个部分和，写入下方 4 MiB 累加器中的一组。B×256 的输出需要 B 拍算完——论文说的”MatrixMultiply 指令耗时 B 个周期”就是这个意思。之所以是 B 而不是 B×256，是因为 256 列的计算是并行的。计算期间权重不需要逐拍读取——它们已经驻留在 PE 里了，后台同时以 256 字节/拍的速率把下一个权重 tile 移入双缓冲，为下一个矩阵乘做准备。

峰值算力的来源此时一目了然——每拍完成 256×256 次乘加，乘上时钟频率：

注意这里的换算：65,536 个 MAC 单元每拍做 65,536 次乘加，每个 MAC 记 2 次运算（乘 + 加），所以 45.9 TeraMAC/s = 91.8 ≈ 92 TOPS。这个”每个 MAC 算 2 次运算”的计数口径在后面 Roofline 分析里还会回来找我们麻烦，先记住。

一个关键问题：为什么 B 拍算完一层，阵列还是会闲着？#

严格说，上述时序成立的前提是权重永远跟得上。权重从 DDR3 到 Weight FIFO 再到 MXU，路径上的带宽是 34 GB/s（DDR3）；而 MXU 消费权重的速率是多少？每个 tile 是 64 KiB，B 拍算完，所以权重的消费速率是：

当 $B = 200$（MLP0 的批大小）时，消费速率约 229 GB/s——是 DDR3 供给带宽（34 GB/s）的 6.7 倍。权重供给严重不足，MXU 必然频繁空转等待。Weight FIFO 只有 4 个 tile 深（4×64 KiB），最多预取 4 个 tile，DRAM 延迟和带宽缺口全靠这 256 KiB 的缓冲来扛，扛不住的部分就变成 stall。这就是”memory-bound”的微观图景，下一节用 Roofline 把它变成精确的数学。

模型定义#

Roofline 模型由 Berkeley 的 Williams、Waterman 和 Patterson 在 2009 年提出（原论文《Roofline: An Insightful Visual Performance Model for Multicore Architectures》），是性能归因的通用工具。它只基于两个硬件参数画图：

• 峰值算力 $P_{peak}$（纵轴天花板，水平线）：硬件每秒最多能做多少次运算；
• 内存带宽 $B_{mem}$（斜线）：硬件每秒最多能从内存取多少字节。

横轴是操作强度（operational intensity）$I$，定义为”每取 1 字节数据能做多少次运算”。实际可达到的性能被两者夹逼：

当 $I \times B_{mem} < P_{peak}$ 时，性能被带宽卡住（斜线区，memory-bound）；反之被算力卡住（水平区，compute-bound）。两条线的交点叫脊点（ridge point）：

脊点左侧的程序提高算力没用，只能提高数据复用率（增大 $I$）或带宽；脊点右侧的程序提高带宽没用，只能提高算力。Roofline 的最大价值是告诉你”优化该往哪个方向使劲”。

为推理定制的 Roofline：操作强度重定义#

TPU 论文对 Roofline 做了一处关键改造：标准 Roofline 的强度分母是”从内存取的所有数据”，而推理场景下权重是只读的、激活和中间结果都存在片上 24 MiB UB 里，数据搬运的大头只有权重。所以论文把操作强度重新定义为：

这个定义下，对一层全连接层，$I = B$（每字节权重被 batch 个样本复用）；对卷积层还要乘上空间复用因子，所以 CNN 的强度远高于 MLP。论文 Table 1 给出的六个模型强度：MLP0 = 200、MLP1 = 168、LSTM0 = 64、LSTM1 = 96、CNN0 = 2888、CNN1 = 1750。

TPU 的 Roofline 图（论文 Figure 5，即下面的图）就是在 log-log 坐标系里画：水平线在 92 TOPS 处，斜线从原点出发，脊点远在右侧 1350 ops/byte 处：

论文原话：“TPU 的 Roofline 有一段很长的斜线部分，这意味着操作强度不高的应用性能受内存带宽限制而非峰值算力限制。六个应用中有五个正好顶着天花板：MLP 和 LSTM 是带宽受限（memory-bound），CNN 是算力受限（compute-bound）”。

用数字验证一遍（顺便验证我们的强度定义是否自洽）。以 MLP0 为例：$I = 200$ MAC/byte，$B_{mem} = 34$ GB/s，则带宽天花板为：

论文实测 MLP0 是 12.3 TOPS——贴着 13.6 TOPS 的天花板跑，剩余差距来自流水线 stall 和激活搬运。CNN0：$I = 2888 > 1350$，落在水平区，天花板就是 92 TOPS，实测 86 TOPS（93% 的峰值）。每一个数字都对得上，这就是 Roofline 作为分析工具的威力——不是定性描述，是精确到一位小数的定量预测。

这里顺带解开一个容易让人困惑的细节：按 $I_{ridge} = P_{peak} / B_{mem} = 92 \text{ TOPS} / 34 \text{ GB/s}$ 算，脊点应该是 2705 ops/byte，为什么论文图里标 1350？因为 92 TOPS 是”每个 MAC 计 2 次运算”的宣传口径，而 Roofline 图里的”运算”按 MAC 计数：46 TeraMAC/s ÷ 34 GB/s ≈ 1353 MAC/byte ≈ 1350。论文在 Table 1 里给的操作强度（MLP0 = 200 等）同样按 MAC 计数。理解这个口径差异很重要——今天各大厂商宣传的 TOPS 数字普遍按”乘+加=2 次运算”计算，比较芯片时稍不留神就会把口径不同的数字直接相除。

三芯片 Roofline 对比：为什么 TPU 赢在”别人够不到天花板”#

论文把三颗芯片的 Roofline 画在同一张 log-log 图上（论文 Figure 8）：

三颗芯片的参数（论文 Table 2）：

K80 的 160 GB/s 是关闭 Boost、开启 SECDED 校验后的有效带宽（标称 240 GB/s）；Haswell 的脊点 13 ops/byte 意味着它的应用只要强度超过 13 就落入算力限制区——而 TPU 的应用强度动辄 200+。换句话说：

• CPU/GPU 的 Roofline 形状”矮而窄”：算力低、带宽高，脊点靠左，但应用实际跑不到自己的屋顶——被延迟约束压着（见下节）；
• TPU 的 Roofline”高而宽”：算力高、带宽低，脊点靠右，MLP/LSTM 这类低强度应用恰好落在长斜线区，虽然达不到峰值，但斜线区的高度本身就比 CPU/GPU 的整个屋顶还高。

这就是论文那句”所有 TPU 的星点都位于其他两条 Roofline 之上”的含义——TPU 哪怕只发挥三成，也够碾压对手的巅峰状态。

延迟约束：CPU/GPU 为什么连自己的屋顶都够不到#

Roofline 给出的是”无约束下的性能上限”，但推理服务还有一个额外的约束——99 百分位响应时间。论文 Table 4 用 MLP0（要求 7 ms 内返回）展示了这个约束的杀伤力：

批大小翻 4 倍，CPU/GPU 吞吐也翻 4 倍——理论上它们都能跑很高，但大批次把延迟顶爆了，应用被迫退回 batch=16，只发挥 37%–42%。TPU 因为单条指令就是一层网络、执行时间确定，敢用 batch=200 还压住延迟在 7.0 ms，发挥 80%。注意 GPU/TPU 的”100%“也并非自己的 Roofline 天花板——论文 Table 4 注明，GPU 和 TPU 的最大 MLP0 吞吐受主机服务器开销限制，它们离屋顶还差着一大截。CPU/GPU 输在两条线上：峰值不够高，且连自己的峰值都用不满。

Table 3：性能计数器的完整账本#

现在回答第二个问题。论文用硬件性能计数器（performance counter）给每个模型记了一本”阵列时间账”，这是整个性能归因的核心数据（论文 Table 3，行百分比为矩阵单元的时间分配或 MAC 利用率）：

第一行”阵列活跃周期”就是那个著名的 28% 平均利用率。按论文的说明，第一、四、五、六行（阵列活跃、权重加载等待、权重移入、非矩阵周期）是对矩阵单元周期活动的完整划分，四行加起来恒等于 100%（以 MLP0 验证：12.7 + 53.9 + 15.9 + 17.5 = 100.0）；其余各行是更细的分解或原因解释，可以重叠。这张表的解读逻辑是：阵列闲着的时间，要么在等权重（stall + shift），要么在处理非矩阵事务（激活、同步），要么在等依赖（RAW）或等输入（PCIe）。下面按三类负载分别讲。

MLP/LSTM：权重带宽卡死的典型 memory-bound#

以 MLP0 为例（占当时部署 TPU 的 61%，它的权重就是平均值的主要贡献者）：阵列只有 12.7% 的时间在算，53.9% 的时间在等权重从 DDR3 搬过来，15.9% 在把新权重块移入阵列——70% 的周期都耗在”喂权重”上。

为什么会等？回到上一节的数值：MLP0 的强度 $I = 200$，带宽天花板 13.6 TOPS，而阵列本身能跑 92 TOPS。阵列的”饭量”是权重组装线的 6.7 倍。具体到一个 64 KiB 权重 tile：

• 从 DDR3 搬一个 tile：$65536 \text{ B} / 34 \text{ GB/s} \approx 1.93\ \mu\text{s}$；
• 阵列算完一个 tile（batch=200）：$200 \text{ 拍} / 700 \text{ MHz} \approx 0.29\ \mu\text{s}$。

搬进来要 1.93 µs，算掉只要 0.29 µs，搬运时间是计算时间的 6.7 倍。Weight FIFO 的 4 个 tile 预取能掩盖部分 DRAM 延迟，但掩盖不了带宽缺口本身——FIFO 很快被抽干，MXU 只能停下来等。LSTM 更极端：强度只有 64–96，权重等待占比高达 58%–62%，实测性能只有 2.8–3.7 TOPS。

CNN0：唯一贴着峰值跑的应用#

CNN0（源自 Inception V2）的强度高达 2888，远超脊点 1350，落在 Roofline 水平区。它的权重加载等待是 0%——因为卷积权重的空间复用让权重消费速度远低于供给速度，权重管线完全够用。阵列 78.2% 的周期在干活，实测 86 TOPS，达到峰值的 93%。CNN0 证明了 MXU 这套设计本身没有毛病：只要喂得饱权重，阵列是能贴着 92 TOPS 跑的。瓶颈不在阵列，在供给端。

CNN1（AlphaGo）：一个值得单独拆的案例#

CNN1 是 DeepMind 的 AlphaGo 网络，强度 1750，理论上应该也是 compute-bound，但实测只有 14.1 TOPS、活跃周期 46.2%、有效 MAC 只有 22.5%（即 65,536 个 MAC 里平均只有约 1.5 万个在干有用的活）。论文归因了三个原因：

1. 4 个全连接层的强度只有 32：AlphaGo 除了卷积还有 4 个 FC 层，这些层没有空间复用、只有 batch 复用，强度退化为批大小 32，远低于脊点 1350——这几层是硬生生的 memory-bound。论文指出 35% 的周期在等这 4 层的权重。
2. 浅特征深度造成”二维内部碎片”：部分卷积层的输入通道数很小（比如第一层只有 3 个 RGB 通道），展开后填充不满 256×256 的阵列，平均只有一半的 MAC 装着有效权重（有效 MAC 22.5% = 活跃周期 46.2% × 平均约一半的有效载荷）。这正是”稀疏”最原始的形态——不是权重稀疏，而是形状不匹配的浪费，TPU 因为时间紧砍掉了稀疏支持，对这类浪费毫无办法。
3. 23% 的 RAW 停顿：层与层之间存在数据依赖（下一层的输入要等上一层算完），脉动阵列 512 拍的流水填充延迟在这种依赖链上无法隐藏，转变成显式的等待。

CNN1 的意义在于：它把”memory-bound”和”compute-bound”的二分法打破了——一个应用可以同时被权重带宽（FC 层）、形状碎片（浅卷积）和依赖延迟（层间 RAW）三座大山压着。这也解释了为什么论文说”我们的计数器无法精确解释阵列空闲的每一拍”：重叠执行的流水线上，多种停顿来源混在一起，只能看到部分账目。

负载构成：平均 28% 是被谁拉低的#

最后把账算圆：六个模型里，MLP0 占部署 TPU 的 61%、LSTM0 占 29%（论文 Table 1），两者加起来 90% 的芯片在跑 memory-bound 负载，利用率 8%–13%；CNN 合计只占 5% 的部署量。所以”28% 的平均利用率”根本不是 MXU 设计的失败，而是负载结构决定的——90% 的 TPU 在跑权重带宽喂不饱的模型。这正是论文用 Roofline + 计数器把每一分性能损失定位到具体原因的原因：不是”芯片不够快”，而是”权重带宽不够宽”。