GPU acceleration of plane-wave density functional theory calculations in Abinit

本文报告了 Abinit 平面波密度泛函理论代码向多 GPU 架构的移植工作，通过算法优化与库函数利用，详细对比了 CPU 与异构节点性能，并重点评估了 LOBPCG 与切比雪夫多项式滤波两种对角化算法在 GPU 上的效率。

要点

这篇论文讲述了一个关于如何让超级计算机“跑”得更快、更省电的故事。

想象一下，科学家们在研究新材料（比如更轻的飞机、更高效的电池）时，需要模拟原子和电子的行为。这就像是在玩一个极其复杂的乐高游戏，但规则是由量子力学决定的。为了算出这些规则，他们使用一种叫Abinit的软件。

过去，这个软件主要靠CPU（电脑的“大脑”）来算。但随着问题越来越复杂（比如要模拟成千上万个电子），CPU 就像是一个勤劳但速度有限的老工匠，算得太慢了。

于是，科学家们决定给 Abinit 装上GPU（图形处理器，原本是用来画游戏画面的“超级画师”）。GPU 拥有成千上万个“小助手”，非常适合处理这种需要大量重复计算的任务。

这篇论文就是他们把 Abinit 成功移植到 GPU 上的经验总结。为了让你更容易理解，我们可以用几个生动的比喻来拆解他们的做法：

1. 核心挑战：从“单兵作战”到“集团军作战”

• 旧模式（CPU）：就像让一个工匠一次只处理一块砖。虽然他很细心，但面对成千上万块砖，效率很低。
• 新模式（GPU）：GPU 有几千个工人。如果还是让他们一次只搬一块砖，那太浪费了。
• 解决方案（批处理/Batch Processing）：
  科学家们发现，与其让每个工人单独干活，不如把砖块打包。
  • 比喻：以前是“一次搬一块砖”，现在是“一次搬一卡车砖”。
  • 在 Abinit 中，他们把电子波函数（计算的核心数据）打包成“批次”。这样，GPU 就能一次性处理大量数据，而不是一个个慢慢算。这就像是用大卡车运货，而不是用自行车。

2. 数据搬家：减少“往返跑”

• 问题：GPU 和 CPU 之间有一条“高速公路”（内存传输）。如果数据在 CPU 和 GPU 之间频繁来回跑，大部分时间都花在“堵车”上，而不是干活上。
• 解决方案（驻留内存）：
  科学家们设计了一个策略，让数据一旦上船，就留在船上。
  • 比喻：以前是每算一步，就把货物从仓库（CPU）搬到工厂（GPU），算完再搬回仓库。现在，他们在开始计算前，一次性把所有原材料搬进工厂，让工人在工厂里把所有活干完，最后只把成品搬出来。
  • 这大大减少了“搬运工”的体力消耗，让 GPU 能专心干活。

3. 两种“解题策略”的比拼

在计算电子结构时，有两个主要的“解题算法”（就像两种不同的解题思路）：

1. LOBPCG：像是一个步步为营的侦探。它需要反复检查、互相验证（正交化），每一步都要停下来确认一下。
   • 缺点：因为它需要频繁地“互相确认”（通信），在 GPU 这种需要大规模并行的环境下，它就像一群人在拥挤的房间里频繁交换纸条，效率不高。
2. 切比雪夫多项式滤波（Chebyshev Filtering）：像是一个拥有透视眼的过滤器。它通过一系列数学变换，直接把“好”的电子找出来，中间不需要太多反复确认。
   • 优点：它更像是在流水线上连续作业，非常适合 GPU 的“大批量”特性。
   • 结果：论文发现，切比雪夫滤波法在 GPU 上表现更好，因为它能更充分地利用 GPU 的算力，而且更省电。

4. 实际效果：快了多少？省了多少？

• 速度：在 NVIDIA 的 GPU 上，使用新方法的计算速度比纯 CPU 快了10 到 17 倍！
  • 比喻：以前需要跑 17 天的路程，现在只要 1 天就能跑完。
• 节能：虽然 GPU 本身耗电，但因为算得太快了，总能耗反而大幅下降。
  • 比喻：就像开一辆大卡车（GPU）跑长途，虽然油耗高，但因为速度极快，跑完同样的路程，总耗油量比开一辆小轿车（CPU）慢慢磨蹭要少得多。
