1. 从矩阵乘法说起为什么AI芯片需要脉动阵列搞AI芯片的人绕不开一个核心问题矩阵乘法怎么算得快、算得省。不管是卷积神经网络里的卷积操作还是Transformer里的注意力机制拆到底层全是矩阵乘加。一个典型的推理任务可能涉及几十亿甚至上百亿次乘加运算如果用通用CPU一个时钟周期一个时钟周期地串行算那基本没法用。我第一次接触脉动阵列这个概念是在研究TPU架构的时候。当时有个疑问GPU里动辄几千个CUDA核心并行度已经很高了为什么Google还要专门搞一个脉动阵列出来后来把数据流捋清楚才明白GPU的并行是空间上的并行——几千个核心同时算不同的数据但每个核心算完的结果要写回寄存器或共享内存下一次计算再从内存里读。这个读写过程消耗的能量远比一次乘加运算本身大得多。这里有个业内公认的数据在28nm工艺下一次32位浮点乘加运算大约消耗3.7pJ的能量而从DRAM读取一个32位数据要消耗640pJ从全局缓冲区读取也要100pJ左右。也就是说数据搬运的能耗是计算本身的几十倍甚至上百倍。所以AI芯片设计的核心矛盾从来不是算得够不够快而是数据搬得够不够少。脉动阵列Systolic Array就是冲着这个矛盾去的。它的核心思想非常朴素让数据像血液在血管里脉动一样有节奏地流过一个个计算单元每个计算单元只负责一次乘加算完就把部分和传给下一个单元而不是写回内存再读出来。数据在阵列内部流动的过程中被反复利用大幅减少了对外部存储的访问。这个思路最早可以追溯到1978年H.T. Kung提出的 systolic array 概念当时是用于信号处理和矩阵运算。但真正让它大放异彩的是Google在2016年发布的TPU v1。TPU v1里那颗256x256的脉动阵列峰值算力达到92 TOPS8位整数而功耗只有40W左右能效比远超同时期的GPU。从那以后脉动阵列就成了AI加速器设计里绕不开的一个选项。你可能会问既然脉动阵列这么好为什么不是所有AI芯片都用它这就涉及到它的适用边界了。脉动阵列最适合的是规则的大规模矩阵乘法尤其是权重固定的推理场景。一旦遇到稀疏矩阵、动态变化的计算图、或者需要频繁分支跳转的任务它的效率就会打折扣。所以理解脉动阵列不能只理解它怎么算更要理解它为什么这样算以及什么情况下不该这样算。2. 脉动阵列的数据流动机制拆解2.1 从一维到二维阵列的基本拓扑要理解脉动阵列先从最简单的一维情况入手。假设我们要算两个向量的点积a[0]*b[0] a[1]*b[1] a[2]*b[2] a[3]*b[3]。传统做法是把a和b都读到寄存器里然后逐个相乘再累加。脉动阵列的做法不同让a的元素从左向右依次流入b的元素从上向下依次流入每个计算单元PE在某个时刻接收一个a和一个b做一次乘法然后把结果累加到从上方传来的部分和里再传给下方。一维阵列只能算向量点积要算矩阵乘法就得扩展到二维。一个NxN的脉动阵列可以同时处理两个NxN矩阵的乘法。具体来说矩阵A的元素按行从左向右流入矩阵B的元素按列从上向下流入每个PE负责计算A的一行和B的一列对应位置的乘积并把部分和沿垂直方向累加。这里有个关键细节数据流入的节奏schedule决定了阵列的利用率。如果A和B的元素同时到达每个PE那PE在每个时钟周期都能做一次有效的乘加。但如果节奏没对齐PE就会空转。所以脉动阵列的设计里数据流的调度和阵列的拓扑结构是同等重要的。我见过一些初学者自己写脉动阵列的仿真最容易犯的错误就是忽略了数据流入的延迟。比如一个4x4的阵列A矩阵的第0行第0列元素在第0个周期进入PE(0,0)但B矩阵的第0列第0行元素也要在第0个周期进入PE(0,0)否则第一个PE就会等。而A的第0行第1列元素要到第1个周期才能进入PE(0,1)因为它在水平方向上要走一步。这种斜对角的数据排布是脉动阵列能高效运转的前提。2.2 权重固定 vs 输出固定两种主流数据流脉动阵列的数据流设计主要分两大流派权重固定Weight Stationary和输出固定Output Stationary。这两个概念在Google的TPU论文里被反复提及也是理解不同AI芯片架构差异的关键。权重固定的思路是把卷积核的权重预先加载到每个PE里然后让输入特征图的数据流过阵列每个PE用自己固定的权重去乘流过的数据部分和沿某个方向累加。这种方式的优势是权重只需要加载一次之后可以反复使用特别适合推理场景——因为推理时权重是固定的而输入数据是不断变化的。TPU v1用的就是权重固定的数据流。输出固定的思路则相反每个PE负责计算一个输出元素输入数据和权重都流过这个PEPE把所有的乘积累加起来最终输出一个完整的结果。这种方式适合输出维度较大、需要频繁写回的场景但在权重复用上不如权重固定高效。还有一种行固定Row Stationary的数据流是权重固定和输出固定的折中方案典型代表是Eyeriss架构。它让每一行的PE共享同一组权重同时让输入数据在行内复用在行间流动。这种设计在卷积神经网络里表现很好因为卷积操作天然具有行方向的局部性。数据流类型权重复用输入复用适用场景典型代表权重固定高中推理、权重固定的矩阵乘TPU v1输出固定低高输出维度大的场景部分GPU张量核心行固定中高卷积神经网络Eyeriss选哪种数据流取决于你的计算模式。如果是Transformer这类权重固定、输入变化的推理任务权重固定通常更优如果是训练场景权重和输入都在变可能就需要更灵活的数据流设计。2.3 部分和的累加路径与流水线设计脉动阵列里最精妙的部分是部分和的累加路径。在一个NxN的阵列里每个PE计算完一次乘积后需要把结果和来自上方PE的部分和相加然后传给下方PE。这个过程形成了一条垂直的累加链。这条链的长度是N意味着一个完整的输出结果需要经过N个PE的累加才能得到。如果每个PE的加法需要1个时钟周期那么从第一个乘积产生到最终结果输出需要N个周期的延迟。这个延迟在阵列规模较大时比如256x256会变得很可观但因为是流水线式的一旦流水线填满每个周期都能输出一个完整的结果。这里有个设计上的权衡阵列越大吞吐量越高但延迟也越大且利用率对数据排布的敏感度越高。TPU v1选择256x256是因为在当时的工艺和功耗约束下这个规模能在吞吐量和延迟之间取得较好的平衡。如果阵列太小比如16x16虽然延迟低但每个周期只能算256次乘加算力上不去如果阵列太大比如1024x1024延迟会达到上千周期而且芯片面积和功耗都会急剧上升。我在实际做架构评估时通常会用一个简单的公式来估算阵列的利用率利用率 有效计算周期数 / 总周期数。对于一个NxN的阵列如果输入矩阵的维度是MxK和KxN那么理想情况下需要MKN次乘加阵列每个周期能做NN次乘加所以需要的周期数是MK/N。但如果M或K不是N的整数倍就会有边缘浪费。比如M100N256那么最后156行的PE就会空转。这也是为什么很多AI芯片在软件层面会做padding把矩阵维度补齐到阵列规模的整数倍。3. 用Python仿真一个4x4脉动阵列光讲原理容易飘咱们直接上手写一个脉动阵列的仿真。我用Python写一个4x4的权重固定脉动阵列模拟矩阵乘法的过程。这个仿真不涉及硬件时序但能帮你把数据流动的每一个周期都看清楚。3.1 仿真框架与数据结构设计先定义PE的结构。每个PE需要存储一个权重值、一个部分和寄存器以及一个输出端口。为了模拟脉动我们用二维数组表示阵列每个元素是一个PE对象。class PE: def __init__(self): self.weight 0 self.partial_sum 0 self.input_val 0 def compute(self): # 每个周期部分和 权重 * 输入 self.partial_sum self.weight * self.input_val return self.partial_sum阵列的输入数据需要按斜对角的方式排布。假设我们要计算A4x4和B4x4的乘积A的行从左向右流入B的列从上向下流入。为了让A[i][k]和B[k][j]在PE(i,j)处相遇A的第i行第k列元素需要在第(ik)个周期进入阵列B的第k行第j列元素需要在第(kj)个周期进入阵列。def systolic_multiply(A, B, N4): # 初始化阵列 array [[PE() for _ in range(N)] for _ in range(N)] # 加载权重这里假设B是权重矩阵B[k][j]加载到PE(k,j) for k in range(N): for j in range(N): array[k][j].weight B[k][j] # 模拟周期 total_cycles 3 * N # 输入流入N周期计算N周期结果流出N周期 results [[0]*N for _ in range(N)] for cycle in range(total_cycles): # 为每个PE准备输入 for i in range(N): for j in range(N): # A的第i行第(cycle - i)列 k cycle - i if 0 k N: array[i][j].input_val A[i][k] else: array[i][j].input_val 0 # 每个PE计算 for i in range(N): for j in range(N): array[i][j].compute() # 收集输出当cycle N-1 j时PE(N-1, j)的部分和就是结果 for j in range(N): if cycle N - 1 j: results[N-1][j] array[N-1][j].partial_sum return results这段代码的核心逻辑是每个周期PE(i,j)从A的第i行取第(cycle-i)列的元素如果索引越界就补0。然后所有PE同时做一次乘加。经过N个周期后PE(N-1,j)的部分和就是最终结果的一列。3.2 逐周期追踪数据在阵列里怎么走跑一下这个仿真把每个周期的状态打印出来你就能直观看到数据是怎么脉动的。假设A和B都是简单的4x4矩阵元素值就是行列索引之和。A [[ij for j in range(4)] for i in range(4)] B [[i*j1 for j in range(4)] for i in range(4)] result systolic_multiply(A, B) print(result)在第一个周期cycle0只有PE(0,0)有有效输入A[0][0]和B[0][0]。其他PE的输入要么是0要么还没到。第二个周期PE(0,1)收到A[0][1]PE(1,0)收到A[1][0]同时PE(0,0)收到A[0][1]因为cycle-i1。你会发现数据像波浪一样从左上角向右下角推进。到第N-1个周期阵列被完全填满所有PE都在做有效计算。这是阵列利用率最高的阶段。之后输入数据逐渐流出部分和继续向下累加直到最后一个周期输出完整结果。这个仿真虽然简单但它揭示了一个重要事实脉动阵列的效率高度依赖于数据排布的精确性。如果A和B的流入节奏差了一个周期整个阵列的计算结果就会错位。在实际硬件里这个节奏是由控制器和FIFO缓冲区来保证的设计难度不小。3.3 仿真结果验证与常见错误排查跑完仿真后一定要和NumPy的矩阵乘法结果对比。我见过不少人写的脉动阵列仿真结果总是差一点排查半天发现是边界条件没处理好。import numpy as np A_np np.array(A) B_np np.array(B) expected A_np B_np print(Expected:\n, expected) print(Systolic:\n, np.array(result))常见的错误有这么几类一是索引越界没补0导致部分和里混入了垃圾数据二是权重加载位置搞反了把B[k][j]加载到了PE(j,k)三是输出收集时机不对在部分和还没累加完就取了结果。排查的时候建议先把阵列规模缩小到2x2手动算一遍每个周期的状态和仿真输出逐周期对比这样最容易定位问题。还有一个容易忽略的点部分和寄存器的位宽。在仿真里我们用Python的无限精度整数不会溢出。但在实际硬件里部分和的位宽是有限的。如果输入是8位整数权重也是8位乘积是16位累加N次后可能需要16log2(N)位。对于256x256的阵列部分和至少需要24位才能保证不溢出。这个位宽计算在硬件设计里是必须做的否则结果会莫名其妙地出错。4. 脉动阵列在真实AI芯片里的工程取舍4.1 TPU v1的256x256阵列为什么是这个规模Google在TPU v1里选择256x256的阵列规模不是拍脑袋决定的。从算力角度看256x256的阵列每个周期能做65536次乘加在700MHz时钟下峰值算力就是65536 * 2 * 700M ≈ 92 TOPS8位整数。这个算力在2016年是非常惊人的足以支撑当时主流的推理任务。但从工程角度看256x256意味着65536个PE每个PE至少需要一个乘法器、一个加法器、一个权重寄存器和一个部分和寄存器。在28nm工艺下这个面积大约占芯片的一半以上。如果阵列再大芯片面积和功耗就会失控如果再小算力又不够看。所以256x256是一个在算力、面积、功耗之间反复权衡后的结果。还有一个关键因素是内存带宽的匹配。TPU v1配备了24MB的片上统一缓冲区带宽达到34GB/s。这个带宽刚好能喂饱256x256的阵列——如果阵列再大缓冲区带宽就跟不上了PE会经常饿肚子。这种计算和存储的匹配设计是AI芯片架构里最考验功力的地方。4.2 稀疏化与结构化稀疏脉动阵列的软肋脉动阵列最大的软肋是对稀疏矩阵的处理效率低。在自然语言处理和推荐系统里很多权重矩阵是稀疏的零元素占比可能超过90%。但脉动阵列的PE不管权重是不是零每个周期都在做乘加。零乘以任何数还是零这些计算完全是浪费。为了解决这个问题业界提出了结构化稀疏的方案。比如NVIDIA的Ampere架构支持2:4的稀疏模式即每4个权重里最多2个非零。硬件可以在加载权重时跳过零元素从而把有效算力翻倍。但这种方式要求稀疏模式是结构化的、可预测的对于非结构化的随机稀疏效果有限。另一种思路是动态门控在PE里加一个判断逻辑如果权重或输入是零就跳过这次乘加直接传递部分和。但这会增加PE的复杂度和面积而且如果稀疏率不够高省下来的计算时间可能还不够弥补判断逻辑带来的开销。我在评估一个稀疏加速方案时通常会算一笔账稀疏带来的算力节省 vs 判断逻辑带来的面积和功耗增加只有当前者明显大于后者时这个方案才值得做。4.3 从推理到训练脉动阵列的适用边界脉动阵列在推理场景下表现优异但在训练场景下就没那么香了。训练需要反向传播权重要不断更新而且梯度计算涉及大量的转置和广播操作。脉动阵列的权重固定数据流在权重频繁变化的场景下需要不断重新加载权重这个开销很大。另外训练时的矩阵维度往往是动态变化的batch size、序列长度都可能变。脉动阵列对固定维度的矩阵乘法效率最高一旦维度变化边缘浪费就会增加。所以很多训练芯片会采用更灵活的架构比如可重构的PE阵列或者基于NoC的众核架构而不是纯粹的脉动阵列。但这不意味着脉动阵列不能用于训练。Google的TPU v2/v3就同时支持训练和推理它们通过双核脉动阵列和高带宽内存的组合在一定程度上解决了训练的需求。只是相比推理训练场景下脉动阵列的利用率会低一些需要软件层面做更多的优化。5. 自己动手评估一个脉动阵列设计5.1 算力、面积、功耗的快速估算方法如果你要评估一个脉动阵列的设计方案不用一上来就跑完整的仿真。先用几个简单的公式做快速估算能帮你快速筛掉不靠谱的方案。算力估算峰值算力 阵列规模^2 * 2 * 时钟频率。比如128x128的阵列在1GHz时钟下峰值算力 1281282*1G 32.7 TOPS。注意这里的2是因为一次乘加算两次操作。面积估算每个PE的面积大约是一个乘法器加一个加法器加若干寄存器的面积。在7nm工艺下一个8位乘法器大约0.001mm²加上寄存器和布线一个PE大约0.002-0.003mm²。128x128的阵列PE总面积大约32-49mm²。再加上缓冲区和控制器整个芯片面积可能在80-100mm²。功耗估算每个PE的动态功耗大约与时钟频率和电压的平方成正比。在0.8V电压、1GHz频率下一个PE的功耗大约0.1-0.2mW。128x128的阵列PE总功耗大约1.6-3.2W。加上内存和控制逻辑整个芯片功耗可能在5-10W。这些数字是粗略的但足够帮你判断一个设计是否在合理的范围内。如果算出来功耗几十瓦、面积几百平方毫米那基本可以判定这个方案在当前工艺下不可行。5.2 用仿真数据反推阵列利用率回到之前的Python仿真我们可以加一些统计代码计算阵列的利用率。利用率 有效乘加次数 / (阵列规模^2 * 总周期数)。def calculate_utilization(A, B, N4): total_cycles 3 * N effective_macs 0 total_macs N * N * total_cycles for cycle in range(total_cycles): for i in range(N): for j in range(N): k cycle - i if 0 k N: effective_macs 1 return effective_macs / total_macs对于4x4的阵列算出来的利用率大约是33%。这个数字看起来很低但这是因为矩阵太小阵列还没填满就结束了。如果矩阵维度远大于阵列规模利用率会趋近于100%。比如256x256的阵列处理1024x1024的矩阵利用率可以到90%以上。这个仿真告诉我们一个重要的工程经验脉动阵列适合处理大矩阵小矩阵用它是浪费。如果你要加速的模型里全是小矩阵乘法那脉动阵列可能不是最优选择反而是一些小规模的向量处理单元更合适。5.3 从仿真到RTL下一步该做什么Python仿真验证了数据流的正确性之后下一步通常是写RTL寄存器传输级代码用Verilog或VHDL实现一个可综合的脉动阵列。这一步的难点在于时序收敛和资源映射。时序收敛方面脉动阵列的关键路径通常是部分和的累加链。如果阵列规模大这条链会很长导致时钟频率上不去。解决办法是插入流水线寄存器把长链打断成多段。但插入寄存器会增加延迟需要重新调整数据排布的节奏。资源映射方面每个PE的乘法器和加法器需要映射到FPGA的DSP slice或ASIC的标准单元。在FPGA上DSP slice的数量是有限的比如Xilinx的UltraScale系列一个芯片可能有几千个DSP。如果阵列规模超过DSP数量就需要用LUT来搭乘法器但LUT乘法器的频率和功耗都不如DSP。我个人的经验是先在Python或C里把数据流和调度验证清楚再用HLS高层次综合快速生成一版RTL评估资源和时序。如果HLS的结果不理想再手动写RTL做优化。这样比一上来就手写RTL效率高得多也不容易在早期陷入时序调试的泥潭。6. 几个容易踩的坑和实战建议6.1 数据排布错位最常见的仿真bug前面提过脉动阵列的数据排布必须精确到周期。我在仿真时遇到最多的bug就是A矩阵和B矩阵的流入节奏差了一个周期。这种错误在结果上表现为整体偏移——比如结果矩阵的每一列都往右移了一位或者部分和里混入了上一轮的数据。排查这种问题最有效的方法是打印每个周期每个PE的输入和部分和然后手动核对前几个周期的数据。如果发现某个PE在应该收到有效数据的周期收到了0或者收到了不该收的数据那就是排布逻辑有问题。建议在仿真代码里加一个verbose模式把每个周期的状态输出到CSV文件用Excel或Python画成热力图一眼就能看出数据流动的波形对不对。6.2 部分和位宽不够导致的溢出这个问题在仿真阶段不容易发现因为Python的整数不会溢出。但一旦转到硬件部分和位宽不够就会导致结果完全错误。计算部分和位宽的公式是位宽 输入位宽 权重位宽 log2(累加次数)。比如输入8位、权重8位、累加256次部分和至少需要88824位。但实际设计中还要留一些余量。因为输入数据可能有符号乘积可能是负数累加过程中可能出现中间值超过最终值的情况。我通常会在理论位宽上加2-4位作为保护带。如果面积紧张可以考虑用饱和截断代替全精度累加但这会引入误差需要评估对模型精度的影响。6.3 阵列规模选择的经验法则选阵列规模不能只看算力。我总结了一个简单的经验法则阵列规模应该约等于你目标模型里最大矩阵维度的平方根。比如你的模型里最大的矩阵乘法是1024x1024那阵列规模选32x32或64x64就比较合适。如果选256x256大部分PE在大部分时间都会空转。另一个考虑因素是内存带宽。阵列每个周期需要从缓冲区读取2N个数据N个输入和N个权重如果缓冲区带宽不够PE就会饿肚子。所以阵列规模要和缓冲区带宽匹配。一个粗略的估算缓冲区带宽字节/周期 2 * N * 数据位宽 / 8。比如N128数据位宽8位那缓冲区带宽至少需要256字节/周期。在1GHz时钟下就是256GB/s。这个带宽要求不低需要在设计初期就考虑进去。6.4 什么时候不该用脉动阵列最后说一个反直觉的建议不是所有AI加速场景都适合脉动阵列。如果你的计算任务里有很多小矩阵乘法、稀疏矩阵、或者动态变化的计算图脉动阵列的效率可能还不如一个简单的向量处理器阵列。我见过一些项目为了追求架构先进性硬上脉动阵列结果因为模型里全是小算子阵列利用率不到20%实际性能还不如用GPU。所以选架构之前一定要先分析你的计算模式矩阵维度有多大稀疏度多高权重是否固定数据复用机会多不多这些问题的答案比架构本身的名气更重要。脉动阵列是一个优雅的设计但优雅不等于万能。理解它的适用边界比盲目崇拜它更有价值。