A Unified Heterogeneous Implementation of Numerical Atomic Orbitals-Based Real-Time TDDFT within the ABACUS Package

本文介绍了 ABACUS 软件包中基于数值原子轨道的实时含时密度泛函理论（RT-TDDFT）的统一异构计算框架，通过三层协同设计显著加速了波函数传播与实空间格点操作，并在从分子到固体的各类体系验证了其高精度与多 GPU 扩展性。

要点

这篇文章介绍了一项名为 ABACUS 的超级计算软件的重大升级。为了让你更容易理解，我们可以把这项研究想象成给一辆赛车换上了“万能引擎”和“超级变速箱”。

1. 背景：为什么要做这个？（赛车的需求）

想象一下，科学家们在研究原子和电子如何运动，特别是当它们受到激光照射时（比如超快电子动力学）。这就像是在观察一场微观世界的赛车比赛。

• 旧问题：以前的软件（赛车）主要是在 CPU（普通引擎）上跑。虽然能跑，但面对极其复杂的计算（比如模拟成千上万个原子的实时反应），速度太慢了，就像用自行车去追 F1 赛车。
• 新挑战：现在的超级计算机都有 GPU（图形处理器，相当于高性能赛车引擎），它们并行计算能力极强。但是，让旧软件直接跑在 GPU 上非常困难，就像强行把自行车的零件装到法拉利上，不仅跑不起来，还容易散架。而且，不同的 GPU 品牌（NVIDIA、AMD 等）就像不同品牌的赛车，零件不通用，导致软件很难移植。

2. 核心创新：三层架构（万能改装方案）

为了解决这个问题，研究团队在 ABACUS 软件中设计了一套**“三层万能改装系统”**，让物理算法和硬件彻底分离：

• 第一层：用户层（驾驶员）
  • 比喻：这是赛车手看到的仪表盘。
  • 作用：用户只需要输入结构、原子轨道等基础数据，不需要关心底层是 CPU 还是 GPU。软件会自动输出光学响应、电荷迁移等物理结果。
• 第二层：算法层（赛车手）
  • 比喻：这是决定赛车如何过弯、加速的策略。
  • 作用：这里负责核心的物理计算，比如电子波函数如何随时间演化。以前这里写死了对 CPU 的依赖，现在它变得“硬件无关”，无论下面挂什么引擎，它都能指挥。
• 第三层：核心抽象层（万能变速箱/传动系统）
  • 比喻：这是最关键的发明。它像是一个智能适配器。
  • 作用：
    1. 统一数据容器：不管数据存在 CPU 内存还是 GPU 显存，都打包成标准的“集装箱”（Tensor），搬运时不用重新包装。
    2. 统一线性代数算子：把复杂的数学运算（如矩阵乘法）封装成标准接口。如果是 CPU，就调用 CPU 的数学库；如果是 GPU，就自动调用 GPU 的加速库。
    3. 统一网格积分：这是处理空间计算的接口，确保在不同硬件上计算结果一致。

简单说：以前写代码要专门为 NVIDIA 显卡写一套，为 AMD 显卡写一套。现在，他们造了一个“万能插头”，物理学家只管写物理公式，软件自动适配任何硬件。

3. 攻克最大难点：速度计与相位（消除“速度瓶颈”）

在模拟光与物质相互作用时，有两种计算方法（规范）：

• 长度规范：适合小分子，但算不了大晶体。
• 速度规范：适合大晶体，但计算量巨大，因为需要处理一种复杂的“相位因子”（可以想象成在计算每个原子周围都有一个快速旋转的陀螺，以前在 CPU 上算这个非常慢，是最大的瓶颈）。

团队的突破：
他们专门为 GPU 设计了一个**“球形网格积分加速器”**。

• 比喻：以前在 CPU 上算这个，就像让一个人拿着计算器，一个一个地数旋转的陀螺，累得半死。
• 现在：在 GPU 上，他们让成千上万个线程同时去数，并且优化了算法，把“数陀螺”变成了“批量处理”。
• 结果：这个最慢的环节，在 GPU 上比 CPU 快了 12 倍！这意味着以前不敢用的“速度规范”现在变得既快又准，科学家可以自由选择最适合物理问题的方法，而不必被计算速度限制。

4. 实际效果：速度与规模（赛车实测）

