第四章：脉动阵列架构——经典Systolic Array原理、数据流模式与Google TPU案例分析
各位同学，今天我们来聊聊脉动阵列。说实话，这个架构在AI加速器领域是个绕不开的话题。我当年第一次接触它，是在做一款图像处理芯片的时候。那时候还没什么AI芯片的概念，但脉动阵列的思路已经非常成熟了。
脉动阵列，英文叫Systolic Array。你想想看，它为什么叫“脉动”？因为数据就像血液一样，在计算单元之间有节奏地流动。每个计算单元只做简单的事——乘加运算，然后把结果传给邻居。这种设计，说白了就是用空间换时间。
4.1 经典Systolic Array原理
脉动阵列的核心思想，是把多个处理单元（PE）排列成二维网格。每个PE只和相邻的PE通信。数据从阵列边缘流入，经过每个PE时完成一次乘加，最后从另一侧流出。
我习惯用一个比喻来理解：想象一条生产线。每个工位只做一道工序，产品从一头进去，经过所有工位后，成品从另一头出来。脉动阵列就是这样，只不过它处理的是矩阵运算。
核心优势：
- 数据复用率高——每个数据被多个PE重复使用
- 通信局部化——PE之间只和邻居通信，没有全局连线
- 高吞吐——流水线式处理，每个时钟周期都能出结果
嗯，这里要注意。脉动阵列虽然好，但它对数据流的组织要求很高。如果数据喂得不对，整个阵列就会空转。我在项目中就吃过这个亏，后面会讲到。
下面这张图展示了经典的二维脉动阵列结构：
4.2 三种数据流模式
脉动阵列的数据流模式，说白了就是决定“谁不动、谁动”。我把它分成三种：Weight Stationary、Input Stationary 和 Output Stationary。每种模式都有自己的适用场景。
4.2.1 Weight Stationary（权重固定）
这种模式下，权重预先加载到每个PE中，然后不动。输入特征图从阵列的一侧流入，部分和在阵列中累加。
我个人习惯用这种模式做卷积加速。为什么？因为权重在推理时是固定的，提前加载好，后面就不用动了。数据复用率很高。
适用场景：卷积神经网络推理。权重固定，输入数据流式处理。
伪代码示意：
// Weight Stationary 数据流
// 假设阵列大小为 M x N
for (每个输入特征图行) {
for (每个输入特征图列) {
for (每个权重) {
// 权重已预加载到PE中
// 输入数据从左向右流动
// 部分和从上向下累加
PE[i][j].mac(input, weight[i][j], partial_sum);
}
}
}4.2.2 Input Stationary（输入固定）
这种模式反过来。输入数据固定在一个PE中，权重在阵列中流动。每个PE持有一份输入数据，权重经过时完成乘加。
嗯，这里要注意。Input Stationary 适合输入数据复用率高的场景。比如全连接层，输入向量被反复使用。
关键区别：
- Weight Stationary：权重不动，输入和部分和流动
- Input Stationary：输入不动，权重和部分和流动
- Output Stationary：部分和不动，输入和权重流动
4.2.3 Output Stationary（输出固定）
这种模式下，每个PE负责计算输出矩阵中的一个元素。输入和权重都从阵列边缘流入，在PE中完成乘加后，结果就留在PE里。
我曾经在一个项目中尝试用Output Stationary做矩阵乘法。效果不错，但有个问题——输出矩阵的尺寸决定了阵列的利用率。如果输出矩阵比阵列小，很多PE就浪费了。
避坑指南：我曾经在设计一款加速器时，选了Output Stationary模式。结果发现，当输出矩阵尺寸不是阵列尺寸的整数倍时，边界处理非常麻烦。后来我改用Weight Stationary，问题就解决了。所以，选哪种模式，一定要看你的实际运算场景。
4.3 Google TPU案例分析
说到脉动阵列，就不得不提Google TPU。TPU v1是2016年发布的，专门用于推理。它里面就有一个256x256的脉动阵列。
我记得第一次看到TPU的架构图时，心里想：这设计太优雅了。它用了Weight Stationary模式，权重预加载到阵列中，输入数据从左边流入，部分和从上往下累加。
TPU的脉动阵列有几个关键设计点：
| 设计特性 | 具体实现 | 我的评价 |
|---|---|---|
| 阵列规模 | 256 x 256 = 65,536个PE | 够大，适合批量处理 |
| 数据流模式 | Weight Stationary | 推理场景最优解 |
| 数值精度 | 8位整数乘加 | 推理够用，功耗低 |
| 累加器 | 32位整数 | 防止精度丢失 |
| 片上存储 | 24MB统一缓冲区 | 减少DRAM访问 |
TPU的脉动阵列是怎么工作的？我简单说一下流程：
- 先把权重从DRAM加载到阵列中。这一步是预加载，权重不动。
- 输入特征图从左边流入阵列。每个时钟周期，数据向右移动一列。
- 每个PE完成一次乘加：output += input * weight。
- 部分和从上往下累加。每个PE把结果传给下方的PE。
- 最后一行PE输出最终结果。
你想想看，这个过程有多高效。每个输入数据被256个PE同时使用，数据复用率极高。而且PE之间只有局部通信，没有全局连线，功耗很低。
TPU性能数据（来自Google论文）：
- 峰值性能：92 TOPS（8位整数）
- 功耗：40-50W（典型负载）
- 能效比：约2 TOPS/W
- 相比CPU/GPU：能效提升30-80倍
嗯，这里要补充一点。TPU的脉动阵列并不是万能的。它擅长的是大矩阵乘法，比如全连接层和卷积层。但对于稀疏运算、不规则运算，它的效率会下降。我后来在做稀疏加速器时，就发现脉动阵列对稀疏数据不太友好——很多PE在空转。
不过话说回来，TPU的成功证明了脉动阵列在AI加速中的价值。它把数据复用和局部通信做到了极致。后来的很多AI芯片，比如华为的昇腾、寒武纪的MLU，都借鉴了脉动阵列的思想。
最后，我给大家一个实用建议。如果你在设计AI加速器，先想清楚你的目标场景：
- 推理场景：优先考虑Weight Stationary，权重固定，数据流式处理
- 训练场景：考虑Output Stationary，因为权重和输入都在变
- 混合场景：可能需要动态切换数据流模式
好了，这一章就到这里。脉动阵列是个好东西，但用得好不好，全看你对数据流的理解。下一章我们聊聊稀疏加速——那又是另一个故事了。
公众号:蓝海资料掘金营，微信deep3321