脉动阵列仿真:用Python理解AI芯片的矩阵乘法内核 简介SystolicArrayDemo是一份用Python演示脉动阵列工作原理的轻量级代码面向对高性能计算、硬件加速器及TPU原理感兴趣的开发者通过软件模拟展现脉动阵列在矩阵乘法中的运行机制。压缩包内共2个文件包含1个Markdown说明文档和1个Python脚本整体仅3KB结构精炼便于快速通读与上手。已有442人学习。代码中可观察到计算单元按二维网格排列数据沿阵列流动以流水线方式完成乘累加操作同时展示队列、缓冲区等数据结构如何调度数据流模拟多单元并行计算的行为。此外项目还涉及性能分析视角帮助读者理解吞吐量、延迟等硬件设计指标并将脉动阵列与Google TPU的应用场景相关联。通过阅读源码和文档开发者既能掌握矩阵乘法的并行分解思路也能为后续接触Verilog、HLS等硬件描述语言打下基础尤为适合作为学习软硬件协同设计的入门参考。1. 脉动阵列是什么一个用 Python 演示 AI 芯片“血液流动”的仿真 Demo第一次听说“脉动阵列”时我以为又是什么玄学架构。直到拿到 SystolicArrayDemo 这份源码才看明白脉动阵列就是让数据像心脏泵血一样一个节拍接一个节拍地在计算单元之间流动。这个项目用不到两百行 Python 代码把 AI 加速芯片里最核心的矩阵乘法部件搬到了命令行里你能亲眼看到矩阵的每一行、每一列在哪个时钟周期进入阵列、在哪一个计算单元相遇相乘、最后怎么累加成完整结果矩阵。适合想做芯片前端仿真、算法到硬件的映射验证、或者单纯想搞懂 NPU 内部结构的开发者反复拆一拆。2. 核心架构拆解PE、数据流与时钟节拍的三角关系2.1 PE一个只做乘加运算的最小计算单元脉动阵列的基础是 PEProcessing Element处理单元。一个 PE 内部没有复杂控制器只有一个乘法器、一个加法器和一个累加寄存器。数据从左边和上边进来在 PE 里做一次乘累加然后把数据继续往右、往下传。在 Python 里PE 可以简化成这样一个类class PE: def __init__(self): self.psum 0 # 部分和累加器最终保存 C[i][j] self.a_val 0 # 水平方向流过的 A 矩阵数据 self.b_val 0 # 垂直方向流过的 B 矩阵数据 def compute(self): self.psum self.a_val * self.b_val注意这里有个关键设计psum 留在 PE 内部不随数据流移动。这对应硬件里的 output-stationary 模式——输出结果固定在一个位置权重和输入从外部流经 PE。每次 compute 执行一次乘加N 次之后psum 里就是 C 矩阵某个元素的内积结果。PE 之间唯一的交互就是相邻 PE 传递 a_val 和 b_val这保证阵列的每个单元只和邻居通信这也是芯片上布线最简单、最省功耗的结构。2.2 三种数据流模式为什么 Demo 选了 output-stationary脉动阵列有几种经典数据流组织方式最常被提到的是这三种数据流模式权重位置输入数据输出数据典型场景Weight-stationary预加载在 PE 内不动激活数据流经阵列部分和流动输出推理芯片的卷积加速Output-stationary输入流经 PE权重和输入都流动部分和留在 PE 内矩阵乘法教学演示Systolic / Row-stationary权重和输入都流动所有数据像流水线一样推进结果在阵列边缘收集通用矩阵乘加速器SystolicArrayDemo 采用的 output-stationary是我认为最适合做教学展示的形式因为计算结果 C[i][j] 一直留在 PE(i,j) 内部跑完时钟周期后直接读取 psum 矩阵就是答案不需要额外处理输出时序验证起来非常直接。而真实芯片里为了让权重复用更常用 weight-stationary你如果后续想对比两种模式对数据复用率的影响可以在第 6 章的改造练习里自己动手切换。2.3 时序公式为什么 A[i][k] 和 B[k][j] 总能撞在同一拍脉动阵列最反直觉的地方就是时序。A 矩阵第 i 行从阵列左边进入B 矩阵第 j 列从阵列上边进入两者的数据都不是同时整行整列流入的而是按行号和列号做了错位延迟。具体规则只有一条A[i][k] 在第 ik 个时钟周期注入阵列第 i 行的最左端然后每个周期向右移动一格。B[k][j] 在第 kj 个时钟周期注入阵列第 j 列的最顶端然后每个周期向下移动一格。那么 A[i][k] 到达 PE(i,j) 的时刻是 (ik) j ikjB[k][j] 到达 PE(i,j) 的时刻是 (kj) i ijk。两个数据从不同的方向出发经过不同长度的路径却在同一个时钟周期到达同一个 PE。这不是巧合而是脉动阵列把空间延迟折算成时间延迟的经典手法——数据不需要额外的路由和仲裁推着走就行。以 3×3 阵列为例子注入调度长这样cycleA 注入行0 / 行1 / 行2B 注入列0 / 列1 / 列20A00 / 0 / 0B00 / 0 / 01A01 / A10 / 0B10 / B01 / 02A02 / A11 / A20B20 / B11 / B0230 / A12 / A210 / B21 / B1240 / 0 / A220 / 0 / B22注意那些空缺的 0是数据没进入阵列时注入的占位值。它们保证 A 和 B 的每一对相乘项都各自错开不会提前或延后拍到错误的 PE 上。我一般会在调试时把这个表打出来对应着看每一拍寄存器里的值比看任何打印都直观。3. 代码逐模块拆解一个矩阵乘法如何流经二维阵列3.1 目录结构与类设计拿到源码后先看文件结构这个 Demo 是典型的小型教学项目布局systolic_demo/ ├── systolic_array.py # 核心模拟器SystolicArray 类含 PE 二维阵列 ├── run_demo.py # 运行入口生成矩阵、注入数据、跑时钟、验证结果 ├── visualize.py # 打印每个时钟周期的寄存器快照 └── requirements.txt # 依赖说明实际只用到一个 numpy整个项目的核心收敛在 systolic_array.py 里的 SystolicArray 类。一个类封装了三个职责二维 PE 寄存器阵列、数据传播逻辑、单周期乘累加。run_demo.py 负责按时序把数据喂进去。visualize.py 是辅助观察工具不改核心逻辑。这种拆分方式很干净你要改算法只动 systolic_array.py要改输入输出只动 run_demo.py互不污染。3.2 数据注入按公式生成每个周期的输入序列注入是脉动阵列最容易出错的一步。前面讲了 A[i][k] 要在第 ik 拍进入第 i 行那么在 run_demo.py 里每个时钟周期要生成一个长度为 size 的注入向量。代码长这样def feed_inputs(sa, A, B): size sa.size total_cycles 3 * size - 1 for t in range(total_cycles): a_in [] for i in range(size): k t - i a_in.append(A[i][k] if 0 k size else 0) # 矩阵左边界注入 b_in [] for j in range(size): k t - j b_in.append(B[k][j] if 0 k size else 0) # 矩阵上边界注入 sa.step(a_in, b_in)理解这段代码的关键是 k t - i。t 是当前时钟周期i 是阵列行号。A[i][k] 要在第 ik 拍进入反解出来就是 k t - i。当 k 超出 [0, size-1] 时说明这个数据还没到该进的时间或者已经全部进完了此时补 0 占位。total_cycles 取 3 * size - 1 也是有讲究的最后一个乘积项 A[size-1][size-1] 乘 B[size-1][size-1] 需要走到最右下角的 PE路径长度正好是 3*size-3 拍多给两拍作为余量。3.3 主循环一个时钟节拍里发生的三件事SystolicArray 的 step 方法是整个模拟器的发动机。每一拍按“先传播后计算”的顺序执行def step(self, a_in, b_in): # 1. 水平传播A 数据从右往左覆盖避免覆盖未传播的值 for i in range(self.size): for j in range(self.size - 1, 0, -1): self.a_reg[i][j] self.a_reg[i][j - 1] self.a_reg[i][0] a_in[i] # 2. 垂直传播B 数据从下往上覆盖 for j in range(self.size): for i in range(self.size - 1, 0, -1): self.b_reg[i][j] self.b_reg[i - 1][j] self.b_reg[0][j] b_in[j] # 3. 所有 PE 并行执行乘累加 for i in range(self.size): for j in range(self.size): self.psum[i][j] self.a_reg[i][j] * self.b_reg[i][j]先说为什么“先传播后计算”如果先算乘法再移动数据刚注入边界的数据要等下一拍才参与运算虽然最终结果也对但时序公式要从 ikj 变成 ikj1调试时心智负担更重。先传播再计算数据进入边界的那一刻立即参与本拍运算和前面表格里的周期完全对齐。再说为什么传播循环要倒着遍历以 A 的水平传播为例a_reg[i][j] 要接收来自 a_reg[i][j-1] 的值。如果从 j0 往 jsize-1 正着遍历上游的值被覆盖前就已经被下游拿走了吗不是的正着遍历时j1 先被改成 a_reg[i][0]等到 j2 时它读到的已经是改过的 a_reg[i][1]等于一个周期内数据连跳两格。这是脉动阵列模拟最常见的 bug。倒着遍历保证每个位置只读取上一拍的值。3.4 结果回收与数值验证所有周期跑完后psum 矩阵就是完整的 C A × B。结果比对代码在 run_demo.py 末尾import numpy as np C_sim np.array(sa.psum, dtypeint) C_ref A B print(systolic result:) print(C_sim) print(numpy reference:) print(C_ref) assert (C_sim C_ref).all(), 结果不一致请检查注入时序这里用整数矩阵做断言相等。为什么强调整数因为浮点加法不满足结合律脉动阵列的累加顺序是按 k 从小到大的顺序逐项相加numpy 矩阵乘法用的是优化过的分块求和两者即使是同一个数学结果浮点数值也可能差最后几位。整数矩阵没有这个问题可以严格断言相等定位问题也更干净。4. 跑起来与调参运行入口、参数表与状态可视化4.1 环境准备与运行入口这个 Demo 对环境要求极低Python 3.8 及以上即可核心模拟器只用标准库。唯一的外部依赖是 numpy作用仅限于最后一步的结果验证如果你想省事用纯 Python 手写一个矩阵乘法做比对也完全可以。项目跑起来很简单在仓库根目录执行python run_demo.py默认会生成一个 3×3 的随机整数矩阵跑完打印脉动阵列的计算结果和 numpy 参考结果断言通过后退出。如果带了可视化可以先跑一遍可视化脚本看数据流动过程python visualize.py --size 4我建议第一次运行先开 visualize看着数据一排排流入阵列再对照第 2 章那张 3×3 注入调度表比直接看结果矩阵更能建立时序直觉。4.2 核心参数说明改哪里能让行为变参数类型默认值作用--sizeint3阵列尺寸同时决定 A、B 矩阵为 size × size--seedint0随机种子固定后每次生成的矩阵可复现--matrixstr无手动传入矩阵字符串如 [1,2;3,4]--cyclesint自动默认 3*size-1一般不用手动改--verboseboolFalse是否打印每个周期的寄存器快照size 是最重要的参数它直接决定阵列的行列数也决定矩阵乘法的维数。改成 4 就是计算 4×4 矩阵乘PE 阵列同时变成 16 个。seed 的作用是让调试可复现跑出错的时候固定 seed每次结果都一样排查起来不用猜。matrix 参数适合你想手工构造边界数据时用比如填一堆 0 或 1 来看传播路径比随机数直观得多。4.3 可视化把每个时钟节拍变成一帧画面visualize.py 的核心是一个打印函数把三个二维数组按拍快照def print_state(sa, cycle): print(f Cycle {cycle} ) print(A_reg:) for row in sa.a_reg: print( .join(f{v:4d} for v in row)) print(B_reg:) for row in sa.b_reg: print( .join(f{v:4d} for v in row)) print(PSUM:) for row in sa.psum: print( .join(f{v:4d} for v in row))读这个输出有个技巧盯着 A_reg 看会发现第 0 行的数据从左边逐个进入同时第 1 行比第 0 行晚一拍开始进第 2 行更晚——这就是错位注入的现场。B_reg 类似只不过方向是从上往下。PSUM 矩阵的变化最有意思它是一排一排“长”出来的最开始只有左上角有非零值随着周期推移右下角的累加项逐渐补齐最终变成完整的结果矩阵。如果你看到某一行的数据在一拍内同时出现在两列那就是碰到了我在避坑章里说的同拍覆盖问题。5. 避坑指南时钟错拍、数据错位与边界掉数的排查记录5.1 现象跑完所有周期后结果矩阵全是零这是个很有意思的翻车现场代码逻辑看起来完全正常每一拍都在传播数据都在乘累加但 psum 纹丝不动全是 0。原因出在 step 方法的执行顺序上。如果你用的是“先计算、后传播”的写法数据在第 t 拍注入边界寄存器后要等到第 t1 拍的 compute 阶段才参与乘法。此时如果 total_cycles 设置得刚好等于 3*size-3最后一个乘积项刚进阵列还没被计算整个结果就差了一拍。更隐蔽的情况是初始化 a_reg 和 b_reg 时误把它们设成 None 或字符串乘法直接报错。解决方法是统一约定“先传播、后计算”并让 total_cycles 3*size-1多留两拍余量。我一般还会在跑完后加一条断言检查右下角 PE 的 psum 非零能提前暴露时序不足的问题。5.2 现象结果矩阵整体向右下方偏移了一个位置特别是 C[0][0] 的位置上出现了 A[0][1] × B[1][0] 这样的乘积项看起来像是整个结果矩阵平移了。原因是注入时没有按行号做延迟直接把 A 的同一列同时注入了所有行。比如有人图省事写成A[i][0]导致第 0 行和第 1 行的数据同一拍进入阵列破坏了 A[i][k] 应该在第 ik 拍进入的时序。两个数据到达 PE 的时刻差不再恒等于 ijk乘积项就会落到相邻 PE 上。解决方法是严格按 k t - i 取元素并且在循环里打印注入序列做核对。这是我调试脉动阵列用得最多的手段把每一拍实际注入的值打出来和纸上推演的调度表并排对比一眼就能看出哪一行提前进、哪一行拖后了拍。5.3 现象最后一列结果恒为 0其他列正常这个 bug 极具迷惑性因为大部分矩阵元素算对了只有最右边一列全是 0而且只出现在 size 大于 3 时。原因是传播时用了正向遍历。以水平传播为例# 错误写法 for j in range(1, self.size): self.a_reg[i][j] self.a_reg[i][j - 1]j1 先被更新成 j0 的值紧接着 j2 又被更新成已经更新过的 j1 的值等于一个周期内数据被连续推了两格。最右边的寄存器在每一拍都收到“被反复覆盖后的值”而且越靠右错得越离谱最终列右侧的 PE 永远等不到正确的数据。解决方法是倒序遍历for j in range(self.size - 1, 0, -1): self.a_reg[i][j] self.a_reg[i][j - 1]同样的原则也适用于垂直传播必须从最下游的 PE 开始覆盖上游的数据才不会被中途改写的值污染。这个坑我踩过一次以后就长记性了凡是涉及一维数组相邻元素搬运统一倒着写。5.4 现象改成浮点矩阵后结果和 numpy 对不上项目改成支持 float 数据后断言C_sim C_ref开始随机失败但不是整块错而是某个元素差 1e-12 这种量级。原因不是算法错而是累加顺序不同。脉动阵列每个 PE 按 k 的递增顺序逐项相加numpy 底层用分块和向量化指令加法的结合顺序不同浮点舍入误差必然不同。这不是 bug是浮点运算的固有属性。解决方法是验证时机分场景调算法时用整数矩阵严格断言相等确认时序正确跑真实浮点数据时换np.allclose(C_sim, C_ref, rtol1e-9, atol1e-9)。要记住脉动阵列对累加顺序敏感这一点在真正做硬件浮点设计时也要考虑。6. 二次开发把 Demo 改造成可配置的矩阵乘法引擎拿到这个项目后最值得做的第一个改造是把 output-stationary 改成 weight-stationary让权重 B 提前预加载进 PEA 数据仍然从左往右流但部分和改为从上往下逐行累加、最终从阵列底部流出。具体做法是在 SystolicArray 里增加一个 weight 二维数组去掉 b_reg 的垂直传播把 step 里的乘法改成self.psum[i][j] self.a_reg[i][j] * self.weight[i][j]再在底部边界把每行的 psum 收集出来。这个改造能让你直观感受到两种模式对数据流动方向的影响。更多可动手的方向一是支持非方阵乘法和多维参数配置把 size 拆成 M、N、K 三个维度注入公式相应改为 k t - i - offset二是加一个“周期计数器”属性把每个 PE 的累加历史记录下来导出成 JSON 再画成时序图三是把 psum 改成从边缘流出的方式模拟真实芯片里结果不回读而是直接进入下一级流水线。每改一步都用随机整数矩阵和 numpy 对比验证一次。有回我为了省事把注入序列的索引直接固定成 0结果整个矩阵错得完全看不出规律排查了半小时才发现是少了按行号延迟。从那以后我每次改数据流调度都会先手画一拍“谁在哪个周期到达哪个 PE”的时序表再动代码。希望帮到你。本文还有配套的精品资源点击获取