团队用真实的材料（如硅晶体、石墨烯等）进行了测试：

• 单卡加速：用一张高端 GPU（NVIDIA A800），比用满 56 个核心的双路 CPU 服务器还要快 3 到 4 倍。
• 大规模并行：当使用 40 张 GPU 一起工作时，软件依然能保持 76% 的极高效率。这意味着如果你把计算任务分给 40 辆车，它们能几乎完美地协同工作，没有太多时间浪费在互相沟通上。
• 适用范围：从单个分子到巨大的晶体，从二维材料到三维固体，结果都非常准确，与之前的权威数据完全吻合。

5. 总结与未来（未来的赛车）

这项工作的意义在于：

1. 快：让模拟超快电子运动（飞秒、阿秒级别）变得可行，以前算一天，现在可能只要一小时。
2. 通用：不再被特定硬件绑定，未来无论是国产芯片（如华为昇腾、海光 DCU）还是其他新型加速器，这套框架都能轻松适配。
3. 开放：这是开源软件 ABACUS 的一部分，全球科学家都可以免费使用，推动材料科学和量子物理的发展。

一句话总结：
作者给 ABACUS 软件装上了一个**“智能硬件翻译器”和一个“超级加速器”**，让原本只能在慢速 CPU 上运行的复杂物理模拟，现在能在 GPU 上全速奔跑，并且能轻松适应各种未来的超级计算机，让科学家能以前所未有的速度和规模去探索微观世界的奥秘。

技术摘要

这是一份关于论文《ABACUS 包中基于数值原子轨道的实时含时密度泛函理论（RT-TDDFT）的统一异构实现》的详细技术总结。

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

背景：
实时含时密度泛函理论（RT-TDDFT）是模拟激发态电子动力学、光与物质相互作用（如光吸收、电荷迁移、高次谐波产生）的关键工具。随着计算需求的增长，利用图形处理器（GPU）加速第一性原理计算已成为趋势。

面临的主要挑战：

1. 实时动力学研究的缺失： 现有的 GPU 加速代码多集中于基态 KSDFT 或线性响应 TDDFT（LR-TDDFT），针对基于**数值原子轨道（NAOs）**的实时动力学研究较少。
2. 速度规范（Velocity Gauge）的瓶颈： 在周期性体系中，速度规范是处理外场的标准方法。然而，在 NAO 基组下，速度规范引入的含时矢量势会导致非局域赝势算符出现位置依赖的相位因子（$e^{-i\mathbf{A}(t)\cdot\mathbf{r}}$）。这使得原本可以通过解析计算的两中心积分失效，必须在全局实空间网格上进行数值积分，计算成本极高，成为主要瓶颈。
3. 硬件依赖性与可移植性： 现有的 GPU 代码往往针对特定架构（如 NVIDIA）深度优化，缺乏硬件无关的抽象层，导致代码维护困难且难以移植到 AMD GPU 或国产加速器（如海光 DCU）上。

2. 方法论与架构 (Methodology)

本文在开源软件 ABACUS 中构建了一个统一的异构计算框架，核心在于设计了三层抽象架构，将物理算法与硬件细节解耦：

A. 统一数据容器层 (Unified Data Containers)

• Tensor 类： 设计了一个运行时多态的容器类，用于管理多维数组（如哈密顿量、重叠矩阵、密度矩阵）。
• RAII 内存管理： 自动处理内存分配与释放，支持 CPU、GPU 和 DCU 后端，无需手动管理指针，消除了内存泄漏风险。
• TensorMap： 提供非拥有（non-owning）的视图，用于无缝集成遗留代码或第三方库。

B. 统一线性代数算子层 (Unified Linear Algebra Operators)

• 多态算子接口： 封装了标准稠密线性代数操作（如 gemm, getrf, getrs 等）。
• 后端自动调度： 根据运行环境自动调用 CPU 上的 BLAS/LAPACK 或 GPU 上的 cuBLAS/cuSOLVER（及对应国产库），确保数值一致性和跨平台兼容性。

C. 统一网格积分接口层 (Unified Grid Integration Interfaces)

