Cascade Pipeline for Leading-Order Matrix Element Evaluation on AMD Versal AI Engine Arrays

本文提出了一种在 AMD Versal AI 引擎阵列上实现的五级级联流水线架构，用于高效计算 $\gamma\gamma \to t\bar{t}g$ 过程的领头阶矩阵元，在保持百万分之一数值精度的同时，实现了每秒 $1.0\times10^6$ 次评估的预估吞吐量，相较于单 CPU 核心获得了 $34\times$ 的加速比和 $7.7\times$ 的能效提升。

要点

想象一下，你正在尝试预测两个微小粒子（如质子）在巨型粒子加速器内部发生的一次大规模、混乱的碰撞结果。为此，物理学家使用一种称为“矩阵元”的复杂数学配方。计算这一配方就像解决一个巨大的多步骤谜题。问题在于，为了获得可靠的答案，他们必须将同一个谜题求解数百万次，每次的初始条件都略有不同。

目前，在标准计算机处理器（CPU）上执行此任务，就像让一个人逐个解决这些谜题。虽然准确，但速度极慢且能耗巨大，尤其是随着粒子加速器变得更为强大时。

本文提出了一种利用一种名为AMD Versal AI 引擎的特殊计算机芯片来解决这些谜题的新方法。作者没有在芯片内部让一个人解决整个谜题，而是构建了一条工厂流水线。

以下是该解决方案的运作方式，分解为简单的概念：

1. “流水线”问题

针对这种特定粒子碰撞（两个胶子转化为一个顶夸克、一个反顶夸克和另一个胶子）的数学配方太大，无法装入芯片上单个微小处理器的内存中。这就像试图将一本 38 页的操作手册塞进一个只能容纳 16 页的口袋。

解决方案： 作者将手册拆分为五个章节。他们构建了一条五阶段流水线。

• 阶段 1： 读取原始原料（碰撞数据）并准备前几步。
• 阶段 2 和 3： 将工作沿流水线传递，为计算添加更多步骤。
• 阶段 4 和 5： 完成最终计算并输出答案。

2. “传送带”（级联流水线）

这五个阶段由一条超高速、专用的传送带连接，称为级联接口。

• 想象一个工厂，工人们不会停下来交谈或等待许可才将盒子传递给下一个人。他们只是瞬间将盒子滑入滑槽。
• 在这块芯片中，“盒子”是被称为令牌的数据块。
• 作者设计了一套严格的规则手册（一种“确定性契约”），以确保工人们永远不会因互相等待而卡住。每个工人都确切知道何时传递盒子、何时接收盒子，因此流水线永远不会堵塞。

3. “超级工厂”（80 条流水线同时运行）

他们使用的芯片（VCK190）就像一个巨大的仓库，里面装有400 名微小工人（称为 Tile）。

• 他们不是只建造一条流水线，而是并排建造了80 条完全相同的流水线。
• 每条流水线有 5 名工人。$80 \text{ 条流水线} \times 5 \text{ 名工人} = 400 \text{ 名工人}$。
• 他们同时工作，并行解决 80 个不同的谜题。

4. 结果：速度与效率

作者将这座“工厂”与两种其他方法进行了测试：标准计算机处理器（CPU）和高端图形处理器（GPU）。

• 速度： 他们的 80 条流水线工厂比单个标准计算机核心快 34 倍。
  • 注： 顶级图形处理器（GPU）在整体速度上仍然更快（比他们的芯片快约 22 倍），但 GPU 是更大、更昂贵的机器。
• 能耗： 这是他们方法大放异彩的地方。由于流水线如此高效且专用，它消耗的电力极少。
  • 为了解决一个谜题，他们的芯片比标准计算机处理器少消耗 7.7 倍的能源。
  • 它的能效不如巨型 GPU，但 GPU 为此消耗了巨大的电力。该芯片的方法是一个“甜蜜点”，适用于你需要速度却无法连接一台耗电巨大的机器的情况。

