从零设计AI加速器:脉动阵列与存储层次实战
1. 为什么我想从零设计一款AI加速器1.1 一个朴素但真实的动机去年我在做一个边缘端视觉推理的小项目模型不大参数量也就几百万级别但部署到目标硬件上之后帧率始终卡在十几帧上不去。我一开始以为是模型的问题换了更轻量的骨干网络做了剪枝和量化效果提升有限。后来用性能分析工具一查发现瓶颈根本不在计算单元上而是在数据搬运——权重和中间特征在存储器和计算阵列之间来回倒腾带宽吃满了计算单元大部分时间在等数据。这件事让我意识到一个很现实的问题我们平时讨论AI芯片张口就是算力多少TOPS、支持什么精度、跑什么模型但真正决定实际性能的往往是那些不那么显眼的东西——存储层次怎么设计、数据怎么复用、控制流怎么调度。这些细节在通用CPU和GPU上被厂商封装好了你不需要关心但一旦你想理解AI加速器到底是怎么回事或者想自己动手做一个这些就是绕不开的核心。所以这篇文章我想把从零设计一款AI加速器这件事完整地拆一遍。不是纸上谈兵地讲概念而是按照一个真实的工程思路从需求定义、架构选型、核心计算单元设计、存储层次规划一直到控制逻辑和验证方法把每个环节的关键决策和背后的理由讲清楚。适合谁看如果你是对计算机体系结构有兴趣的软件工程师或者是刚接触芯片设计、想搞明白NPU内部到底怎么运转的学生再或者你只是好奇矩阵乘法为什么能加速这件事这篇文章应该都能给你一些实在的东西。1.2 先搞清楚AI加速器到底在加速什么在动手之前得先把问题定义清楚。AI加速器不是万能芯片它之所以能比CPU快是因为它针对特定类型的计算做了极致的优化。那这个特定类型的计算是什么答案就是矩阵运算尤其是矩阵乘法。无论是卷积神经网络里的卷积操作还是Transformer里的注意力机制拆到最底层本质上都是矩阵乘加运算。一个卷积层可以展开成im2col形式的矩阵乘法一个全连接层本身就是矩阵乘法注意力里的Q乘K转置也是矩阵乘法。所以AI加速器的核心任务就是高效地做矩阵乘法。这里就引出一个关键概念乘加运算MACMultiply-Accumulate。一次MAC就是a乘以b再加上c即c c a * b。AI加速器里通常会有大量的MAC单元并行工作比如一个256x256的脉动阵列里面就有65536个MAC单元每个时钟周期能完成65536次乘加。相比之下CPU虽然单核频率高但并行度有限一个周期能做的乘加次数远不及专用的阵列结构。但光有计算单元还不够。矩阵乘法有个特点数据复用率高。比如C A * BA的每一行会和B的每一列做点积A中的元素会被复用多次B中的元素也是。如果每次计算都从外部存储器重新读数据带宽根本扛不住。所以AI加速器的设计很大程度上是在解决怎么让数据在正确的时间出现在正确的位置这个问题。1.3 设计目标的量化定义在开始画架构图之前我习惯先把设计目标量化。不然做到一半很容易迷失方向。对于一款AI加速器我通常会定义以下几个维度的指标峰值算力单位是TOPS每秒万亿次操作或GOPS。这个数字由MAC单元数量和时钟频率决定。比如1024个MAC单元跑在1GHz理论峰值就是1024 * 2 * 1e9 2.048 TOPS乘和加各算一次操作。能效比单位是TOPS/W即每瓦功耗能提供多少算力。这个指标对边缘设备尤其重要。存储带宽包括片上缓存带宽和外部存储带宽。片上带宽通常远高于外部带宽设计时要尽量让数据在片上复用。支持的精度INT8、INT4、FP16还是混合精度。精度越低计算单元面积越小能效越高但模型精度可能受影响。灵活性是只支持特定算子还是能通过指令编程支持多种网络结构。这几个指标之间是相互制约的。想要高算力就得堆MAC单元面积和功耗就上去了想要高能效就得降低精度或者减少数据搬运想要灵活就得增加控制逻辑的开销。所以设计的过程本质上是在这些约束之间找平衡。我给自己定的目标是一款面向边缘推理的AI加速器峰值算力在1 TOPS左右INT8能效比尽量做到5 TOPS/W以上支持INT8和INT4两种精度能跑常见的卷积网络和轻量级Transformer。这个目标不算激进但足够覆盖大部分边缘场景也足够让我把整个设计流程走通。2. 架构选型从CPU、GPU到NPU的思考路径2.1 为什么不用CPU或GPU直接做有人可能会问既然CPU和GPU都能做矩阵运算为什么还要专门设计AI加速器这个问题值得认真回答因为它直接决定了你的架构设计方向。CPU的设计目标是通用性。它有复杂的指令流水线、分支预测、乱序执行、多级缓存这些机制让它能高效地处理各种不同类型的任务但代价是大量的面积和功耗花在了控制逻辑和缓存一致性上真正用于计算的ALU算术逻辑单元占比很小。一个典型的CPU核心可能只有不到10%的面积是真正的计算单元。做矩阵乘法时CPU的向量指令比如AVX系列虽然能一次处理多个数据但并行度仍然有限而且数据在寄存器和缓存之间的搬运开销很大。GPU的设计目标是高吞吐并行计算。它有成千上万个简单的计算核心通过SIMT单指令多线程的方式并行执行。GPU做矩阵乘法确实比CPU快很多但它的问题在于第一功耗高高端GPU动辄几百瓦边缘设备根本扛不住第二GPU的存储层次复杂编程模型也复杂对于固定的推理任务来说很多灵活性是浪费的第三GPU的调度开销和内核启动开销在边缘场景下不可忽略。NPU神经网络处理单元的设计思路则完全不同它放弃通用性专门为神经网络的计算模式做优化。具体来说就是针对矩阵乘法的高数据复用特性设计专用的计算阵列和存储层次让数据流动的路径最短、复用次数最多、控制逻辑最简。这样就能在相同的面积和功耗下获得远高于CPU和GPU的算力。2.2 三种主流架构范式的对比在NPU的设计中有几种主流的架构范式我在选型时做了详细对比第一种是脉动阵列Systolic Array。这是Google TPU采用的结构。它的核心思想是让数据像水流一样在计算单元阵列中有节奏地流动每个计算单元从相邻单元接收数据计算后再传给下一个单元。这种结构的优点是数据复用率极高控制逻辑简单适合大规模的矩阵乘法。缺点是灵活性较差对于非规则的计算模式比如稀疏矩阵、动态形状支持不好。第二种是乘加树MAC Tree或广播结构。这种结构中输入数据通过广播的方式同时发送给多个计算单元计算单元并行计算后通过加法树汇总结果。它的灵活性比脉动阵列好但数据复用率相对低一些因为广播本身消耗带宽。第三种是SIMD/SIMT风格的向量处理单元。这种结构更接近GPU的思路通过宽向量指令来并行处理数据。灵活性最高但能效比通常不如前两种。我最终选择了脉动阵列作为核心计算结构理由有三第一我的目标场景是边缘推理计算模式相对固定以卷积和矩阵乘法为主脉动阵列的高复用率正好匹配第二脉动阵列的控制逻辑简单可以把更多面积留给计算单元和缓存第三脉动阵列的设计方法成熟有大量的文献和开源实现可以参考对于从零开始的设计来说风险可控。2.3 整体架构的顶层规划确定了核心计算结构之后我开始规划整体架构。一款AI加速器的完整架构通常包含以下几个部分计算核心由多个处理单元PE组成的脉动阵列负责执行矩阵乘加运算。片上缓存包括输入缓存、权重缓存和输出缓存用于暂存参与计算的数据减少对外部存储器的访问。控制单元负责解析指令、调度数据流、控制计算阵列的时序。数据搬运引擎负责在外部存储器和片上缓存之间搬运数据通常支持DMA直接内存访问。接口包括与主机的通信接口如PCIe或AXI和与外部存储器的接口如DDR控制器。这些部分之间的关系可以用一个简单的类比来理解计算核心是工厂里的生产线片上缓存是生产线旁边的物料架控制单元是调度员数据搬运引擎是叉车接口是工厂的大门。生产线的效率不仅取决于机器本身的速度还取决于物料能不能及时送到、成品能不能及时运走。在我的设计中计算核心是一个16x16的脉动阵列共256个MAC单元。每个MAC单元支持INT8乘加也支持INT4模式下的双倍吞吐。片上缓存分为三块输入缓存256KB、权重缓存256KB和输出缓存128KB。控制单元采用简单的状态机加微指令的方式支持基本的矩阵乘法、卷积和池化操作。数据搬运引擎支持双缓冲机制可以在计算当前数据块的同时预取下一块数据。3. 核心计算单元脉动阵列的设计细节3.1 脉动阵列的工作原理脉动阵列的核心思想是让数据在计算单元之间脉动式地流动。以最简单的矩阵乘法C A * B为例假设A是MxK的矩阵B是KxN的矩阵C是MxN的矩阵。在脉动阵列中A的元素从左侧流入B的元素从上方流入每个计算单元在每个时钟周期完成一次乘加并将结果累加。具体来说假设我们有一个NxN的脉动阵列要计算两个NxN矩阵的乘积。在第一个时钟周期A的第一行第一个元素和B的第一列第一个元素进入左上角的计算单元完成一次乘加。第二个时钟周期A的第一个元素向右移动B的第一个元素向下移动同时A的第二个元素和B的第二个元素进入阵列。这样每个计算单元在每个周期都在做有用的计算数据复用率极高。这种结构的精妙之处在于A的每个元素只从外部读取一次然后在阵列中向右流动被N个计算单元依次使用B的每个元素也只读取一次向下流动被N个计算单元使用。整个计算过程中外部存储器的访问次数大大减少带宽压力自然就小了。3.2 MAC单元的具体实现MAC单元是脉动阵列的基本构件。一个典型的MAC单元包含以下部分乘法器负责计算a * b。对于INT8乘法可以用8x8的乘法器对于INT4可以用4x4的乘法器面积更小。累加器负责将乘积累加到之前的结果上。通常是一个位宽较宽的寄存器比如32位以防止溢出。寄存器用于暂存输入数据和部分和支持数据的流动。控制逻辑根据模式选择信号决定当前是INT8模式还是INT4模式以及是否清零累加器等。在设计MAC单元时有几个关键决策第一乘法器的位宽和精度。INT8乘法器输出16位结果INT4乘法器输出8位结果。为了支持两种精度我设计了一个可配置的乘法器在INT4模式下将两个INT4乘法器合并使用实现双倍吞吐。这样虽然增加了一些控制逻辑但灵活性提升明显。第二累加器的位宽。累加器的位宽决定了能累加多少个乘积而不溢出。对于INT8乘法每个乘积最大是127*12716129如果用32位累加器最多可以累加约26万个乘积而不溢出。对于典型的卷积层这个范围足够了。如果担心溢出可以在累加过程中做饱和处理或者定期将部分和写回缓存。第三数据流动的方向。在脉动阵列中数据流动的方向决定了阵列的形状和连接方式。我选择了经典的权重固定、输入流动、部分和流动的方式权重在计算开始前加载到阵列中并保持不动输入数据从左侧流入部分和从上方流入、下方流出。这种方式适合卷积操作因为卷积核权重在同一个特征图上滑动时是固定的。3.3 阵列规模的权衡脉动阵列的规模即NxN中的N是一个需要仔细权衡的参数。N越大峰值算力越高但面积和功耗也越大而且对缓存容量的要求也越高。我做了简单的估算假设每个MAC单元的面积约为0.01平方毫米在28nm工艺下一个16x16的阵列就是256个MAC单元面积约2.56平方毫米。如果换成32x32就是1024个MAC单元面积约10.24平方毫米。对于边缘设备来说10平方毫米的计算核心已经相当大了加上缓存和控制逻辑整个芯片面积可能超过20平方毫米成本会明显上升。另一方面N越大数据复用的收益也越大。16x16的阵列每个权重元素可以被复用16次32x32的阵列复用次数翻倍到32次。但复用收益是递减的从16到32的提升并不足以抵消面积和功耗的增加。最终我选择了16x16的规模峰值算力在1GHz时钟下为256 * 2 * 1e9 512 GOPSINT8如果算上INT4模式的双倍吞吐可以达到1 TOPS左右。这个算力对于边缘推理来说已经相当充裕了。3.4 数据流控制与时序设计脉动阵列的时序设计是另一个关键点。数据必须在正确的时钟周期出现在正确的位置否则计算结果就会出错。这需要精确的控制逻辑。我采用了一种简单的波前控制方式在计算开始前先将权重加载到阵列中然后在每个时钟周期从左侧和上方分别注入输入数据和部分和经过N个周期后第一个输出结果从右下角流出再经过N个周期所有结果都流出。为了处理不同大小的矩阵我设计了分块计算的机制。如果矩阵的维度大于阵列的规模就将矩阵分成NxN的块逐块计算块与块之间的部分和通过输出缓存暂存和累加。这样无论矩阵多大都可以用固定规模的阵列来处理。注意脉动阵列的时序设计非常容易出错建议在RTL仿真阶段就用小规模的矩阵比如4x4做验证确保数据流动的时序完全正确后再扩展到大规模。4. 存储层次比计算更关键的战场4.1 为什么存储是AI加速器的真正瓶颈在AI加速器的设计中有一个经常被忽视的事实计算单元很少成为瓶颈存储才是。这不是说计算不重要而是说在大多数实际工作负载中数据搬运的时间和能耗远超过计算本身。有一个经典的估算在28nm工艺下一次32位浮点乘加的能耗约为3.7皮焦耳而从外部DDR读取32位数据的能耗约为640皮焦耳是计算的170多倍。即使是从片上缓存读取能耗也有几十皮焦耳。这意味着如果数据复用率不够高加速器的能效会被存储访问严重拖累。这就是为什么AI加速器的存储层次设计如此重要。目标很简单让数据尽可能在离计算单元最近的地方被复用减少对远距离存储的访问。4.2 三级存储层次的设计我设计了一个三级存储层次第一级是寄存器文件位于每个MAC单元内部。它暂存当前参与计算的数据和部分和访问延迟最低能耗最小但容量也最小。在我的设计中每个MAC单元有少量的寄存器用于暂存输入和部分和。第二级是片上缓存包括输入缓存、权重缓存和输出缓存。这一级用SRAM实现容量在几百KB级别访问延迟在几个时钟周期能耗远低于外部存储。输入缓存暂存从外部读取的特征图数据权重缓存暂存卷积核参数输出缓存暂存计算结果。第三级是外部存储器通常是DDR。容量大但访问延迟高、能耗大。设计的目标是尽量减少对这一级的访问。三级存储之间的数据搬运由数据搬运引擎控制。我采用了双缓冲机制当计算单元在处理当前数据块时数据搬运引擎在后台预取下一块数据到另一个缓冲区。这样计算和数据搬运可以重叠进行减少等待时间。4.3 缓存容量的计算与分配缓存容量的分配需要根据目标工作负载来定。我以典型的卷积层为例做了估算假设输入特征图是224x224x3比如MobileNet的第一层卷积核是3x3x3x32。如果每次处理一行输入224x3个元素需要的输入缓存约为22431字节672字节。权重缓存需要存储33332864个权重约864字节。输出缓存需要存储22432个输出约7168字节。但这只是单层的情况。实际上为了充分利用脉动阵列通常需要一次处理多个输入通道和输出通道。如果一次处理16个输入通道和16个输出通道输入缓存需要224163584字节权重缓存需要3316162304字节输出缓存需要224*163584字节。考虑到需要双缓冲以及支持更大的特征图和更多的通道数我将输入缓存和权重缓存各设为256KB输出缓存设为128KB。这个容量在边缘设备上是可以接受的SRAM的面积大约在几平方毫米级别。4.4 数据复用策略的优化存储层次设计好了接下来就是怎么最大化数据复用。我采用了以下几种策略权重复用权重在脉动阵列中保持不动被所有输入数据复用。这是脉动阵列的天然优势。输入复用同一块输入数据可能与多个不同的卷积核做卷积产生多个输出通道。我将输入数据缓存在输入缓存中供多个输出通道的计算复用。输出复用在分块计算中部分和需要在块与块之间累加。我将部分和暂存在输出缓存中避免频繁写回外部存储器。通道分组对于通道数很多的层我将通道分组每组内的数据在片上完成所有计算后再写回。这样可以减少中间结果的搬运。这些策略的组合使用可以将外部存储器的访问量降低一到两个数量级从而大幅提升能效比。5. 控制逻辑与指令集设计5.1 控制单元的角色控制单元是AI加速器的大脑负责解析指令、调度数据流、控制计算阵列的时序。它的设计目标是在保证功能正确的前提下尽量简单因为复杂的控制逻辑会消耗面积和功耗。我采用了一种微指令状态机的控制方式。微指令是一条条简单的控制信号组合比如从输入缓存读取一行数据、启动脉动阵列计算、将结果写入输出缓存等。状态机根据当前执行的阶段依次发出这些微指令。5.2 指令集的设计指令集是控制单元与软件之间的接口。我设计了一套简单的指令集包含以下几类指令配置指令设置计算模式INT8/INT4、矩阵维度、缓存地址等参数。数据搬运指令控制DMA在外部存储器和片上缓存之间搬运数据。计算指令启动脉动阵列执行矩阵乘法或卷积。同步指令等待某个操作完成或者设置屏障。指令的格式我设计得比较紧凑每条指令32位包含操作码、操作数和一些控制位。指令序列由编译器或手动编写存储在指令缓存中由控制单元依次读取和执行。5.3 卷积操作的映射卷积是AI加速器最常执行的操作之一。将卷积映射到脉动阵列上需要一些技巧。最常用的方法是im2col将输入特征图按照卷积核的感受野展开成矩阵然后做矩阵乘法。比如一个3x3的卷积核在5x5的输入上滑动可以展开成一个9x9的矩阵每个位置取9个元素然后与9x1的卷积核向量做矩阵乘法。但im2col会增加存储开销因为输入数据被重复存储了。另一种方法是直接做卷积让脉动阵列的输入数据按照卷积的滑动窗口方式流动。这种方式更节省存储但控制逻辑更复杂。我选择了im2col的方式因为它实现简单而且可以复用矩阵乘法的硬件。虽然存储开销大一些但通过合理的分块和缓存策略可以控制在可接受的范围内。5.4 控制流的优化控制流的优化目标是减少空闲周期让计算单元尽可能保持忙碌。我采用了以下几种方法流水线化将数据搬运、计算和结果写回分成三个阶段让它们重叠执行。当计算单元在处理第i块数据时DMA在搬运第i1块数据同时第i-1块的结果在写回。预取根据指令序列提前将下一块需要的数据从外部存储器预取到片上缓存。乱序执行对于没有数据依赖的指令允许乱序执行以提高并行度。但这会增加控制逻辑的复杂度我在第一版设计中暂时没有采用。提示控制逻辑的验证是设计中最耗时的部分之一。建议在RTL仿真中构造各种边界情况比如空矩阵、单元素矩阵、维度不匹配的矩阵等确保控制逻辑在所有情况下都能正确工作。6. 验证方法与常见问题排查6.1 功能验证的基本流程AI加速器的验证是一个系统工程需要从模块级到系统级逐步验证。模块级验证对每个模块MAC单元、缓存、DMA、控制单元单独验证。比如MAC单元需要验证INT8和INT4模式下的乘法结果是否正确累加器是否会溢出清零信号是否有效等。集成验证将各个模块连接起来验证它们之间的接口和时序。比如验证DMA能否正确地将数据写入缓存控制单元能否正确地启动计算阵列等。系统级验证在完整的加速器上运行真实的神经网络模型比如MobileNet或ResNet的一个层将结果与CPU上的参考实现对比确保数值一致。我使用Verilog编写RTL代码用Verilator做仿真用Python生成测试向量和参考结果。这种组合灵活且开源适合个人开发者。6.2 常见问题与排查技巧在设计过程中我遇到了不少问题这里整理成速查表问题现象可能原因排查方法解决方案计算结果全为零权重未正确加载检查权重加载时序确保权重在计算前已写入阵列计算结果部分错误数据流动时序错误用小规模矩阵仿真调整数据注入的时钟周期累加结果溢出累加器位宽不足检查累加器位宽和累加次数增加位宽或定期写回部分和能效比低于预期外部存储访问过多统计DDR访问次数优化数据复用策略增加缓存容量控制逻辑死锁状态机状态转换错误检查状态机的所有转换路径增加超时机制或修复状态转换仿真速度过慢测试向量过大减小测试规模先用小规模验证再逐步扩大6.3 性能评估与迭代设计完成后需要对性能做评估。我主要关注以下几个指标峰值算力利用率实际算力与峰值算力的比值。如果利用率低于50%说明数据供应或控制流有问题。能效比实际功耗与算力的比值。我通过仿真估算动态功耗包括计算单元、缓存和DMA的功耗。面积通过综合工具估算芯片面积确保在目标工艺下可接受。根据评估结果我做了几轮迭代第一轮优化了数据复用策略将DDR访问量降低了60%第二轮优化了控制流将空闲周期减少了30%第三轮调整了缓存容量在面积和性能之间找到了更好的平衡点。6.4 从设计到实现的注意事项如果你也想从零设计一款AI加速器有几个坑我想提前提醒第一不要一开始就追求大而全。先做一个能跑通的小规模设计比如4x4的阵列验证整个流程然后再逐步扩展。我见过太多人一开始就设计32x32的阵列结果在验证阶段卡住最后不了了之。第二存储层次的设计要尽早考虑。不要等到计算单元设计完了才想缓存的事。存储和计算是紧密耦合的需要协同设计。第三验证的重要性不亚于设计。一个没有经过充分验证的设计即使逻辑看起来正确也可能在实际运行中出错。建议把至少40%的时间花在验证上。第四善用开源工具和参考设计。有很多开源的NPU项目和论文可以参考比如NVDLA、Gemmini等。不要从零造轮子站在巨人的肩膀上效率更高。第五关注工艺和功耗。设计出来的东西最终要流片或部署到FPGA上工艺决定了面积和功耗的上限。在设计的早期阶段就要对目标工艺有清晰的认知。7. 一些个人体会与后续扩展方向做这个项目的过程中我最大的体会是AI加速器的设计本质上是一个数据搬运的问题而不是计算的问题。计算单元的设计相对直接无非是乘法器和加法器的组合但怎么让数据高效地流动、复用、暂存才是真正考验设计功力的地方。脉动阵列之所以经典就是因为它用简单的结构解决了数据复用的问题。另一个体会是软硬件协同设计非常重要。加速器的指令集、数据布局、算子映射方式都会影响最终的性能。如果只关注硬件而忽略软件设计出来的加速器可能很难用如果只关注软件而不了解硬件也无法充分发挥硬件的潜力。后续如果继续深入我会考虑几个方向一是支持稀疏计算利用神经网络的稀疏性进一步提升能效二是支持动态形状适应不同输入尺寸的模型三是探索存内计算等新型存储技术从根本上解决数据搬运的瓶颈四是完善编译器和工具链让开发者能更方便地把模型部署到加速器上。这个项目让我对计算机体系结构有了更深的理解也让我意识到很多看似复杂的技术拆解开来都是一些朴素的思想加上大量的工程细节。如果你也对这方面感兴趣建议从一个小规模的设计开始动手做一遍比看十篇文章都有用。