GPU Acceleration and Portability of the TRIMEG Code for Gyrokinetic Plasma Simulations using OpenMP

本文介绍了将 TRIMEG 回旋动力学等离子体模拟代码通过可移植的 OpenMP API 成功移植到 GPU 架构的工作，通过混合 MPI-OpenMP 并行化以及离子温度梯度模模拟，证明了显著的加速效果并验证了物理准确性。

要点

大局观：烹饪一场宇宙风暴

想象一下，试图预测恒星内部的天气。在现实世界中，我们无法直接把温度计插进太阳或核聚变反应堆内部；那里太热、太混乱了。因此，科学家们使用超级计算机来运行等离子体（超热、带电的气体）的“虚拟模拟”。

TRIMEG 代码是一种非常复杂且精密的模拟等离子体的“食谱”。它追踪着数以十亿计的微小粒子（就像风暴中每一颗细小的沙粒），观察它们如何旋转、碰撞并产生湍流。问题在于？这个食谱极其沉重。在标准计算机（CPU）上运行它，就像是用一把勺子去搬动一座大山，效率太低，耗时太长。

目标： 作者 Giorgio Daneri 希望通过使用 GPU（图形处理器）来加速这一过程。把 CPU 想象成一位极其聪明但一次只能切一种蔬菜的“主厨”；而 GPU 则像是一个拥有 1 万名“副厨”的厨房，他们可以同时切菜。这篇论文的研究重点在于：如何让这位主厨的食谱完美地适配这 1 万名副厨组成的军队，并且让这套方案能同时适用于两种不同品牌的厨房（NVIDIA 和 AMD）。

挑战：“通用翻译官”难题

作者选择了一个名为 OpenMP 的工具来进行这种“翻译”。你可以把 OpenMP 想象成一个通用翻译官，它告诉计算机：“嘿，把食谱的这一部分交给 GPU 去处理。”

然而，作者遇到了两个主要的障碍：

1. “编译器”故障： 将代码转化为机器语言的软件（编译器）并不完美。这就像是使用一个有时会忘记如何表达“盐”或“热量”的通用翻译官。作者不得不重写部分代码以适应翻译器的特性。例如，原代码使用了高级的“多态性”（一种让对象可以改变形状或身份的高级编程方式），但 GPU 的翻译器无法理解这种“变形”，因此作者必须将这些形状“压平”成固定的方块，才能让它们正常工作。
2. “交通堵塞”： 在主计算机（CPU）和 GPU（副厨们）之间移动数据是非常缓慢的。如果你不停地停下来传递食材，副厨们就会处于闲置状态。作者必须重构代码，使得所有食材在开始时就一次性全部移交给 GPU，而不是在两者之间频繁地往返运输。

解决方案：重构厨房

为了让代码能在 NVIDIA 和 AMD 的 GPU 上都能运行，作者对 TRIMEG 代码进行了一些“手术”：

• 压平地图： 代码使用了一张复杂的地图来寻找粒子的位置。这张地图就像一个杂乱无章的档案柜。作者将其压平为一个单一的、笔直的列表，这样 GPU 就能瞬间读取，而不会迷失方向。
• 修复“竞态”： 有时，当成千上万名副厨试图同时在同一块白板上书写时，他们会互相覆盖彼此的字迹（即“竞态条件”）。作者找到了代码中发生这种情况的地方，并进行了修复，确保每个人都在自己的赛道内书写。
• “一刀切”的折中方案： 由于两种 GPU 品牌（NVIDIA 和 AMD）使用的语言略有不同，作者创建了一个能同时兼容两者的单一代码版本，即便这意味着要使用一些“权宜之计”（例如使用一种对两者都通用的内存分配方式，即使这对其中某一个品牌来说并不是绝对最快的）。

结果：奏效了吗？

作者使用两个著名的“测试用例”（就像是新车的标准驾驶测试）将新的 GPU 版本与旧的 CPU 版本进行了对比：

1. 旋风案例（The Cyclone Case）： 一个简化的等离子体湍流模拟。
2. TCV-X21 案例： 一个涉及等离子体边缘的更复杂、更真实的模拟。

结论：

• 速度： GPU 版本明显更快。在某些测试中，单机运行速度比 CPU 版本快了近 30 倍。
• 准确性： GPU 的结果与 CPU 的结果几乎完全吻合。其“天气模式”（能量增长和湍流结构）看起来是一致的。
• 移植性： 代码成功地在 NVIDIA 和 AMD 硬件上运行，无需为每种硬件重新编写代码。

