商汤 HPC 团队携手 DeepModeling，实现燃烧 CFD 全隐式求解器全流程 GPU 计算

商汤高性能计算 (HPC) 团队秉承开源之心，与北京科学智能研究院 (AI for Science Institute, Beijing)的 DeepFlame团队合作，助力 DeepModeling 社区打造整合下一代异构超级计算和人工智能加速基础设施(如GPU和FPGA)的燃烧流体力学计算开源软件。

前言

DeepFlame是面向AI for Science时代开发的燃烧流体计算开源平台，旨在突破传统计算流体力学（Computational Fluid Dynamics，简称CFD）在燃烧领域长期以来的应用落地难题。

商汤科技HPC团队多年来深耕计算优化技术，深入研究多种底层硬件架构，自研高性能神经网络算子库、图像处理算子库和代数计算库等核心计算引擎，并面向社区开源了高性能深度学习推理平台OpenPPL，以及端到端部署工具集。

商汤科技HPC团队携手DeepFlame团队，共同探索燃烧反应流体力学计算中的HPC异构加速。此前，在2023年6月底发布的DeepFlame版本1.2中，我们实现了有限体积法求解Navier-Stokes方程的GPU加速。此次，DeepFlame新版本1.3于2024年1月2日发布，历经半年的开发和测试，我们完成了低马赫数全隐式求解器dfLowMachFoam的全流程GPU实现与优化。

低马赫数全隐式求解器的全流程GPU实现

相比于显式求解方法，隐式求解方法具有较强的数值稳定性同时可以适应大步长的时间推进。因此，隐式方法在大部分仿真场景尤其是低马赫数场景下具有更高的计算效率，在工程应用领域中得到了广泛的应用。DeepFlame的低马赫数求解器dfLowMachFoam采用了全隐式求解方法，在零维均质点火、一维层流预混火焰、二维三叉火焰等简单流动场景，以及二维各项同性湍流预混火焰、二维/三维泰勒格林涡火焰等复杂流动场景下，该求解器都展现出了优异的计算性能和仿真精度。
为了实现全隐式求解器的全流程GPU加速，我们对全隐式求解器的工作流进行了精心设计，如下图所示。一旦初始化数据被传输到GPU，时间步推进所涉及的所有计算都在GPU上进行。具体而言，主要包括三类：1）密度方程（rhoEqn）、动量方程（UEqn）、组分方程（YEqn）、能量方程（EEqn）以及压力方程（pEqn）的全隐式PDE求解，2）化学源项的神经网络推理，3）热物性求解。其中，基于有限体积方法的隐式离散、热物性和输运参数的计算均通过自研CUDA kernel实现；神经网络推理和线性方程组的求解则借助第三方库libTorch和AmgX完成。

低马赫数全隐式求解器的GPU优化

为了充分发挥GPU的能力，我们进行了一系列的性能优化与显存优化。在性能方面，通过改进kernel实现和kernel launch过程来提升计算效率。在存储方面，通过静态数据复用和动态数据分配来减少GPU上的显存占用。

1.性能优化

a.kernel优化

算法大体上可以分为访存约束和计算约束两类。在本项目中，基于有限体积法的隐式离散通常显示出访存约束的特性，热物性和输运参数相关的环节则表现出更高的计算强度。

对于访存约束的kernel，最重要的优化策略是全局内存的合并访问。项目中广泛存在向量数组和张量数组，每个数组元素包含多个分量。在CPU代码中沿着分量方向存储，这有助于连续数据访问，然而它不利于GPU上的数据局部性。因此，我们合并相同分量，改为存储分量数组，从而提高了GPU带宽效率。

对于计算约束的kernel，我们充分利用多级存储发挥数据复用的潜力，例如采用shared memory存储质量和摩尔分数，采用constant memory存储常系数等。我们还对一些任务繁重的kernel进行计算划分，有效地减少了每个kernel的寄存器用量，进而提高GPU占用率。此外，我们还对sqrt()和pow()等长延迟计算指令进行了等效替代优化。

b.kernel launch优化

GPU kernel的每次执行都需要在CPU上执行一次kernel launch操作，从而产生额外的开销。隐式离散中边界计算的特征是kernel数量多，kernel计算量小，这导致kernel执行时间法充分掩盖kernel launch开销。我们通过CUDA Graph机制进行了针对性的优化。考虑到隐式离散过程中每个时间步上的kernel顺序乃至kernel参数都是固定的，这些kernel可以封装在CUDA Graph中，从而可以将一个时间步的一系列kernel一次性提交到GPU。这种方法显著提高了隐式离散过程中的kernel launch效率。

