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

This paper presents a unified heterogeneous computing framework within the ABACUS package that accelerates real-time time-dependent density functional theory simulations based on numerical atomic orbitals through co-designed abstraction layers, demonstrating significant speedups on single GPUs and high parallel efficiency across multiple GPUs for large-scale electron dynamics studies.

The Gist

The Big Picture: Simulating the "Dance" of Electrons

Imagine you are trying to film a high-speed dance battle between millions of tiny dancers (electrons) inside a material like silicon or a molecule. You want to see exactly how they move when hit by a flash of light (like a laser).

This is what Real-Time Time-Dependent Density Functional Theory (RT-TDDFT) does. It's a super-complex mathematical movie camera that simulates how electrons react to light in real-time.

However, there's a problem: The movie is too heavy.
Running these simulations on standard computer processors (CPUs) is like trying to film that dance battle with a single, slow-moving camera. It takes days or weeks to render just a few seconds of the movie.

The Solution: The authors of this paper built a new engine for the ABACUS software (a popular tool for these simulations) that runs on GPUs (Graphics Processing Units). Think of GPUs as a stadium filled with thousands of tiny, super-fast cameras working in perfect unison.

---

The Three-Layer "Smart Factory"

The authors didn't just throw the old code onto a GPU; they completely redesigned the factory floor. They built a 3-Level System to make sure the work gets done efficiently, no matter what kind of hardware you have.

1. The User Layer (The "Customer")

• What it is: This is where scientists type in their instructions (e.g., "Simulate a silicon crystal").
• The Analogy: Imagine a customer walking into a restaurant. They just look at the menu and order a burger. They don't need to know how the grill works, who the chef is, or if the kitchen uses gas or electric stoves. They just want the burger.
• The Benefit: Scientists can use the software without needing to be computer experts.

2. The Algorithm Layer (The "Head Chef")

• What it is: This is the logic that decides what needs to be calculated (e.g., "Move the electrons forward in time," "Calculate the forces").
• The Analogy: The Head Chef looks at the order and says, "Okay, we need to chop onions, grill the patty, and melt the cheese." The Chef doesn't care who does the chopping or which stove is used; they just manage the flow of the recipe.
• The Benefit: The physics logic stays the same. The scientists can focus on the science, not the computer code.

3. The Core Layer (The "Universal Kitchen Staff")

• What it is: This is the magic part. It's a "translator" that takes the Chef's orders and assigns them to the right workers (CPUs or GPUs).
• The Analogy: Imagine a kitchen where you have both human chefs (CPUs) and a swarm of robot arms (GPUs). Usually, you'd have to write two different recipes: one for humans and one for robots.
  • The Innovation: This paper created a Universal Translator. It takes the Chef's order ("Chop onions") and instantly figures out: "Oh, the robot arm is free, let's give it to the robot!" or "The human is free, let's give it to the human!"
  • The Result: The same code runs perfectly on an Intel CPU, an NVIDIA GPU, or even a Chinese DCU chip, without rewriting the recipe.

---

The "Speed Trap" and How They Fixed It

There was a specific problem with simulating light-matter interactions called the "Velocity Gauge."

• The Problem: In the old way of doing this, calculating how the electrons move under a specific type of light field was like trying to count grains of sand on a beach by picking them up one by one with tweezers. It was incredibly slow and became a "bottleneck" that stopped the whole simulation.
• The Fix: The authors built a specialized GPU tool (a "Spherical Grid Integration") that acts like a giant vacuum cleaner. Instead of picking up grains one by one, it sucks up the whole beach in seconds.
• The Result: This specific step became 12 times faster on the GPU. It removed the "speed trap," allowing scientists to use the most accurate physics methods without waiting forever.

---

The Results: From Days to Hours

The team tested their new system on everything from tiny molecules to huge chunks of silicon.

1. Speed: On a single powerful GPU, their system was 3 to 4 times faster than a massive, fully-loaded computer server with 56 CPU cores.
2. Efficiency: When they used 40 GPUs working together (like a team of 40 robots), the system didn't slow down due to communication issues. It kept working at 76% efficiency.
3. Accuracy: They checked their math against known benchmarks (like comparing their movie to a famous, award-winning documentary). Their results matched perfectly.