• 对比：NVIDIA 的显卡（A100 等）表现非常出色，而 AMD 的显卡虽然也能用，但在某些特定步骤（如“雷利 - 里兹”步骤，可以理解为最后整理答案的步骤）上，效率不如 NVIDIA。

5. 总结与未来

这篇论文告诉我们，仅仅把旧代码搬到新硬件上是不够的。

• 关键启示：必须重新设计算法，让计算任务更适合 GPU 的“大批量、高并发”特性。
• 未来展望：他们建议未来可以进一步减少那些需要“互相确认”的步骤（比如减少正交化次数），或者开发新的方法（如“频谱切片”），让 GPU 能处理更大规模的材料模拟。

一句话总结：
科学家们给 Abinit 软件换上了“超级引擎”（GPU），并重新设计了“驾驶规则”（算法），让它在计算材料科学难题时，不仅快如闪电，而且更加省油，为未来设计新材料打开了新的大门。

技术摘要

这是一份关于 Abinit 软件包在 GPU 架构上加速平面波密度泛函理论（DFT）计算的详细技术总结。

1. 研究背景与问题 (Problem)

• 计算挑战：大规模材料科学的量子模拟需要计算数千个电子态，这涉及求解在大量平面波基组上离散的 Kohn-Sham 方程。随着电子态数量的增加，传统 CPU 架构难以高效利用，导致模拟时间过长。
• 硬件趋势：现代高性能计算（HPC）系统广泛采用图形处理单元（GPU）作为加速器。然而，将现有的 DFT 代码移植到 GPU 面临巨大挑战，不仅需要利用厂商库，还需要对核心算法（特别是迭代对角化过程）进行重构，以识别并优化适合 GPU 的数学操作（如线性代数、FFT）。
• 现有局限：早期的 Abinit GPU 移植尝试（如 2009-2012 年）主要基于自定义 CUDA 内核或仅针对特定硬件，且未充分利用现代混合编程模型。此外，现有的对角化算法（如 LOBPCG）在 GPU 上可能因通信开销和内存带宽限制而效率低下。

2. 方法论与技术架构 (Methodology)

本文介绍了 Abinit 在 2024-2025 年间从零开始重新实现的 GPU 版本，主要采用以下策略：

A. 编程模型与移植框架

• 混合编程模型：采用 MPI + OpenMP Offloading 模型（OpenMP 5.0+），利用指令（Directives）简化代码移植，同时保持跨平台（NVIDIA CUDA 和 AMD ROCm）的可移植性。
• 批处理（Batch Processing）：
  • 核心思想是将电子波函数（$\Psi$）视为数据块，而非单个向量。
  • 将原本针对单个能带（band）的循环操作重组为针对多个能带的批处理操作。
  • 利用 Level-3 BLAS（矩阵 - 矩阵乘法）替代 Level-2 BLAS，并使用批处理 FFT（Batched FFTs）来同时处理多个平面波系数向量，从而最大化 GPU 吞吐量并减少内核启动开销。
• 内存管理：
  • GPU 驻留（GPU-resident）：在自洽场（SCF）迭代开始时将波函数从 CPU 传输到 GPU，并在整个对角化过程中保留在 GPU 显存中，仅在计算电子密度（需要特定内存布局）时进行必要的 CPU-GPU 数据交换。
  • 通信优化：使用 GPU 感知（GPU-aware）MPI 实现进程间通信，直接在 GPU 之间传输数据，避免不必要的显存 - 主机内存拷贝。

B. 迭代对角化算法对比

文章重点比较了两种用于求解 Kohn-Sham 问题的迭代算法：

1. LOBPCG (Locally Optimal Block Preconditioned Conjugate Gradient)：
   • 基于向量导向，通过块预条件共轭梯度法构建子空间。
   • 缺点：在每次子空间迭代中都需要进行正交化和 Rayleigh-Ritz 步骤，导致频繁的 MPI 通信（All-to-All）和内存密集型操作。
2. Chebyshev 多项式滤波 (Chebyshev Filtering)：
   • 基于谱导向，利用切比雪夫多项式放大目标谱区间。
   • 优点：消除了块间依赖，仅需在滤波开始前和结束后进行全局通信。
   • 数学特性：在相同的内迭代次数下，Chebyshev 滤波的算术强度（Arithmetic Intensity）是 LOBPCG 的 $k$ 倍（$k$ 为多项式阶数），更适合 GPU 的高计算密度特性。

