来源论文: https://arxiv.org/abs/2603.21835v1 生成时间: Mar 24, 2026 17:54

迈向超大规模超快动力学模拟：ABACUS 软件包中基于数值原子轨道（NAO）的实时 TDDFT 异构并行实现深度解析

0. 执行摘要

实时时间依赖密度泛函理论（RT-TDDFT）是研究强场超快电子动力学、光吸收谱以及非平衡态电子-离子耦合过程的核心方法。然而，RT-TDDFT 的高计算复杂度限制了其在超大体系（数千原子）中的应用。本文深入探讨了近期由 Taoni Bao 等人发表的一项突破性工作：在国产开源电子结构计算软件包 ABACUS（Atomic-orbital Based Ab-initio Computation at USTC）中实现了一套统一的异构计算框架。该框架采用了三层架构设计（用户层、算法开发层、核心底层异构抽象层），成功地将基于数值原子轨道（NAO）的 RT-TDDFT 算法移植到了 GPU 等硬件加速器上。通过引入硬件无关的 Tensor 抽象和针对速度规范（Velocity Gauge）下的非定域赝势投影算子优化的 GPU 核函数，该工作在处理 bulk silicon 等体系时实现了 3×–4× 的整体壁钟时间加速，关键算子加速甚至超过 12×。这一进展标志着大规模（>1000 原子）、长时标（fs 级）的非线性电子动力学模拟进入了高性能异构并行时代。

1. 核心科学问题，理论基础，技术难点，方法细节

1.1 核心科学问题：为何需要 RT-TDDFT 的异构加速？

电子结构动力学的模拟通常有两种路径：线性响应 TDDFT（LR-TDDFT）和实时 TDDFT（RT-TDDFT）。LR-TDDFT 在处理弱场激励时非常高效，但在处理非线性过程（如高次谐波产生 HHG、多光子电离等）或极短脉冲激励下的复杂电子迁移时，RT-TDDFT 具有天然优势。然而，RT-TDDFT 要求在每一个极细微的时间步（通常在 atto-second 数量级）更新波函数和 Hamiltonian。在 NAO 基组下，这意味着需要频繁进行矩阵乘法、求逆以及复杂的实空间网格积分。随着体系规模的增加，计算量呈立方级（$O(N^3)$）增长，传统的 CPU 集群已难以满足大规模、长寿命动力学模拟的需求。因此，如何构建一套既能保持代码可维护性，又能最大化利用 GPU 算力的异构框架，是当前量子化学软件面临的核心技术挑战。

1.2 理论基础：NAO 基组下的 TDKS 方程

在 ABACUS 中，体系的含时波函数 $\psi_{nk}(\mathbf{r}, t)$ 被展开在 Bloch-like NAO 基组下：

$$\psi_{nk}(\mathbf{r}) = \sum_{\mathbf{R}} \sum_{\mu} c_{n\mu,k} e^{i\mathbf{k}\cdot\mathbf{R}}\phi_{\mu}(\mathbf{r} - \tau_I - \mathbf{R})$$

其中 $\phi_{\mu}$ 是具有径向截断性质的数值原子轨道。在 RT-TDDFT 模拟中，我们需要求解含时 Kohn-Sham (TDKS) 方程：

$$H\psi_{nk}(\mathbf{r}, t) = i \frac{\partial}{\partial t}\psi_{nk}(\mathbf{r}, t)$$

将其投影到非正交的 NAO 基组上，得到矩阵形式：

$$S_k^{-1} H_k C_k(t) = i \frac{\partial}{\partial t}C_k(t)$$

这里 $S_k$ 是重叠矩阵，$H_k$ 是 Hamiltonian。波函数的随时间演化通过传播子 $U_k$ 实现：

$$C_k(t + \Delta t) = U_k(t + \Delta t, t) C_k(t)$$

ABACUS 实现了多种传播子方案，包括 Crank-Nicolson、4阶泰勒展开以及强制时间反演对称（ETRS）方法。其中 Crank-Nicolson 方法因其无条件稳定性而在本文中被广泛使用，其核心在于求解如下线性方程组：

