手把手图解：用Python模拟一个3x3的Systolic Array（脉动阵列）计算卷积

手把手图解用Python模拟3x3脉动阵列的卷积计算奥秘在AI加速器设计中有一种被称为计算引擎心脏的架构正悄然改变着深度学习推理的效能边界。想象一下当你在手机相册中使用人像虚化功能时背后可能有数百亿次卷积运算正在发生——而让这些计算高效完成的核心技术之一就是脉动阵列(Systolic Array)。不同于传统CPU的通用计算模式这种由多个处理单元(PE)规则排列组成的计算矩阵通过精妙的数据流动实现了惊人的能效比。我们将用Python代码揭开这个数据流水线芭蕾的神秘面纱。1. 脉动阵列的生物学灵感与计算哲学早在上世纪80年代计算机科学家H.T. Kung就从人体血液循环系统获得灵感——就像心脏有节奏地泵送血液数据在计算单元间也能脉动传输。这种架构有三个典型特征数据复用每个输入数据元素会流经多个PE被重复利用局部通信PE只与相邻单元直接交互避免全局数据搬运流水并行不同PE同时处理数据流的不同阶段class ProcessingElement: def __init__(self): self.weight 0 # 当前存储的权重 self.psum 0 # 部分累加和 self.input_reg 0 # 输入寄存器 def compute(self, input_data): self.psum self.weight * input_data return self.psum在卷积运算场景下3x3脉动阵列的表现尤为亮眼。假设我们有以下卷积核和输入特征图卷积核 (W)输入特征图 (X)w11 w12 w13x11 x12 x13 x14w21 w22 w23x21 x22 x23 x24w31 w32 w33x31 x32 x33 x34传统计算需要9次乘加操作才能得到一个输出点而脉动阵列通过数据流动可以实现权重预加载到各PE特征图元素按特定节奏流入阵列每个时钟周期完成部分乘积的累加2. 权重驻留(Weight Stationary)模式实现Google TPU采用的设计范式是典型的权重驻留策略。让我们用NumPy构建一个3x3阵列import numpy as np def initialize_systolic_array(size3): 初始化3x3脉动阵列 return { weights: np.random.randn(size, size), psum: np.zeros((size, size)), input_registers: np.zeros((size, size)) }数据流动的关键时序如下表所示周期PE(0,0)PE(0,1)PE(1,0)说明1x11*w11--第一数据进入左上PE2x12*w11x11*w12x21*w11数据向右下方传播3x13*w11x12*w12x22*w11同时计算多个乘积............持续脉动计算可视化数据流动的代码示例def visualize_dataflow(cycle): 模拟数据在阵列中的流动 fig, ax plt.subplots() for i in range(3): for j in range(3): if (ij) cycle (ij3): ax.add_patch(plt.Rectangle((j, 2-i), 1, 1, colorskyblue)) ax.text(j0.5, 2-i0.5, fPE({i},{j}), hacenter) ax.set_xlim(0, 3) ax.set_ylim(0, 3)注意权重驻留模式的优势在于减少权重搬运能耗特别适合卷积核不变的推理场景3. 输出驻留(Output Stationary)策略对比另一种有趣的设计是让部分结果驻留在PE中而权重和输入数据相向流动。这种模式更适合需要累加多个部分结果的场景def output_stationary_pe(): 输出驻留型PE单元 pe { output: 0, # 累积的输出结果 weight_reg: 0, # 临时存储流动权重 input_reg: 0 # 临时存储流动输入 } return pe两种模式的性能对比如下指标权重驻留输出驻留权重带宽需求低中输入带宽需求高中控制复杂度简单中等PE利用率75%~90%50%~70%实现双向数据流动的关键代码def bidirectional_dataflow(pe_array): 双向数据流动模拟 # 权重从左向右移动 for i in range(3): for j in range(2, 0, -1): pe_array[i][j][weight] pe_array[i][j-1][weight] # 输入从右向左移动 for i in range(3): for j in range(2): pe_array[i][j][input] pe_array[i][j1][input] # 垂直方向的部分和传递 for i in range(2): for j in range(3): pe_array[i][j][psum] pe_array[i1][j][psum]4. 完整卷积计算的时空调度要实现完整的3x3卷积需要精心设计数据输入的时空关系。以下是关键步骤权重预加载阶段def load_weights(array, kernel): 将卷积核加载到PE阵列 for i in range(3): for j in range(3): array[i][j][weight] kernel[i][j]特征图输入编排每行数据需要错开1个周期输入边缘填充需要特殊处理输出收集策略每个输出点需要9个周期完成计算有效输出从第5个周期开始产生时空图表示例每个格子代表PE在特定周期的活动PE坐标周期1周期2周期3周期4周期5(0,0)x11w11x12w11x13w11x14w11-(0,1)-x11w12x12w12x13w12x14w12(1,0)-x21w11x22w11x23w11x24w11完整的卷积模拟函数def simulate_convolution(input_map, kernel): 完整的3x3卷积模拟 array [[ProcessingElement() for _ in range(3)] for _ in range(3)] load_weights(array, kernel) output np.zeros((input_map.shape[0]-2, input_map.shape[1]-2)) for cycle in range(2*3-1 input_map.shape[1]-2): # 数据输入逻辑 if cycle input_map.shape[1]: array[0][0].input_reg input_map[0][cycle] if 1 cycle input_map.shape[1]1: array[0][1].input_reg input_map[0][cycle-1] array[1][0].input_reg input_map[1][cycle-1] # ... 其他PE的输入控制 # 计算并传递部分和 for i in range(3): for j in range(3): if has_data(array[i][j]): array[i][j].compute() if i 2: array[i][j].psum array[i1][j].psum # 输出收集 if cycle 4: # 第一个完整输出产生的周期 output_row (cycle -4) // 1 if output_row output.shape[0]: output[output_row, 0] array[0][2].psum # ... 其他输出位置 return output在实际项目中调试脉动阵列时最耗时的往往是确定数据输入的精确时序。有次为了定位一个边界错误我不得不为每个PE添加了周期级的日志输出最终发现是权重加载时序比理论模型多了一个周期延迟。这也印证了硬件设计中魔鬼在细节中的真理。5. 现代AI加速器的脉动变体虽然基础原理相同但现代AI芯片对经典脉动阵列做了诸多优化数据重利用增强Tesla Dojo的2D脉动网格灵活数据流Graphcore的IPU支持可编程数据流稀疏计算SambaNova的稀疏脉动架构以下是一个支持可配置数据流方向的PE改进版class AdvancedPE(ProcessingElement): def __init__(self): super().__init__() self.dataflow_config { weight_dir: right, # 权重流动方向 input_dir: down, # 输入流动方向 psum_dir: left # 部分和流动方向 } def set_dataflow(self, config): 配置数据流方向 self.dataflow_config.update(config)对于想深入研究的开发者建议从以下几个方向入手优化尝试实现Winograd变换与脉动阵列的结合探索不同数据流策略的能耗模型添加量化支持来模拟真实芯片行为在Jupyter Notebook中实时观察数据流动的小技巧是使用IPython.display.clear_output()配合matplotlib动画这比静态图表直观十倍。当我第一次看到自己实现的脉动阵列像活物一样呼吸计算时那种顿悟的喜悦至今难忘。