C. 底层实现

• 利用厂商库（cuBLAS, rocBLAS, cuFFT, rocFFT, cuSOLVER, rocSOLVER, ELPA）处理核心线性代数运算。
• 实现了抽象层，统一处理波函数对象、MPI 分布和底层库调用，屏蔽了硬件差异。
• 针对非局域算符（Non-local operator）进行了代码重构，利用 gemmBatchedStrided 内核高效处理不同原子类型的投影子。

3. 主要贡献 (Key Contributions)

1. 完整的 GPU 移植策略：提供了从算法设计（批处理、内存布局）到编程模型（OpenMP Offloading）及底层库集成的完整方案。
2. 算法性能理论分析：
   • 建立了 LOBPCG 和 Chebyshev 滤波的理论成本模型（FLOP 计数、通信量、算术强度）。
   • 证明了 Chebyshev 滤波在 GPU 上具有更高的算术强度，且通信开销随系统规模增长更慢。
3. 多厂商性能基准测试：
   • 在 NVIDIA (A100, H100, H200) 和 AMD (MI250X, MI300) GPU 上进行了详细测试。
   • 对比了纯 CPU 节点与 CPU+GPU 混合节点的性能和能耗。
4. 算法优化建议：指出 Rayleigh-Ritz 步骤是主要瓶颈，建议通过增加 Chebyshev 滤波的多项式阶数来减少 SCF 迭代次数，从而在 GPU 上获得最佳整体性能。

4. 实验结果 (Results)

• 加速比 (Speedup)：
  • NVIDIA GPU：表现优异。在 2 个 GPU 节点上即可超越 8 个 CPU 节点的性能（加速比约 13-17 倍）。
  • AMD GPU：也有显著加速，但略低于 NVIDIA，主要受限于 Rayleigh-Ritz 步骤中 hegvd 库在 ROCm 上的性能。
• 能耗效率：
  • GPU 节点显著降低了能耗。在 Jean Zay (NVIDIA) 上，能耗节省因子高达 11-15 倍；在 Adastra (AMD) 上约为 3-5 倍。
  • 随着节点数增加，CPU 能耗线性增长，而 GPU 能耗增长较缓（NVIDIA）或接近线性（AMD），显示出更好的扩展性。
• 算法对比：
  • Chebyshev 滤波 vs. LOBPCG：Chebyshev 滤波在 GPU 上表现更佳。增加多项式阶数虽然增加了单次 SCF 步的计算量，但显著提高了本征矢量的收敛精度，减少了总的 SCF 迭代次数，从而在 GPU 上实现了更短的总运行时间。
  • 瓶颈分析：Roofline 模型显示，哈密顿量应用（Hamiltonian application）和滤波步骤是计算密集型（Compute-bound），而 Rayleigh-Ritz 和正交化步骤是内存/通信密集型（Memory/Communication-bound）。
• 扩展性：随着 GPU 节点数增加，Rayleigh-Ritz 步骤的加速效果不明显，成为扩展性的瓶颈。

5. 意义与展望 (Significance)

• 技术突破：该工作证明了通过算法重构（特别是从 LOBPCG 转向 Chebyshev 滤波）和批处理策略，可以充分发挥现代 GPU 架构在大规模平面波 DFT 计算中的潜力。
• 能源效率：GPU 加速不仅大幅缩短了计算时间，还显著降低了单位计算的能耗，符合绿色计算的趋势。
• 未来方向：
  • 针对 Rayleigh-Ritz 步骤的瓶颈，建议采用 Spectrum Slicing（谱切片） 方法，将大规模对角化分解为多个并行的小规模问题，以进一步减少通信和内存压力。
  • 该 GPU 版本已集成到 Abinit v10.6 及更高版本中，为材料科学界提供了强大的开源工具。

总结：本文不仅成功将 Abinit 移植到 GPU，更重要的是通过深入的理论分析和实验验证，确立了 Chebyshev 多项式滤波 作为 GPU 上平面波 DFT 对角化的首选算法，并提供了详细的性能优化指南，为未来在异构计算架构上进行大规模电子结构计算奠定了基础。