AI-Assisted GPU optimization of the Stochastic Variational Method for Few-Body Boson System

AI-Assisted GPU optimization of the Stochastic Variational Method for Few-Body Boson System Noritaka Shimizu∗ , Shigeyoshi Aoyama† , Yusuke Tsunoda∗ , Akira Nukada∗ ∗ Center for Computational Sciences, University of Tsukuba, Tsukuba, Japan

{shimizu, ytsunoda}@nucl.ph.tsukuba.ac.jp, nukada@cs.tsukuba.ac.jp

† Computing Research Center, High Energy Accelerator Research Organization (KEK), Tsukuba, Japan

arXiv:2610.02745v1 [physics.comp-ph] 2 Oct 2026

shigeyoshi.aoyama@kek.jp

Abstract—The stochastic variational method (SVM) is one of the most powerful methods to solve quantum few-body systems precisely in various fields, such as nuclear and atomic physics. To the best of our knowledge, no SVM code optimized for GPUs has been reported. We developed a parallel SVM code for few-body clusters of 4 He atoms and optimized it for two GPU architectures, NVIDIA GH200 and AMD Instinct MI300A, with the entire code written by an AI coding agent. On the algorithmic side, we introduce the secular-equation method with the Gu–Eisenstat prescription for solving the generalized eigenvalue problem within the SVM framework. The main part of the code tuning is to batch the many small matrices so that the GPUs are used efficiently, together with an array layout in memory optimized for coalesced addressing. For the five-body system with 4800 SVM basis states, the tuned SVM code runs 16.2 (14.7) times faster on the NVIDIA GH200 (AMD MI300A) than the same tuned code on the Intel Xeon Max CPU system, and 168 (153) times faster than the first working CPU implementation. The achieved performance brings 6-body and larger cluster systems within reach. Index Terms—Quantum many-body problem, stochastic variational method, GPU acceleration

I. Introduction 1

Quantum few-body problems, namely calculating ground-state and low-lying excited states of atomic, nuclear, and exotic clusters, are fundamental in computational physics. The Stochastic Variational Method (SVM) with a correlated Gaussian basis [1], [2] is one of the most successful methods for obtaining high-precision solutions, particularly for problems with strong correlations, such as halo nuclei or weakly bound atomic clusters like helium dimers and trimers. In the present work, we adopt the SVM to evaluate a few-body system consisting of 4 He atoms. The helium-4 cluster system has attracted much interest in connection with the Efimov physics. Much effort has been devoted to elucidating this physics with the Gaussian expansion method [3], density functional theories, and other approaches [4]. Precise quantum many-body 1 (c) 2026 IEEE. Personal use of this material is permitted. Permission from IEEE must be obtained for all other uses, in any current or future media, including reprinting/republishing this material for advertising or promotional purposes, creating new collective works, for resale or redistribution to servers or lists, or reuse of any copyrighted component of this work in other works.

calculations beyond the five-body system remain difficult because they have to treat the strong short-range repulsion and are limited by available computational resources. In the SVM, the wave function is written as a linear combination of correlated Gaussian basis states, which are selected stochastically from many candidates so that the variational energy is systematically lowered as the basis space is enlarged. In the original SVM the selection is performed sequentially: candidates are generated and their energies evaluated one by one, and a candidate is adopted if a certain energy gain is achieved. Such a sequential selection, however, hampers efficient parallel computation, and “Parallel SVM” was proposed as a parallel-computing-based extension of the SVM [5]. In this method a fixed number of candidate basis states are evaluated in parallel, which exploits the performance of modern massively parallel computers efficiently, and the best one is chosen among these candidates incrementally. In the present work, we follow this strategy which is expected to be beneficial also for GPU computing. Most of the latest supercomputers are equipped with Graphics Processing Units (GPUs), which accelerate scientific computations by an order of magnitude over CPU systems. Although the SVM algorithm is well suited to parallel computation, GPU programming remains a demanding task for physicists. The recent growth of AI coding agents drastically reduces this burden: performance tuning is particularly well suited to them, because the correctness of every code modification can be verified automatically against reference results and the target is clearly defined, namely the measured wall time. In the present work, we develop a highly tuned SVM code employing AI-coding and discuss its performance on GPU machines. This paper is organized as follows: the present Hamiltonian and the framework of the SVM are described in Sect. II, Sect. III is devoted to the algorithm and code improvements, and the energy convergence and the code performance are discussed in Sect. IV. Sect. V summarizes this work. II. Method We present the Hamiltonian of the few-body 4 He boson system and the SVM framework used to obtain its