Why Does This Matter?

Think of this like upgrading from a flip phone to a smartphone.

• Before: Scientists could only simulate small, simple systems or very short moments in time. It was like trying to watch a movie on a flip phone—pixelated and slow.
• Now: With this new framework, scientists can simulate huge materials (like entire computer chips) and watch ultra-fast events (like electrons moving in femtoseconds) in high definition.

This opens the door to designing better solar cells, faster computer chips, and new medical materials by understanding exactly how electrons behave when hit by light, all without needing a supercomputer the size of a building.

Summary in One Sentence

The authors built a "universal translator" for scientific software that lets complex electron simulations run 12 times faster on modern graphics cards, turning a task that used to take days into one that takes hours, all while keeping the math accurate and the code easy to use.

Technical Summary

1. Problem Statement

Real-Time Time-Dependent Density Functional Theory (RT-TDDFT) is essential for simulating ultrafast electron dynamics, non-equilibrium phenomena, and light-matter interactions. However, implementing RT-TDDFT efficiently on modern heterogeneous computing architectures (CPUs and GPUs) faces three major challenges, particularly when using Numerical Atomic Orbitals (NAOs):

1. Lack of GPU Support for Real-Time Dynamics: Most existing GPU-accelerated electronic structure codes focus on ground-state Kohn-Sham DFT or Linear-Response TDDFT. Real-time propagation, which requires explicit time-stepping and self-consistent loops, remains under-explored in the NAO basis.
2. The Velocity Gauge Bottleneck: In periodic systems, the velocity gauge is theoretically preferred but introduces a severe computational bottleneck in NAO codes. The nonlocal pseudopotential operator acquires a position-dependent phase factor ($e^{-i\mathbf{A}(t)\cdot\mathbf{r}}$), destroying the two-center integral structure. This forces expensive numerical integration on real-space grids at every time step, which is difficult to parallelize efficiently.
3. Software Portability and Maintainability: Developing GPU code often requires architecture-specific tuning (e.g., explicit memory management, kernel optimization), leading to code duplication and poor portability across different hardware (NVIDIA, AMD, Hygon DCUs). A unified, hardware-agnostic framework is needed to decouple physical algorithms from hardware specifics.

2. Methodology and Architecture

The authors implemented a unified heterogeneous computing framework within the ABACUS package, structured into three logical layers:

A. Unified Data Abstraction (The "Tensor" Layer)

• Tensor Container: A runtime-polymorphic data container for multi-dimensional arrays (Hamiltonians, overlap matrices, density matrices). It encapsulates memory management (RAII principle) and device affinity (CPU/GPU/DCU), allowing a single codebase to execute on diverse hardware without source-level specialization.
• TensorMap: A non-owning "view" for integrating legacy modules or external buffers into the unified framework without invasive refactoring.

B. Unified Linear Algebra Operators

• A polymorphic interface for standard dense linear algebra operations (e.g., gemm, getrf, getrs).
• It automatically dispatches to optimized backends: BLAS/LAPACK for CPUs and cuBLAS/cuSOLVER (or equivalents for AMD/Hygon) for accelerators.
• This layer handles the heavy lifting of wavefunction propagation (solving linear systems for the Crank-Nicolson propagator) and Energy-Density Matrix (EDM) construction.

C. Unified Grid Integration Interfaces

• Uniform Grid Integration (Gint Module): Handles charge density, local potentials, and forces. It uses a two-level grid decomposition (MeshGrid and BigGrid) and employs batched GEMM kernels on GPUs to transform irregular memory accesses into efficient matrix operations.
• Spherical Grid Integration (Velocity Gauge Solution): A dedicated, GPU-accelerated module for the nonlocal pseudopotential overlaps in the velocity gauge.
  • Algorithm: Instead of processing neighbors sequentially, it uses an atom-level batching scheme where a single GPU kernel processes all neighboring orbitals of a central atom simultaneously.
  • Optimization: It utilizes Lebedev-Laikov (angular) and Gauss-Legendre (radial) quadrature grids. Reductions are performed using warp-level primitives (__shfl_down_sync) and shared memory to avoid atomic operations, achieving massive parallelism.