$$(S_k + \frac{i\Delta t}{2} H_k) C_k(t + \Delta t) = (S_k - \frac{i\Delta t}{2} H_k) C_k(t)$$

1.3 技术难点：规范选择（Gauge Selection）与投影算子瓶颈

在外部电磁场作用下，规范的选择极大地影响计算开销：

  1. 长度规范 (Length Gauge)：通过 $\mathbf{E}(t) \cdot \mathbf{r}$ 项引入。优点是公式简单，但缺点是破坏了周期性边界条件，仅适用于分子体系。
  2. 速度规范 (Velocity Gauge)：通过矢量势 $\mathbf{A}(t)$ 引入。这是处理周期性固体体系的标准方法。然而，在 NAO 基组下，速度规范会导致非定域赝势投影算子 $V_{NL}$ 获得一个随位置相关的相位因子： $$V_{NL}(t) = e^{-i\mathbf{A}(t)\cdot\mathbf{r}} V_{NL} e^{i\mathbf{A}(t)\cdot\mathbf{r}}$$ 这种相位因子破坏了 NAO 基组中原本高效的“双中心积分”特性，迫使程序必须在实空间网格上进行数值积分。由于每一个时间步都要重新计算，这成为了 RT-TDDFT 模拟中最大的性能瓶颈。

1.4 方法细节：三层架构设计

为了解决异构开发中的代码冗余和硬件依赖问题，作者设计了三层逻辑架构：

  • (a) Users Layer (用户层)：处理输入输出（STRU, .orb, .upf 等文件）和最终物理性质（光响应、电荷迁移等）的呈现。用户无需感知底层硬件切换。
  • (b) Algorithm Developers Layer (算法开发层)：实现了 RT-TDDFT 的工作流，包括 Hamiltonian 构建、波函数传播、电荷混合及 Post-processing。该层通过调用底层抽象接口来实现算力透明。
  • (c) Core Underlying Heterogeneous Abstraction Layer (底层异构抽象层)：
    • Unified Data Containers (Tensor)：基于 RAII 原则，封装了 CPU 和 GPU 上的内存管理。Tensor 不仅存储数据，还记录了设备属性（CPU/GPU）和数据排布，实现了“一次编写，处处运行”。
    • Unified Linear Algebra Operators：对 BLAS 和 LAPACK 库（如 MKL, cuBLAS, cuSOLVER）进行了高度抽象，使得算法层可以像写公式一样调用 getrf (LU分解) 和 getrs (线性系统求解)。
    • Unified Grid Integration Interfaces (Gint)：这是该工作的核心加速点，专门针对 NAO 轨道和非定域投影算子的实空间积分进行了 GPU 核心优化（如原子级 batching、warp-level reduction 等）。

2. 关键 benchmark 体系，计算所得数据，性能数据

2.1 物理准确性验证（Accuracy Validation）

作者首先在多种维度（0D 到 3D）的体系上验证了异构 RT-TDDFT 的数值精度：

  1. Anthracene (蒽分子)：在长度、速度和混合规范下，QZTP 基组得到的吸收谱高度重合，且与 DGDFT 软件包的基准数据完美吻合。验证了 CPU 与 GPU 实现之间的一致性（误差达到数值精度上限）。
  2. (CdSe)6 团簇：计算得到的吸收谱成功复现了 1.5-4.0 eV 范围内的特征三峰结构。相比于 CP2K 采用的紧凑型基组，ABACUS 的 QZTP 基组表现出更好的高能态描述能力（约 0.05 eV 的红移）。
  3. 一维氢链 (1D H-chain)：模拟了 Peierls 畸变下的电子动力学，成功复现了 Qbox 在有限频率区域的主峰（2.50 eV）。
  4. 二维 h-BN 单层：计算得到的含时复电导率 $\text{Re}\sigma(\omega)$ 表现出明显的双峰结构（5.6 eV 和 14.5 eV），与 Octopus 及 PySCF 的基准数据高度一致。
  5. 三维大块硅 (Bulk Silicon)：重点展示了“混合规范（Hybrid Gauge）”在修正速度规范带来的低频发散问题上的卓越表现。这证明了异构实现在处理固体介电函数时的稳健性。