5. 准确性检查

他们确保自己的“流水线”没有出错。他们将芯片得出的答案与“金标准”双精度计算结果进行了比较。

• 结果几乎完美匹配。差异极小（约为百万分之一），对于他们正在进行的物理计算而言，这种差异可忽略不计。

总结

简而言之，作者将一项对单块计算机芯片来说过于庞大的复杂物理计算，切分为五个可管理的部分，并构建了 80 条并行流水线以同时解决所有问题。这种方法创造了高速度与低能耗的“甜蜜点”，为运行理解大型强子对撞机（LHC）宇宙所需的模拟提供了一种强大的替代方案。

技术摘要

以下是论文《在 AMD Versal AI 引擎阵列上评估领头阶矩阵元级联流水线》的详细技术总结。

1. 问题陈述

现代高能物理（HEP）事件生成器，如 MadGraph5_aMC@NLO (MG5aMC)，在评估粒子碰撞的**矩阵元（$|M|^2$）**方面面临严重的计算瓶颈。随着大型强子对撞机（LHC）进入高亮度阶段，对这些计算的需求呈非线性增长，而 CPU 的扩展能力却有限。

• 瓶颈： 矩阵元评估占事件生成总时间的 30–40%，特别是涉及额外实辐射的多喷注过程。
• 挑战： 现有的 GPU 解决方案（如 CUDACPP）提供高吞吐量，但功耗巨大。现场可编程门阵列（FPGA）具有能效优势，但在现代 AI 引擎阵列上，其面临**每片 16 kB 的程序存储器（PM）**限制。复杂过程（如 $gg \to t\bar{t}g$）的单体实现超出了此内存限制，导致无法直接映射到单个 Tile 上。

2. 方法论

作者提出了一种在AMD Versal AI 引擎（AIE）阵列（具体为 VCK190 平台）上的级联流水线架构，以克服内存限制并最大化并行性。

A. 目标平台与架构

• 硬件： AMD Versal XCVC1902 ACAP，包含 400 个 AI 引擎 Tile，排列成 $50 \times 8$ 的网格，时钟频率为 1.25 GHz。
• 流水线分解： $gg \to t\bar{t}g$ 过程（涉及 16 个费曼图和 10 个不同的 HELAS 函数）被分解为五级流水线。
  • 第 1 阶段： 波函数生成（外部旋量/矢量）和令牌初始化。
  • 第 2 和第 3 阶段： 费米子 - 矢量顶点评估（拆分 12 个图以平衡内存）。
  • 第 4 阶段： 三胶子顶点评估（包括延迟的离壳玻色子生成器）。
  • 第 5 阶段： 四胶子接触项和色矩阵约化。
• 内存管理： 为了适应 16 kB 的限制，作者采用了程序内存分区和延迟评估。例如，离壳玻色子生成器（$FFV1P0\_3$）从第 1 阶段移至第 4 阶段，将第 1 阶段的内存使用量从 17.8 kB 减少至 15.5 kB。

B. Tile 间通信（级联协议）

• 机制： 各阶段通过384 位单向级联接口进行通信（带宽 60 GB/s）。
• 令牌协议： 采用确定性、无死锁的协议，各阶段交换包含波函数和部分振幅的结构化“令牌”。
  • 扩展令牌（第 1–4 阶段）： 携带 5 个外部波函数、3 个预计算传播子和 6 个色流振幅（每螺旋度 18 次传输）。
  • 精简令牌（第 4–5 阶段）： 在本地评估三胶子顶点后，仅携带 5 个波函数和振幅（每螺旋度 12 次传输）。
• 确定性： 系统强制执行“级联契约”，具有相同的循环结构、无条件写入和静态匹配的令牌计数，以确保零开销、无流控制的运行。

C. 软件适配

