基于Kokkos的MOOSE框架GPU多物理场仿真：静态多态性设计与性能可移植实现

《SoftwareX》：Kokkos-MOOSE: GPU-based multiphysics simulation in the MOOSE framework

　　

Kokkos-MOOSE扩展了MOOSE（Multiphysics Object-Oriented Simulation Environment，多物理场面向对象仿真环境），为其增加了基于Kokkos构建的GPU兼容有限元组件。该方法采用静态多态性技术，保留了MOOSE高度依赖动态多态性的熟悉对象模型。Kokkos-MOOSE对MOOSE设计理念和对象模型的紧密遵循，最大限度地减少了开发人员移植现有MOOSE应用程序所需的工作量，同时通过GPU加速实现了显著的性能提升。此外，Kokkos提供的跨架构可移植性，结合MOOSE的面向对象架构及其丰富的可复用框架组件集合，消除了为异构计算平台重写整个应用程序的必要性。

一、研究背景与问题

MOOSE是一个用于构建领域应用（尤其在核工程领域）的有限元多物理场框架，其底层依赖libMesh提供有限元计算机制，并通过PETSc实现可扩展的非线性与线性求解器。迄今为止，MOOSE主要依靠消息传递接口在多核CPU上实现大规模问题的并行扩展。

然而，大规模多物理场仿真越来越多地运行在异构高性能计算系统上，其计算性能主要来自GPU。这给MOOSE这类成熟框架带来了设计挑战：一方面，必须向GPU暴露足够的并行度；另一方面，不能抛弃现有用户所依赖的应用接口，特别是支持多态性的灵活扩展点和架构模式。这些通常通过模板方法模式实现——一种行为设计模式，在基类中定义算法的骨架，允许派生类在不改变整体结构的情况下覆盖某些步骤。难点在于对动态多态性的重度依赖，这在GPU后端中存在限制和性能开销。

二、研究内容与结论

Kokkos-MOOSE通过一种静态多态性设计模式解决了上述问题，该模式保留了熟悉的MOOSE编程接口，降低了采用门槛，并最大程度减少了移植现有应用的工作量。保留原始设计至关重要，因为它使应用开发者能够以最小干扰继续使用框架的广泛基础设施，同时获得GPU加速效果。此外，通过构建在Kokkos之上——一种用于跨不同硬件实现性能可移植代码的C++编程模型——Kokkos-MOOSE避免了针对不同架构维护独立代码库，从而降低了长期维护成本，提高了对当前和未来HPC系统的可移植性。

三、关键技术方法

1. 1.
   分离式编译系统：Kokkos-MOOSE源文件通过专用*.K扩展名标识并使用GPU工具链编译，而其他MOOSE源码沿用标准CPU工具链。系统定义了预处理器保护宏MOOSE_KOKKOS_ENABLED、MOOSE_ENABLE_KOKKOS_GPU和MOOSE_KOKKOS_SCOPE，分别用于条件性暴露Kokkos功能、GPU专属功能和隔离Kokkos符号。这种模型将Kokkos特定代码局限在专用源文件中，使其影响局部化，确保新的GPU实现能与庞大的遗留CPU代码库共存。
2. 2.
   数据容器体系：基础容器Moose::Kokkos::Array替代Kokkos Views，存储对两个编译作用域均可见的普通类型数据，同时支持嵌套数组和显式的宿主/设备数据拷贝操作。在此基础上构建了Moose::Kokkos::JaggedArray和Moose::Kokkos::Map等高级容器。OpenMP后端模拟相同的宿主-设备模型，使CPU和GPU构建在行为上保持一致，便于在CPU调试和测试期间暴露设备内存错误和同步缺失问题。
3. 3.
   值绑定机制：Kokkos-MOOSE倾向于使用可拷贝的值、轻量级句柄或包装器作为设备端成员，而非宿主引用。Moose::Kokkos::ReferenceWrapper在需要引用时提供支持。这些句柄和包装器利用其拷贝构造函数作为钩子来同步宿主和设备数据，利用Kokkos在每次并行调度前拷贝函子的特性，使设备代码无需用户显式同步即可访问最新数据。
4. 4.
   静态多态性调度架构：基类拥有组装算法，派生类仅提供弱形式评估。公共钩子函数标记为KOKKOS_FUNCTION，线程私有索引和上下文作为参数传入。基类入口函数通过重载的operator()按操作类型标记并以具体对象类型Derived为模板参数，实现对派生类方法的静态绑定。注册宏将分派器注册到分派器注册表中，分派器作为Kokkos函子，其parallelFor()启动并行循环，通过类型擦除基类的虚函数接口供宿主代码调用。

四、研究结果

4.1 瞬态热传导示例

研究人员展示了一个5MW_th微反应器的瞬态热传导仿真，模型采用1/12对称性，包含190,744个单元和1,394,008个自由度（DOFs）。温度场使用二阶拉格朗日基函数，热管由800K体温度的対流边界条件表示，外部表面施加300K环境的对流和辐射边界条件。空间变化的热物理性质通过材料赋值。瞬态运行40,000秒，共61个时间步，最终达到稳态。

问题采用牛顿-克雷洛夫方法求解：J(u_n)δu_n+1 = -r(u_n)，u_n+1 = u_n + δu_n+1。残差和雅可比矩阵由MOOSE内核和边界条件组装，修正项通过PETSc的克雷洛夫求解器迭代求解。

4.2 性能对比

CPU计算使用两颗Intel Xeon Platinum 8480+处理器（112核），GPU计算使用单块NVIDIA H200 GPU。残差评估实现10.0倍加速，雅可比评估实现27.4倍加速，展示了Kokkos-MOOSE以最小移植代价有效加速MOOSE有限元组装的能力。求解阶段加速比较低（3.3倍），限制了整体性能（总加速比5.9倍）。GPU计算所需的线性迭代次数约为CPU的2.5倍，归因于CPU与GPU实现的HYPRE BoomerAMG预处理器之间的差异。

在与原始MOOSE的CPU对比中，Kokkos-MOOSE在112个CPU核上运行速度约慢20%，主要由于雅可比评估耗时增加50%，且内存使用略有增加。

五、总结与讨论

Kokkos-MOOSE通过静态多态性设计模式成功保留了MOOSE的编程接口，使应用开发者能以最小干扰实现GPU加速。该工作证明了在保持大型框架对象模型完整性的前提下，通过精心设计的编译分离、数据容器和调度机制，可以实现有效的GPU加速。研究结果表明，残差和雅可比组装的加速效果显著，但求解器性能成为整体加速的主要瓶颈。因此，未来的开发重点将放在调整PETSc/HYPRE的GPU参数以及优化与GPU版PETSc/HYPRE的集成上。值得注意的是，当前版本的Kokkos-MOOSE尚未提供所有MOOSE功能的完整移植，高级特性如网格位移与自适应、不连续伽辽金（DG）内核和接触有限元尚未支持，将逐步添加。