eigenstates. Since the SVM is a variational method, the energy, namely the expectation value of the Hamiltonian, has to be evaluated efficiently. The Hamiltonian is given in Subsect. II-A, and the SVM framework together with the formulae for evaluating the energy in Subsect. II-B. The detailed formulae to compute the matrix elements are described in Subsect. II-C. The AI-assisted coding process is described in Subsect. II-D. A. Hamiltonian of 4 He atomic cluster We begin with the Hamiltonian of N -body system, including the two-body potential as H =T +V =−

NJ N X ℏ2 2 X ∇a + V (rij ), 2ma a=1 i<j

(1)

where ∇2a and ma are the Laplacian with respect to the Jacobi coordinate a and its reduced mass, respectively. The symbol NJ is the number of the Jacobi coordinates, namely NJ = N −1. The symbol rij is the relative distance between two particles i and j. The mass of the 4 He atom, m, is taken as ℏ2 /(mkB ) = 12.12 KÅ2 [6]. In the present work, the two-body potential is approximated as a tworange Gaussian function of the relative distance r as V (r) =

2 X

vν e−(r/µν )

2

(2)

ν=1

with v1 = 2000 K, v2 = −8.28233 K, µ1 = 1Å, and µ2 = 3.48813 Å. These parameters were determined to reproduce the scattering length and the dimer energy of the LM2M2 potential [7], and this potential was adopted in some previous studies [8], [9]. The first term represents a strong repulsive force caused by the Pauli principle, and the second a shallow attractive van der Waals force. This strong repulsion makes it difficult for many other methods to compute the wave function precisely. B. Stochastic Variational Method The SVM wave function is a linear combination of the correlated Gaussian basis states as ΨNb (x) =

Nb X

ck SϕAk (x),

(3)

k=1

where S is the symmetrization operator for the Bose statistics. The symbol x = (⃗r1 , ⃗r2 , · · · , ⃗rNJ ) is a set of Jacobi coordinates for an N particle system as shown in Fig. 1 with the number of the coordinates being NJ = N −1. The Jacobi coordinates are defined as relative coordinates of the particles so that the center-of-mass motion is subtracted exactly. Only one of the possible sets is shown; any rearrangement of the Jacobi coordinates is absorbed into the matrix A introduced below. In the present work we consider only the J π = 0+ state, so that the global vector for the angular momentum is not required and the functional form is kept simple [2].

Fig. 1. Jacobi coordinates, ⃗ ra , of five-body system (N = 5 and NJ = 4).

The correlated Gaussian basis ϕ(x) is defined as     NJ X 1 1 ϕA (x) = exp − xT Ax = exp − Aab ⃗ra · ⃗rb  , 2 2 a,b=1

(4) where the symbol A is a symmetric positive definite matrix of size NJ ×NJ and thus contains NJ (NJ +1)/2 parameters for each matrix. In the SVM framework, the matrices A are generated randomly, and the best one to minimize the energy is selected based on the variational principle. The eigenenergies and the coefficients ck in Eq. (3) are determined by solving the Hill–Wheeler equation, or the generalized eigenvalue problem, written here in the bra-ket notation with ϕA (x) = ⟨x|A⟩, as X X Hkk′ ck′ = E Nkk′ ck′ . (5) k′

Hkk′ = ⟨Ak |HS|Ak′ ⟩,

k′

Nkk′ = ⟨Ak |S|Ak′ ⟩.

(6)

Thus, the wave function is written as a linear combination of the symmetrized correlated Gaussian basis, whose matrix elements can be computed efficiently employing the Gaussian integral as discussed in Subsect. II-C. Algorithm 1 shows the parallel SVM. At each step, where K basis states have been accepted, Nt candidates are evaluated in parallel and the one that lowers the energy most is accepted. Algorithm 1 Algorithm of the parallel SVM [5]. 1: for K = 0, 1, ..., Nb − 1 do 2: prepare Nt random matrices A(1) , . . . , A(Nt ) 3: for t = 1, . . . , Nt do 4: set AK+1 = A(t) 5: compute H1:K+1,K+1 , N1:K+1,K+1 6: solve Eq. (5) with dimension K +1 → energy E (t) 7: end for 8: t ← arg min E (t) 9: AK+1 ← A(t) and reconstruct H and N 10: end for

C. Matrix elements in SVM The most time-consuming part of the SVM method is to evaluate the Hamiltonian and norm matrix elements in Eq. (6). For identical Bose particles, it is necessary to symmetrize the wave function under the exchange of particles. A particle permutation acts on Jacobi coordinates as an orthogonal matrix P , and transforms a Gaussian into another Gaussian with matrix P T AP . The symmetrized basis element is X (p) ⟨Ak |HS|Ak′ ⟩ = ⟨Ak |H|Ak′ ⟩, (7) p (p)

where Ak′ = PpT Ak′ Pp and Pp is an NJ × NJ matrix representing the exchange of particles corresponding to the permutation p. Since the Hamiltonian commutes with S, it is sufficient to symmetrize the wave function on one side. Consequently, the computation of every matrix element demands a sum over the N ! permutations, which makes the computational cost scale as O(N !). The norm matrix element Nkk′ is computed in a similar way. The overlap between two different basis states is calculated efficiently by employing the Gaussian integral as  (p) ⟨Ak |Ak′ ⟩ =

(2π)NJ det(B)

3/2 (p)

, B = Ak + Ak ′

(8)

where the determinant is calculated by the Cholesky decomposition of the matrix B. The norm matrix element is calculated as the summation concerning p. Hereafter, we omit the index (p) for simplicity. The matrix elements of the kinetic term are computed as ⟨Ak |T |Ak′ ⟩

=

3 ⟨Ak |Ak′ ⟩ Tr Ak B −1 Ak′ M , (9) 2

where M is the diagonal matrix whose elements are Mab = ℏ2 δab /ma and ma is the reduced mass of the Jacobi coordinate a. Thus, the computation cost of Eq. (9) is O(NJ3 ) due to the Cholesky decomposition of B to construct B −1 . For the potential term, the relative position vector ⃗rij of the particle pair (i, j) is written as a linear combination of the Jacobi coordinates as ⃗rij

=

2 rij

=

NJ X

wa(i,j) ⃗ra

a=1 T (i,j)

x w

(w(i,j) )T x,

(10)

where the constant coefficient vector w(i,j) ∈ RNJ is given by the inverse Jacobi transformation. Thus, the potential is rewritten in terms of the Jacobi coordinates as X (11) V (rij ) = vν exp(−xT w(i,j) (w(i,j) )T x/µ2ν ). ν

Thanks to the matrix determinant lemma, the matrix elements of the potential term are calculated efficiently as X vν ⟨Ak |V |Ak′ ⟩ = ⟨Ak |Ak′ ⟩ . (ij) (1 + 2s /µ2ν )3/2 ν,i,j s(ij)

=

(w(i,j) )T B −1 w(i,j) .

(12)

While s(ij) is independent of ν, it depends on k, k ′ , p, i, and j. The pair index (i, j) runs over N (N −1)/2 values, so that the computational cost per matrix element is O(N 4 ), which is heavier than the kinetic term in Eq. (9) and is the bottleneck part of the whole SVM computation. A typical SVM calculation builds the variational basis incrementally: at each step, hundreds to thousands of trial Gaussian basis states are evaluated, and the best one, which has the lowest energy eigenvalue, is added to the basis. For every candidate, the generalized eigenvalue problem in Eq. (5) is solved to obtain the energy. For physically interesting problems the basis size reaches Nb = 100–10000 or more, depending on the number of particles, and the costs of the matrix-element evaluation and solving the eigenvalue problems become prohibitive on a single CPU. The former cost can be reduced through GPU-oriented code optimization, while the latter can be reduced by introducing the secular-equation method. These improvements are described in the next section. D. AI-assisted coding By using AI agents, we developed the whole SVM code from scratch. We gave many instructions to the AI agents (Claude Opus mainly) with natural language, the code was developed gradually in a few months. We do not touch the code directly except for a few trivial modifications. Concerning the GPU optimization, we asked the AI agent to measure the profile data, and it measured and analyzed the data suggesting us how to optimize the code. At the early stage, the code was developed by the one of the authors (N.S.) who has little experience in the CUDA developments, and wrote the ”baseline” code including the technique discussed in Subsect. III-A. After that, the discussions with High-Performance-Computing experts accelerates the CUDA optimization, indicating that the domain knowledge of the CUDA coding and GPU hardware is essential even for AI-coding. The code development was completed after more than 1000 brief interactions with AI and many trial-and-errors including the technique discussed in Subsect. III-B. The production code consist of 13,682 lines of Fortran, CUDA and HIP codes. III. Performance improvements In this section, we focus on two improvements which dominate the gain over a straightforward implementation, and they act on different components. The secularequation treatment of the generalized eigenvalue problem removes the O(Nb3 ) cost from the per-candidate loop,

which is discussed in Subsect. III-A. The transposition of the memory layout makes the matrix-element kernel use the memory system efficiently, which is discussed in Subsect. III-B. The computations for Nt candidates are mutually independent, so that the Nt candidates are distributed evenly over the Nproc MPI processes, each of which handles Nt /Nproc candidates. Within a process the workload is divided between GPU and CPU: the matrix elements are computed on GPU when available, while the eigenvalue problems of the candidates are solved on the CPU with OpenMP threads. The candidates of a process are further split into chunks of the permutations, and the chunks are pipelined so that the CPU-side processing of one chunk overlaps with the GPU execution of the following ones. A. Solving the generalized eigenvalue problem with the secular equation At each SVM step, a (K + 1)-dimensional generalized eigenvalue problem must be solved for each of the Nt candidates. Solving it from scratch costs O(K 3 ), so that this stage alone would cost O(Nt K 3 ) per step and O(Nt Nb4 ) for the whole run, which overtakes the cost of the matrix elements as soon as Nb becomes large, even though its efficient GPU implementation is feasible [10]. In Ref. [2], the secular-equation technique was introduced in the SVM to avoid this problem. Let the eigenpairs of the K-dimensional problem be already solved, with the eigenvectors Û normalized with respect to the norm matrix, Û T NK Û = I and Û T HK Û = diag(εi ). Appending a candidate as the (K +1)-th basis state adds one row and one column to the K ×K Hamiltonian and norm matrices. Transforming the new row and column into the eigenbasis of the K × K problem gives vk =

K X k′ =1

Ûk′ k Hk′ ,K+1 ,

uk =

K X

Ûk′ k Nk′ ,K+1 (13)

k′ =1

for k = 1, . . . , K, together with the diagonal elements HK+1,K+1 and NK+1,K+1 . By applying the GramSchmidt orthogonalization so that the new basis state is orthogonal to the fixed K basis states, the generalized eigenvalue problem is rewritten as a standard eigenvalue problem of the symmetric arrowhead-structure matrix   ε1 β1  ..  .. . .  b = (14) H ,   εK β K  β1 · · · βK α with γ = NK+1,K+1 − uT u and P HK+1,K+1 − 2 u·v + k εk u2k α= , γ

βk =

v k − εk uk , √ γ (15)

where the empty matrix elements are zero. The K + 1 eigenvalues of Eq. (14) are the roots of the secular equation det(Ĥ − EI) = 0, namely f (E) = (α − E) −

K X k=1

βk2 = 0. εk − E

(16)

If γ is too small or non-positive, the candidate is linearly dependent on the accepted basis, and the candidate is rejected. Thanks to Cauchy’s interlacing theorem, the roots interlace the eigenvalues εk , so that each of them is bracketed in advance and is obtained by the bisection method in a numerically stable manner. In the candidate loop only a few of the lowest eigenvalues are required. The basis transformation of Û costs O(K 2 ) and each evaluation of f (E) costs O(K), so that the candidate energy is evaluated in O(K 2 ) instead of O(K 3 ), and no eigenvector is computed at this stage. While this Û transformation is achieved as a matrix-vector product for each candidate, we collect the vectors of many candidates into a chunk and perform the products as a matrix-matrix product for further acceleration. This is the “Secular equation” component of Fig. 3, which is executed Nt times per step. The eigenvectors are needed only once per step, for the candidate that is actually accepted, because the (K + 1)th step requires all the eigenpairs of the K × K problem. They are constructed from the roots of Eq. (16) by the Gu–Eisenstat prescription [11]: building the eigenvectors directly from the βk loses their orthogonality when two roots lie close to each other, and the prescription instead reconstructs the couplings from the computed roots, which restores orthogonality to working accuracy. The cost of this update is O(K 3 ), but it is dominated by a single matrix–matrix product, the back transformation of the eigenvectors into the original basis, which runs far faster than a general eigenvalue solver such as dsygvd. This is the “Incremental update” component, executed once per step. This step is parallelized using CPU OpenMP threads, with GPU acceleration enabled for K ≥ 2000. Since each incremental update uses the eigenpairs of the preceding step as exact input, rounding errors accumulate over the steps. We therefore reconstruct the eigenpairs from scratch every 100 basis states, which is called “Eigensolver” component in Tab. III and is executed 48 times in a run with Nb = 4800 performed on GPU (cusolverDnDsygvd / hipsolverDsygvd) if available. The cost per step is thus Nt ×O(K 2 )+O(K 3 ) instead of Nt × O(K 3 ). The benchmark of Fig. 3 shows the effect of this improvement: on the GH200 the per-candidate part costs 10.1 s although it is executed Nt = 1200 times more often than the once-per-step update, which costs 14.6 s, while the 48 library solves cost only 3.2 s. The MPI communication follows the same division. Only the row of the winning candidate have to be exchanged, which is O(K) data per step; the K × K eigenvector matrix is

