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

This paper presents the GPU acceleration of the Abinit code for plane-wave DFT calculations, detailing the algorithmic revisions and vendor library integrations required for multi-GPU scalability while comparing the performance of Locally Optimal Block Preconditioned Conjugate Gradient and Chebyshev polynomial filtering diagonalization methods on heterogeneous CPU-GPU architectures.

The Gist

Imagine you are trying to solve a massive, incredibly complex jigsaw puzzle. This isn't just any puzzle; it's the puzzle of understanding how atoms and electrons behave inside materials (like the titanium in your phone or the silicon in a computer chip). This is the job of a software program called Abinit.

For years, Abinit has been running on supercomputers made of thousands of standard computer processors (CPUs). But recently, the world of computing has shifted. We now have GPUs (Graphics Processing Units)—the same chips that power video games and AI—which are like having thousands of tiny, super-fast workers who can all do the same simple task at the exact same time.

This paper is the story of how the team behind Abinit moved their puzzle-solving operation from a team of slow, careful workers (CPUs) to a stadium full of lightning-fast, synchronized workers (GPUs).

Here is the breakdown of their journey, using simple analogies:

1. The Problem: Too Many Pieces, Too Slow

In the world of quantum physics, the "puzzle pieces" are called electronic wave functions. To solve the puzzle, the computer has to do a massive amount of math to figure out where these electrons are.

• The Old Way (CPU): Imagine a single librarian trying to sort a million books. They do it one by one, very carefully. It's accurate, but it takes forever.
• The New Way (GPU): Imagine a stadium with 10,000 librarians. If you give them a simple instruction like "Sort all the red books," they can do it instantly. The challenge is that the old Abinit code was written for the single librarian, not the stadium.

2. The Strategy: "Batching" the Work

The biggest mistake you can make with a stadium of workers is giving them one book at a time. They spend all their time waiting for the next book.

• The Analogy: Instead of handing a worker one book, you hand them a whole stack.
• The Fix: The team changed Abinit to use Batch Processing. Instead of calculating the math for one electron at a time, they group thousands of electrons together and feed them to the GPU all at once. This keeps the "stadium" busy and eliminates the waiting time.

3. The Traffic Jam: Moving Data

GPUs are like a high-speed race track, but the data lives in the CPU's garage. Moving data back and forth is slow and causes traffic jams.

• The Analogy: Imagine the workers (GPU) are in a factory, but the raw materials (data) are in a warehouse (CPU). If you have to drive a truck back and forth for every single brick, the factory sits idle.
• The Fix: The team decided to move the entire pile of raw materials to the factory floor at the start of the day. They keep the data on the GPU as long as possible, only moving it back to the CPU when absolutely necessary. This keeps the race track clear.

4. The Two Main Algorithms: The Sprinter vs. The Marathoner

To solve the puzzle, Abinit uses two different mathematical strategies (algorithms). The paper compares them like two different types of athletes:

• Algorithm A: LOBPCG (The Sprinter)

  • How it works: It takes a step, stops to check its position (communicates with other workers), takes another step, and stops again.
  • The Flaw: It stops a lot. Every time it stops to check, it has to talk to other workers across the network. This "talking" (communication) is slow. On a GPU, where speed is everything, stopping to chat kills performance.
  • Verdict: Good for small jobs, but gets bogged down on massive puzzles.

• Algorithm B: Chebyshev Filtering (The Marathoner)

  • How it works: It runs a long, continuous stretch of work without stopping to check its position. It does a huge amount of math in one go, then checks once at the very end.
  • The Win: Because it keeps running without stopping to talk, it utilizes the GPU's massive speed perfectly. It does more work per "stop."
  • Verdict: This is the winner for GPUs. It turns the GPU into a powerhouse.

5. The Results: Speed and Energy Savings

The team tested this new setup on real supercomputers using both NVIDIA (the "gold standard" for GPUs) and AMD chips.

