计算单元设计：脉动阵列与乘加器

各位同学，今天我们来聊聊深度学习加速器里最核心的模块——计算单元。说实话，这部分内容我做了十几年，每次回看都觉得有意思。脉动阵列这个概念，最早是80年代提出的，但真正在AI芯片里大放异彩，也就是这几年的事。

脉动阵列（Systolic Array）原理

脉动阵列，说白了就是让数据像血液一样，在计算单元之间规律地流动。每个计算节点只做一件事：乘加运算。数据从边界流入，计算结果从另一边流出。

我刚开始接触这个概念时，觉得它特别像工厂里的流水线。每个工人只拧一个螺丝，但整条线跑起来，效率惊人。脉动阵列也是这个道理。

核心思想：数据复用 + 局部通信。每个PE（处理单元）只和邻居通信，避免了全局互联的带宽瓶颈。

为什么会这样设计？你想想看，深度学习模型里，卷积和矩阵乘占了90%以上的计算量。这些运算有个特点：数据可以被反复使用。脉动阵列正好利用了这个特性。

举个例子，一个256x256的脉动阵列，理论上可以同时进行65536次乘加运算。但实际项目中，我遇到过数据喂不饱阵列的情况——这就是所谓的“内存墙”问题。

脉动阵列数据流示意图 权重输入 特征输入 PE PE PE PE PE PE PE PE PE PE PE PE PE PE PE PE 部分和输出 权重数据流 特征数据流 部分和流

乘加器（MAC）单元设计

每个PE里最核心的东西，就是乘加器。一个MAC单元 = 一个乘法器 + 一个加法器 + 一个累加寄存器。听起来简单，但设计起来门道不少。

我个人习惯把MAC设计分成三个层次：

  • 位宽选择：INT8、INT16还是FP16？这直接影响芯片面积和功耗。我做过一个项目，用INT8比FP16省了4倍面积。
  • 流水线级数：2级还是3级？级数越多频率越高，但延迟也越大。这是个trade-off。
  • 累加器精度：这个容易被忽略。我曾经吃过亏——累加器精度不够，导致大模型训练时梯度爆炸。

我的经验：对于推理芯片，INT8 MAC足够用。但训练芯片，至少需要FP16，累加器建议用FP32。

来看一个简单的MAC Verilog代码示例：

module mac_unit #(
    parameter DATA_WIDTH = 8,
    parameter ACC_WIDTH  = 32
)(
    input  clk, rst_n,
    input  [DATA_WIDTH-1:0] a, b,
    input  valid_in,
    output reg [ACC_WIDTH-1:0] result,
    output reg valid_out
);
    reg [ACC_WIDTH-1:0] accumulator;
    wire [DATA_WIDTH*2-1:0] product = a * b;
    
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            accumulator <= 0;
            valid_out <= 0;
        end else if (valid_in) begin
            accumulator <= accumulator + product;
            valid_out <= 1;
        end else begin
            valid_out <= 0;
        end
    end
    
    always @(*) begin
        result = accumulator;
    end
endmodule

数据流优化：三种主流策略

数据流怎么安排，直接决定了阵列的利用率。我见过太多设计，阵列利用率不到50%——说白了就是浪费硅片。

目前主流的有三种数据流策略：

策略 核心思想 适用场景 优缺点
Weight Stationary 权重固定，特征图流动 卷积层（权重复用度高） 权重搬运少，但特征图带宽要求高
Output Stationary 输出固定，输入流动 全连接层（输出维度大） 减少部分和搬运，但控制复杂
Row Stationary 行数据固定，按行流动 通用矩阵乘 平衡了前两者的优缺点

Weight Stationary（权重固定）

这是最直观的策略。权重加载到PE后就不动了，特征图数据在阵列里流动。每个PE计算一个权重和多个特征图的乘积。

我在做第一代AI芯片时用的就是这种策略。优点是控制简单，但有个坑——如果卷积核很大，权重加载时间会很长。

关键指标：权重复用度 = 特征图尺寸 / 步长。复用度越高，这种策略越划算。

Output Stationary（输出固定）

这种策略反过来——输出部分和固定在PE里，权重和特征图都流动。每个PE负责计算输出特征图的一个像素点。

嗯，这里要注意。Output Stationary对部分和的累加特别友好，因为结果就在本地。但控制逻辑比Weight Stationary复杂不少。

Row Stationary（行固定）

这是我个人比较推崇的策略。它把权重矩阵的一行和特征图的一行固定下来，按行进行流水计算。说白了，就是前两种策略的折中方案。

Google的TPU用的就是类似思路。我记得在某个项目里，用Row Stationary把阵列利用率从60%提到了85%。

避坑指南：我曾经在一个项目里，盲目追求高利用率，把阵列设计得特别大。结果发现数据搬运成了瓶颈，阵列经常空转。后来我学乖了——阵列大小要和片内缓存带宽匹配。

实际设计中的取舍

讲完理论，说说实际设计中的几个关键决策点：

  1. 阵列大小：16x16还是128x128？我建议从应用场景倒推。如果是边缘设备，16x16就够了；云端芯片，至少64x64起步。
  2. 数据精度：混合精度越来越流行。我最近做的芯片，支持INT4/INT8/FP16动态切换。
  3. 稀疏性支持：模型里很多权重是0。如果跳过这些0计算，能省一半功耗。但硬件支持稀疏性，控制逻辑会复杂很多。

最后说一句，设计计算单元没有银弹。每种策略都有适用场景，关键是要理解你的目标模型和部署场景。我见过太多人，拿着论文里的架构直接套用，结果流片回来发现根本跑不动目标模型。

好了，这一章就到这里。记住，好的计算单元设计，是数据流、存储和计算三者完美配合的结果。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →