在 AI 硬件加速领域,systolic 阵列因其高效的数据复用和并行计算能力而备受关注。然而,当阵列规模缩小到 2×2 的最小配置时,微架构设计面临独特的挑战:如何在极有限的硬件资源下实现高效的数据流编排?本文以 TinyTinyTPU 项目为案例,深入分析 2×2 systolic 阵列的微架构设计与数据流优化策略。
1. 2×2 systolic 阵列的微架构设计挑战
传统的 systolic 阵列设计通常面向大规模配置(如 Google TPU 的 256×256 阵列),其优化重点在于最大化计算单元利用率。但对于 2×2 最小阵列,设计目标发生了根本性转变:
核心挑战:
- 数据流编排复杂度不降反增:小规模阵列需要更精细的时序控制来维持正确的计算波前
- 内存带宽与计算单元平衡:每个 PE(Processing Element)都需要及时获得数据,但总线带宽有限
- 流水线深度优化:需要在有限的硬件资源下实现完整的后 MAC 处理流水线
TinyTinyTPU 的设计体现了这些挑战的应对策略。该项目采用 SystemVerilog 实现,包含完整的处理流水线:MMU(2×2 矩阵乘法单元)→ 累加器(对齐 + 双缓冲)→ 激活 + 归一化 + 损失 → 量化 → 统一缓冲区。这种完整流水线的设计确保了即使是最小阵列也能执行完整的神经网络推理任务。
2. 对角线波前权重加载的数据流编排
在 systolic 阵列中,正确的权重加载时序至关重要。对于 2×2 阵列,TinyTinyTPU 采用了3 周期交错权重加载方案,确保权重以对角线波前模式在阵列中传播。
具体实现机制:
// 控制信号设计
en_weight_pass // 广播到所有PE,在权重加载阶段启用psum直通
en_capture_col0 // 第0列PE的捕获使能
en_capture_col1 // 第1列PE的捕获使能(交错)
权重加载时序表(以权重矩阵 W = [[1, 2], [3, 4]] 为例):
| 周期 | col0_out | col1_out | 第 0 列捕获 | 第 1 列捕获 |
|---|---|---|---|---|
| 0 | 3 | 0(交错) | 否 | 否 |
| 1 | 1 | 4 | 是 | 否 |
| 2 | (保持) | 2 | 否 | 是 |
最终权重分布:
- PE00: W[0,0] = 1
- PE01: W[0,1] = 2
- PE10: W[1,0] = 3
- PE11: W[1,1] = 4
这种交错捕获机制确保了权重在阵列中的正确传播方向。第 1 列的捕获使能比第 0 列延迟 1 个周期,形成了对角线波前。正如相关研究指出,systolic 阵列的性能优化需要精细的时序控制,特别是在小规模配置中,时序错误的容错空间更小。
3. 交错列权重 FIFO 的内存带宽优化
内存带宽是 systolic 阵列设计的关键瓶颈。对于 2×2 阵列,TinyTinyTPU 通过dual_weight_fifo.sv模块实现了创新的带宽优化策略。
设计要点:
- 共享数据总线:两个独立的 4 项队列共享一个数据总线,同时填充两个 MMU 列
- 交错读取时序:
- 第 0 列:组合逻辑读取输出(无延迟)
- 第 1 列:寄存器输出,带 1 周期交错用于对角线波前
- 同步推进:单个
pop信号同时推进两个读指针
带宽优化效果:
- 减少 50% 的权重加载总线访问
- 通过交错时序避免总线冲突
- 保持对角线波前传播的正确性
这种设计体现了小规模阵列特有的优化机会:当阵列规模较小时,可以通过更精细的时序控制来优化带宽利用率,而大规模阵列通常采用更粗粒度的数据分块策略。
4. 数据流编排的可落地参数清单
基于 TinyTinyTPU 的设计经验,以下是 2×2 systolic 阵列数据流优化的关键参数与配置清单:
4.1 时序参数配置
权重加载阶段:3周期
- 周期0:初始化交错
- 周期1:第0列捕获
- 周期2:第1列捕获
计算阶段:3周期
- 激活流带行交错
首次结果输出:5周期(从计算开始)
结果间隔:1周期(连续有效输出间)
4.2 内存带宽优化参数
权重FIFO深度:4项/列
总线共享策略:单总线双列填充
读取交错:1周期延迟(第1列相对第0列)
带宽节省:~50%(相比独立总线)
4.3 流水线深度配置
完整流水线阶段:5级
1. MMU(矩阵乘法单元)
2. 累加器对齐
3. 激活函数(ReLU/ReLU6)
4. 归一化(增益/偏置/移位)
5. 量化(int8饱和)
双缓冲配置:累加器内存双缓冲
对齐逻辑:列输出解交错
4.4 控制信号设计规范
权重加载控制:
- en_weight_pass:全局直通使能
- en_capture_col0:第0列定时捕获
- en_capture_col1:第1列定时捕获(交错)
数据流控制:
- 激活水平流动(向右)
- 部分和垂直流动(向下)
- 权重PE内固定
5. 可扩展性设计与局限性分析
虽然 2×2 阵列规模有限,但其设计原则具有重要的教育意义和可扩展性参考价值。
可扩展设计原则:
- 模块化 PE 设计:每个 PE 独立且接口一致,便于阵列扩展
- 分层控制结构:局部控制(PE 级)与全局控制(阵列级)分离
- 参数化配置:关键时序和尺寸参数可配置,适应不同规模
实际应用局限性:
- 计算吞吐量有限:2×2 阵列仅适合极小规模矩阵运算或教育演示
- 内存带宽瓶颈:即使优化后,带宽仍是限制因素
- 实际部署价值:主要用于 FPGA 原型设计和教学,而非生产环境
然而,正如相关研究指出,systolic 阵列的分区优化可以显著提升推理吞吐量。对于更大规模的阵列,可以采用 1D 分区策略,将网络层切片分配到不同的阵列分区中,通过流水线方式提高整体利用率。
6. 监控与调试要点
对于 2×2 systolic 阵列的开发和调试,以下监控点至关重要:
关键监控信号:
- 权重加载时序:验证对角线波前是否正确建立
- 数据流方向:确保激活水平流动、部分和垂直流动
- FIFO 状态:监控 dual_weight_fifo 的填充和读取状态
- 流水线停顿:检测各阶段的数据依赖和停顿情况
调试工具链:
- Verilator 仿真环境
- cocotb Python 测试框架
- VCD 波形查看(GTKWave/Surfer)
- 完整的测试套件覆盖所有核心模块
结论
2×2 最小 systolic 阵列的微架构设计展示了在极端资源约束下的优化艺术。通过对角线波前权重加载、交错列 FIFO 带宽优化、精细的时序控制等策略,即使在最小规模的配置下也能实现正确的 systolic 计算语义。
虽然 2×2 阵列的实际计算能力有限,但其设计原则为更大规模 systolic 阵列的优化提供了重要参考:模块化设计、分层控制、参数化配置、以及针对特定规模的数据流优化策略。这些经验对于理解 systolic 阵列的本质、开发 FPGA 加速器原型、以及教学演示都具有重要价值。
在 AI 硬件加速快速发展的今天,从小规模原型中积累的设计经验,往往能为大规模生产系统的优化提供关键洞察。2×2 systolic 阵列的设计之旅,正是这种从微观到宏观、从原理到实践的技术演进缩影。
资料来源:
- TinyTinyTPU GitHub 仓库:https://github.com/alanma23/TinyTinyTPU
- Systolic Array 分区优化研究:https://nachiket.github.io/publications/part-systolic_fpt-2019.pdf