1. 从矩阵乘法说起为什么AI芯片偏偏看上了脉动阵列搞AI芯片的人绕不开一个词——矩阵乘法。不管你是做CNN卷积、Transformer注意力还是全连接层推理底层拆开来看八成以上的计算量都压在矩阵乘加这一件事上。问题在于矩阵乘法这东西计算量大、数据复用率高、对内存带宽极度饥渴。你用CPU去跑算力还没喂饱内存带宽先跪了你用GPU去跑虽然并行度高但功耗和面积又上去了。所以做专用AI加速器的人脑子里想的都是同一件事怎么用最少的功耗、最小的面积把矩阵乘法的吞吐量拉满。脉动阵列Systolic Array就是在这个背景下被反复拿出来讨论的方案。它最早不是为AI设计的上世纪七八十年代H.T. Kung提出这个概念的时候目标是用规整的PEProcessing Element阵列做高效的信号处理和矩阵运算。但有意思的是几十年后TPU把它推到了AI芯片舞台的中央。脉动阵列的核心思路用一句话概括让数据像心跳一样有节奏地在PE阵列中流动每个PE只做最简单的乘加数据从一边进、从另一边出中间不做全局广播不做复杂的地址计算。这种设计的好处是数据复用率极高控制逻辑极简功耗效率非常好。你可以把它想象成一条流水线每个工位只负责拧一颗螺丝零件从上一站传过来拧完传给下一站不需要每个工位都跑去仓库取零件。但脉动阵列也不是万能药。它的规整性意味着灵活性差遇到稀疏矩阵、非规则计算、动态shape的时候效率会打折扣。而且阵列尺寸一旦定下来就很难改这对算法快速迭代的AI领域来说是个不小的风险。所以理解脉动阵列不能只看它“快”还要看它“为什么快”、“在什么条件下快”、“什么情况下会变慢”。这篇文章我会从基本原理讲起拆解PE的工作机制、数据流的组织方式、权重 stationary 和 output stationary 的区别然后聊到实际芯片设计中的取舍包括阵列尺寸怎么定、数据位宽怎么选、如何应对稀疏性最后分享一些在仿真和RTL实现中容易踩的坑。适合做AI加速器架构、数字IC设计、以及想搞清楚TPU内部到底怎么跑的人。2. 脉动阵列的底层工作机制PE、数据流与节拍控制2.1 一个PE到底在做什么脉动阵列的基本单元叫PEProcessing Element。别看名字唬人它干的事极其简单一个乘法器加一个累加器再加几个寄存器。每个时钟周期PE从左边或上边接收数据做一次乘加把结果往右边或下边传。就这么简单。但正是这种简单让它能跑得很快。因为PE内部没有复杂的控制逻辑没有指令译码没有分支预测关键路径短时钟频率可以拉得很高。而且所有PE结构一模一样布局布线非常规整后端实现的时候面积和功耗都好控制。一个典型的PE有三个端口输入数据、输入权重、部分和输出。在权重stationary模式下权重提前加载到PE内部寄存器里输入数据从左侧流入部分和从上往下流动。每个周期PE执行partial_sum_out partial_sum_in (data_in * weight_reg)数据进来乘一下加上从上面来的部分和再往下面送。权重不动数据横向流动部分和纵向流动。这就是最经典的脉动阵列数据流。2.2 数据流动的“节拍感”从哪来脉动阵列之所以叫“脉动”是因为数据在阵列中的流动是有节奏的像心跳一样。每个PE在每个周期都做一次操作没有空闲没有等待。这种节奏感来自于数据的错位排列。举个例子假设你要算一个2x2矩阵乘以2x2矩阵。输入矩阵A的行从左向右流动矩阵B的列从上向下流动。为了让每个PE在正确的时刻拿到正确的数据你需要在数据进入阵列之前做skew也就是错开排列。第一行数据第0周期进第二行数据第1周期进第三行第2周期进以此类推。这样当数据流到某个PE的时候正好和从另一个方向流过来的权重相遇完成乘加。这个skew操作看起来简单但在硬件实现的时候很关键。如果skew的周期数算错了整个阵列的计算结果就全乱了。而且skew的深度和阵列尺寸直接相关NxN的阵列skew深度就是N-1个周期。2.3 权重stationary vs 输出stationary脉动阵列的数据流组织方式决定了它的适用场景。最常见的两种是权重stationary和输出stationary。权重stationary的意思是权重提前加载到PE里在整个计算过程中不动输入数据流动部分和流动。这种方式适合权重复用率高的场景比如卷积神经网络里同一个卷积核要在输入特征图上滑动很多次权重不变输入变。TPU用的就是这种方式。输出stationary的意思是部分和固定在PE里不动输入数据和权重都流动。这种方式适合输出复用率高的场景比如全连接层里输出神经元的累加结果需要保留很久。但这种方式对PE的寄存器压力更大因为部分和要一直存着。实际芯片设计里很少纯用一种模式通常是混合的。比如Google TPU的MXUMatrix Multiply Unit就是权重stationary但它在系统层面做了很多优化让数据复用最大化。数据流类型固定不动流动数据适用场景寄存器压力权重stationary权重输入、部分和卷积、权重复用高中输出stationary部分和输入、权重全连接、输出复用高高输入stationary输入权重、部分和输入复用高低2.4 为什么脉动阵列能省带宽脉动阵列最吸引人的地方是它的数据复用能力。在一个NxN的阵列里每个输入数据被复用N次每个权重被复用N次。这意味着你从内存里读一次数据就能在阵列里用N次。对于矩阵乘法这种计算密度极高的操作这能大幅降低对内存带宽的需求。举个例子一个256x256的脉动阵列做一次矩阵乘法输入数据只需要从内存读256次权重读256次但实际完成了256x256x256次乘加。计算量是内存访问量的几万倍。这就是为什么TPU能用相对较低的内存带宽跑出很高的吞吐量。但这里有个前提数据必须能规整地流入阵列。如果矩阵是稀疏的或者shape不规整很多PE就会空转复用率下降带宽优势就没了。所以脉动阵列的高效率是建立在“稠密、规整、大矩阵”这个假设上的。3. 阵列尺寸、位宽与稀疏性设计中的三个硬约束3.1 阵列尺寸怎么定不是越大越好做架构设计的时候第一个要回答的问题是阵列做多大128x128256x256还是512x512阵列越大理论峰值算力越高因为算力是N^2级别的。但阵列越大面积也越大功耗越高而且良率会下降。更重要的是阵列越大对数据供给的要求越高。如果内存带宽喂不饱大阵列就是浪费。实际设计里阵列尺寸的确定要综合考虑几个因素目标模型的矩阵维度、内存带宽、功耗预算、面积预算。比如TPU v1用的是256x256的MXU因为当时的目标模型主要是CNN卷积核和特征图的维度大概在这个量级。如果做的是Transformer加速器矩阵维度可能更大但也要考虑注意力矩阵的动态性。还有一个容易被忽略的点阵列尺寸和skew深度直接相关。NxN的阵列skew深度是N-1。如果N很大skew的寄存器开销也不小。而且skew逻辑会增加控制复杂度影响时序收敛。实操心得在架构探索阶段不要一上来就定死阵列尺寸。先用仿真模型跑一遍目标负载看看不同尺寸下的利用率和带宽需求再结合面积和功耗估算做取舍。我见过不少项目阵列做大了结果发现大部分时间利用率不到30%算力全浪费了。3.2 数据位宽8bit、16bit还是混合精度AI芯片的位宽选择直接影响到算力、功耗和精度。早期TPU用的是8bit整数因为推理场景对精度要求不高8bit够用而且面积和功耗都小。后来训练场景需要更高的精度16bit浮点和bfloat16开始流行。脉动阵列的位宽设计有个特点乘法器的面积和位宽的平方成正比。8bit乘法器面积大概是16bit的四分之一。所以位宽翻倍算力不一定翻倍但面积肯定翻倍。这就是为什么很多AI芯片主打8bit或4bit推理用低位宽换高吞吐。但低位宽有个问题累加器的位宽不能太低。即使输入是8bit累加结果也可能很大需要32bit甚至更宽的累加器。所以PE内部通常是“窄乘法、宽累加”的结构。乘法器做8x8累加器做32bit这样既能保证精度又能控制面积。混合精度是另一个趋势。有些芯片支持多种位宽模式比如8bit模式下阵列跑满16bit模式下阵列拆成两半用。这种灵活性对算法迭代很友好但控制逻辑会复杂不少。位宽模式乘法器面积典型算力适用场景精度风险4bit最小最高量化推理高8bit小高主流推理中16bit中中训练、高精度推理低32bit大低科学计算极低3.3 稀疏性脉动阵列的“天敌”脉动阵列最怕稀疏矩阵。因为它的设计前提是每个PE每个周期都有活干一旦矩阵稀疏很多PE就在做无用功乘零加零功耗照跑算力浪费。稀疏性在AI模型里其实很常见。剪枝后的模型、ReLU后的激活、注意力矩阵都有大量零。如果直接丢进稠密脉动阵列效率会掉得很厉害。有研究显示50%稀疏度的矩阵在稠密阵列上利用率可能只有50%甚至更低。应对稀疏性有几种思路。一种是硬件上支持稀疏跳过PE检测到零输入就跳过计算省功耗。但这需要额外的控制逻辑而且会破坏脉动阵列的规整性。另一种是在数据流层面做压缩把稀疏矩阵压缩成稠密格式再送进阵列但解压缩的开销也不小。实际产品里纯稀疏加速器很少见更多是在稠密阵列基础上做一定程度的稀疏优化。比如支持结构化稀疏2:4稀疏每四个元素里有两个零硬件可以固定跳过。这种方式对阵列改动小收益也比较确定。4. 从RTL到系统脉动阵列落地时的真实挑战4.1 数据供给阵列再快喂不饱也是白搭脉动阵列的算力是N^2级别的但数据供给是N级别的。这意味着阵列越大对数据带宽的要求越高。一个256x256的阵列每个周期需要256个输入数据和256个权重如果时钟跑1GHz带宽需求就是256x2x4Bytex1GHz大概2TB/s。这个带宽在芯片内部还能靠SRAM撑一撑但要从DRAM读基本不可能。所以实际设计里脉动阵列前面通常有一级甚至多级缓存。权重先加载到片上SRAM输入数据也先缓存在SRAM里然后以阵列需要的节奏喂进去。这个缓存的设计很关键缓存太小阵列会饿缓存太大面积和功耗又上去了。还有一个细节权重的加载时间。权重stationary模式下权重加载是一次性的但加载本身需要时间。如果矩阵很小加载时间可能比计算时间还长这时候阵列利用率就很低。所以有些设计会做权重预加载或者支持权重流式加载减少等待。4.2 部分和的累加与溢出处理脉动阵列的部分和是纵向流动的从阵列顶部流到底部最后输出。这个过程中部分和的位宽会逐渐增加。如果累加器位宽不够就会溢出结果就错了。假设输入是8bit权重是8bit乘积是16bit。一个256x256的阵列最底下的PE要累加256个乘积结果可能达到16824bit。所以累加器至少要24bit通常做32bit留余量。但32bit累加器面积不小每个PE都要一个256x256就是65536个面积很可观。有些设计会用分段累加比如每16个PE做一次局部累加然后再把局部结果加起来。这样每个PE的累加器可以小一点但需要额外的加法树。取舍点在于面积和时序的平衡。注意部分和的溢出是静默错误不会报错但结果全错。仿真的时候一定要用边界数据测比如全1输入、全最大值输入看看累加器会不会溢出。4.3 控制逻辑简单不等于没有脉动阵列的控制逻辑确实比CPU/GPU简单得多但也不是没有。你需要控制权重的加载、数据的skew、阵列的启动和排空、部分和的输出。这些控制信号要精确到周期错一个周期结果就错。特别是排空阶段。阵列算完之后部分和还在阵列里流动需要额外的时间排空。排空的周期数也是N-1。如果控制逻辑没算对排空时间最后几个结果就丢了。还有一个容易踩的坑阵列的启动和停止。启动的时候数据要逐渐填满阵列前N-1个周期阵列是没满的算力在爬坡。停止的时候阵列逐渐排空后N-1个周期算力在下降。对于大阵列这个爬坡和排空的开销占比不大但对于小阵列可能占很大比例。所以小阵列适合小矩阵大阵列适合大矩阵不能一概而论。4.4 时序收敛与物理实现脉动阵列的规整性对后端实现是好事但也不是没有挑战。最大的挑战是时钟树。NxN的阵列时钟要送到每个PE时钟树很长skew很难控制。如果时钟skew太大PE之间的数据传递就会出错。另一个挑战是布线拥塞。PE之间的数据线是横向和纵向的如果阵列很大布线资源会很紧张。特别是部分和的纵向连线每个PE都要往下传线宽和间距都要仔细规划。实际项目中脉动阵列的物理实现通常会用定制布局PE做硬核阵列做规整排列电源和时钟网格专门设计。这也是为什么脉动阵列适合ASIC不太适合FPGA。FPGA上做脉动阵列布线资源和时钟资源都受限效率会打折扣。5. 脉动阵列在AI芯片中的实际表现与边界5.1 TPU的实战数据说明了什么Google TPU v1的MXU是256x256的8bit脉动阵列峰值算力92 TOPS功耗大概40W。这个能效比在当时是非常惊人的。但实际跑模型的时候利用率并不是100%。根据公开资料TPU v1在跑CNN推理的时候MXU利用率大概在20%到60%之间取决于模型和batch size。为什么利用率上不去主要原因是数据供给和模型shape的匹配问题。如果卷积核太小或者特征图通道数不是256的倍数阵列就会有空转。还有batch size太小的时候矩阵维度不够阵列填不满。TPU v2之后做了很多改进比如支持更大的batch、更好的数据复用、更高的带宽。但脉动阵列的基本架构没变说明这个方向是对的只是需要系统层面的配合。5.2 脉动阵列不适合什么脉动阵列不是万能的。以下几种情况它表现不好稀疏矩阵前面说过稀疏性会让PE空转效率骤降。动态shape脉动阵列的skew和控制逻辑是固定的如果矩阵维度动态变化需要额外的控制逻辑来适配灵活性差。小矩阵矩阵太小阵列填不满启动和排空开销占比高。非矩阵运算脉动阵列是为矩阵乘加设计的遇到激活函数、归一化、池化这些操作需要额外的硬件单元。所以实际AI芯片里脉动阵列通常只是其中一个模块周围还有向量单元、标量单元、特殊函数单元。脉动阵列负责最重的矩阵乘法其他操作交给别的模块。5.3 和其他架构的对比和脉动阵列竞争的主要是SIMD/SIMT架构比如GPU和可重构架构比如FPGA。GPU的优点是灵活什么都能算但能效比不如专用阵列。FPGA的优点是可重构适合算法快速迭代但峰值算力低功耗高。脉动阵列的优点是能效比高适合稠密矩阵乘法但灵活性差。实际选择的时候要看目标场景。如果是云端推理模型固定、batch大、追求能效比脉动阵列是很好的选择。如果是边缘设备模型多变、功耗受限可能需要更灵活的架构。如果是训练对精度和灵活性要求高GPU可能更合适。架构类型能效比灵活性适用场景典型代表脉动阵列高低稠密矩阵推理TPUSIMT中高训练、通用计算GPU可重构低高算法迭代、小批量FPGA存内计算极高低低功耗推理新型AI芯片6. 仿真与验证脉动阵列开发中的踩坑记录6.1 功能仿真skew逻辑最容易出错我第一次做脉动阵列仿真的时候skew逻辑写错了结果整个阵列的输出全是乱的。排查了半天发现是skew的周期数算错了。NxN的阵列第i行数据应该在第i个周期进入我写成了第i1个周期导致所有数据都错位了一拍。这个坑很典型。skew逻辑看起来简单但周期数、数据排列、控制信号的配合任何一个细节错了结果就不对。而且这种错误在仿真波形上不容易看出来因为数据都在动只是错位了。最好的排查方法是用小矩阵比如2x2做仿真手动算一遍预期结果然后对比波形。如果2x2对了再放大到4x4、8x8逐步验证。6.2 时序仿真部分和的建立时间部分和的纵向流动是组合逻辑路径从上面的PE传到下面的PE中间可能经过多个PE。如果阵列很大这条路径会很长时序可能不收敛。解决办法是在PE之间插入寄存器把部分和的传递做成流水线。但这样会增加延迟而且控制逻辑要相应调整。另一种办法是分段累加每几个PE做一次寄存器打拍减少关键路径长度。实际项目中时序收敛通常要迭代好几轮。建议在架构设计阶段就考虑时序不要等到RTL写完再改。比如阵列尺寸不要一味求大位宽不要一味求宽给时序留余量。6.3 验证覆盖率怎么确保测全了脉动阵列的验证最难的是覆盖率。因为数据流是并行的状态空间很大。你不能像验证CPU那样跑几个指令序列就完了。你需要验证各种数据模式、各种矩阵维度、各种边界条件。我的经验是验证要分层次。先验证单个PE的功能再验证一行PE的联动再验证整个阵列。数据模式要覆盖全零、全一、最大值、最小值、随机值、稀疏值。矩阵维度要覆盖等于阵列尺寸、小于阵列尺寸、大于阵列尺寸、非方阵。还有一个容易漏的点阵列的启动和排空。很多验证只测稳态不测启动和排空结果流片后发现边界情况出错。启动和排空的测试用例一定要单独写确保控制逻辑正确。6.4 功耗仿真动态功耗的估算脉动阵列的功耗主要来自动态功耗也就是PE的翻转。每个周期PE的乘法器和累加器都在翻转功耗和阵列尺寸、时钟频率、数据翻转率都相关。估算功耗的时候不能只看峰值。实际跑模型的时候数据翻转率是变化的。比如全零输入的时候翻转率很低功耗也低。随机输入的时候翻转率大概50%功耗最高。所以功耗估算要用实际负载的翻转率不能用理论峰值。还有一个经验脉动阵列的时钟树功耗占比很高。大阵列的时钟树可能占整个阵列功耗的30%以上。所以时钟门控很重要空闲的PE要把时钟关掉省功耗。7. 写在最后一些个人体会脉动阵列这个东西原理不复杂但做好不容易。它的优势在于规整和高效劣势在于死板和局限。做AI芯片架构不能因为TPU用了脉动阵列就盲目跟风要看自己的目标场景和数据特征。我在实际项目里最大的体会是脉动阵列的效率不取决于阵列本身而取决于数据供给和系统配合。阵列做得再大数据喂不饱就是浪费。所以架构设计要从系统层面考虑内存、缓存、数据流、控制逻辑每一环都要匹配。还有一个建议如果你刚开始做脉动阵列先用Python或C写一个周期精确的仿真模型把数据流和控制逻辑跑通再写RTL。这样能省很多调试时间。RTL仿真慢波形难看先用高级语言验证逻辑效率高得多。最后脉动阵列只是AI芯片的一种选择不是唯一选择。存内计算、近存计算、可重构架构都在快速发展。保持开放心态根据实际需求选方案比迷信某一种架构更重要。