• Speed: They found that using GPUs made the calculations 13 to 17 times faster than using just CPUs. In some cases, 4 GPU nodes did the work of 128 CPU nodes!
• Energy: Because the GPUs finish the job so much faster, they use less total electricity. It's like driving a sports car that finishes a race in 2 minutes versus a truck that takes 2 hours; even if the car burns more gas per minute, it uses far less total fuel to finish the race.
• The Catch: The "Rayleigh-Ritz" step (a specific part of the math where they organize the final pieces) is still a bit slow on GPUs, especially on AMD chips. It's like the one part of the factory where the workers still have to stop and chat. The team is working on fixing this next.

The Bottom Line

This paper is a success story of modernizing old software. By rethinking how the math is done (batching data) and choosing the right strategy (Chebyshev filtering over LOBPCG), the team turned Abinit into a GPU monster.

Why does this matter?
Scientists can now simulate larger, more complex materials in a fraction of the time. This means we can design better batteries, more efficient solar panels, and new drugs much faster than before. They didn't just buy faster computers; they taught the computers how to run a better race.

Technical Summary

1. Problem Statement

Large-scale electronic structure calculations using Kohn-Sham Density Functional Theory (DFT) with plane-wave basis sets are computationally intensive. The primary bottleneck lies in solving the Kohn-Sham eigenvalue problem to determine electronic wave functions, which involves:

• Iterative diagonalization of large Hamiltonian matrices.
• Massive linear algebra operations (matrix-matrix multiplications) and Fast Fourier Transforms (FFTs).
• Handling thousands of electronic states and plane-wave basis functions.

While High-Performance Computing (HPC) has shifted toward heterogeneous architectures (CPU + GPU), porting legacy codes like Abinit to GPUs is challenging. Previous attempts were limited by:

• Inefficient data movement between host (CPU) and device (GPU).
• Algorithmic structures not optimized for GPU parallelism (e.g., reliance on small matrix operations or frequent synchronization).
• The need to balance the efficiency of different iterative diagonalization algorithms (specifically LOBPCG vs. Chebyshev filtering) on GPU hardware.

2. Methodology and Porting Strategy

The authors present a comprehensive re-implementation of Abinit (v10.6+) for GPU architectures, moving away from custom CUDA kernels to a modern, portable approach.

A. Programming Model and Architecture

• Model: The code utilizes OpenMP offloading (version 5.0+) for the CPU-GPU interface. This allows for directive-based programming, minimizing code rewriting and ensuring portability across NVIDIA (CUDA) and AMD (HIP/ROCm) hardware.
• Abstraction Layer: A new software layer was developed to manage distributed matrices and vectors, hiding the complexity of underlying GPU libraries (cuBLAS, rocBLAS, cuSOLVER, rocSOLVER, cuFFT, rocFFT) from the main application logic.
• Batch Processing: A core optimization strategy involves batching. Instead of processing electronic bands one by one, data is regrouped into larger chunks. This transforms Level-2 BLAS operations into Level-3 (matrix-matrix) operations and allows for batched FFTs, significantly increasing throughput and reducing kernel launch overhead.

B. Memory Management

• GPU-Resident Wave Functions: To minimize costly Host-Device (H2D) transfers, the wave function is transferred to the GPU once at the start of a Self-Consistent Field (SCF) iteration and kept resident in GPU memory throughout the iterative diagonalization process.
• Layout Optimization: The code manages two memory layouts:
  1. Row-distributed: Optimized for linear algebra solvers (eigensolvers).
  2. Column-distributed (Waveform layout): Optimized for Hamiltonian application and FFTs.
  • Switching between these layouts is handled via GPU-aware MPI all-to-all communications, decoupling data movement from computation.

C. Algorithmic Revisions

The paper focuses on two iterative diagonalization algorithms, analyzing their suitability for GPUs:

1. LOBPCG (Locally Optimal Block Preconditioned Conjugate Gradient): Relies on block-wise orthogonalization and frequent Rayleigh-Ritz (RR) procedures.
2. Chebyshev Polynomial Filtering: Uses spectral filtering to amplify desired eigenvalues. It applies the Hamiltonian operator $k$ times in a sequence before a single global RR step.

