第四节：脉动阵列架构（二）：经典Systolic Array设计（Google TPU v1）

好，咱们接着聊脉动阵列。上一节我们把脉动阵列的基本概念和原理讲清楚了，这一节我带你看看一个真正在工业界大放异彩的经典案例——Google TPU v1 的脉动阵列设计。

说实话，TPU v1 刚出来那会儿，我第一反应是：这玩意儿也太「朴实无华」了吧？没有复杂的乱序执行，没有超标量，就是一个大号的矩阵乘法器。但后来仔细一琢磨，才明白什么叫「大道至简」。

TPU v1 的整体架构概览

TPU v1 是 Google 在 2015 年部署的专用加速器，专门用于推理任务。它的核心就是一个 256×256 的脉动阵列。你想想看，256×256 意味着什么？意味着一次可以完成 65536 个 MAC 操作。

我个人习惯把 TPU v1 的架构分成三个关键部分：

• 脉动阵列核心：256×256 的 MAC 单元矩阵
• 权重缓冲区：从片外 HBM 读取权重，然后「喂」给脉动阵列
• 累加器：处理部分和，完成最终的矩阵乘法结果

嗯，这里要注意，TPU v1 的脉动阵列是权重固定（Weight Stationary）的数据流方式。什么意思呢？就是权重数据待在 MAC 单元里不动，输入数据在阵列里流动，部分和也在流动。

数据复用与带宽优化

做芯片设计的人都知道，带宽永远是瓶颈。我曾经在一个项目里，计算单元利用率只有 30%，为什么？因为数据喂不进去。TPU v1 是怎么解决这个问题的？靠的就是数据复用。

咱们来算一笔账：

数据复用类型	复用次数	带宽节省比例
权重复用	256 次（一个权重被 256 个输入复用）	256x
输入复用	256 次（一个输入被 256 个权重复用）	256x
部分和复用	256 次（在累加过程中）	256x

说白了，TPU v1 的脉动阵列通过三种数据复用方式，把片外带宽需求降低了三个数量级。这就是为什么它可以用相对较低的带宽（DDR3-2133，约 34 GB/s）支撑起 92 TOPS 的算力。

权重固定数据流详解

咱们具体看看权重固定是怎么工作的。假设我们有一个 3×3 的脉动阵列：

// 伪代码描述权重固定数据流
// 权重预先加载到每个 PE 中
for (int i = 0; i < 256; i++) {
    for (int j = 0; j < 256; j++) {
        // 输入数据从左侧流入，从上到下流动
        PE[i][j].weight = W[i][j];  // 权重固定
        PE[i][j].input = I[i][j];   // 输入流动
        PE[i][j].partial_sum += PE[i][j].weight * PE[i][j].input;
    }
}

你可能会问：为什么选择权重固定？我记得在 2018 年做的一个推理芯片项目里，我们对比过几种数据流方式。对于推理场景，权重是固定的（模型训练好后就不变了），所以把权重留在 PE 里是最合理的——加载一次，用到底。

脉动阵列的时序设计

嗯，这里有个容易踩坑的地方——时序。脉动阵列的时序设计比想象中要复杂得多。

我曾经在一个项目里，把脉动阵列的时钟频率设得太高，结果数据在阵列里传播时，不同路径的延迟差异导致数据到达时间不一致。说白了就是「数据打架」了。

TPU v1 是怎么解决这个问题的？它采用了流水线寄存器的设计：

• 每个 PE 之间都有寄存器
• 数据在每个时钟周期只移动一个 PE
• 所有 PE 同步工作

这样做的好处是：时序收敛变得非常简单。你不需要担心长路径的延迟问题，因为每个 PE 之间的路径长度是固定的、可控的。

带宽优化的具体策略

TPU v1 的带宽优化不仅仅是靠数据复用，还有几个关键设计：

1. 双缓冲机制：权重缓冲区采用双缓冲设计，一个缓冲区在计算，另一个在加载下一批权重。这样计算和加载可以重叠。
2. 数据广播：输入数据通过广播方式同时送到阵列的每一行，而不是逐行传输。
3. 部分和累加：部分和在阵列内部累加，只有最终结果才写回片外存储器。

你想想看，如果没有这些优化，TPU v1 的带宽需求会是多少？我们来算一下：

// 假设矩阵乘法 C = A × B
// A 的大小：256 × 256
// B 的大小：256 × 256
// C 的大小：256 × 256

// 如果不做数据复用：
// 需要读取 A 的 256×256 = 65536 个元素
// 需要读取 B 的 256×256 = 65536 个元素
// 总共需要读取 131072 个元素
// 每个元素 8 位（INT8），总共 1 MB 数据

// 如果做数据复用（权重固定）：
// 权重 B 只需要加载一次：65536 个元素
// 输入 A 按行流入：256 个元素/周期
// 总共只需要读取 65536 + 256 = 65792 个元素
// 节省了将近一半的带宽

这个计算虽然简化了，但道理是通的。数据复用带来的带宽节省是实实在在的。

实际设计中的权衡

做脉动阵列设计，说白了就是在做权衡。TPU v1 的设计团队做了几个关键决策：

设计决策	优点	缺点
256×256 阵列大小	算力高，适合大矩阵	小矩阵利用率低
INT8 精度	推理场景足够，功耗低	训练场景精度不够
权重固定数据流	推理场景带宽效率高	训练场景需要频繁更新权重
无本地缓存	设计简单，面积小	灵活性差

我个人觉得，TPU v1 最聪明的地方在于：它知道自己要解决什么问题。它不是为了做一个通用的 AI 芯片，而是专门为 Google 的推理场景（比如 RankBrain、语音识别）设计的。所以它的设计决策都非常「极端」——把推理场景的性能做到极致，其他场景？抱歉，不是我的菜。

下一节，我们会继续深入脉动阵列的变体设计，看看除了权重固定之外，还有哪些数据流方式，以及它们各自适用的场景。