Hotdry.
归档
ai-systems

TinyTinyTPU：2×2脉动阵列TPU风格矩阵乘法单元在FPGA上的实现

深入分析最小化2×2脉动阵列TPU风格矩阵乘法单元的SystemVerilog实现，涵盖对角线波前权重加载、完整后MAC流水线及多层MLP推理支持。

在 AI 硬件加速领域，脉动阵列因其高效的数据复用和并行计算能力而备受关注。Google 的 TPU（Tensor Processing Unit）正是基于脉动阵列架构的成功案例。然而，对于 FPGA 开发者而言，理解并实现一个完整的脉动阵列系统仍面临诸多挑战。TinyTinyTPU 项目提供了一个极简但完整的 2×2 脉动阵列 TPU 风格矩阵乘法单元实现，为学习和研究脉动阵列架构提供了绝佳的起点。

脉动阵列的核心优势与设计挑战

脉动阵列是一种规则排列的处理单元（PE）网格，数据在网格中以 "脉动" 的方式流动 —— 激活值水平流动，部分和垂直流动。这种设计最大限度地减少了数据移动，提高了计算密度和能效。然而，在 FPGA 上实现脉动阵列需要考虑多个工程因素：资源利用率、时序收敛、数据流控制以及权重加载机制。

TinyTinyTPU 采用 2×2 的最小规模设计，这并非性能最优的选择，而是为了教学和原型验证的目的。正如研究指出，FPGA 上脉动阵列的资源利用率是一个关键挑战，未充分利用的 FPGA 资源会显著影响推理吞吐量。

2×2 脉动阵列的架构设计

处理单元（PE）设计

PE 是脉动阵列的基本计算单元，每个 PE 执行乘累加（MAC）操作：psum_out = psum_in + (in_act × weight)。TinyTinyTPU 的 PE 设计采用单周期寄存器输出，确保时序可预测性。

关键控制信号包括：

  • en_weight_pass：在权重加载阶段控制部分和直通
  • en_weight_capture：触发权重寄存器锁存

这种分离的控制机制允许在权重加载期间保持计算流水线的连续性，是脉动阵列高效运行的基础。

矩阵乘法单元（MMU）组织

2×2 MMU 的组织结构如下：

PE00 → PE01    激活值水平流动（向右）
  ↓       ↓     
PE10 → PE11    部分和垂直流动（向下）
  ↓       ↓
acc0    acc1    输出到累加器

MMU 负责：

  • 将激活值馈送到行（row0 直接输入，row1 带偏斜寄存器）
  • 通过垂直部分和路径加载权重，每列具有独立的捕获时序
  • 向累加器输出两个部分和列

对角线波前权重加载机制

权重加载是脉动阵列设计中最复杂的部分之一。TinyTinyTPU 采用 3 周期交错的权重加载方案，确保权重在阵列中以正确的对角线波前模式传播。

对于权重矩阵 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

这种交错加载机制通过dual_weight_fifo.sv模块实现，该模块包含两个独立的 4 项队列，共享一个数据总线来填充两个 MMU 列。列 0 提供组合逻辑读取输出（无延迟），列 1 提供带 1 周期偏斜的寄存器输出，以支持对角线波前。

完整的后 MAC 流水线设计

TinyTinyTPU 不仅实现了核心的矩阵乘法，还包含了完整的后处理流水线，模拟了真实 TPU 的数据流：

累加器子系统

累加器捕获 MMU 输出，对齐交错的列，并支持累加 / 覆盖模式的双缓冲：

  • accumulator_align.sv：消除列输出偏斜
  • accumulator_mem.sv：双缓冲 32 位存储

激活流水线

activation_pipeline.sv整合了所有后累加器处理阶段：

  1. 激活函数：支持 ReLU、ReLU6 或直通
  2. 归一化：增益 / 偏置 / 移位操作
  3. 并行损失计算：L1 损失计算
  4. 仿射 int8 量化：带饱和的量化

统一缓冲区

unified_buffer.sv提供字节宽度的同步 FIFO，支持就绪 / 有效背压机制，确保输出数据的稳定传输。

多层 MLP 推理支持

TinyTinyTPU 的一个重要特性是支持多层神经网络推理。test_mlp_integration.py演示了多层感知机（MLP）的推理流程：

权重FIFO → MMU（脉动） → 累加器 → 激活流水线 → 统一缓冲区
     ^                                                      |
     +------------------ 反馈（下一层） --------------------+

2 层 MLP 示例

考虑一个简单的 2 层 MLP：

  • 输入：A = [[5, 6], [7, 8]]
  • 第一层权重：W1 = [[1, 2], [3, 4]]
  • 第二层权重：W2 = [[1, 1], [1, 1]]