2.2 性能分析（Performance Metrics）

性能测试是在 A800 GPU 和 Intel Xeon Gold 6348 CPU 上进行的：

2.2.1 系统规模扩展性 (System-size Scaling)

对于 48 到 1200 个原子的硅超原胞，总的计算时间显示出：

  • GPU 加速比：在 1200 个原子规模时，单张 GPU 相比于完全负载的双路 CPU 节点（56 核）实现了 3×–4× 的整体加速。
  • 算子级加速：波函数演化（evolve_k）在 GPU 上比单 MPI 进程的 CPU 实现快了 12× 以上。这主要归功于 GPU 对密集线性代数运算（LU 分解和替代）的极高吞吐量。

2.2.2 关键瓶颈的突破：Gint 模块

速度规范下的球面网格积分（snap_psibeta）是速度规范的“致命伤”。在 CPU 上，该步骤消耗了小体系模拟 40%-80% 的时间。通过 GPU 上的原子级批处理（Atom-level batching）和消除冗余核函数启动：

  • 加速比：GPU 版的球面积分比优化后的 OpenMP 并行 CPU 版快了 12× 甚至更高（对于 1200 原子体系）。
  • 影响：使得“规范惩罚（Gauge Penalty）”几乎被消除，科研人员可以根据物理需求自由选择规范，而不必过多担心计算成本。

2.2.3 多 GPU 强扩展性 (Strong Scaling)

  • 测试规模：在 40 张 GPU 上运行 1728 原子的硅体系。
  • 效率：从 16 张 GPU 扩展到 40 张 GPU，并行效率保持在 76% 左右。尽管存在受限的 ScaLAPACK 通信开销（目前分布式线性求解器主要依赖 NVIDIA 的 cuBLASMp），但整体扩展性依然能够支持大规模生产任务。

3.1 关键类实现：Tensor 与 TensorMap

在代码实现中，Tensor 类是所有操作的基石。其不使用模版类，而是采用运行时多态和显式枚举（Data type, Device type）：

// 示例逻辑展示，非直接源代码
Tensor A(DataType::ComplexDouble, DeviceType::GPU, {N, N});
// 自动在 GPU 内存中分配空间，并在析构时释放 (RAII)

对于现有的 CPU 遗留数据排布，作者引入了 TensorMap，这是一种“非所有”的视图（View），允许将 CPU 内存中的 raw pointer 包装成 Tensor 参与统一接口运算，极大地降低了代码重构的成本。

3.2 线性代数接口抽象

为了解耦物理算法与底层数学库，ABACUS 内部实现了一套 Dispatch 系统：

  • 如果 Tensor 在 CPU 上，getrf 会调用 mkl_dgetrf 或 lapack_dgetrf。
  • 如果 Tensor 在 GPU 上，则透明地调用 cusolverDnDgetrf。 这种设计确保了 evolve_k.cpp 等物理逻辑文件中完全不出现任何特定的 vendor 库调用。

3.3 关键算子：snap_psibeta_atom_batch_gpu

这是针对非定域赝势投影算子相位因子处理的专门核函数。其核心思想是：

  1. 原子级调度：一个线程块（thread block）处理一个中心原子的所有轨道。
  2. 共享内存优化：利用 GPU 的 Shared Memory 存储 Lebedev-Laikov 球面网格权重和投影函数。
  3. 向量化规约：使用 __shfl_down_sync 等 warp 指令进行高效的球面角度积分规约，避免了昂贵的原子操作（Atomic Add）。

3.4 复现指南与开源链接

  • 软件包名称：ABACUS (Atomic-orbital Based Ab-initio Computation at USTC)
  • 版本要求：v3.9.0.24 及以上版本。
  • 编译需求：需要支持 CUDA 的编译器（nvcc）以及 cuBLAS, cuSOLVER, cuFFT 库。分布式加速建议安装 cuBLASMp 和 cuSOLVERMp。
  • 开源仓库：
  • 复现步骤：
    1. 克隆 abacus-develop 的 develop 分支。
    2. 使用 CMake 配置，开启 -DENABLE_CUDA=ON。
    3. 运行测试算例（位于 examples/rt_tddft），比较 CPU 和 GPU 版本的波函数演化轨迹。