• 均匀网格积分 (Gint)： 针对电荷密度、局域势等常规量，采用模块化设计，利用批处理（Batched）策略将稀疏的原子轨道计算转化为高效的矩阵运算，并在 GPU 上加速。
• 球面网格积分 (Spherical Grid Integration)： 这是本文的核心创新之一。 针对速度规范下的非局域赝势相位因子积分，开发了专用的 GPU 加速方案：
  • 使用原子中心的球面网格（Lebedev-Laikov 角度积分 + Gauss-Legendre 径向积分）。
  • 设计了 snap_psibeta_atom_batch_gpu 内核，采用原子级批处理策略，利用 Warp 原语进行无原子操作的归约，极大提升了并行效率。

D. 物理实现细节

• 规范选择： 支持长度规范（有限体系）、速度规范（周期性体系）和混合规范（Hybrid Gauge）。混合规范通过引入局部相位因子解决了 NAO 基组在速度规范下的低频发散问题。
• Ehrenfest 动力学： 实现了包含 Hellmann-Feynman 力、Pulay 力、非正交性修正力及 Ewald 力的完整力计算，并支持离子运动与电子传播的耦合。

3. 关键贡献 (Key Contributions)

1. 统一的异构框架： 首次在基于 NAO 的 ABACUS 包中实现了完整的 RT-TDDFT 异构计算，通过 Tensor 抽象层实现了“一次编写，多架构运行”（CPU/GPU/DCU）。
2. 消除速度规范瓶颈： 通过专用的 GPU 球面网格积分方案，成功解决了速度规范下非局域赝势积分的计算难题，使其计算效率与混合规范相当，消除了算法上的“速度规范惩罚”。
3. 高性能多 GPU 扩展： 针对多 GPU 环境，采用了基于 QR 分解的并行求解策略（替代传统的 LU 分解），解决了多右端项（多能带）并行求解的瓶颈，实现了大规模强扩展。

4. 实验结果 (Results)

物理验证

• 测试体系： 涵盖了从 0D（蒽分子、CdSe 团簇）、1D（氢链）、2D（h-BN 单层）到 3D（块体硅）的多种体系。
• 精度对比： 计算的光吸收谱、介电函数等物理量与 DGDFT、CP2K、Qbox、Octopus 等主流代码的基准结果高度一致。
• 规范一致性： 验证了长度、速度和混合规范在各自适用体系下的一致性，特别是混合规范成功消除了速度规范在周期性体系中的低频发散伪影。

性能表现

• 单节点加速比： 在单块 NVIDIA A800 GPU 上，相比双路 56 核 Intel CPU 节点，整体墙钟时间加速了 3-4 倍。
• 关键内核加速：
  • 波函数传播（evolve_k）：加速约 6-7 倍。
  • 球面网格积分（速度规范瓶颈）：加速超过 12 倍，彻底消除了该步骤的时间占比（从 CPU 的 40-80% 降至 GPU 的 <5%）。
• 多 GPU 强扩展： 在 40 张 GPU 上对 1728 个原子的硅体系进行测试，实现了 76% 的并行效率，证明了代码处理大规模、长时程模拟的能力。
• 通信开销： 数据传输开销占总时间的比例较低（约 9%），且随着系统增大，计算密集型操作掩盖了通信延迟。

5. 意义与展望 (Significance & Outlook)

• 科学意义： 该工作为研究超快电子动力学（飞秒至阿秒尺度）提供了一个高性能、可移植的第一性原理平台。它使得在复杂材料中进行大规模、长时程的 RT-TDDFT 模拟成为可能，特别是解决了 NAO 基组在速度规范下的效率难题。
• 技术意义： 提出的“统一异构抽象层”设计模式为科学计算软件的未来发展提供了范本，有效平衡了物理算法的复杂性与硬件架构的多样性。
• 未来方向：
  • 提升多设备加速的跨平台能力（目前多 GPU 强依赖 NVIDIA 库，需适配更多国产/AMD 硬件的分布式库）。
  • 探索混合精度计算（Mixed Precision），利用 Tensor Core 进一步提升吞吐量。
  • 将框架扩展至更高级的含时多体方法（如 GW-BSE）。

总结：
本文通过构建统一的异构计算架构，成功在 ABACUS 中实现了基于数值原子轨道的高效 RT-TDDFT 模拟。其核心突破在于利用 GPU 加速解决了速度规范下的计算瓶颈，并实现了优异的跨平台性能和大规模并行扩展能力，为超快电子动力学的大规模模拟奠定了坚实基础。