The matrix-element computation of the K-th basis decomposes into N ! × K × Nt mutually independent tasks, each of which builds an NJ × NJ matrix B, factorizes it by the Cholesky decomposition, and computes the matrix elements in Eqs. (9) and (12). With the O(N 4 ) cost of one such task and the Nb (Nb + 1)/2 matrix elements of the whole run, the total cost is O(N 4 N !Nb2 Nt ). For the fivebody system with Nt = 1, 200 already N !Nt = 144, 000 tasks are available even at Nb = 1, and thus the computation is embarrassingly parallel and maps naturally onto a GPU: one thread block is assigned to a candidate basis and its threads stride over the permutations, that is, one thread processes one permutation. Since the number of particles is fixed through a whole computation, the kernel is instantiated as a template over NJ = N − 1, all inner loops are unrolled, and every intermediate matrix is held in registers. (p) The permuted matrices Ak′ = PpT Ak′ Pp do not depend on the index of the accepted basis state, so they are computed once per candidate and reused for all Nb pairs. How they are stored turns out to be decisive. In the natural array-of-structures (AoS) layout, the NJ2 elements of one permutation are contiguous, so that adjacent threads, which handle adjacent permutations, read addresses NJ2 doubles apart. Profiling with NVIDIA Nsight Compute showed that each 32-byte memory sector then delivered only 8 useful bytes, that the L1 cache throughput was the saturated resource at 67%, and that the DRAM utilization was negligible at 0.02%: the kernel was limited by the number of memory transactions, not by the bandwidth. The remedy is to transpose the array indices into a structure-of-arrays (SoA) layout in which the permutation index is the innermost and contiguous one. Namely, (p) (1) (2) (3) the matrix Ai,j is stored as A1,1 , A1,1 , A1,1 · · · , so that adjacent threads read adjacent doubles and the loads of a warp coalesce into full sectors. The arithmetic is untouched, so that the energies are bit-identical to those of the previous layout, while the matrix-element stage becomes 1.3–1.6 times faster. The same transformation applies to the CPU code, where the vector lanes play the role of the threads. Vectorizing inside one small matrix fails, because the innermost contraction has length NJ = 4 and fills only half of the eight lanes of a 512-bit register, and because the Cholesky factorization is sequential within one matrix. Instead the permutation loop is moved from the outermost to the innermost position and tiled, with the same structure-ofarrays layout, so that each vector instruction advances the same elementary step of eight independent factorizations. Non-positive pivots cannot break out of a vector loop and

IV. Results We show the energy convergence of the SVM method in Subsect. IV-A, and the code performance in Subsect. IV-B. A brief discussion of the 6-body system is shown in Subsect. IV-C. A. Convergence of the SVM method. Nt=300 Nt=600 Nt=1200 Nt=2400

-1 E (K)

B. Memory layout of the matrix-element kernel

are handled by lane masking, the software counterpart of the divergence handling of the GPU. The correspondence is one to one — GPU thread = SIMD lane, coalesced access = unit stride — and reflects a single requirement, namely that the parallel dimension be contiguous in memory.

-1.25

Nt=4800 Nt=9600 Nt=14400

(a)

-1

10 E - E0(K)

replicated on every rank and updated redundantly, so that no O(K 2 ) matrix is ever communicated. The measured “MPI synchronization” is accordingly at the percent level, 1.2 s out of 70.8 s.

-2

10

-3

10

(b) 0

2400

4800

7200

9600

Fig. 2. (a) Convergence of the ground-state energies of the five-body system against the number of the basis states Nb . The grey, blue, green, orange, red, purple and brown lines denote the SVM results with Nt = 300, 600, 1200, 2400, 4800, 9600, and 14400, respectively. (b) Same as (a), but the vertical axis shows the energy subtracted by the estimated exact energy, ∆E = E − E0 , on a logarithmic scale.

As an example, we show the energy convergence of the five-body system against the number of the basis states Nb in Fig. 2 (a). As Nb increases, the energy decreases monotonically and is nearly converged by Nb ∼ 4800 for any Nt , although the converged value depends slightly on Nt . Figure 2 (b) shows the energy difference ∆E from the estimated exact energy on a logarithmic scale. The exact value is estimated as E0 = −1.30851(1) K employing the energy-variance extrapolation [12], in which the energy is extrapolated to zero energy variance. The details of this extrapolation method will be reported elsewhere. The energy decreases with Nb , and the convergence behavior improves as Nt increases. The gain becomes progressively smaller: that from Nt = 9600 to Nt = 14400 is only 0.09 mK, indicating convergence with respect to Nt . Table I lists the wall time as a function of Nt . It is approximately linear in Nt as Wall time = 181.5 + 0.1664 Nt .

