首页 / 资讯中心 / 文章详情

AI芯片脉动阵列原理与实战:从数据复用到4x4时序推演

AI芯片脉动阵列原理与实战:从数据复用到4x4时序推演 ★ FEATURED ARTICLE
1. 为什么一颗AI芯片里最值钱的往往不是算力数字我第一次真正把脉动阵列Systolic Array这四个字和实际芯片对上号是在看一份TPU的架构资料时。当时有个疑问一直绕不过去为什么谷歌第一代TPU峰值才92 TOPS却敢说在推理任务上吊打同时期一堆几百TOPS的GPU后来把数据流捋清楚才明白决定AI芯片效率的从来不是那个最大的数字而是数据在计算单元之间怎么流动、搬了多少次、搬了多远。脉动阵列就是干这件事的。它本质上是一种让数据像心跳一样有节奏地在计算单元阵列中流动的硬件结构。名字里的脉动systolic来自心脏泵血——数据从边缘输入一拍一拍地往内部推进每个计算单元在每个时钟周期做一次乘加然后把部分结果传给下一个单元。整个过程没有大规模的数据回读权重可以预先驻留在阵列里激活值像水流一样穿过。这篇文章我想聊的不是教科书上那段定义而是作为一个实际接触过AI加速器设计的人把脉动阵列的基本原理、为什么这么设计、怎么算尺寸、怎么处理边界、实际踩过哪些坑讲透。适合正在学AI芯片架构的学生、想搞懂TPU底层逻辑的工程师以及做边缘推理芯片选型的从业者。看完你至少能自己动手推一遍4x4脉动阵列的时序知道为什么它省带宽、为什么它怕稀疏、为什么现代芯片又开始给它加各种补丁。2. 脉动阵列到底解决什么问题从冯诺依曼瓶颈说起2.1 传统计算架构在矩阵乘法上的三个致命伤要理解脉动阵列得先知道它要对抗什么。AI推理和训练里超过80%的运算量集中在矩阵乘法和卷积上而这两者本质都是乘加运算的堆叠。用最朴素的CPU思路做矩阵乘法问题立刻暴露访存墙每次乘加都要从内存取两个操作数、写回一个结果。算一个CA×BA和B的每个元素会被反复读取N次内存带宽瞬间成为瓶颈。算力再高喂不饱也是白搭。功耗墙数据搬运的能耗远高于计算本身。在28nm工艺下一次32位浮点乘加大约消耗几十皮焦而从DRAM取一个数可能消耗几百甚至上千皮焦差了一到两个数量级。芯片功耗大头其实花在搬数据上。面积墙如果给每个乘法器都配独立的取数通路互连线会占据大量芯片面积布线拥塞直接拖垮频率。我见过不少初版AI加速器设计算力标得很漂亮实测利用率不到20%问题几乎都出在数据供给上。脉动阵列的核心价值就是用规则的数据复用把访存次数压到最低。2.2 数据复用脉动阵列的省钱逻辑矩阵乘法C[i][j] Σ A[i][k] × B[k][j]这里存在天然的复用机会A的同一行元素会被C的整行复用B的同一列元素会被C的整列复用。脉动阵列把这个复用做到了硬件层面——每个数据一旦进入阵列就会被沿途的多个计算单元共享。举个直观的例子。一个N×N的脉动阵列做N×N矩阵乘法理想情况下每个A元素从左侧进入向右流动被同一行的N个PE依次使用每个B元素从上方进入向下流动被同一列的N个PE依次使用每个PE只做一次乘加把部分和向下传递。结果是每个输入数据只从外部存储器读取一次就完成了N次复用。相比朴素实现访存次数直接降到1/N量级。这就是为什么脉动阵列在能效上能拉开差距——它把搬数据这件事的边际成本摊薄了。2.3 和SIMD、GPU的路线差异有人会问GPU不也是大量并行计算单元吗区别在于数据流组织方式。GPU走的是SIMT路线靠大量线程和共享内存做复用灵活但控制开销大寄存器文件和共享内存的读写能耗不低。脉动阵列走的是空间数据流路线控制逻辑极简几乎不需要指令译码数据流动本身就是程序。代价是灵活性差。脉动阵列天生适合稠密、规则的矩阵运算遇到稀疏、不规则的计算就抓瞎。这也是为什么后来的AI芯片大多采用脉动阵列其他单元的混合架构而不是纯脉动阵列打天下。3. 脉动阵列基本原理拆解一拍一拍把数据推过去3.1 一个PE里到底有什么脉动阵列的基本单元叫PEProcessing Element处理单元。别看名字唬人一个最基础的PE内部结构简单到令人发指一个乘法器一个加法器累加器若干寄存器用来暂存从左边和上边来的数据几条数据通路把结果往右下传。它每个时钟周期就干一件事把左边来的激活值乘以从上边来的权重加到自己的累加器里然后把激活值往右传、权重往下传。没有分支预测没有乱序执行没有缓存一致性协议。极简所以能效高。3.2 权重固定Weight Stationary数据流脉动阵列有多种数据流变体最常见的是权重固定。逻辑是这样的权重B[k][j]预先加载到对应PE里在整个计算过程中不动。激活值A[i][k]从左侧按行输入每拍向右推进一格。部分和从上方往下流动逐列累加。以4×4阵列算4×4矩阵乘法为例时序大致是时钟周期输入动作阵列状态T0加载权重B到各PE权重就位T1A[0][0]从左侧进入第0行只有PE(0,0)有效T2A[0][1]进入A[0][0]右移PE(0,0)、PE(0,1)工作T3A[0][2]进入前两个右移第0行前三个PE工作T4第0行填满A[1][0]开始进入第1行阵列逐渐被填满T7阵列满载所有PE都在乘加T8起结果从底部逐拍流出流水线排空可以看到阵列有个**填充fill和排空drain**的过程。对于N×N阵列算N×N矩阵总周期数约为3N-2拍其中只有N拍是满载的。这就是脉动阵列的固有开销——阵列越大填充排空占比越小效率越高但大阵列的利用率和灵活性又会下降这是个需要权衡的点。3.3 为什么是脉动而不是广播有人会想为什么不直接把A的整行广播给所有PE非要让它一拍一拍地传原因有两个第一广播的扇出fanout太大。一行N个PE同时接收同一个数据驱动能力要求高互连线延迟和功耗都会上去。脉动传递每次只传给相邻PE扇出为1布线规整频率容易做高。第二脉动传递天然对齐了计算时序。A[i][k]到达PE(i,j)的时刻正好是它需要和B[k][j]相乘的时刻不需要额外的同步逻辑。这种数据流动即调度的设计把控制复杂度降到了最低。实操心得很多初学者画脉动阵列时序图时容易忽略填充阶段。实际做性能建模时一定要把3N-2这个周期数算进去否则算出来的利用率和实测会差一大截。4. 尺寸怎么定从矩阵规模反推阵列大小4.1 阵列尺寸与矩阵维度的匹配关系脉动阵列的尺寸不是随便定的它直接决定了能一次处理多大的矩阵块。假设阵列是N×N那么它一次能处理的最大矩阵乘法是N×N乘N×N。如果实际矩阵更大就要做分块tiling把大矩阵切成N×N的小块分多次喂进阵列。这里有个关键参数叫阵列利用率。对于M×K乘K×N的矩阵乘法用N×N阵列处理需要的周期数约为总周期 ≈ (M/N) × (K/N) × (N N - 1) 填充排空开销当M、K、N都远大于阵列尺寸时填充排空可以忽略利用率趋近于1。但如果矩阵维度不能被阵列尺寸整除就会产生尾部浪费——最后一块没填满PE空转。4.2 一个具体的尺寸计算案例假设你要设计一个推理芯片目标模型的主要卷积层是224×224×64的输入3×3卷积核输出64通道。把卷积展开成矩阵乘法后大致是M224×22450176K3×3×64576N64。如果阵列选32×32M方向需要50176/32≈1568块K方向需要576/3218块N方向需要64/322块每块计算周期约3232-163拍总周期≈1568×18×2×63≈3.56M拍。如果阵列选64×64M方向784块K方向9块N方向1块每块周期约127拍总周期≈784×9×1×127≈0.90M拍。看起来大阵列快很多但64×64阵列的PE数量是32×32的4倍面积和功耗也上去了。而且N64时64×64阵列在N方向只用了一半存在浪费。实际选型要在吞吐、面积、利用率之间找平衡通常会让阵列尺寸和主流模型的通道数对齐比如64、128、256这些2的幂次。4.3 分块策略与数据复用层级大矩阵切块后数据复用分几个层级阵列内复用PE之间的脉动传递这是最省的一层片上缓存复用同一块权重被多个激活块复用权重可以缓存在片上不用反复从DRAM取DRAM复用最外层尽量让每个数据从DRAM只取一次。好的脉动阵列设计会把权重固定在片上让激活值流过这样权重带宽需求极低。这也是为什么脉动阵列特别适合权重密集型的推理场景——权重加载一次可以服务大量输入。注意分块时K方向的切分要特别小心。如果K方向切得太碎每次都要重新加载部分和累加开销会吃掉收益。一般建议K方向的块尽量大让部分和在片上累加完再写回。5. 实操手推一个4×4脉动阵列的完整时序5.1 设定与初始状态光看原理容易飘我们动手推一遍。目标用4×4脉动阵列计算CA×B其中A和B都是4×4矩阵。约定A的元素从左侧按行输入第i行的元素A[i][0..3]依次从第i行左端进入B的元素从上方按列输入第j列的元素B[0..3][j]依次从第j列顶端进入每个PE(i,j)计算A[i][k]×B[k][j]的累加部分和向下传递权重固定方案B预先加载A脉动流动。为简化假设B已经加载完毕我们只看A的流动和计算。5.2 逐拍时序推演设A矩阵为A [a00 a01 a02 a03] [a10 a11 a12 a13] [a20 a21 a22 a23] [a30 a31 a32 a33]B矩阵为B [b00 b01 b02 b03] [b10 b11 b12 b13] [b20 b21 b22 b23] [b30 b31 b32 b33]时序如下只列关键拍T1a00从第0行左端进入到达PE(0,0)。PE(0,0)计算a00×b00存入累加器。T2a00右移到PE(0,1)a01进入PE(0,0)。PE(0,1)计算a00×b10注意这里b的下标因为PE(0,1)固定的是B的第1列元素b10PE(0,0)计算a01×b00。T3a00到PE(0,2)a01到PE(0,1)a02进入PE(0,0)。三个PE同时工作。T4第0行填满a00到PE(0,3)同时a10从第1行左端进入PE(1,0)。此时第0行4个PE全在工作第1行开始填充。T5-T7第1、2、3行依次填满。到T7整个4×4阵列满载16个PE同时做乘加。T8a00到达PE(0,3)后完成它的使命开始从右侧流出。同时第0行的部分和开始向下传递到第1行对应PE。T9-T10部分和逐行向下累加最终从第3行底部流出形成C的各列结果。T11最后一个结果流出计算完成。总周期数填充4拍满载4拍排空3拍≈11拍符合3N-210的估算考虑边界略有出入。5.3 从时序看效率这个例子里4×4阵列算4×4矩阵理论最少需要4拍如果无限并行实际用了约11拍效率约36%。如果换成32×32阵列算32×32矩阵总周期约95拍效率约34%。阵列越大填充排空占比越小但单次计算的绝对开销也越大。真正提升效率的办法是连续计算多个矩阵块让流水线一直保持满载填充排空只在最开始和最后发生一次。这就是为什么实际芯片都采用批量处理策略把多个输入拼成一个大batch让阵列持续工作。实操心得做性能建模时我习惯把有效计算周期/总周期作为核心指标。如果这个比值低于70%说明分块策略或batch size有问题需要调整。6. 常见问题与排查脉动阵列落地时的那些坑6.1 稀疏矩阵为什么让脉动阵列很难受脉动阵列的设计前提是稠密、规则的计算。一旦遇到稀疏矩阵比如剪枝后的模型大量PE会因为对应的权重或激活为零而空转但时钟还在走功耗照耗。这就是所谓的稀疏利用率问题。解决办法有几种结构化稀疏强制剪枝成规则模式比如每4个权重保留2个硬件可以跳过零块动态门控检测到零输入时关闭对应PE的时钟省动态功耗混合架构稀疏部分交给专门的稀疏引擎稠密部分交给脉动阵列。我实测过非结构化稀疏在纯脉动阵列上收益很有限甚至因为索引开销反而变慢。结构化稀疏才是和脉动阵列匹配的方案。6.2 边界处理与padding的坑卷积运算涉及边界padding展开成矩阵乘法后边界块往往填不满。如果处理不当会出现阵列部分PE处理无效数据浪费周期部分和累加时把padding的零也算进去虽然结果对但功耗白费分块边界不对齐导致额外的数据搬运。我的经验是在数据预处理阶段就把padding做掉让进入阵列的数据都是有效数据边界块用掩码控制哪些PE参与计算。这样虽然增加了一点控制逻辑但避免了阵列空转。6.3 常见问题速查表问题现象可能原因排查方向阵列利用率低于50%分块太小或batch不足增大batch调整tiling功耗异常高PE空转或时钟未门控检查稀疏处理和门控逻辑结果累加错误部分和传递时序错位核对填充排空周期数频率上不去互连线延迟或扇出过大检查脉动传递的布线权重带宽瓶颈权重未固定或反复加载确认权重驻留策略6.4 一个真实的调试案例之前调一个卷积加速器发现实测吞吐只有理论值的40%。查了半天最后定位到权重加载策略原本设计是每个输入块都重新加载权重导致权重带宽成了瓶颈。改成权重常驻片上、只加载一次后吞吐直接翻倍。这个坑很典型——脉动阵列的威力在于数据复用如果复用没做起来它就退化成一个普通的乘法器阵列。7. 脉动阵列的演进与适用边界7.1 从TPU到现代AI芯片的演化谷歌TPU v1是脉动阵列最著名的落地案例256×256的阵列权重固定数据流专为推理设计。之后的TPU v2/v3加入了训练能力阵列结构也有调整支持浮点和更灵活的数据流。现代AI芯片里脉动阵列往往不是唯一单元而是和向量单元、标量单元、专用激活函数单元组合使用。演进方向主要有几个支持多种数据流权重固定、输出固定、行固定可切换适应不同层类型混合精度INT8、FP16、BF16甚至FP8不同精度用不同阵列配置稀疏加速加入零跳过逻辑提升稀疏模型效率可重构阵列可以动态划分成多个小阵列处理不同规模的运算。7.2 什么场景适合什么场景别硬上脉动阵列的甜点区是稠密、规则、权重可复用的计算大矩阵乘法标准卷积全连接层。不适合的场景稀疏、不规则的计算图神经网络、稀疏注意力控制流复杂的任务小规模、低并行的运算阵列填不满效率极低。选型时我的建议是先看模型的计算模式再决定要不要用脉动阵列。如果模型以稠密卷积和矩阵乘法为主脉动阵列几乎是能效最优解如果模型稀疏度高、结构不规则硬上脉动阵列只会事倍功半。7.3 给初学者的学习路径如果你想深入脉动阵列我建议按这个顺序走先把矩阵乘法的数据复用逻辑搞清楚理解为什么访存是瓶颈手推一个4×4脉动阵列的时序画清楚每一拍的数据位置用Python或Verilog写一个简单的脉动阵列模型跑通矩阵乘法研究TPU论文看工业级设计怎么处理分块、精度、稀疏尝试做性能建模算清楚利用率、带宽、功耗的账。最后分享一个我自己常用的验证方法用一个小矩阵比如8×8手工算一遍结果再和脉动阵列模型的输出对比。如果对不上八成是时序对齐或者部分和传递出了问题。这个方法帮我抓过好几次隐蔽的时序bug比看波形图还快。
阅读完成 · 觉得有帮助?
咨询建站