从零设计AI加速器：矩阵乘加阵列与存储层次实战

1. 为什么我想从零设计一款AI加速器1.1 一个朴素但真实的动机三年前我在做一个边缘侧视觉检测项目模型不大参数量不到两百万但推理延迟始终卡在40毫秒下不来。我试过换更快的CPU、把模型量化到INT8、甚至把后处理全部搬到GPU上效果都不理想。问题不在于算力不够而在于数据在存储器与计算单元之间来回搬运的能耗和延迟远远超过了计算本身。那段时间我翻了很多资料看到一句话印象很深矩阵乘法本身不贵贵的是把矩阵搬来搬去。这句话直接把我推向了AI加速器设计这个方向。所谓AI加速器说白了就是专门为神经网络里那些重复度极高的运算主要是矩阵乘加定制的硬件。它不像CPU那样什么都能干但什么都不算特别快也不像GPU那样靠海量线程堆吞吐而是把神经网络的计算模式吃透之后用最省的方式把矩阵运算做掉。NPU、TPU、各种端侧推理芯片本质上都是这个思路的不同实现。这篇文章不是学术论文也不是芯片设计教材的读书笔记。我想做的是把“从零设计一款AI加速器”这件事按照一个从业者真实的思考路径拆开先想清楚要解决什么问题再决定架构长什么样然后落到具体的矩阵运算单元、存储层次、数据流调度最后聊聊验证和踩过的坑。如果你是对硬件加速感兴趣的软件工程师或者是刚入门的数字IC设计者又或者只是好奇NPU里面到底在干什么这篇内容应该都能给你一些可以直接参考的东西。1.2 先明确边界我要设计的到底是什么“AI加速器”这个词太宽了。从数据中心里几百瓦的推理卡到手机SoC里几平方毫米的NPU再到单片机旁边挂的一颗几毛钱的协处理器都叫AI加速器。如果不先把边界画清楚后面的架构讨论就是空中楼阁。我给自己定的目标是面向边缘侧推理的定点矩阵加速器。具体约束如下目标场景卷积神经网络和轻量级Transformer的推理不做训练数据精度INT8为主支持INT16累加不碰浮点算力目标在200MHz时钟下做到1TOPS左右的等效算力存储约束片上SRAM不超过256KB不接外部DRAM接口AXI4-Lite做控制AXI4-Stream做数据输入输出为什么这么定因为边缘侧最缺的不是峰值算力而是能效比和面积效率。你在一块智能摄像头的主控旁边塞一颗几百毫瓦的加速器如果还要外挂DDR那整体功耗和成本就失控了。把数据留在片上、把精度压到INT8、把矩阵运算做成流水线这才是边缘加速器该有的样子。提示很多初学者一上来就想做“通用AI加速器”既支持卷积又支持全连接还支持各种激活函数结果架构越做越复杂最后连仿真都跑不通。先把一个核心运算做扎实比什么都重要。2. 架构设计从矩阵乘法出发倒推硬件结构2.1 为什么矩阵乘法是核心中的核心神经网络推理拆到最底层绝大部分计算量都落在矩阵乘法上。卷积可以展开成矩阵乘im2col全连接本身就是矩阵乘Transformer里的注意力机制核心也是Q、K、V三个矩阵的乘加。所以只要把矩阵乘法加速做好就覆盖了90%以上的计算需求。一个典型的矩阵乘是C A × B其中A是M×KB是K×NC是M×N。在CPU上这件事靠循环嵌套加SIMD指令来做在GPU上靠成千上万个线程并行而在AI加速器上我们要做的是设计一个二维的乘加阵列让A的行和B的列在硬件里直接相遇一个周期内完成一批乘加。这里有个关键概念叫数据复用。矩阵乘里每个A的元素会被复用N次每个B的元素会被复用M次。如果每次乘加都从存储器里重新读数据那带宽根本扛不住。所以加速器的核心设计目标就是让数据在计算单元附近尽可能久地停留减少对存储器的访问次数。2.2 脉动阵列一个经典但不过时的选择说到矩阵加速绕不开脉动阵列Systolic Array。我第一次看到这个概念的时候觉得挺玄乎后来想明白了它本质上就是一个让数据像心跳一样有节奏地在计算单元之间流动的结构。具体来说假设我有一个8×8的乘加单元PE阵列。A矩阵的元素从左侧流入B矩阵的元素从上方流入每个PE负责一个乘加操作结果向下或向右传递。数据在阵列里“脉动”前进每个周期都有新的数据进来也有算完的数据出去。这样做的好处是每个PE只需要和相邻的PE通信布线短频率容易做高数据复用率极高A的每个元素进入阵列后会被整行PE用到控制逻辑简单不需要复杂的调度器当然脉动阵列也不是没有缺点。它的灵活性较差适合规则的矩阵运算遇到稀疏矩阵或者不规则计算就效率骤降。但对于边缘推理这种以规则卷积为主的场景它依然是最优解之一。我最终选择的是一种可配置的脉动阵列PE阵列大小固定为16×16但支持三种工作模式——标准矩阵乘模式、深度卷积模式、以及向量乘加模式。这样既能覆盖常规卷积也能处理MobileNet里的深度可分离卷积。2.3 存储层次比计算更值得花心思的地方前面说了矩阵乘本身不贵贵的是搬数据。所以在设计加速器的时候我在存储层次上花的精力比计算单元还多。整个存储层次从外到内分为四级层级容量带宽用途外部接口-AXI4-Stream输入特征图和权重全局SRAM128KB256bit/cycle缓存输入、权重、输出行缓冲4KB512bit/cycle为PE阵列提供对齐数据寄存器文件每个PE 32B本地PE内部暂存为什么这么分因为PE阵列每个周期要吃掉16×16256个乘加操作对应需要16个A元素和16个B元素。如果直接从全局SRAM读带宽根本不够。所以我在PE阵列旁边放了行缓冲提前把数据准备好PE阵列只管算不管等。这里有个经验行缓冲的深度要至少能装下两行数据这样才能做双缓冲一边算当前行一边预取下一行。我一开始只做了一行缓冲结果PE阵列有30%的时间在等数据算力利用率直接掉到70%以下。后来加了双缓冲利用率才回到95%以上。2.4 数据流调度让每个周期都不浪费数据流调度是加速器设计里最像“艺术”的部分。同样的硬件调度策略不同实际性能可能差一倍。我采用的是输出 stationary 加权重 stationary 的混合策略。具体来说权重在计算开始前一次性加载到PE阵列的本地寄存器里整个推理过程中不变输入特征图按行流入每行计算完后输出部分和部分和在PE阵列内部累加直到所有K维度处理完才写出这样做的好处是权重只需要读一次输入特征图也只读一次输出部分和只在最后写一次。整个过程中全局SRAM的访问次数被压到了最低。注意输出 stationary 策略对片上存储的要求比较高因为部分和要一直留在PE阵列里。如果K维度太大部分和寄存器会不够用。我的做法是把K维度分块每块算完先写到全局SRAM最后再累加。虽然多了一次读写但总比寄存器溢出强。3. 核心模块实现从PE到控制器的完整拆解3.1 乘加单元PE的设计细节PE是整个加速器里最基础的单元它的设计直接决定了加速器的面积、功耗和频率。一个PE的核心就是一个乘法器加一个累加器。但要做得好有几个细节必须注意第一乘法器的位宽选择。INT8乘INT8得到INT16这是最自然的做法。但如果你要做INT8乘INT8加INT16累加那累加器就要至少24位防止多次累加后溢出。我一开始用了16位累加器结果在跑ResNet的一个中间层时直接溢出了输出全是错的。后来改成24位问题解决。第二流水线级数。乘法器本身有延迟如果PE里不加流水线时钟频率上不去。我在PE里插了一级流水线把乘法和累加分到两个周期做。代价是控制逻辑稍微复杂一点但频率从100MHz提到了200MHz算力直接翻倍。第三清零和使能逻辑。PE需要支持三种操作清零累加器、累加当前乘积、输出结果。这些控制信号要做得干净不能有毛刺。我用了同步复位加时钟使能的方案实测下来很稳。下面是一个PE的简化Verilog代码你可以直接参考module pe #(parameter DW 8, parameter AW 24) ( input wire clk, input wire rst_n, input wire en, input wire clr, input wire signed [DW-1:0] a, input wire signed [DW-1:0] b, output reg signed [AW-1:0] out ); wire signed [2*DW-1:0] prod; assign prod a * b; always (posedge clk or negedge rst_n) begin if (!rst_n) out 0; else if (en) begin if (clr) out {{(AW-2*DW){prod[2*DW-1]}}, prod}; else out out {{(AW-2*DW){prod[2*DW-1]}}, prod}; end end endmodule这段代码里有个细节符号扩展。因为prod是16位有符号数累加器是24位所以要把prod的符号位扩展到24位再相加。如果忘了这一步负数乘法就会出错。3.2 行缓冲与数据对齐行缓冲的作用是把全局SRAM里的数据整理成PE阵列需要的格式。听起来简单做起来坑很多。最大的坑是数据对齐。卷积运算里输入特征图的一个3×3窗口要对应9个权重但这9个数据在SRAM里不是连续存放的。如果每次都要重新计算地址那控制逻辑会非常复杂。我的做法是在行缓冲里做一个滑动窗口每次新进来一行数据窗口就往下滑一行同时把最旧的一行挤出去。这样PE阵列永远看到的是一个对齐好的3×3窗口不需要关心地址计算。另一个坑是边界处理。卷积在特征图边缘需要补零如果补零逻辑放在PE阵列里做会浪费计算周期。我把它放在行缓冲的写入端当写入地址超出特征图范围时直接写零。这样PE阵列完全感知不到边界的存在计算效率不受影响。3.3 控制器与状态机控制器是整个加速器的“大脑”它要协调数据加载、计算、输出三个阶段的时序。我设计了一个五状态的状态机IDLE等待启动信号LOAD_WEIGHT从全局SRAM加载权重到PE阵列LOAD_INPUT加载输入特征图到行缓冲COMPUTEPE阵列计算输出部分和WRITE_OUT把结果写回全局SRAM状态之间的切换条件要仔细设计。比如从LOAD_INPUT到COMPUTE必须等行缓冲填满至少两行才能开始否则PE阵列会饿死。从COMPUTE到WRITE_OUT必须等所有部分和累加完毕。实操心得状态机的调试建议先用仿真跑一遍完整的推理流程把每个状态的持续时间打印出来。我当时发现COMPUTE状态占了总时间的85%但其中有20%是在等行缓冲。后来调整了LOAD_INPUT的预取策略把这20%压到了5%以下。3.4 量化与反量化模块INT8推理离不开量化。模型训练时是浮点部署时要转成定点这中间有一个缩放因子scale和零点zero point。我的加速器里集成了一个简单的量化模块支持每层独立的scale和zero point。具体做法是权重在加载时就已经量化好了直接存INT8输入特征图在写入行缓冲前做量化浮点转INT8输出部分和在写回SRAM后做反量化INT24转浮点这里有个经验scale和zero point最好用2的幂次这样量化就是简单的移位操作不需要乘法器。虽然精度会损失一点但在边缘场景下完全够用而且省下来的面积和功耗非常可观。4. 验证与调试那些只有踩过才知道的坑4.1 功能验证从单元测试到系统级仿真加速器设计最怕的就是流片回来发现算错了。所以在RTL阶段就要把验证做扎实。我的验证策略分三层第一层是PE级单元测试。用随机生成的A和B矩阵在PE阵列上跑一遍和Python的numpy结果对比。这一步主要验证乘加逻辑和累加器位宽。第二层是模块级测试。把行缓冲、控制器、PE阵列连起来跑一个完整的卷积层。输入用真实的特征图数据输出和PyTorch的CPU结果对比。这一步主要验证数据流和状态机。第三层是系统级仿真。把整个加速器挂到一个小型的SoC模型上跑一个完整的MobileNet推理。这一步主要验证接口和中断逻辑。三层验证跑下来基本能覆盖95%以上的功能bug。剩下的5%往往是时序相关的需要上板调试。4.2 时序收敛那些让人头疼的路径200MHz的目标频率在28nm工艺下不算高但也不是随便就能达到的。我遇到的最大问题是PE阵列的布线延迟。16×16的PE阵列每个PE都要和上下左右的PE通信布线资源非常紧张。如果综合工具把PE摆得太散关键路径就会出现在PE之间的连线上。我的解决办法是用手动布局约束把PE阵列约束成一个紧凑的矩形在PE之间的连线上插入寄存器把长路径打断对时钟树做平衡确保所有PE的时钟偏斜在可接受范围内这些手段用下来时序余量从-0.3ns变成了0.15ns算是勉强收敛。4.3 常见问题速查表问题现象可能原因排查方法解决办法输出全零PE阵列使能信号没拉高检查控制器的en信号修正状态机切换条件输出溢出累加器位宽不够打印中间部分和增加累加器位宽算力利用率低行缓冲带宽不足统计PE等待周期加双缓冲或加宽SRAM位宽时序不收敛PE间布线太长看时序报告的关键路径手动布局加插入寄存器量化后精度掉太多scale粒度太粗逐层对比浮点和定点输出改成每通道独立scale4.4 上板调试从仿真到现实的鸿沟仿真跑通不代表上板能跑。我第一次上板的时候加速器直接不工作读出来的全是0xDEADBEEF。排查了两天才发现是AXI4-Stream的握手信号有问题仿真里ready和valid的时序是理想的但实际硬件里ready信号有延迟导致数据丢失。后来我养成了一个习惯所有跨时钟域的信号都加同步器所有握手信号都做超时保护。虽然多了一点面积但稳定性提升了一个数量级。提示上板调试建议先用一个最简单的测试用例比如4×4的矩阵乘确认数据通路是通的。然后再逐步加大规模不要一上来就跑完整模型。5. 性能评估与优化方向5.1 实测性能数据在200MHz时钟下16×16的PE阵列理论峰值算力是16 × 16 × 2 × 200M 102.4 GOPS实测下来跑一个标准的3×3卷积层算力利用率在92%左右等效算力约94 GOPS。跑深度可分离卷积时利用率会掉到75%左右因为深度卷积的复用率低PE阵列有一部分在空转。功耗方面在28nm工艺下整个加速器的动态功耗约180mW静态功耗约15mW。能效比大约是0.5 TOPS/W在边缘场景下算是中规中矩。5.2 下一步可以怎么优化如果要把这个设计再往前推一步我觉得有几个方向值得尝试第一支持稀疏计算。很多模型剪枝后权重里有大量零如果PE阵列能跳过零权重算力可以再提升30%以上。做法是在PE里加一个零检测逻辑遇到零直接跳过乘加。第二增加多核扩展。单核16×16的算力有限如果做成4核每个核负责不同的输出通道算力可以线性扩展。代价是核间通信和任务调度的复杂度上升。第三优化数据流支持Transformer。现在的架构对卷积很友好但对注意力机制里的矩阵转置和softmax支持不够。如果要做端侧大模型这部分必须补上。我个人在实际操作中的体会是AI加速器设计最难的不是某个模块的实现而是在面积、功耗、性能、灵活性之间找到平衡点。你不可能什么都想要必须根据目标场景做取舍。边缘侧就老老实实做INT8、做小阵列、做高复用数据中心才需要考虑浮点、大阵列、高带宽。想清楚这一点后面的设计决策就会顺畅很多。最后再分享一个小技巧如果你也在做类似的加速器设计建议先把Python的行为级模型写出来用numpy把整个推理流程跑通再开始写RTL。这样你在调试RTL的时候永远有一个“标准答案”可以对比效率会高很多。我一开始跳过这一步直接写Verilog结果在调试累加器溢出的时候花了整整一周后来补了Python模型半天就定位到了问题。