Key Theoretical Insight:

• Arithmetic Intensity: Chebyshev filtering achieves $k$ times higher arithmetic intensity than LOBPCG per MPI rank because it performs $k$ Hamiltonian applications between communication steps, whereas LOBPCG communicates after every block update.
• Communication Overhead: LOBPCG requires $k \times b$ (blocks) all-to-all communications per iteration, while Chebyshev filtering requires only 2 global communications per iteration.

3. Key Contributions

1. Full GPU Port of Abinit: A complete re-implementation supporting NVIDIA and AMD GPUs using OpenMP offloading and vendor libraries, eliminating the need for custom kernels.
2. Algorithmic Comparison: A rigorous theoretical and empirical comparison of LOBPCG and Chebyshev filtering on GPUs. The authors demonstrate that Chebyshev filtering is superior for GPU architectures due to higher arithmetic intensity and lower communication frequency.
3. Performance Metrics: Introduction of metrics based on arithmetic intensity and roofline models to classify operations as compute-bound (Hamiltonian application, FFTs) or memory/communication-bound (orthogonalization, Rayleigh-Ritz).
4. Energy Efficiency Analysis: A "node-to-node" comparison showing that GPU-accelerated nodes significantly reduce both execution time and energy consumption compared to CPU-only nodes.

4. Results and Performance

The authors benchmarked the code on supercomputers Jean Zay (NVIDIA A100/H100) and Adastra (AMD MI250X).

• Speedup:
  • NVIDIA GPUs: Achieved excellent speedups. For a 255-atom Titanium system, 2 GPU nodes outperformed 8 CPU nodes (approx. 17x speedup for the filtering step).
  • AMD GPUs: Showed good speedup but were less efficient than NVIDIA, particularly in the Rayleigh-Ritz step, attributed to performance limitations in the ROCm hegvd routine.
• Scaling:
  • Strong Scaling: GPU scaling efficiency lags behind CPU as node count increases (consistent with Amdahl's law), primarily due to the Rayleigh-Ritz step becoming a bottleneck.
  • Chebyshev vs. LOBPCG: Increasing the polynomial degree in Chebyshev filtering improves eigenvector accuracy without degrading performance, whereas increasing minimization lines in LOBPCG yields diminishing returns. Chebyshev filtering allows for fewer total SCF iterations, leading to faster convergence.
• Energy Efficiency:
  • GPU nodes demonstrated significant energy savings. On Jean Zay (NVIDIA), energy savings were substantial (up to 15x compared to CPU-only scaling), whereas AMD nodes showed linear scaling in energy consumption with node count.
• Roofline Analysis:
  • Compute-Bound: Hamiltonian application (FFT and non-local potential) and spectral filtering operate near peak performance (Tensor cores).
  • Memory-Bound: The Rayleigh-Ritz procedure and orthogonalization are memory-bound and communication-bound, limiting scalability for large numbers of bands.

5. Significance and Future Perspectives

• Modernization: This work establishes Abinit as a state-of-the-art tool for exascale computing, capable of leveraging heterogeneous architectures efficiently.
• Algorithm Selection: The paper provides a clear guideline for users: Chebyshev filtering is the preferred algorithm for GPU-accelerated DFT calculations due to its high arithmetic intensity and lower communication overhead.
• Future Work: The authors suggest that to further overcome the Rayleigh-Ritz bottleneck (which scales poorly with band count), the Spectrum Slicing method could be adopted. This would distribute the eigenvalue problem across GPUs, avoiding global diagonalization of large matrices.

In conclusion, the paper demonstrates that successful GPU acceleration of DFT requires not just hardware porting, but fundamental algorithmic restructuring to maximize arithmetic intensity and minimize data movement, with Chebyshev filtering emerging as the optimal strategy for modern GPU clusters.