计算过程：

  1. 第一层：H = ReLU (A × W1) = [[23, 34], [31, 46]]
  2. 第二层：Y = ReLU (H × W2) = [[57, 57], [77, 77]]

这种反馈机制使得 TinyTinyTPU 能够处理实际的神经网络推理任务，而不仅仅是孤立的矩阵乘法。

SystemVerilog 迁移与测试框架

从 Verilog 到 SystemVerilog

项目从原始的 Verilog 实现迁移到 SystemVerilog，带来了多项改进：

  1. 类型系统增强：

    • input wire / output reg → input logic / output logic
    • reg / wire → logic

  2. 过程块语义化：

    • always @(posedge clk) → always_ff @(posedge clk)
    • always @(*) → always_comb

  3. 数组语法简化：

    • [0:DEPTH-1] → [DEPTH]

  4. 明确的 case 语句：

    • case → unique case（适当情况下）

这种迁移不仅提高了代码的可读性和可维护性，还增强了综合兼容性。

基于 cocotb 的测试框架

Verilog 测试平台被 cocotb Python 测试平台取代，带来以下优势：

  • 更可维护和可读的测试代码
  • 使用 Python 工具进行更轻松的调试
  • 通过 Verilator 生成 VCD 波形
  • Pytest 集成用于测试发现和报告

测试覆盖包括所有核心模块：PE、MMU、权重 FIFO、累加器、激活函数、归一化器、激活流水线和 MLP 集成测试。

FPGA 部署的工程考虑

流水线时序总结

阶段 周期数 描述
权重加载 3 周期 带对角线波前的交错列捕获
计算 3 周期 带行偏斜的激活值流式传输
首个结果 5 周期 从计算开始到第一个累加器输出
结果间隔 1 周期 连续有效累加器输出之间的间隔

资源利用率优化

虽然 TinyTinyTPU 是 2×2 的最小设计，但实际部署时需要考虑资源优化策略。研究表明，通过将 FPGA 划分为多个方形脉动阵列，并将这些阵列的放置问题表述为 2D 背包问题，可以显著提高资源利用率。

对于大型神经网络，可以采用分区架构，在不同区域部署不同大小的脉动阵列，以匹配各层的计算需求。这种方法的吞吐量增益可达 7-22.7 倍，尽管会有 1.2-7.6 倍的延迟牺牲。

频率与布局考虑

脉动阵列的工作频率受布局位置影响显著。实验数据显示，2×2 脉动阵列在不同 FPGA 区域的工作频率分布呈现明显的空间相关性。优化布局可以确保时序收敛并最大化工作频率。

实际应用与扩展方向

教育价值

TinyTinyTPU 作为一个最小化但完整的实现，具有重要的教育价值：

  1. 理解脉动阵列基本原理：通过简单的 2×2 设计，清晰展示数据流动模式
  2. 学习权重加载机制：对角线波前加载是理解更大规模阵列的基础
  3. 掌握完整 AI 加速器流水线：从矩阵乘法到后处理的完整流程

原型开发平台

对于希望开发自定义 AI 加速器的团队，TinyTinyTPU 可以作为：

  1. 验证架构概念的起点：快速验证新的数据流或计算模式
  2. 测试工具链的基础：集成新的编译器或调度器
  3. 性能建模的参考：建立小规模模型的性能预测

扩展方向

基于 TinyTinyTPU，可以探索多个扩展方向：

  1. 规模扩展：将 2×2 阵列扩展到更大规模（如 8×8 或 16×16）
  2. 精度支持：增加对混合精度（int4、fp16 等）的支持
  3. 稀疏性利用：集成稀疏矩阵乘法优化
  4. 动态重构：利用 FPGA 的动态部分重构特性

结论

TinyTinyTPU 项目展示了如何在 FPGA 上实现一个完整但最小化的 TPU 风格脉动阵列。通过 2×2 的设计规模，它降低了学习曲线，同时保持了架构的完整性。对角线波前权重加载机制、完整的后 MAC 流水线以及多层 MLP 推理支持，使其成为一个有价值的教学工具和原型开发平台。

对于 FPGA 开发者而言，理解脉动阵列的实现细节是开发现代 AI 加速器的关键一步。TinyTinyTPU 提供了一个绝佳的起点，帮助开发者掌握从基本 PE 设计到完整系统集成的全流程技能。随着 AI 硬件加速需求的不断增长，这类开源项目将在培养下一代硬件工程师方面发挥重要作用。

资料来源

  1. TinyTinyTPU GitHub 仓库：https://github.com/alanma23/TinyTinyTPU
  2. "Implementing FPGA-optimized Systolic Arrays using 2D Knapsack and Evolutionary Algorithms" - 关于 FPGA 脉动阵列优化的学术研究
  3. "A TPU-like Design for FPGA Benchmarking" - FPGA 上 TPU 风格设计的基准测试研究
查看归档