(17)

The constant term is the Nt -independent cost (eigensolver, incremental update and I/O), while the slope reflects the

TABLE I SVM energies of the five-body 4 He cluster with respect to the number of trial bases per basis function, Nt . Runs were performed with the tuned code and Nb = 9600 on 4 nodes of Miyabi-G.

−1.303765 −1.305901 −1.306942 −1.307563 −1.307869 −1.308062 −1.308144 −1.30851(1)

∆E (mK)

Wall time (s)

4.74 2.60 1.57 0.95 0.64 0.45 0.36

231.8 277.8 381.7 577.1 987.3 1782.7 2574.2

The exact energy, E0 , is estimated with the energy-variance extrapolation. ∆E is the energy relative to E0 .

Total 1147.5 s

1000

139.8 s

80

600 400

B. Code Performance We present the wall time of the tuned SVM code for the five-body system with MPI parallelization. The SVM calculations were performed with Nb = 4800 and Nt = 1200 and their wall times were measured. Four devices are used on every platform, namely 4 nodes for Miyabi-G, 2 nodes (4 CPUs) for Miyabi-C, and 1 node (4 APUs) for Sirius supercomputer, which are summarized in Tab. II. Table III shows the wall time decomposed into six components, corresponding to the steps of the SVM algorithm described in Sect. II. The times are measured on the three devices with AI-tuned code and with the first-implemented code (baseline) as a comparison. The performance tuning on the CPU substantially accelerates the code, achieving a speedup of up to a factor of 10; the port to the GH200 reaches 168. Figure 3 also shows the measured time as breakdown bars. “Matrix elements” is the evaluation of the couplings of every candidate to the accepted basis states, Eqs. (9) and (12), each of which carries the sum over the N ! permutations. “Secular equation” is the energy evaluation of each candidate from those couplings, “Incremental update” is the update of all eigenpairs by the secular-equation method once the winning candidate has been accepted, and “Eigensolver” denotes the full solution of the generalized eigenvalue problem using the dsygvd library call or its variants. “MPI synchronization” is the collective communication that selects the globally best candidate, and “Other” collects the remaining host-side work such as the candidate generation and the file I/O. On the Sirius supercomputer we found large synchronization delays among the MPI processes when 24 OpenMP threads were used per APU, even though each APU has 24 CPU cores. Running with 23 threads per APU gives better performance and more stable wall times than with 24 threads.

6.5 s 5.2 s

4.3 s 3.2 s 10.1 s

13.0 s 3.8 s

14.6 s

40

15.3 s

939.1 s

20

37.5 s

34.1 s

Miyabi-G (4 GH200) 4 MPI × 72 OMP

Sirius MI300A (4 APU) 4 MPI × 23 OMP

200 0

0

Miyabi-C (4 Xeon Max 9480) 4 MPI × 56 OMP

Component (stack order, top

Other MPI sync

matrix-element evaluation, which scales linearly with the number of trials. Since the computational cost also scales as Nb2 , Nt and Nb have to be chosen as a trade-off under a given wall-time budget.

Total 77.9 s

Total 70.8 s

60 800

Wall time (s)

300 600 1200 2400 4800 9600 14400 E0 (est.)

E (K)

Wall time (s)

Nt

1200

Eigensolver Secular eq

bottom)

Incremental upd Matrix elements

Fig. 3. Wall-time breakdown of the tuned SVM code on Miyabi-C with Intel Xeon CPU (left), Miyabi-G with GH200 (middle), and Sirius with MI300A (right). The wall time of the SVM computation of the five-body system with Nb = 4800 and Nt = 1200 for 4 devices with MPI parallel. The components are the matrix-element computation (blue), the incremental update (orange), the secular equation (light red), the eigensolver (dark red), the MPI synchronization (purple), and the others (grey). Note that the vertical scale of the left panel differs from that of the other two.

