Acceleration of Atomistic NEGF: Algorithms, Parallelization, and Machine Learning

TL;DR

Problem: Ab-initio quantum transport with NEGF is computationally demanding, limiting realistic device sizes and inclusion of scattering effects. Approach: The authors develop a GPU-accelerated Serinv/RGF solver and a DFT+NEGF workflow within QuaTrEx to enable scalable simulations that incorporate electron-phonon and electron-electron interactions. Contributions: They demonstrate 80% weak scaling on Frontier for a ~25k-atom silicon nanoribbon with GW, validate current conservation, and show that an EGNN can predict Hamiltonian entries with ~2 meV accuracy to reproduce transmission curves. Impact: The framework expands the reach of atomistic quantum transport toward realistic devices and points to ML-assisted acceleration as a viable path to practical, ab-initio simulations.

Abstract

The Non-equilibrium Green's function (NEGF) formalism is a particularly powerful method to simulate the quantum transport properties of nanoscale devices such as transistors, photo-diodes, or memory cells, in the ballistic limit of transport or in the presence of various scattering sources such as electronphonon, electron-photon, or even electron-electron interactions. The inclusion of all these mechanisms has been first demonstrated in small systems, composed of a few atoms, before being scaled up to larger structures made of thousands of atoms. Also, the accuracy of the models has kept improving, from empirical to fully ab-initio ones, e.g., density functional theory (DFT). This paper summarizes key (algorithmic) achievements that have allowed us to bring DFT+NEGF simulations closer to the dimensions and functionality of realistic systems. The possibility of leveraging graph neural networks and machine learning to speed up ab-initio device simulations is discussed as well.

Figures (6)

• Figure 1: Iterative scheme to compute the non-equilibrium Green's functions (NEGF) for electrons (retarded: $\mathbf{G}^R$) and lesser/greater: $\mathbf{G}^{\lessgtr}$ in the presence of electron-phonon and/or electron-electron interactions. To determine the corresponding scattering self-energies ($\mathbf{\Sigma}^{R,\lessgtr}$), the phonon ($\mathbf{D}^{R,\lessgtr}$) and screened Coulomb ($\mathbf{W}^{R,\lessgtr}$) Green's functions must be evaluated, which requires the knowledge of the phonon self-energy ($\mathbf{\Pi}^{R,\lessgtr}$) and/or of the polarization function ($\mathbf{P}^{R,\lessgtr}$). All these quantities are matrices $\mathbf{M}$ that depend on the energy ($E$) or frequency ($\omega$). Their entries ($M_{ij}$) represent correlations between two points located at $R_i$ and $R_{j}$. The calculation of all these matrices involves solving linear systems of equations and open boundary conditions (LSE+OBC, red blocks) or performing energy convolutions through fast Fourier transforms (FFT, blue blocks). In a parallel implementation, it is convenient to store all $M_{ij}$ entries for one or a few $E$/$\omega$ points when dealing with LSE+OBC, while it is preferable to have access to a few $M_{ij}$ for multiple $E$/$\omega$ when computing the $\mathbf{\Sigma}^{R,\lessgtr}$, $\mathbf{\Pi}^{R,\lessgtr}$, or $\mathbf{P}^{R,\lessgtr}$ terms with FFT. Data can be "transposed" with all-to-all communication operations from one representation ($E$,$\omega$;$i,j$) to the other ($i,j$;$E$,$\omega$).
• Figure 2: Simplified representation of the Serinv algorithm serinv for the parallel calculation of the NEGF equations (retarded and lesser/greater). It relies on the repeated application of Schur complements. Starting from a block tri-diagonal input matrix, $P$ partitions are created, each of them being attributed to different computing units (CPUs or GPUs). After local processing, a reduced system of equations arises, which must be solved sequentially before reconstructing selected entries of the Green's functions locally. Note the mapping between the large and reduced systems (blocks with the same color).
• Figure 3: (a) Schematic view of a Si nano-ribbon with a height of 1.5 nm, a width of 5 nm, and a varying length $L_{tot}$. The surface of the nano-ribbon is passivated with hydrogen atoms. (b) Weak scaling results for the Si nano-ribbon in (a) with $L_{tot}$=52.1 nm and a total of $N_A$=25,344 atoms, in the presence of electron-electron interactions within the GW approximation, on the Frontier supercomputer. The experiment starts by measuring the execution time on 1 node of the machine with $N_E$=2 energy points, each of them being handled by 4 GPUs. Subsequently, $N_E$ increases proportionally with the number of nodes, reaching 18,800 on 9,400 nodes. The total time is decomposed into its communication and computation components.
• Figure 4: Simulation results for the Si nano-ribbon in Fig. \ref{['fig4']}(a) with $L_{tot}$=21.7 nm. A linear potential drop is applied with a voltage difference of 0.2 V between both ribbon extremities. (a) Local density-of-states extracted at $x$=0 in the ballistic limit of transport (thin blue line) and in the presence of electron-electron interactions within the GW approximation (thick red curve). (b) Spectral current distribution extracted at $x$=2.2, 10.9, and 19.5 nm. The inset shows the integral of this quantity (electrical current) as a function of $x$, demonstrating current conservation.
• Figure 5: Illustration on how machine learning can be inserted to accelerate ab-initio quantum transport simulations by predicting the entries of Hamiltonian matrices $H$ based on training from DFT. A single device structure (left), here a valence change memory (VCM) cell, is divided into several partitions and fed to an equivariant graph neural network (EGNN). Once trained with DFT data from the original VCM on multiple GPUs, this EGNN can accurately predict the Hamiltonian matrix of different VCM configurations with various oxygen vacancy distributions. During prediction, the average error of the $H$ entries does not exceed $\sim$2 meV with respect to DFT icml. Note that EGNNs scale with $O(N)$, where $N$ is the number of atoms, and DFT with $O(N^3)$.