DALES大气模型GPU加速:OpenACC实现与优化策略
1. 荷兰大气大涡模拟(DALES)模型概述
荷兰大气大涡模拟(DALES)模型是一个高分辨率的大气数值模拟工具,专门用于研究大气边界层中的湍流、对流和云物理过程。作为大气科学研究的重要工具,DALES能够模拟水平分辨率10-200米、垂直分辨率10-50米尺度的大气现象,这使得它能够显式地模拟云团和对流上升气流等精细结构。
DALES模型采用Fortran 90编写,核心架构基于MPI的域分解并行策略。整个代码库约7.5万行,分为物理过程、数值计算和工具模块三大类。模型使用三维数组作为基础数据结构,在初始化时一次性分配内存,因为计算网格是固定的。特别值得注意的是,数组的ijk索引中,k方向(垂直方向)不是连续存储的,这种内存布局对后续的GPU优化提出了挑战。
提示:DALES采用z-pencil(垂直于地面的方向)的域分解方式,生成笛卡尔子域,其水平方向的大小与MPI进程数成反比。每个分区在水平方向添加halo cells(通常1-3层)来确保跨MPI进程的数据一致性,通过非阻塞的发送/接收操作进行数据交换。
2. GPU加速的必要性与挑战
2.1 大气模拟的计算需求
大气数值模型是计算最为密集的物理模型之一。其计算复杂性来源于气流动力学的混沌特性,以及与各种传输过程(如降水、化学反应、太阳和热红外辐射、海洋/陆地相互作用)的耦合。模拟的准确性和可靠性强烈依赖于模型离散化的空间和时间分辨率,反过来又依赖于对大量计算资源的高效利用。
传统上,大气数值模型依赖于在强大超级计算机上的并行实现。过去十年中,GPU加速器在高性能计算(HPC)中的大规模使用改变了这一格局。目前,GPU占据了Top500超级计算机中四分之三的计算能力。气候学界因此探索了多种有效利用GPU的途径。
2.2 GPU加速路径选择
早期工作集中在将选定的计算密集型内核转换为CUDA,实现了单个内核最高20倍的加速。随后的努力旨在加速整个求解器,导致了NIM-CUDA、GALES或MicroHH等求解器的开发,后两者引入了先前Fortran代码库的C++/CUDA重写。然而,这种激进的方法需要大量的代码开发投入,对用户来说意味着颠覆性的改变。
近年来,基于指令的加速方法(如OpenACC)已成功应用于天气和气候Fortran代码库,在实现GPU加速的同时,大幅减少了代码修改量。考虑到DALES代码库的规模(约7.5万行)和模型参数调优与验证的长期性,完全用更现代、对GPU更友好的编程语言(如C++)重写求解器并不可行。此外,DALES被研究人员和学生广泛使用,后者通常计算机科学背景有限,这使得Fortran的命令式编程风格成为一个合理的选择。
3. OpenACC加速实现策略
3.1 数据管理优化
DALES中的数据容器分布在物理、数值和工具Fortran模块中,每个模块管理其计算所需的多维数组的动态内存分配和释放。在早期GPU(如NVIDIA Kepler及更早版本)上,有限的设备内存需要仔细管理加载到设备内存中的数据,数据在设备和主机之间的传输被发现是影响性能的主要瓶颈。
利用现代GPU更大的内存,DALES采用了将整个CPU内存镜像到设备上的策略。每个模块在CPU上初始化数据后使用enter data copyin子句一次性将数据传输到设备,在CPU上释放数组前使用exit data delete子句清除设备内存。这样,在启动计算内核时,数据总是假定已经存在于设备上,只有当执行诊断或检查点操作时才发生从设备到主机内存的数据拷贝。
3.2 核心移植策略
DALES典型的计算子程序可分为两类:
- 嵌套的独立ijk或ijkn循环,遍历三个空间维度和可能的第四个标量索引空间维度
- 类似的嵌套循环,但在对应于垂直方向的k索引上存在依赖关系
考虑到DALES典型用例中ijk循环大小的广泛变化,我们没有针对特定案例大小过早优化OpenACC编译指示。相反,默认方法是在尽可能多的紧密嵌套循环上使用collapse子句,以向GPU暴露最大程度的并行性,同时保持可移植性。
对于包含地表层(k=1)特殊处理的ijk循环,为避免设备代码中的分支和由此产生的线程发散,k=1的情况被分离到单独的内核中,使用async子句来隐藏启动延迟。对于带有k依赖的ijk循环,当依赖严格是加法性时使用原子操作(如在云沉降内核中,云滴在相邻k层之间转移),或者引入临时数组。
3.3 代码重构关键技术
为了能够collapse尽可能多的循环,我们对多个计算内核进行了局部重构,将仅依赖k的计算移到最内层循环,以重复计算为代价。在微物理和热力学部分需要更广泛的重构,因为局部状态的差异会导致显著的线程发散。两个典型的重构例子:
雨微物理计算优化:DALES中的所有雨微物理计算都使用掩码,以利用雨滴通常存在于有限的k层范围内并在展向方向上"阵雨"中聚集的观察结果。最初的CPU实现基于k-only循环结合Fortran数组语法和exit语句,在GPU上表现不佳。我们改用完全collapse的ijk循环上的归约操作,但需要为CPU和GPU维护不同的代码版本,因为GPU方法在CPU上被证明更昂贵。
热力学饱和调整方案:DALES在热力学例程中使用饱和调整方案来诊断云液态水含量。为相应调整液态水位温,需要Newton-Raphson算法。每个计算单元中的局部条件可能导致达到所需精度所需的Newton迭代次数不同,引入线程发散。通过展开固定数量的迭代(精度足够)替换迭代算法,缓解了发散问题,在CPU和GPU上都得到了更高效的内核。
4. 外部库的GPU加速
4.1 辐射传输计算
辐射传输计算使用RTE-RRTMGP库执行。该库由两部分组成,均用Fortran实现并通过OpenACC指令支持GPU卸载:
- RRTMGP:通过插值表格数据从大气输入状态计算分子和云的辐射特性
- RTE:解决辐射问题并输出跨空间域辐射通量
由于内存限制,实际上不可能同时运行所有列计算。总列数因此被分成若干批次。最大化每批列数(或最小化批次数)很重要,因为多次调用RRTMGP和RTE库会带来与内核启动开销相关的额外计算成本。
4.2 泊松求解器的FFT
为了在GPU上执行泊松求解器所需的3D FFT,我们依赖cuFFT库。傅里叶变换在一个空间方向接一个方向地执行。因为傅里叶变换是非局部的,它需要访问沿执行方向的整个场数据,跨越域分区。这需要通过执行转置来重新组织3D场内存布局,这需要密集的全通信。
值得注意的是,cuFFT不支持DALES中用于N大小实数到复数变换的FFTW"半复数"编号格式。我们采用了CaNS中使用的方法,对cuFFT的输入/输出进行最小化的前/后处理步骤,以及修改波数的重新排序。
5. 性能评估与优化
5.1 测试案例与验证
我们选择Cloud Botany模拟集合中的一个案例作为GPU移植项目的测试案例。Cloud Botany是一组103个DALES模拟,研究浅积云在亚热带海洋上的行为,探索云及其中尺度组织在不同物理条件下的变化。模拟域在两个水平方向上宽150公里,足够大以允许在更大尺度(称为中尺度)上的组织。
为确保OpenACC实现的正确性,我们比较了CPU和GPU实现的模型输出。由于系统的混沌性质,舍入误差会迅速导致不同的模拟轨迹。因此,我们基于空间和时间平均统计数据进行比较,这些数据应该保持接近。结果显示,对于一阶和二阶统计量(湍流动能),GPU和CPU输出表现出极好的一致性,仅在降水平均液滴直径上观察到局部最大相对差异约3%。
5.2 单节点性能
我们在荷兰国家超级计算机Snellius上比较了不同GPU硬件的性能:
- CPU节点:192个AMD Genoa 9654核心
- GPU节点:4个NVIDIA A100
- GPU节点:4个NVIDIA H100
测试调整了水平方向的网格点数以适应不同的MPI空间域分区,确保填充GPU内存。性能测量在模型时间10分钟的短模拟中进行,仅包括时间步进循环中花费的时间,实验重复5次。
结果显示,在节点基础上,DALES在A100上的整体加速比约为4倍,在H100上约为11.5倍。各代码组件的耗时比例保持接近其CPU值,除了A100上的泊松求解器性能低于平均水平,加速比接近2。进一步测试表明,这与NVHPC 24.5(CUDA 12.1.1)中的cuFFT求解器效率较低有关。
5.3 弱扩展性分析
为了评估DALES随研究需求扩展的能力,A100案例从1个GPU扩展到64个GPU(0.25到16个计算节点),利用测试案例的周期性在水平方向复制初始数据以保持每个GPU的工作量不变。结果显示DALES的并行效率下降较快,主要由泊松求解器的较差扩展性能驱动,而其他组件由于几乎不需要通信,保持接近1的效率。
值得注意的是,本研究未考虑强扩展,因为加速的目标是处理更大的空间域而非缩短现有问题的计算时间。
6. 内核优化与自动调优
6.1 内核调优的必要性
第2.3节介绍的加速策略能够快速加速整个代码库,在核启动参数方面给予编译器完全自由,并确保相对良好的可移植性。然而,少量内核表现出低占用率(<50%),需要更仔细的考虑。我们采用两种策略来提高这些内核的性能:
- 手动修改内核以改善内存访问、缓存使用和寄存器压力
- 使用Kernel Tuner(KT)自动调整内核启动参数
6.2 Kernel Tuner简介
Kernel Tuner(KT)是一个通用的GPU应用程序自动调优框架,用Python编写,支持调优用CUDA、HIP和OpenCL编写的计算内核,最近还扩展到处理基于指令的语言如OpenACC和OpenMP。
KT提供了一组辅助函数,允许用户从代码其余部分提取可调优内核以及所有必要的C预处理器宏,并负责所用数据结构的分配和初始化。用户通过向源代码添加特定的KT编译指示来指导可调优内核的选择和数据结构的识别,这些编译指示是对原始代码的唯一必要更改。
6.3 目标内核优化
优化的计算内核是平流和次网格尺度扩散算子中的基于模板的内核。这些内核合计占总计算时间的20%。特别是标量次网格扩散内核diffcsv(类似于列表2但包括标量的外部循环和4维Fortran数组)和sources内核(评估交叉导数并显著重用相邻数据)是优化的主要目标。
手动优化diffcsv内核采用了三种策略:
- 使用cache指令将各向异性网格间距数据加载到软件管理的缓存中,而非多次读取1D数组
- 内核分裂以减少寄存器压力(从108降到62+56)
- 用ijk collapse循环和n sequential循环替换ijkn collapse循环
结果显示,结合策略1和3,次网格模块的整体性能可提高近5%。而通过分裂长内核来减少寄存器压力导致性能下降,因为生成的内核实际上是顺序运行的。
使用KT对diffcsv内核进行自动调优,在A100上实现了最高35%的性能提升(n=10的情况),在H100上仅获得最高10%的提升。有趣的是,H100上的时间分布不是A100的简单平移,两台GPU的最佳参数集略有不同。然而,在H100上获得最佳性能的参数在A100上也能获得接近最佳的性能(差距在5%以内),反之亦然。
对于sources内核,当仅使用collapse子句时,KT未能找到比默认参数更好的启动参数。转而使用tile(tx,ty,tz)子句和手动分块两种方法。结果显示,使用OpenACC tile在A100和H100上都获得了最高15%的性能提升,最佳性能使用tx=64、ty=4和tz=2获得。相比之下,手动引入分块最多导致性能小幅下降,大多数情况下性能更差。
7. 实际应用建议
7.1 移植策略选择
对于类似DALES的大型传统科学计算代码,我们推荐采用渐进式的GPU移植策略:
- 首先使用OpenACC指令实现基础版本,利用编译器的自动优化能力
- 通过性能分析识别热点内核
- 对关键内核进行针对性优化,包括:
- 数据局部性优化
- 循环结构调整
- 使用高级OpenACC特性如cache、tile等
- 考虑使用自动调优工具如Kernel Tuner寻找最优参数
7.2 性能优化技巧
基于DALES项目的经验,我们总结了以下性能优化建议:
- 内存管理:现代GPU具有足够的内存容量,建议采用"镜像"策略,一次性将所需数据传输到GPU,减少主机与设备间的数据传输
- 循环优化:
- 优先使用collapse暴露最大并行度
- 对于存在条件分支的循环,考虑分离特殊条件到独立内核
- 对于存在垂直依赖的循环,使用原子操作或临时数组
- 内核融合:将连续独立的内核融合,减少内核启动开销
- 外部库使用:优先选择已有GPU支持的库,注意内存布局兼容性
7.3 多GPU扩展建议
对于需要多GPU扩展的应用,我们建议:
- 首先优化单GPU性能,确保每个GPU的计算效率最大化
- 分析通信模式,识别瓶颈(如DALES中的泊松求解器)
- 考虑采用混合并行策略,结合MPI和OpenACC
- 对于弱扩展场景,确保每个GPU的工作负载均衡
8. 总结与展望
通过OpenACC指令集实现的DALES GPU加速项目展示了传统Fortran科学计算代码在现代GPU架构上获得显著性能提升的可行路径。在不完全重写代码的情况下,我们在NVIDIA A100上实现了约4倍的加速,在更新的H100上实现了11.5倍的加速。这一成果使得研究人员能够在相同时间内模拟更大的空间域或更高分辨率的大气过程,为大气科学研究提供了更强大的工具。
特别值得注意的是,自动调优工具如Kernel Tuner在优化特定内核方面表现出色,特别是在较旧的A100架构上可实现最高35%的额外性能提升。然而,随着GPU架构的演进(如从A100到H100),自动优化的收益有所降低,这可能表明新一代GPU的编译器优化已经相当成熟。
未来工作可以集中在以下几个方向:
- 改进多GPU扩展性能,特别是泊松求解器的通信模式
- 探索更多内核优化机会,特别是内存访问模式的进一步优化
- 将经验应用到其他大气和气候模型中
- 评估在新一代GPU架构(如NVIDIA Blackwell)上的性能表现
从工程实践角度看,DALES项目证实了OpenACC作为传统科学代码GPU移植工具的价值,它平衡了实现工作量和性能收益,使研究团队能够在不完全重写代码库的情况下利用现代GPU的计算能力。这一经验对于众多类似规模和历史代码库的科学计算项目具有重要参考价值。