• HELAS 库移植： 将标准的 MG5aMC HELAS 库（最初为标量双精度 C++）移植到AI 引擎向量内部函数（单精度 float32）。
• 优化措施：
  • 向量化： 波函数映射到 8 宽 SIMD 向量。
  • 复数除法： 用单个硬件倒数指令替代 Smith 方法（2 次除法）。
  • 螺旋度缓存： 预计算 32 种螺旋度配置的 10 个波函数，通过位索引查找进行选择，将评估次数减少 16 倍。
  • 色约化： 将色归一化除法折叠为编译时常量。

D. 系统部署

• 规模： 80 个独立流水线被映射到 400 个可用 Tile 上（每个流水线 5 个 Tile）。
• I/O： 可编程逻辑（PL）中的分组交换架构将相空间点分发给流水线并收集结果。

3. 主要贡献

1. 内存驱动的流水线架构： 引入了一种新颖的五级级联流水线，成功将复杂的多图矩阵元计算分区到多个 AI 引擎 Tile 上，克服了 16 kB PM 限制。
2. 确定性级联契约： 开发了一种使用波函数令牌和静态循环结构的无死锁通信协议，消除了对复杂流控制硬件的需求。
3. 完整的 HELAS 移植： 成功将完整的 HELAS 振幅库移植到 AI 引擎向量内部函数，结合了复杂的优化措施，如二进制索引螺旋度缓存和简化的复数除法。
4. 可扩展部署： 展示了利用 VCK190 100% AI 引擎计算资源的 80 条流水线的理论部署。

4. 结果

• 吞吐量： 80 条流水线阵列的预计吞吐量为每秒 $1.0 \times 10^6$ 次矩阵元评估（ME/s）。
  • 这比单线程 CPU 核心（Intel i5-10600）快34 倍。
  • 虽然低于 NVIDIA A100 GPU（$2.18 \times 10^7$ ME/s），但 AI 引擎解决方案的能效显著更高。
• 能效：
  • AI 引擎： 54.8 µJ/ME（在 54.8 W 的 AIE 域功耗下）。
  • CPU： 422 µJ/ME。
  • GPU： 7.3 µJ/ME（但在 159 W 功耗下）。
  • 改进： 与 CPU 基线相比，AI 引擎的能效提高了7.7 倍。
• 精度： 经 MG5aMC 双精度参考验证。
  • 平均相对误差：1.43 ppm（百万分之一）。
  • 最大相对误差：168 ppm。
  • 这种精度水平被认为足以满足领头阶（LO）计算，因为物理不确定性（尺度变化、部分子分布函数 PDFs）主导了数值误差。
• 资源利用率：
  • 程序内存： 第 1 阶段是瓶颈，利用率为 94.7%（15,514 字节）。
  • 可编程逻辑： 使用量适中（4.72% LUTs，2.87% 寄存器），为额外逻辑留有余地。

5. 意义与未来工作

• 意义： 这项工作证明了AI 引擎阵列适用于高通量、高能效的 HEP 事件生成，特别是在 GPU 功耗包络不可行的功耗受限环境（如 LHC 的在线触发系统或边缘计算）中。它建立了一种在受限内存 Tile 上分区复杂物理内核的系统化方法。
• 局限性： 当前实现仅限于领头阶（LO）过程。延迟数据基于循环近似模拟器，而非整个阵列的直接硬件时序。
• 未来方向：
  • 螺旋度过滤： 预计算有效的螺旋度掩码以减少内层循环迭代，可能使吞吐量翻倍。
  • 更高多重性： 扩展流水线深度以处理更复杂的过程（如 $t\bar{t}ggg$）。
  • NLO 集成： 调整架构以支持涉及环路积分的次领头阶（NLO）计算。
  • 硬件演进： 利用具有更大阵列或更高时钟频率的下一代 Versal 设备。

总之，该论文提出了一种针对特定 HEP 工作负载的稳健、高能效的 GPU 加速替代方案，利用 AMD Versal AI 引擎独特的级联能力，解决了复杂矩阵元计算中固有的内存分区挑战。