The computation of the matrix elements is the largest component of the wall time, and increases rapidly with the number of particles. Thus, the optimization of this part plays a key role in increasing the number of particles and the GPU acceleration is most effective as discussed in Subsect. III-B. In this part, the MI300A is 10% faster than the GH200, although the theoretical peak of its vector arithmetic logic unit (VALU) is 80% higher. Measuring the FLOP count with the hardware counters of the GH200, the FP64 throughput of the matrix-element computation reaches 16.4 TFLOP/s, 48% of the theoretical peak of the VALU of the GH200. The corresponding throughput on the MI300A reaches 18.0 TFLOP/s, namely 29% of its VALU peak. Note that the tensor cores are not used on either system for the present code since the size of the matrix A is too small. The “Eigensolver” part is the wall time of solving the eigenvalue problem with the dsygvd subroutine, provided by cuSOLVER 13.0 on the GH200 and by rocSOLVER 7.2.0 on the MI300A every 100 basis states to avoid the accumulation of the numerical errors. The comparison indicates a higher level of maturity of the cuSOLVER implementation. C. 6-body system We briefly discuss the possibility of extending the calculation to the six-body system. The five-body system with Nt = 14400 and Nb = 9600 was computed in 2574 seconds with 4 nodes, namely 2.9 node×hours on MiyabiG (GH200). With the same Nt and Nb , the computational cost scales as O(N 4 N !), and thus the six-body calculation is expected to be 12.4 times more expensive than the five-

TABLE II Computing platforms for the benchmark test.

Miyabi-C (CPU)

Miyabi-G (GPU)

Sirius (GPU)

Processor Intel Xeon Max 9480 NVIDIA GH200 (Grace 72c + H100) AMD Instinct MI300A Devices used 4 CPUs (224 cores) 4 GH200 4 APUs FP64 peak / device 3.4 TFLOP/s 34 TFLOP/s (vector) 61.3 TFLOP/s (vector) Compilers ifort 2021.10.0 + MKL nvcc CUDA 13.0, nvfortran ROCm 7.2.0 + gfortran 11.5.0 Numerical libraries MKL 2023.2.0 cuBLAS / cuSOLVER hipBLAS / rocSOLVER Parallelization 4 MPI × 56 OpenMP 4 MPI × 72 OpenMP 4 MPI × 23 OpenMPa Mem BW / device 1.6 TB/s 4.0 TB/s 5.3 TB/s a Although the APU contains 24 CPU cores, 23 OpenMP threads are invoked to avoid OS noise. TABLE III Wall-time breakdown for the five-body SVM run (K = 4800 basis states, Nt = 1200 trials per step; four devices on every platform; seconds). “CPU baseline” is the pre-optimization code with the original flat-MPI decomposition (112 MPI × 2 OpenMP); the tuned columns use the configurations of Table II. The CPU-to-GPU step combines the port and the hardware: the software-only gain on identical CPU hardware is the 10.4× of the second column. “Eigensolver” is the full generalized eigensolve refreshed every 100 basis steps (48 calls): QR-based dsygv in the CPU baseline, divide-and-conquer dsygvd (MKL / cuSOLVER / rocSOLVER) in the tuned columns; the per-step rank-one update is the separate “incremental update” row.

Component

CPU baseline

CPU tuned

GH200 tuned

MI300A tuned

Matrix elements Incremental update Secular equation Eigensolver (full solve) MPI synchronization Other

7435.9 3217.0 52.4 417.3 780.0 13.5

939.1 139.8 29.0 22.6 12.5 4.5

37.5 14.6 10.1 3.2 1.2 4.3

34.1 15.3 3.8 13.0 5.2 6.5

Total Speedup vs. CPU baseline

11916.1 1

1147.6 10.4×

70.8 168×

77.9 153×

body case. As a preliminary benchmark, we measured the wall time for the six-body case with the same conditions, and found that it is only 10.2 times larger. Although it is difficult to predict the values of Nt and Nb required for the six-body calculation, even if both Nt and Nb are doubled, the computational cost is estimated to be 237 node×hours, which is feasible with current computational resources. V. Summary We presented the algorithm and code developments of the SVM for the five-body 4 He atomic system. The present code was highly tuned with the aid of AI coding. Employing the secular equation for low-lying eigenvalues significantly accelerates the SVM computation and makes a large number of basis states feasible, which is directly relevant to the capability of the SVM. We achieved a 168fold acceleration on the GH200 over the baseline CPU code, and a 16-fold acceleration over the tuned CPU code. The loop exchange, which enables the efficient use of SIMD on CPUs and of coalesced addressing on GPUs, has a significant impact. This code makes 6-body calculations and beyond computationally feasible. Since the whole code was developed by an AI agent, the present study shows the ability of AI to write a performant scientific code for current GPU architectures, while a working knowledge of the hardware is required. In particular, the secular-