局限性（不足之处）

作者坦诚地说明了局限性：

• “翻译官”尚未完美： 这些 GPU 的编译器（将代码转化为机器语言的软件）仍在成熟阶段。有时它们产生的数学结果与 CPU 略有不同，这可能会导致随时间推移而产生的微小误差。
• 硬件不匹配： 如果你的计算机拥有大量的 CPU 核心但只有一个 GPU，当你尝试向 GPU 输入过多任务时，GPU 可能会不堪重负。作者发现，为了获得最佳效果，你需要平衡“厨师”（MPI 进程）的数量与可用的“副厨”（GPU 线程）的数量。
• 并非“万灵药”： 虽然粒子移动部分的计算速度得到了巨大提升，但模拟的其他部分（例如求解磁场方程的部分）仍然在 CPU 上运行，因为将这些特定部分转移到 GPU 的工具尚未准备就绪。

总结

简而言之，这篇论文讲述的是一个关于工程智慧的故事。作者将一个沉重、缓慢且复杂的模拟代码，成功地教会了如何在现代、强大的显卡上运行。他们穿越了软件漏洞和编译器限制的雷区，创造出了一个能在两种不同硬件上运行的版本，证明了我们可以在不损失准确性的前提下，更快速地模拟聚变等离子体。这是朝着使核聚变研究更高效迈出的关键一步，尽管实现完全自动化、完美翻译的旅程尚未结束。

技术摘要

技术摘要：利用 OpenMP 实现 TRIMEG 代码在 Gyrokinetic 等离子体模拟中的 GPU 加速与可移植性

问题陈述
等离子体物理模拟，特别是用于研究托卡马克聚变装置中不稳定性与湍流的回旋动力学（gyrokinetic）模型，具有极高的计算强度。TRIMEG 代码是一种利用非结构化三角形网格上的 C1 有限元方法的高精度粒子-网格（PIC）求解器，由于实现现实模拟通常需要海量粒子（通常为 $10^7$ 至 $10^8$ 个），面临着显著的执行时间挑战。虽然该代码已经采用了 MPI 进行多节点并行化，但粒子推进（particle pushing）和网格-粒子（G2P）操作构成了主要的性能瓶颈，占据了高达 80% 的总执行时间。挑战在于如何使用图形处理器（GPU）加速这些特定的“热点”区域，同时保持代码复杂的面向对象结构（包括多态性和派生类型），并确保在不同硬件架构（特别是 NVIDIA 和 AMD）之间的可移植性。

方法论
本研究重点在于使用 OpenMP Offloading API（版本 4.0 及更高版本）将 TRIMEG 代码移植到 GPU 架构。研究方法包括：

1. 目标选择： 确定粒子推进内核及相关的 G2P 操作（回溯、密度计算以及分布函数插值）为主要的加速目标，因为它们具有高算术强度且缺乏粒子间的依赖关系。
2. 面向可移植性的代码重构： 为了克服 amdflang (AMD) 和 nvfortran (NVIDIA) 编译器的限制，需要进行大量的重构。主要挑战包括：
   • 多态性： 两种编译器在 GPU target 区域内处理 class() 派生类型和类型绑定过程（type-bound procedures）时都表现不佳。解决方案是尽可能将代码重构为使用非多态的 type() 声明，并利用基类/扩展类层次结构以及通过 Fortran INCLUDE 指令复制函数体的方法，来解决粒子类与场类之间的循环依赖问题。
   • 动态数组： 代码使用了一个模仿 C++ vector 的自定义库来实现动态数组。由于 GPU 内核难以处理动态分配或复杂的指针间接寻址，因此将绑定框与网格三角形之间的映射结构从“结构体数组”展平为“一维数组”，以实现高效的内存传输。
   • 内存管理： 在初始化阶段实施预先内存分配，以最大限度地减少运行时延迟。对于 AMD 平台，尽可能利用统一共享内存（USM）；而对于不支持 USM 的 NVIDIA 平台，则使用显式的 enter data、update 和 exit data 指令。
   • 数值一致性： 为确保 GPU 结果与 CPU 参考值匹配，禁用了会改变浮点运算的编译器优化（如 Fused-Multiply-Add 指令），例如 AMD 使用 -ffp-contract=off，NVIDIA 使用 -Mnofma。通过将外部 B-spline 库从共享对象成员切换为局部声明的自动数组，解决了竞态条件问题。