D. Theoretical Framework

• Gauge Choices: The code supports Length, Velocity, and Hybrid Gauges. The Hybrid gauge (incorporating local phase factors) is used to avoid low-frequency divergence in periodic systems while maintaining compatibility with localized bases.
• Dynamics: Supports Ehrenfest dynamics for electron-ion coupling, calculating forces (Hellmann-Feynman, Pulay, Orthogonality, Ewald) and propagating ions via the velocity Verlet algorithm.

3. Key Contributions

1. First Unified Heterogeneous RT-TDDFT in NAOs: The paper presents the first fully hardware-agnostic implementation of RT-TDDFT using numerical atomic orbitals, capable of running on CPUs, NVIDIA GPUs, AMD GPUs, and Hygon DCUs.
2. Elimination of the Velocity Gauge Penalty: The development of a specialized GPU-accelerated spherical grid integration kernel effectively removes the computational bottleneck associated with the velocity gauge in NAO-based codes.
3. High-Performance Multi-GPU Scaling: Implementation of a distributed linear solver strategy (using QR factorization via cusolverMp) to solve the Crank-Nicolson equation for thousands of bands across multiple GPUs, overcoming limitations of single-right-hand-side solvers.
4. Robust Software Architecture: The "Tensor" abstraction layer successfully decouples physics algorithms from hardware details, ensuring long-term maintainability and portability.

4. Results and Performance

Physical Validation

The implementation was validated against established benchmarks (DGDFT, CP2K, Qbox, Octopus, SIESTA) across diverse systems:

• 0D: Anthracene molecule and (CdSe)₆ nanocluster.
• 1D: Hydrogen chain.
• 2D: Hexagonal Boron Nitride (h-BN) monolayer.
• 3D: Bulk Silicon.
• Outcome: Excellent agreement in optical absorption spectra, dielectric functions, and current densities. The hybrid gauge successfully corrected the low-frequency divergence seen in standard velocity gauge implementations.

Performance Benchmarks (Bulk Silicon)

• Single-GPU vs. CPU: On a single NVIDIA A800 GPU, the code achieves a 3×–4× wall-clock speedup over a fully utilized dual-socket 56-core Intel Ice Lake CPU node.
• Kernel Acceleration:
  • Wavefunction Propagation (evolve_k): ~6–7× speedup on GPU vs. best CPU config; >12× vs. single-MPI CPU.
  • Spherical Grid Integration: >12× speedup on GPU vs. CPU, effectively eliminating the velocity gauge overhead.
  • Uniform Grid Integration: GPU slightly outperforms the full CPU node, maintaining linear scaling.
• Strong Scaling (Multi-GPU):
  • Tested on systems up to 1728 atoms across 40 GPUs.
  • Achieved 76% parallel efficiency when scaling from 16 to 40 GPUs.
  • The wavefunction propagation module scales well (near-ideal for large systems), while grid integration shows saturation due to linear scaling ($O(N)$) and communication overhead, but this does not dominate the total runtime for large systems.

5. Significance

This work establishes a high-performance, portable platform for large-scale first-principles simulations of ultrafast electron dynamics. By solving the specific computational bottlenecks of the velocity gauge in NAO-based methods and providing a unified hardware-agnostic framework, it enables researchers to:

• Simulate complex materials (from molecules to bulk solids) with longer time scales and larger system sizes than previously possible.
• Choose the most physically appropriate gauge (including the velocity gauge for periodic systems) without being constrained by computational cost.
• Future-proof their simulations against the rapid evolution of HPC hardware, as the code can easily adapt to new accelerators (e.g., Hygon DCUs) through the abstraction layer.

The paper represents a significant step forward in making RT-TDDFT a practical tool for studying non-equilibrium phenomena in realistic materials using modern supercomputing resources.