equation solver for massively parallel computation and the Gu–Eisenstat prescription that preserves the orthogonality of the eigenvectors are known to be efficient, but implementing them is a burdensome task for physicists. The recent growth of AI coding enables us to implement them in a short time and to exploit the latest HPC environments efficiently. As a future work, further improvements of the SVM algorithm are anticipated, such as the gradient-based optimization [13] and Bayesian optimization [14]. Acknowledgment N. S., S. A., and Y. T. acknowledge the KAKENHI budgets (25H01268) and the JST ERATO project (JPMJER2304). N. S. thanks also the KAKENHI budgets (25K00995, 25K07294, 25K07330, 24H00239). The numerical calculations were performed at the Miyabi [15] and Sirius [16] supercomputers under the MCRP program of Center for Computational Sciences, University of Tsukuba. The coding was performed mainly by Anthropic Claude Opus and Fable. The draft, figure and tables were prepared with the aid of Claude Opus. References [1] K. Varga and Y. Suzuki, ”Precise solution of few-body problems with the stochastic variational method on a correlated Gaussian basis”, Phys. Rev. C 52, (1995) 2885.

[2] Y. Suzuki and K. Varga, ”Stochastic variational approach to quantum-mechanical few-body problems”, Springer (2002). [3] E. Hiyama, Y. Kino, and M. Kamimura, ”Gaussian expansion method for few-body systems”, Prog. Part. Nucl. Phys. 51, (2003) 223. [4] M. Gattobigio, A. Kievsky, and M. Viviani, ”Spectra of helium clusters with up to six atoms using soft-core potentials”, Phys. Rev. A 84, (2011) 052503. [5] S. Aoyama, N. Shimizu and Y. Tsunoda, ”Parallel computingbased extension of the stochastic variational method”, PDAA workshop, CANDAR-2026. [6] T. Pang, ”The breakup of a helium cluster after removing attractive interaction among a significant number of atoms in the cluster”, Scientific Reports, 10, (2020) 5767. [7] R. A. Aziz and M. J. Slaman, ”An examination of ab initio results for the helium potential energy curve”, J. Chem. Phys. 94, (1991) pp. 8047-8053. [8] D. Yoshida, E. Hiyama, T. Imamura, I. Kanamori, and H. Matsufuru, ”Solving large-scale eigen problem in quantum fewbody system on massive parallel computer”, Proc. HPC ASIA 2026 Workshops, (2026) 133. [9] S. Aoyama, Y. Tsunoda, N. Shimizu, T. Myo, H. Matsufuru, and I. Kanamori, ”Development and validation of the computational code GEM for simulating quantum few-body systems through optimization on the Fugaku computer”, Proc. ICCSA 2026 Workshops, (2026) 484. [10] H. Tokura, T. Honda, Y. Ito, K. Nakano, M. Nishino, Y. Hirota, and M. Saeki, ”An efficient GPU implementation of bulk computation of the eigenvalue problem for many small real non-symmetric matrices”, Int. J. Netw. Comput. 7, (2017) 227. [11] M. Gu and S. C. Eisenstat, ”A stable and efficient algorithm for the rank-one modification of the symmetric eigenproblem”, SIAM J. Matrix Anal. Appl. 15, (1994) 1266. [12] M. Imada and T. Kashima, ”Path-integral renormalization group method for numerical study of strongly correlated electron systems”, J. Phys. Soc. Jpn. 69, (2000) 2723. [13] P. Recchia, D. Basu, M. Gattobigio, C. Miniatura, and S. Bressan, ”The steepest slope toward a quantum few-body solution: gradient variational methods for the quantum fewbody problem”, Few-Body Syst. 65, (2024) 102. [14] S. Aoyama, ”Bayesian variational method for precision few-body calculations”, arXiv:2607.25265 [nucl-th] [15] Miyabi supercomputer, JCAHPC, https://www.jcahpc.jp/ supercomputer/miyabi.html [16] Sirius supercomputer, Center for Computational Sciences, University of Tsukuba, https://www.ccs.tsukuba.ac.jp/ supercomputer/

Related documents

• AI-Assisted GPU optimization of the Stochastic Variational Method for Few-Body Boson System#1187059
• Decoupling the Magnetic Field Operator as an Independent Operator Class in Stochastic Series Expansion Quantum Monte Carlo#1050167
• Large-scale quantum simulations of dissipative spin-1/2 Heisenberg chains#1110495
• Unconventional quantum spin mixing of two spin-1 atoms beyond the single-mode approximation#1355032
• Adaptive Algorithmic Simulation for Nonlinear Eigenvalue Problems in Mathematical Physics#250530