3. 性能评估： 实现方案在 Viper 集群 (AMD MI300A)、Raven (NVIDIA A100) 和 Pitagora (NVIDIA H100) 集群上进行了测试。性能评估通过以下方式进行：
   • 内核剖析（Profiling）： 使用 rocprof-compute 和 nsys 来分析资源占用率、内存带宽和指令组合。
   • 可扩展性研究： 进行强扩展性测试，以评估混合 MPI-OpenMP offloading 的效率，特别是在考察 GPU 过度订阅（oversubscribing）多个 MPI 进程时的影响。
   • 网格规模探索： 通过调整 OpenMP team 的数量和每个 team 中的线程数来最大化硬件利用率。

核心贡献

• 首次跨厂商移植： 本工作展示了通过 OpenMP offloading，将复杂的面向对象 Fortran 回旋动力学代码同时移植到 NVIDIA 和 AMD GPU 上的开创性尝试。
• 编译器规避方案： 本论文记录了针对编译器在多态、动态数组和过程指针方面局限性的特定且非平凡的规避方法。同时也强调了 nvfortran 和 amdflang 在这些功能方面缺乏全面文档的情况。
• 混合并行化分析： 研究对混合 MPI-OpenMP offloading 的权衡进行了详细分析，表明虽然 GPU 加速对粒子推进器非常有效，但由于原代码缺乏 OpenMP 多线程支持，必须通过过度订阅来实现，这可能会限制在具有高核心数但 GPU 资源有限的节点上的扩展性。
• 数值验证： 实现过程包含了严格的验证流程，通过对比能量增长率和 2D 模态结构与 CPU 结果，确认了 GPU 版本能够高保真地重现物理特性，尽管由于编译器特定的浮点处理可能存在微小的数值偏差。

结果

• 加速比： 对于 $32 \times 10^6$ 个电子的现实工作负载，与 TOK 集群上使用 GCC 编译的 CPU 版本相比，GPU 实现方案在 AMD Viper 节点上实现了约 14.8 倍 的加速，在 NVIDIA Pitagora 节点上实现了约 29.6 倍 的加速。
• 内核效率： 粒子推进内核占据了运行时间的大部分。在 AMD MI300A 上的剖析显示，该内核具有高算术强度，L1/L2 缓存命中率超过 80%，但仅有 18% 的内存访问是合并（coalesced）的。
• 可扩展性限制： 强扩展性测试显示，虽然 GPU 加速部分扩展良好，但整体应用加速受到非加速部分（如使用 PETSc 的场求解器）以及 GPU 过度订阅开销的制约。在 NVIDIA Pitagora 集群上发现，由于测试使用的编译器版本（nvfortran 24.9）不支持通过 OpenMP 进行多 GPU 支持，导致无法同时利用节点上的所有可用 GPU。
• 正确性： 对 Cyclone 情形（ITG 模态）和 TCV-X21 情形（非线性 ITG 不稳定性）的模拟证实，GPU 版本能够正确重现 CPU 版本观察到的能量增长率和模态结构，其差异归因于随机数生成器的初始化以及编译器特定的浮点运算差异，而非算法错误。

意义与主张
本文主张，虽然 OpenMP offloading 为不同 HPC 架构间的可移植性提供了一条充满希望的路径，但它对于复杂的遗留代码而言并非“无缝”的解决方案。这项工作证明，要实现一个高性能的 GPU 版本，需要进行深入的编译器探索和大规模的代码重构，以绕过当前编译器对高级 Fortran 特性支持的局限性。

作者强调，这种可移植性的成功很大程度上取决于具体的编译器版本，而非仅仅取决于编程范式。他们得出结论：TRIMEG GPU 实现是一个功能完备且准确的工具，能够为计算量最大的部分提供实质性的加速。然而，他们也谦逊地指出，由于编译器对多设备 offloading 支持尚不成熟，以及底层 CPU 代码结构缺乏 OpenMP 多线程支持，硬件的全部潜力（尤其是多 GPU 节点）目前仍受到限制。这项工作为其他试图将复杂 Fortran 代码移植到异构架构的研究者提供了实用的指南和“替代性文档”。