脉动阵列原理与Python仿真:从矩阵乘法到数据流硬件设计 简介这份Python演示项目面向对硬件加速器、脉动阵列架构感兴趣的开发者通过软件模拟帮助理解TPU等芯片中矩阵乘法加速的核心原理。资源为zip压缩包共2个文件包含1个Python脚本模拟脉动阵列的数据流与运算逻辑和1个Markdown说明文档梳理设计思路与使用方法整体仅3KB轻量易读。包内脚本演示了数据如何在二维计算单元间流动、流水线式乘累加操作以及并行调度逻辑适合作为学习高性能计算与硬件仿真的入门示例。已有442人浏览下载适合想从零了解张量处理单元内部机制或为后续硬件设计打基础的读者。1. 脉动阵列为什么一份Python仿真比硬件框图更能讲清它的底层机制脉动阵列这个名词最劝退的不是数学而是“脉动”这两个字。第一次看到加速器内部框图你面对的是一排排整齐的格子每根连线都在做数据的搬运但整张图是静态的。你会觉得它像某种黑匣子知道它能算矩阵乘法却说不清数据到底怎么流成了结果。这份用Python写成的SystolicArrayDemo做了一件硬件框图做不到的事——它把“脉动”拆成了一拍一拍的寄存器快照每过一个时钟周期你都能看到哪个数据进入了哪个处理单元哪个部分和被累加到了哪一行。搞懂脉动阵列不需要先把整本体系结构啃完但需要回答这样几个问题为什么每个处理单元不从内存里直接取数为什么数据要错开一拍才能算对为什么芯片内部宁可铺一大堆乘法器也不愿意把乘法交给一条简单循环处理这些问题在阅读文字时很容易被跳过落到一份逐时钟仿真的代码里就成了纸面遮不住的具体行为。你能直接观察到数据从左向右移动、部分和从上向下累积看到一个矩阵乘法怎样被“时间化”成一个二维网格上的接力。把这条链路走通之后再回去看任何加速器文档里的Wiring图那些格子就都动起来了。2. 从三重循环到一张网格矩阵乘法如何在脉动阵列里被重新编排先用最朴素的视角看矩阵乘法。CA×B任何一行输出都来自A的某一行与B的某一列做内积。软件里的写法是三重循环外层遍历输出的行和列内层遍历求和的下标。脉动阵列没有把这三重循环原样变成三段流水而是做了一个更激进的映射输出矩阵的每个位置对应网格里的一个处理单元内层那个逐步累加的k循环被拉伸成了时钟轴上的时间序列。从这一刻起循环的秩序不见了。代码里的顺序访问变成了空间上的邻居通信循环变量变成了数据流动的方向。理解脉动阵列的难点不是每个周期发生了什么而是这个“循环到空间、k到时间”的改写为什么能保持计算结果不变同时换来巨大的带宽收益。2.1 循环体的三个自由度哪一层展开成空间哪一层被时间化回到那份教科书式的三重循环代码for i in range(M): for j in range(N): for k in range(K): C[i][j] A[i][k] * B[k][j]在这段代码里i和j是两个可以自由互换次序的维度因为它们各自只决定一次计算结果的位置k才是真正做累加的维度。软件版本里i、j、k都被翻译成内存访问的地址循环每走一步都要从DRAM里把数据取回来。硬件如果想要加速最直接的做法是把i和j展开成M×N条并行计算单元每个单元负责一个C[i][j]的累加工作。问题马上出现如果每个单元都独立从内存取A[i][0..K-1]和B[0..K-1][j]读到的全都是重复数据。M×N个单元里同一份A[i][k]会被N个不同的单元各取一次同一份B[k][j]会被M个单元各取一次。这种重复取数的代价在芯片上非常致命。脉动阵列的改法是不让每个单元都拥有访问全局内存的权利而是让数据从阵列边缘进入在每个单元里用一次然后直接传给下一个邻居。k方向的计算被压到了时间轴上每个时钟周期完成一次乘加后一个时钟周期算下一项。这样阵列内部不再需要任何访问内存的“单元内存接口”只有最外圈与DRAM通信。这是脉动阵列最核心的设计动机用邻居通信代替全局广播用数据流动代替重复取指。2.2 数据驻留策略权重不动、输入不动还是输出不动确定了“每个PE算一个输出位置”之后还要决定谁留在原处谁在网格间流动。三种经典设计分别叫权重驻留WS、输入驻留IS、输出驻留OS在仿真demo里都能通过改变数据注入方向来复现。设计驻留在PE里的数据在网格间流动的数据部分和的走向典型场景权重驻留WS权重矩阵输入激活自下而上或自上而下累加推理阶段权重固定激活流式进入输入驻留IS输入激活权重在相邻PE间传递卷积滑动窗口复用输入输出驻留OS)无输出就地累加输入和权重同时穿行尽头输出不进累加器长卷积核减少部分和搬移WS是示范里最容易看懂的方案先把B矩阵的权重烧录进每个PE再把A矩阵的元素从左侧连续送入每个PE乘完自己的权重后把结果加进路过的部分和最后从阵列底部读出C。你在Demo里看到PE的data_in寄存器每拍向右移动一格那就是A在流动看到partial_sum每拍向下传递一次那就是C在累积。实际的硬件实现里这三种策略通常要混用比如分块卷积时在不同维度上同时使用IS和WS。但初学阶段WS一个就足够讲通所有核心概念。这份demo之所以把WS作为默认路径正是因为权重预先放好之后输入数据的运动规律非常纯粹便于打印和观察。2.3 时间差是脉动阵列的第一性原理软件循环里没有“时间”概念一次add指令完成之后下一次add无脑继续。硬件阵列里数据从左边进入第一个PE再移动到第二个PE本来就需要一个时钟周期。这意味着为了使同一个输出位置上的多个乘积项能依次相遇不同数据之间必须人为制造时间差。一个三维直觉值得反复品味把整个阵列想象成一张网数据从左边进入后每一拍向右移动一格。如果同一行的两个相邻PE各自负责不同k值的乘积源数据必须相隔一拍出发才能在同一个累加链路上按顺序到达。这种时间差让所有PE在这一拍里都在计算“同一个k对应的子问题”谁也不多算一步谁也不空等一拍。仿真器里这个时间差体现在数据缓冲区的设计上。你看代码时如果发现给第j列喂的数据总比第j-1列晚一拍那不是在预加延迟而是在还原硬件里数据从左边邻居移动到当前PE所花费的那一拍。搞懂了这一点你就已经拿下了脉动阵列的一半。3. 源码级仿真一个PE、一个周期、一个乘加操作理解了映射原则之后可以拆开这份Demo的实现了。它的结构基本可以看成三层的洋葱最内层是一个处理单元处理单元只描述“一个时钟周期内发生什么事”中间层是二维的PE网格负责数据注入和结果收集最外层是一个驱动循环像一个节拍器每循环一次就模拟一个硬件时钟周期。这种分层建模方式本身就是对硬件的真实模仿。硬件里没有“run一次乘法”的函数只有时钟边沿触发后组合逻辑根据当前寄存器值算出新值、再在下一个边沿写入寄存器的过程。仿真器把“时刻”抽象成循环变量逐拍推进就能复刻出同步时序电路的行为。3.1 PE类一个寄存器、一个乘加器、一条转发路径一个最基本的PE只需要三个成员变量保存当前流入数据的data_in、预存的权重weight、以及不断累积的partial_sum。tick方法代表一个时钟周期class ProcessingElement: def __init__(self): self.data_in 0.0 self.weight 0.0 self.partial_sum 0.0 def load_weight(self, w): self.weight w def tick(self, data_in, psum_in): # 本周期完成一次乘加MAC self.partial_sum psum_in self.data_in * self.weight # 锁存新进入的数据供下一周期向右邻居传递 self.data_in data_in return self.partial_sum看这段代码需要抓住两点。第一乘法的两个操作数并不是同一个周期到达的权重是预先加载的data_in是上一个周期锁存下来的。如果你想在tick里把data_in改成“本周期刚传入的新数据”整个阵列的时序会乱。第二partial_sum不是从自身的历史值累加而是由外部传入的psum_in决定这对应硬件里部分和沿着垂直方向在PE间传递的走线。为什么要这么绕因为每个PE内部只有一个乘加器它不能既做当前k的乘法又保存上一k的加法结果。部分和需要一个外部寄存器链承载在阵列里形成纵向的传递通道。Demo里让你打印partial_sum打印出的其实是这条传递通道在当前位置的值而不是PE私有的状态。3.2 阵列类时钟驱动、数据路由、边界补零PE之上是SystolicArray类。它把二维PE网格组装起来并在每个时钟周期决定“谁把数据传给谁”。这里有一个仿真实战里最容易误解的地方仿真循环是顺序执行的但硬件阵列是并发执行的。因此不能在一个tick里直接读邻居的最新值必须提前取出邻居上一周期的输出。class SystolicArray: def __init__(self, rows, cols, weights): self.rows rows self.cols cols self.pe [[ProcessingElement() for _ in range(cols)] for _ in range(rows)] for r in range(rows): for c in range(cols): self.pe[r][c].load_weight(weights[r][c]) self.psum_buffer [[0.0] * cols for _ in range(rows 1)] self.data_buffer [[0.0] * (cols 1) for _ in range(rows)] def step(self, input_column, psum_top): # 先读上一周期的旧值避免写后读竞态 old_data [[p.data_in for p in row] for row in self.pe] old_psum [row[:] for row in self.psum_buffer] outputs [] for r in range(self.rows): for c in range(self.cols): data old_data[r][c - 1] if c 0 else input_column[r] psum old_psum[r][c] if r 0 else psum_top[c] out self.pe[r][c].tick(data, psum) self.psum_buffer[r 1][c] out self.data_buffer[r][c 1] self.pe[r][c].data_in outputs.append(self.psum_buffer[r 1][c]) return outputs注意data取的是old_data[r][c-1]也就是左邻居上一周期锁存的data_in若c等于0则从阵列左侧的input_column取新数据。同理psum取自上一行缓冲区的旧值。这样每个PE在本周期只能看到邻居上一周期的结果和真实硬件在时钟同步下的行为一致。边界补零也很重要input_column和psum_top在没有数据可喂的时候必须显式补零。硬件里这些端口可以视为悬空为零但在Python的嵌套列表里默认值是None还是0会让后面乘出来的结果完全不同。这份Demo把“没有信号”显式建模成0.0正是为了让部分和链在阵列边缘能正常闭合。3.3 运行一个示例从载入权重到矩阵结果边跑边打印把上面两个类拼起来跑一个2×2的小矩阵乘打印每个周期结束后每个PE的data_in、weight和partial_sum三张表并排对照数据流动立刻清晰起来。# 权重矩阵 W [[1.0, 2.0], [3.0, 4.0]] # 输入激活按脉动节奏逐列送入 activations [ [1.0, 0.0], [2.0, 0.0], [0.0, 1.0], [0.0, 2.0], ] sa SystolicArray(2, 2, W) for clk in range(6): col activations[clk] if clk len(activations) else [0.0, 0.0] sa.step(col, [0.0, 0.0]) # 打印每个周期结束后所有PE的累加值 for r in range(2): print(fclk{clk} row{r} psum, [p.partial_sum for p in sa.pe[r]])参数那里没有魔法rows和cols是PE网格形状weights是预存的权重矩阵activations按列送入每列代表输入矩阵的一列数据。时钟循环次数只要大于数据总列数与阵列行数之和就能保证最后一批数据完全流出阵列。这份Demo解决的核心问题不是让矩阵乘法跑得更快而是让你亲眼看到“乘法发生在数据到达的那一拍”。跑完一遍你会理解为什么硬件加速器里没有“for k in range(K)”循环——它已经把k循环展开成了时间每一拍找一个不同的k来乘。那些preload、input_column、psum_buffer看似琐碎实际上每一个都是在模拟真实芯片里的寄存器和走线。4. 数据斜移Skewing为什么数据必须“歪着”流入阵列第一次接触脉动阵列的人十有八九会发出灵魂拷问数据为什么不能整整齐齐地一行行送进去非要斜着排答案藏在阵列计算的节奏里如果数据同时到达所有PE同一个PE在同一拍会拿到好几个本该由不同时钟周期处理的项累加器瞬间被“撑爆”语义。斜移不是给数据流的装饰而是脉动阵列能正确工作的时间——空间映射。它用“出发时间不同”来补偿“传输路径长度不同”让本该在同一个周期相遇的数据正好在同一个PE内部碰面。这部分是整个SystolicArrayDemo里最值得一行行看的地方。4.1 不歪着送数据会发生什么考虑一个只有一行的点积阵列三个PE横向排开权重w0、w1、w2已经预存。如果所有输入a0、a1、a2在同一拍从左侧同时涌入并且每个PE都能接收全部三个数据那第一个PE会把a0w0、a1w1、a2w2的全部乘积当场累加出来第二个PE也会做完全相同的事——因为它看到的输入数据集合是一样的。问题来了三个PE算出的结果彼此完全相同没有一个PE能只负责其中一项。想让PE0只算a0w0、PE1只算a1w1、PE2只算a2w2就必须让a0比a1早一个时钟周期到达a1再比a2早一个周期到达。数据在时间轴上“错开”每个PE才会在它应当负责的那一项出现的那个周期动手。这正是“脉动”中流动部分和“歪斜”的直接体现。如果阵列是二维的错误会升级成部分和之间的竞争多个数据在同一拍抵达不同行时累加链路会在同一周期收到多个本应不同k的中间结果输出的C矩阵会出现整块的乱序。早期做这类仿真的人最常遇到的就是“对角线永远是对的但右下角的数据像是被洗过牌”一般就是漏掉了行方向的斜移。4.2 斜移的本质把下标差翻译成时间差斜移在数学上做的是一件很干净的事给数据矩阵的不同行添加不同长度的延迟使原来的矩形变成平行四边形。下面的代码展示了这个padding过程def skew(matrix, step1): 把每一行数据依次向右错开step列形成斜向输入。 rows, cols matrix.shape skewed np.zeros((rows, cols step * (rows - 1))) for i in range(rows): skewed[i, i * step:i * step cols] matrix[i] return skewed对每一行来说它仍然保留原来的数据顺序但行与行之间整体延迟了i拍。这样送到阵列后第i行的数据与第i-1行的数据不会同时抵达同一个PE而是错开了一个时钟周期。硬件实现里这些“延迟”来自布线本身第1行的数据直接从左边进入第0列第2行的数据要绕过一个寄存器才能进入下一列。仿真器里看不到这些物理线路但skew函数的每一行赋值都在模拟那根多绕的线。你在Demo的注释里如果看到“align data to systolic rhythm”这样的描述它指的就是这个把矩形变平行四边形的过程。4.3 从寄存器快照看“脉动”究竟在动什么打印一组脉动阵列的寄存器快照你会看到一个很有感染力的现象每个时钟周期整列data_in的值像多米诺骨牌一样同时向右移动一格。第4章开头那个“脉动”的意象在这里落地成代码可见的行为数据在网格中整体平移每个PE在数据经过自己时执行一次乘加然后把累加值沿纵向传递下去。可以试试在step()里把每个PE的data_in单独打印出来def snapshot(self, clk): for r in range(self.rows): states , .join(fd{p.data_in:.1f} p{p.partial_sum:.1f} for p in self.pe[r]) print(fclk{clk} row{r}: {states})你会看到第0拍的输入数据在第二拍整体移到第1列在第三拍移到第2列与此同时每个PE的partial_sum也在沿着列方向递增。纯静态的框图里你只能用箭头想象这个过程在这个Demo里它是打印出来的客观事实。斜移之所以让人困惑是因为它的效果隐藏在时间维度只有把时间切成显式的快照才能看出“歪着进”和“算得对”之间的因果关系。5. 避坑手册Python仿真脉动阵列最容易翻车的五个细节把脉动阵列写成Python仿真最大的风险不是硬件概念难而是Python本身的便利性在悄悄掩盖模块级时序错误。下面这些坑都是反复踩过的每一条都用“现象→原因→解决”的方式记录下来希望能帮你少走弯路。5.1 Python负索引回卷把越界行为变成“看似正确的数据”现象斜移和边界补零的代码里写了matrix[r][c - delay]当delay太大导致下标变成负数时程序没有抛异常而是静默取到了数组末尾的数据。结果矩阵依然形状完整、数值看起来也合理完全不像出了错。原因C语言里数组越界直接崩溃Python的列表负索引却在合法语法范围内——buffer[-1]返回的是最后一行。很多初学者把边界处理写成了“不处理”恰好就踩中了这个回卷语义。更隐蔽的是由于脉动阵列本身就在大量使用“延迟”操作负索引产生的错位在视觉上很难和故意设计的延迟区分开。解决养成写显式边界条件的习惯不要依赖Python的自动回卷。可以在一开始就把边界外区域pad成0并让所有下标访问都走一个带判断的辅助函数idx大于等于0才取数据否则返回0.0。宁可多写两行判断也不要让自己面对一个“看起来对、跑起来错”的结果。def safe_get(buffer, row, col): # Python的负索引在这里是陷阱必须显式判负数 if col 0 or row 0: return 0.0 return buffer[row][col]提示在涉及阵列边界的仿真代码里永远不要直接用负下标当作“取不到数据”的策略。5.2 用0填充“暂无数据”掩盖了调度错误现象pe.tick()里如果data_in是None乘法会直接TypeError。为了省事有人把初始化数据全改成0.0程序不报错了但原本应该在某个时钟周期“还没有数据到达”的位置被误当成“数据值为0”调度时序错误被悄悄吞掉。原因None和0在语义上天差地别。None表达“这个信号还没有到达”0是一个真实参与运算的数值。用0去填充所有未就绪位置会让未到达的数据和真的0值无法区分错误就会悄悄混进累积结果。解决调试阶段用float(nan)作为未到达信号。任何正确路径都不应该把nan带进最终结果一旦输出出现nan说明某个PE在错误的时间读取了尚未到达的数据。定位到具体周期后再决定是把这一拍的数据当作0还是修正输入节奏。5.3 吞吐量和延迟混为一谈测试周期给少了还自以为正确现象跑完Demo后把cycles设得比理论延迟小很多输出结果依然“正确”。后来仔细一看数据不是从左边邻居逐拍传过来的而是每个周期直接从输入矩阵里抽了一行塞给所有PE相当于把并发访问当成了流水省下的周期全是假的。原因硬件脉动阵列的延迟是固定的数据从最左边进入、流到最右边至少要经过cols个周期加上K次累加总延迟远比单纯的计算周期数大。仿真代码里如果图方便从全局矩阵直接索引就会破坏这种有序流动导致理论上不可能在这么短周期内完成的计算却在打印时给出了正确结果。解决把每个PE的data_in来源打印成一张二维表纵轴是时钟横轴是PE编号检查每个周期到底有几个新数据从阵列边缘进入。正常情况下一拍只应该有一个新数据列进入最左列如果看到某一拍所有PE的data_in同时变化基本可以断定“并发取数”混进来了。硬件里的吞吐量是每周期算完一个子矩阵乘法延迟是把这个子矩阵完全推出阵列要等的周期数两者别混着用时序计算。5.4 同一拍内“先更新再读取”把同步电路模拟成了顺序执行现象tick()里先更新左边PE的data_in随后右边PE马上读取左边邻居的data_in结果数据在同一拍里多移动了一格——每个时钟周期后数据的位置都比预期超前了一个PE。输出顺序全乱但累加值看起来还挺连续。原因硬件里寄存器输出在时钟沿之后才更新同一时钟周期内所有PE看到的都是上一周期的旧值。Python脚本是顺序执行的写在前面的一条赋值语句会对后续语句立即可见模拟的是组合逻辑的实时传播而不是寄存器在时钟边沿的锁存。解决把step()拆成两个阶段第一个阶段遍历所有PE读走它们上一周期的输出第二个阶段才调用tick执行更新。或者像本Demo的SystolicArray实现那样先复制old_data和old_psum后续计算全部基于快照而不是直接读写邻居属性。这个约束是仿真脉动阵列的黄金法则。5.5 长序列累加用浮点随机误差干扰正确性判断现象K1024的长矩阵乘跑完把仿真结果和numpy的标准矩阵乘法一对比最后几十个浮点尾数对不上。有人立刻怀疑调度逻辑错了开始检查数据流浪费了大量时间。原因每个PE的partial_sum在一拍里累加一个乘积软件里float64的舍入误差在长序列里缓慢累积和硬件里的定点累加行为本来就不完全一致。数值上的微小差异不是架构错误而是运算精度的本质不同。解决在做调度正确性验证时先用小整数矩阵跑逻辑或者把浮点乘法替换成定点整数运算等数据流验证无误后再换回浮点对照最终数值。精度差异应该用明确容差去判断而不是拿全等比较在卡正确性。6. 设计取舍当你在软件里模拟硬件寄存器你其实选了哪条路线写这份SystolicArrayDemo的过程本质上是在做一个选择到底要把仿真的粒度卡在“能算出正确矩阵”这一层还是逼真到“每个时钟周期每个寄存器的状态都可观测”。前者用numpy一行矩阵乘法就够了后者却要求你逐周期、逐PE地推进状态。这也是所有软件数据流模型都会遇到的分岔路。我倾向于把仿真的粒度卡在时钟周期这一层哪怕它慢。因为脉动阵列之所以能工作靠的不是“乘法器足够多”而是“每个乘法器在精确的时刻拿到精确的数据”。如果你把PE之间的数据流动压进一个巨大矩阵乘法里速度上去了但“精确时刻”这个信息完全丢失阵列里真正有趣的行为也就消失了。逐周期仿真等于是在软件里支付了时间成本买回硬件的时序可见性。另一个取舍在于边界处理和数据流描述的直白程度。显式补0、双缓冲寄存、每个周期一个step()方法这些都比直接写循环要啰嗦但正是这些啰嗦保证了模型的同步语义不被Python的顺序执行破坏。仿真器的价值是作为理解工具而不是作为性能基线。我调试这类数据流模型时最深的教训是不要关掉snapshot打印去“相信一次求解”。最初我也尝试省略中间状态只输出最终矩阵节省了大量打印时间可一旦结果有偏差排查时间直接翻了几倍因为没有任何周期级证据可以对照。从那以后我每次搭新仿真器都会强制先跑一遍空转打印每个PE的data_in和partial_sum确认边界节奏完全正确后才会关掉中间态继续往下做别的优化。每一步都以一个能解释每个周期行为的最小算例作为锚点再慢慢放大。这个习惯帮我避开了无数个“结果像模像样、调度错到离谱”的深夜也希望帮到你。本文还有配套的精品资源点击获取