2.显存优化

a.静态数据复用

我们对求解器全流程用到的数据进行分类，以最大化数据复用。网格信息、热物性和输运参数等全局数据可以跨控制方程共享，AmgX求解过程中的矩阵A、向量b以及AmgX handle等也实现了跨控制方程的复用。此外，我们在各个隐式离散环节内部提高了离散化操作的粒度，消除了大量中间数据的存储。这些策略有效地减轻了GPU显存负担。

b.动态数据分配

基于CUDA流分配器和CUDA Graph机制，我们进一步实现了动态的显存分配和释放。它避免了传统的内存分配接口带来的昂贵的设备同步开销，从而以极低的性能代价实现动态内存管理。每个动态更新的场量会在每个时间步内不断地异步创建、异步使用、异步销毁，并通过CUDA Graph批量提交。该方案可以最小化各个场量的生命周期，大大降低GPU显存开销。

成果展示

本版本发展的全流程GPU框架在准DNS（直接数值模拟）的三维泰勒格林涡火焰模拟中进行了性能测试和显存测试。其中，CPU模拟在一个32核CPU芯片上进行（32个MPI进程），GPU模拟在4卡V100上进行，这也是当前异构计算服务器的典型配置（1CPU+4GPU，或者2CPU+8GPU）。

如下图(a)所示，将化学源项计算转换为GPU神经网络推理后，整体加速达到6.25倍。进一步地将所有剩余的操作（PDE求解、热物性求解）移植到GPU后，整体加速提升到近35倍。下图(b)详细说明了PDE过程（主要包括隐式离散化和求解线性系统）的计算成本，每一组数据的左栏和右栏分别表示CPU和GPU的时间。PDE过程的总体加速约10倍。其中，隐式离散化是由项目自研的CUDA kernel实现的，获得了75倍的加速。相比之下，线性系统求解由AmgX库提供支持，显示出相对较低的加速，后续需要对该环节进行进一步优化。

显存测试情况如下表所示。对于256 256 256的三维泰勒格林涡火焰模拟，每张卡上占用的显存总量大约为11GB。 CFD数据（包括网格和场信息）约占总使用量的44%——它已经通过动态数据复用减少了2.3倍。跨方程复用的AmgX句柄和CUDA句柄分别消耗了35%和17%的GPU内存——它们主要由编译后的kernel指令、CUDA上下文和其他CUDA资源组成，因此，随着问题规模的增大，它们会表现出最小的增长。DNN推理过程在优化批处理大小和数值精度后仅占用不足4%的显存。

本版本发展的全流程GPU框架还进行了剑桥燃烧器的LES（大涡模拟）仿真，它是众所周知的分层预混旋流燃烧的基准测试，可以帮助研究人员更好地了解燃烧中的物理和化学过程，为实际的燃烧器设计和优化提供依据。测试网格规模从250万至1000万，这涵盖了LES研究的典型范围。250网格规模的模拟可以在单个GPU上进行，1000万网格规模使用4个GPU卡完成。该测例进一步证明了DeepFlame低马赫流求解器的全流程GPU框架的实际适用性。

剑桥分层预混旋流燃烧仿真温度云图

上述一系列测试也意味着，对于具有上千万个网格的准DNS或者LES场景，可以在一个小型工作站内，使用4张GPU卡，以低得多的功耗，提供与具有数千个CPU核的超级计算机相当的性能。

下一步计划

1.实现几何多重网格求解（GAMG）的GPU加速。

当前，DeepFlame的线性系统求解由Nvidia AmgX库提供支持，显示出相对较低的加速比，特别是在求解压力泊松方程时。这主要是由于GPU侧的AmgX库只能使用代数多重网格方法（AMG），而CPU侧采用了更高效的几何多重网格方法（GAMG）。为了获得更好的整体加速效果，后续将对该环节进行进一步优化。

2.探索神经网络模型压缩和推理加速。

在DeepFlame 1.0版本中，神经网络替代化学源项ODE求解之后，计算热点转移到了PDE离散系数矩阵构造和线性系统求解上；此次DeepFlame 1.3版本的更新实现了PDE的全流程GPU加速，计算热点再次回到了化学源项，也就是神经网络的推理之上。因此，模型压缩和推理加速将是一下阶段的重要目标。而DeepFlame的开发正式进入了AI与HPC双驱动模式，未来的优化将是有趣的“此消彼长”相竞状态。