High-performance Computation of Kubo Formula with Vectorization of Batched Linear Algebra Operation

One of the main challenges in data-driven material development is the small data problem. It is not uncommon for data scientific methodologies to fail to work due to a lack of sufficient data in a target domain. For overcoming this problem, computational approaches based on high-throughput first-principles calculations have attracted much attention and are actively used in data generation [29, 24, 25, 41]. Compared with experimental approaches, the computational approach has advantages in terms of data production, allowing us to construct a large-scale dataset or to run an exhaustive virtual screening.

Currently, almost all high-throughput calculation projects stand on the Kohn–Sham density functional theory (KS-DFT), which is an established and generic method for predicting the electronic structure of materials without assuming empirical parameters. The DFT fundamentally characterizes the ground state of electrons. It enables the computation of static quantities such as electron density, total energy, force fields, and magnetic moments, thus providing useful insights into structural stability, chemical reactions, and various practical applications[31].

External-field responses such as transport properties and optical properties not only give helpful knowledge of materials under a driving forces but also allow direct comparison between calculations and measurements. However, unlike the static quantities, these properties are not included in the standard DFT procedure because they must gather electron excitation, which is out of the range of DFT. Computing response functions in the first-principles level, we have to consider some theoretical extensions of DFT such as time-dependent DFT or additional post-process calculation depending on specific tasks.

Recently, a post-process scheme based on the Kubo formula with an effective tight-binding model has been widely used in the context of high-throughput calculations. The Kubo formula is a general formula for linear response phenomena, which can be evaluated in a single framework regardless of the system and target [30]. Usually, the electronic structure obtained by DFT is often downfolded to a tight-binding (TB) model such as Wannier TB to make its evaluation easy [32]. On the basis of such a scheme, several studies have successfully produced datasets of transport properties, for example, the charge conductivity, the thermoelectric conductivity, and the spin conductivity[38, 44, 39, 43].

Although the Kubo formula is an attractive method for high-throughput calculations due to its versatility and stability, it is computationally expensive, limiting the yield of workflow even if combined with TB models. Roughly estimated, the evaluation of Kubo formula typically costs 10 times as much as that of DFT calculation. Moreover, in contrast to the DFT core program, for which implementations utilizing a general-purpose graphics processing unit (GPGPU) exist, speeding up of Kubo formula remains a major challenge because it involves overcoming multiple computational bottlenecks. Modern high-performance computing (HPC) implementations of the Kubo formula can play a vital role in increasing the yield of data generation.

In this study, we develop an efficient method of evaluating the Kubo formula by optimizing for vector processors. As a practical example, we perform benchmark tests in computation of the anomalous Hall conductivity (AHC) [33] and show the results of performance comparisons between the vectorized method and an existing method of CPU implementation.

To examine the performance of the vector implementation, we conduct a demonstration on a vector machine NEC SX-Aurora TSUBASA where the vector length is 256 [36]. For comparison, the baseline implementation is also tested on a CPU machine, Xeon Gold, with 32-process parallelization by flat MPI. We are aware that comparison between different architectures is not straightforward, but here we compare single node performances as a compromise measure. The computing and compiling environments are listed in Table 1.

In this demonstration, we measure the computation times to evaluate AHC of the bcc-Fe and its supercell structures. Wannier TB models are constructed from the DFT calculation using the Quantum Espresso code [28] and the Wannier90 code [37]. We refer to the WANNIER-LINEAR-RESPONSE code [42] as the baseline implementation of Kubo formula and vectorize it by the method introduced in Sec.2.3.

Let us begin with the vector processing rate, a performance indicator defined as the proportion of vectorized processing in total processing. According to the profiler, our vectorized code achieves sufficiently high scores reaching 98-99 % in all demonstrations.

Fig.2 shows the computation times as a function of $N_{k}$. Here, we use the [1,1,2]-supercell where $N_{d}=48$ as it has 4 atoms with 2 spins and 6 Wannier projections per atom. From this data, we can see that the results on the SX-Aurora are about twice as fast as those of 32-process parallel execution on Xeon Gold, indicating a significant acceleration due to the vectorized implementation. This trend becomes more pronounced as $N_{k}$ increases. We again emphasize that this comparison is only a single aspect and does not lead general superiority of the vector processor over the CPU. Rather, this is clear evidence of the usefulness of the vectorization technique on this task.

Next, Fig.3 shows the computation times on the $N_{k}=150^{3}$ case as a function of $N_{d}$. Although the performance ratios varied by matrix size, in all cases, the SX-Aurora is more than 1.7 times faster than the Xeon Gold and 2.17 times at maximum. Note that the performance ratio on $N_{d}=192$ is lower because it is a special condition for better performance of the Xeon Gold rather than defects of the SX-Aurora. Together with the results in Fig.2, we claim that vectorized implementations have shown significant speedups in a wide range of tasks on small-scale systems, which is a desirable feature in high-throughput computations.

Finally, we examine the effect of the padding on the performance. Fig.4 compares the vectorized implementation with and without the padding, showing that the padding improves the performance by 1.5 times at maximum. The improvement becomes notable where $N_{d}\geq 96$ (8 atoms), that is, the memory-bank conflict causes serious performance degradation except in very small systems. Hence, padding is critically important in practical use.

In this study, we have proposed a method to accelerate the computation of Kubo formula by using batched matrix operations optimized to vector processors. Our findings have led to the following significant insights:

First, the introduction of vectorization was demonstrated to accelerate the computation of Kubo formula. Through benchmark comparisons between the vector-based NEC SX-Aurora TSUBASA and the scalar-based Xeon machines, we verified that the SX-Aurora has up to approximately 2.2 times faster node performance. This result indicates that our proposed method is effective and leads to increased data production in the context of high-throughput calculations.

The computational performance improvement due to padding was observed. This result underscores the importance of avoiding the memory-bank conflict to improve computational efficiency.

Furthermore, our approach is not limited to Kubo formula; it can be extended to various similar tasks relying on batched linear algebra operations. The vectorization approach could become a powerful alternative to the MPI in such cases. In addition, the same strategy can be applied to GPGPU as long as memory consumption is acceptable. The GPGPU implementation is an interesting future work.

Y.Y and T.K. are grateful to Dr. Taizo Shibuya at NEC Corporation for his valuable suggestions.

[sorting=count]