4. 关键引用文献，以及你对这项工作局限性的评论

4.1 关键引用文献

  1. Runge-Gross (1984) [1]：RT-TDDFT 的理论基石。
  2. Yabana & Bertsch (1996) [2]：定义了实空间实时间演化的标准方法。
  3. Chen et al. (2010) [40]：ABACUS 软件包中数值原子轨道的系统化构建方法。
  4. Zhao & He (2025) [46]：提出了解决速度规范低频发散的 Hybrid Gauge 方法，本文 GPU 加速方案对其进行了深度融合。
  5. Soler et al. (2002) [47]：SIESTA 方法的经典文献，定义了基于 NAO 的 Hamiltonian 构建范式。

4.2 局限性评论

尽管这项工作在异构加速方面取得了巨大成功，但从技术角度看，仍存在以下局限：

  1. 分布式线性求解器的依赖性：目前多 GPU 并行依赖于 NVIDIA 特有的 cuBLASMp 和 cuSOLVERMp。这使得在非 NVIDIA 显卡（如 AMD, Intel 或国产天数智芯、壁仞）上的大规模分布式运行仍面临挑战。作者虽然提出了基于 QR 分解的替代方案（Section 3.2.2），但在极其稀疏的算例中性能尚未完全超越 LU 分解。
  2. 非绝热效应的简化：当前的动力学基于 Ehrenfest 近似，这是一种平均场处理方法。它无法描述波函数的坍缩（Decoherence）或多势能面间的非绝热跳转（Surface Hopping）。对于某些涉及长程关联或强激子效应的超快过程，可能需要更高级的非绝热动力学模型。
  3. 内存带宽限制：实空间网格积分（Gint）虽然获得了 12× 加速，但由于其典型的带宽受限特性，在大体系下受制于 PCIe 数据传输速率。未来的混合精度计算（Mixed Precision）可能是进一步提升吞吐量的关键。

5. 其他补充：未来愿景与开发者建议

5.1 异构框架的“长久主义”

这项工作最值得称道的是其“三层架构”设计。对于高性能计算软件开发者来说，直接写 CUDA kernel 虽然能在短期内获得极致性能，但往往会导致代码腐烂。ABACUS 采用的 Tensor 抽象层为未来的自动微分（Automatic Differentiation）和 AI-assisted DFT 铺平了道路。随着深度学习势能面（Deep Potential）等技术的发展，RT-TDDFT 与机器学习的结合（如利用 ML 模型加速电荷密度预测）将是未来的热点。

5.2 速度规范与长度规范的博弈

长期以来，周期性体系的 RT-TDDFT 计算被视为“昂贵”的代名词。本文通过 GPU 加速球面积分，向领域展示了一个重要信号：速度规范下的相位因子不再是不可逾越的障碍。这对于研究拓扑绝缘体的非线性光学效应、晶体中的 HHG 以及强场物理具有重要意义。

5.3 混合精度计算的潜力

在 RT-TDDFT 的迭代传播中，由于时间步长极小，波函数的微小演化是否一定需要双精度（FP64）？作者在结论部分提到，Tensor 架构已经支持单精度（FP32）或混合精度。通过在非核心物理量计算中引入 FP32，有望在不损失物理准确性的前提下，利用 NVIDIA Tensor Cores 进一步压榨 GPU 性能，将模拟体系规模推向万原子级别。

5.4 结语

ABACUS 的这项统一异构实现不仅是代码的迁移，更是对 NAO 基组下实时含时演化算法的一次深刻重构。它不仅展示了中国国产电子结构软件在高性能计算领域的强大竞争力，也为全球量子化学社区提供了一个高扩展、可维护的科研平台。对于致力于材料科学、光伏技术及 atto-second 物理的科研人员来说，ABACUS 的这一新特性无疑是手中最锐利的武器之一。