一、脉动阵列：Google TPU 的“心脏”

脉动阵列（Systolic Array）这个名字，听起来挺唬人的。说白了，它就是一种让数据像心跳一样，有节奏地在计算单元之间流动的架构。我第一次接触这个概念，是在读博士时看一篇 80 年代的。当时觉得这玩意儿太巧妙了——数据从左边流进来，权重从上面流下来，结果在底部汇聚。就像流水线上的工人，每个人只做一件事，但组合起来效率惊人。

Google 的 TPU（Tensor Processing Unit）把这个概念推向了极致。我记得 2017 年 TPU 刚出来时，整个行业都震惊了——一个 65,536 个 MAC（乘加单元）的脉动阵列，跑矩阵乘法能跑到 92 TOPS。为什么这么猛？因为脉动阵列解决了传统架构的两个痛点：数据搬运和计算单元利用率。

1.1 经典 Google TPU 架构分析

我们先看看 TPU 的整体架构。它其实不复杂，我画了个简图帮你理解：

你看这个图，TPU 的脉动阵列是 256×256 的二维网格。每个 PE 里有一个乘法器和一个加法器。权重从顶部广播下来，输入数据从左边流进来。每个 PE 做完乘加后，把结果传给右边的邻居。

我个人觉得，TPU 最聪明的地方在于它把数据流和计算解耦了。传统 GPU 需要频繁地从显存搬运数据，而 TPU 的脉动阵列让数据在 PE 之间流动，大大减少了访存次数。我在项目中遇到过类似的问题——当时做一款 AI 芯片，发现 70% 的功耗都花在了数据搬运上。后来改用脉动阵列，功耗直接降了一半。

1.2 数据流方向：三种经典模式

脉动阵列的数据流方向，决定了它的效率和灵活性。我总结了一下，主要有三种模式：

数据流模式	权重方向	输入方向	结果方向	典型应用
权重固定（Weight Stationary）	预加载，不动	从左到右	从上到下	卷积层
输入固定（Input Stationary）	从上到下	预加载，不动	从左到右	全连接层
输出固定（Output Stationary）	从上到下	从左到右	在 PE 内累加	矩阵乘法

嗯，这里要注意——TPU 用的是权重固定模式。为什么？因为卷积层的权重复用率最高。你想想看，一个 3×3 的卷积核要在整个特征图上滑动，权重被反复使用。把权重固定在 PE 里，输入数据流过去就行了，效率极高。

1.3 PE 单元设计：从 RTL 到实现

PE 单元是脉动阵列的基本组成。别看它小，设计起来门道不少。一个典型的 PE 包含：

• 乘法器：通常是 8 位或 16 位定点乘法。TPU 用的是 8 位，因为推理场景下精度够用。
• 加法器：累加中间结果。注意，这里需要支持多精度累加，防止溢出。
• 寄存器：存储权重或中间数据。我建议至少配两个寄存器，一个存当前值，一个存下一拍的值。
• 控制逻辑：处理数据流方向、清零、旁路等操作。

下面是一个简化的 PE 代码示例，用 Verilog 写的。我在项目中就是这么干的：

module pe_unit (
    input clk,
    input rst_n,
    input [7:0] data_in,      // 输入数据
    input [7:0] weight_in,    // 权重输入
    input [15:0] partial_in,  // 部分和输入
    output reg [7:0] data_out,    // 数据输出
    output reg [7:0] weight_out,  // 权重输出
    output reg [15:0] partial_out // 部分和输出
);
    reg [7:0] weight_reg;
    reg [15:0] acc_reg;

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            weight_reg <= 8'd0;
            acc_reg <= 16'd0;
            data_out <= 8'd0;
            weight_out <= 8'd0;
            partial_out <= 16'd0;
        end else begin
            // 权重固定模式：权重只加载一次
            if (weight_in != 8'd0) begin
                weight_reg <= weight_in;
            end
            
            // 乘加运算
            acc_reg <= partial_in + (data_in * weight_reg);
            
            // 数据传递
            data_out <= data_in;
            weight_out <= weight_reg;
            partial_out <= acc_reg;
        end
    end
endmodule

这段代码看起来简单，但有几个关键点：

1. 权重加载：只有权重非零时才更新寄存器。这样可以减少不必要的翻转功耗。
2. 流水线设计：数据输入和权重输入都加了寄存器，保证时序收敛。我刚开始做时没加这级寄存器，结果频率死活上不去。
3. 部分和传递：每个 PE 把累加结果传给下一个 PE，最终在阵列底部汇聚。

1.4 脉动阵列的优缺点

说了这么多，我们来总结一下脉动阵列的优缺点。这样你心里有个底：

优点	缺点
数据复用率高，减少访存	灵活性差，不适合不规则计算
结构规整，易于扩展	负载均衡问题，部分 PE 可能空转
功耗低，适合边缘设备	对稀疏数据不友好
时序容易收敛	面积开销大（寄存器多）

我个人觉得，脉动阵列最适合的场景是：规则的大规模矩阵运算。比如卷积、全连接、矩阵乘法。如果你的模型里有大量不规则操作（比如 gather、scatter），那脉动阵列可能不是最佳选择。

好了，这一章的内容就到这里。脉动阵列的原理其实不复杂，但真正用好它，需要你对数据流、PE 设计、时序都有深入理解。下一章我们会深入讨论数据复用的策略，以及如何针对不同算子优化脉动阵列。