PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Noritaka Shimizu, Shigeyoshi Aoyama, Yusuke Tsunoda, Akira Nukada

arXiv:2610.02745v1physics.comp-phcs.DCnucl-th

TL;DR

Matrix-element evaluation is the bottleneck in the whole SVM computation, with computational cost per matrix element scaling as O(N 4). The code uses a secular-equation method with the Gu–Eisenstat prescription for eigenvectors and a memory layout that makes adjacent GPU threads access adjacent doubles. The five-body SVM code achieved 168-fold acceleration on the GH200 over the baseline CPU code and 16-fold acceleration over the tuned CPU code.

  • Problem

    Matrix-element evaluation is the bottleneck in the whole SVM computation, with computational cost per matrix element scaling as O(N 4).

    Source: C. Matrix elements in SVM · p. 3

  • Method

    The code uses a secular-equation method with the Gu–Eisenstat prescription for eigenvectors and a memory layout that makes adjacent GPU threads access adjacent doubles.

    Source: A. Solving the generalized eigenvalue… · p. 4, B. Memory layout of the matrix-element… · p. 5

  • Results

    The five-body SVM code achieved 168-fold acceleration on the GH200 over the baseline CPU code and 16-fold acceleration over the tuned CPU code.

    Source: §V · p. 7

  • Takeaways & Limitations

    The authors conclude that the code makes six-body calculations and beyond computationally feasible.

    Source: §V · p. 7

Abstract

from arXiv · show

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.

II. Method

The method evaluates eigenstates of a few-body 4He system using the SVM framework. Its Hamiltonian includes a two-body potential modeled by a two-range Gaussian whose parameters reproduce the LM2M2 potential’s scattering length and dimer energy.

II. Method

  • The SVM framework is used to obtain eigenstates of the few-body 4He boson system by evaluating the energy, the Hamiltonian’s expectation value.The method section identifies the Hamiltonian, SVM framework, matrix-element formulae, and AI-assisted coding process as its components.Source: §II · pp. 1–2

A. Hamiltonian of 4He atomic cluster

  • The 4He atomic-cluster Hamiltonian includes a two-body potential approximated by a two-range Gaussian function of relative distance.Its parameters were determined to reproduce the LM2M2 potential’s scattering length and dimer energy.Source: A. Hamiltonian of 4He atomic cluster · p. 2
  • The two-range potential combines a strong Pauli-principle repulsion with a shallow attractive van der Waals force.The strong repulsion makes precise wave-function computation difficult for many other methods.Source: A. Hamiltonian of 4He atomic cluster · p. 2

B. Stochastic Variational Method

The SVM represents the wave function with symmetrized correlated Gaussian basis states and determines energies and coefficients through a generalized eigenvalue problem. In the parallel variant, candidate states are evaluated together and the one yielding the lowest energy is selected.

B. Stochastic Variational Method

  • The SVM wave function is a linear combination of symmetrized correlated Gaussian basis states, with Bose symmetry and Jπ = 0+ specified for this work.The basis uses Jacobi coordinates with the center-of-mass motion removed; the figure shows N = 5 and N_J = 4.Source: B. Stochastic Variational Method · p. 2, Figure 1 · p. 2
  • The generalized eigenvalue problem determines the basis coefficients and eigenenergies, and the resulting wave function uses matrix elements evaluated with Gaussian integrals.The problem is also called the Hill–Wheeler equation.Source: B. Stochastic Variational Method · p. 2
  • At each basis-growth step, the parallel SVM evaluates N_t candidates and accepts the one that lowers the energy most.The algorithm prepares random matrices for candidates and solves the generalized eigenvalue problem for each candidate’s energy.Source: B. Stochastic Variational Method · p. 2

C. Matrix elements in SVM

SVM matrix elements require Bose symmetrization over particle permutations, while the potential-term evaluation is a computational bottleneck. The method uses Gaussian integrals and matrix decompositions to compute overlaps and kinetic and potential contributions.

C. Matrix elements in SVM

  • For identical bosons, each matrix element sums over N! particle permutations, although symmetrizing the wave function on one side is sufficient because the Hamiltonian commutes with the symmetrization operator.A particle permutation acts on Jacobi coordinates through an orthogonal matrix and transforms a Gaussian by A → P^T A P.Source: C. Matrix elements in SVM · p. 3
  • The overlap is evaluated using a Gaussian integral, with the determinant obtained by Cholesky decomposition; the kinetic matrix-element cost is O(N_J^3).The kinetic-term cost follows from the Cholesky decomposition used to construct B^−1.Source: C. Matrix elements in SVM · p. 3
  • The potential-term computation costs O(N^4) per matrix element and is the bottleneck of the SVM calculation.The pair index spans N(N−1)/2 particle pairs, and s^(ij) depends on k, k′, p, i, and j.Source: C. Matrix elements in SVM · p. 3
  • A typical calculation evaluates hundreds to thousands of trial basis states at each step, and basis sizes can reach N_b = 100–10000 or more.The matrix-element and eigenvalue-solve costs can become prohibitive on a single CPU; the text identifies GPU-oriented optimization and the secular-equation method as ways to reduce these costs.Source: C. Matrix elements in SVM · p. 3

D. AI-assisted coding

The SVM code was developed from scratch with AI coding agents, mainly Claude Opus, and optimized through profiling and iterative interactions. The account also notes that CUDA and GPU-hardware expertise informed the optimization process.

D. AI-assisted coding

  • AI coding agents developed the SVM code from scratch, with the authors giving natural-language instructions and making only a few trivial code modifications.The code was developed gradually over a few months, mainly using Claude Opus.Source: D. AI-assisted coding · p. 3
  • GPU optimization used AI-agent profiling, but discussions with HPC experts highlighted the importance of CUDA and GPU-hardware domain knowledge.The initial baseline code was written by an author with little CUDA-development experience.Source: D. AI-assisted coding · p. 3
  • Development involved more than 1000 brief AI interactions and produced 13,682 lines of Fortran, CUDA, and HIP code.The process included many trials and errors.Source: D. AI-assisted coding · p. 3

III. Performance improvements

Two improvements dominate performance gains: a secular-equation treatment reduces per-candidate eigenvalue-solver work, while transposed memory layout improves matrix-element kernel efficiency.

III. Performance improvements

  • The secular-equation treatment removes the O(K^3) cost from the per-candidate loop.The memory-layout transposition improves memory-system efficiency in the matrix-element kernel.Source: §III · pp. 3–4
  • Candidate evaluations are independent and distributed across MPI processes, with CPU work overlapped with GPU execution across permutation chunks.Matrix elements run on GPUs when available; candidate eigenvalue problems are solved on CPUs with OpenMP threads.Source: §III · p. 4

A. Solving the generalized eigenvalue problem with the secular equation

The secular-equation method evaluates candidate energies without solving each enlarged generalized eigenvalue problem from scratch. Gu–Eisenstat eigenvector reconstruction and periodic full solves support accurate incremental updates.

A. Solving the generalized eigenvalue problem with the secular equation

  • The method transforms the appended candidate's matrix row and column into the accepted basis eigenbasis, then rewrites the generalized problem as a symmetric arrowhead eigenproblem.A candidate linearly dependent on the accepted basis is rejected; interlacing brackets the roots for stable bisection.Source: A. Solving the generalized eigenvalue… · p. 4
  • The candidate loop costs O(K^2) per candidate rather than O(K^3), replacing repeated full solves with secular-equation energy evaluation.The generalized eigenvalue solve would otherwise cost O(NtK^3) per step; the revised cost is Nt×O(K^2)+O(K^3).Source: A. Solving the generalized eigenvalue… · p. 4
  • Only the accepted candidate needs eigenvectors; Gu–Eisenstat reconstruction from computed roots restores orthogonality to working accuracy.This once-per-step update is dominated by a matrix–matrix back transformation and is parallelized with CPU OpenMP, with GPU acceleration enabled for K ≥2000.Source: A. Solving the generalized eigenvalue… · p. 4
  • Eigenpairs are rebuilt from scratch every 100 basis states to control accumulated rounding errors; a run with Nb = 4800 therefore performs 48 full solves when GPU libraries are available.On GH200, the per-candidate work costs 10.1 s, the once-per-step update 14.6 s, and the 48 library solves 3.2 s; MPI synchronization is 1.2 s of 70.8 s.Source: A. Solving the generalized eigenvalue… · pp. 4–5

B. Memory layout of the matrix-element kernel

Matrix-element work comprises many independent permutation tasks, making it amenable to GPU parallelism. A structure-of-arrays layout coalesces accesses, while corresponding CPU vectorization uses the same contiguous parallel dimension.

B. Memory layout of the matrix-element kernel

  • For five bodies with Nt = 1,200, 144,000 independent permutation tasks are available even at Nb = 1, mapping naturally to GPU threads.Each thread block handles a candidate, with threads striding over permutations.Source: B. Memory layout of the matrix-element… · p. 5
  • The natural array-of-structures layout made each 32-byte memory sector deliver only 8 useful bytes, with L1 throughput saturated at 67% and DRAM utilization at 0.02%.The kernel was limited by memory-transaction count, not bandwidth.Source: B. Memory layout of the matrix-element… · p. 5
  • Transposing to a structure-of-arrays layout makes permutation indices contiguous, coalesces warp loads, and makes the matrix-element stage 1.3–1.6 times faster without changing energies.The energies are bit-identical to those from the previous layout.Source: B. Memory layout of the matrix-element… · p. 5
  • CPU vectorization also tiles independent factorizations across lanes, using lane masking for non-positive pivots and contiguous memory for the parallel dimension.Vectorizing within one small matrix is inefficient because its contraction length is NJ = 4 and Cholesky factorization is sequential.Source: B. Memory layout of the matrix-element… · p. 5

IV. Results

The results section presents five-body energy convergence and the associated computational trade-off between trial count and wall time.

IV. Results

  • The five-body ground-state energy is nearly converged by Nb ∼4800 for any Nt, while increasing Nt improves convergence with diminishing gains.The gain from Nt = 9600 to Nt = 14400 is only 0.09 mK; Fig. 2 plots energy against Nb and, in panel (b), the difference from the estimated exact energy on a logarithmic scale.Source: A. Convergence of the SVM method · p. 5, Figure 2 · p. 5
  • The estimated exact energy is E0 = −1.30851(1) K, obtained by extrapolating energy to zero energy variance.The details of the extrapolation method are reported elsewhere.Source: A. Convergence of the SVM method · p. 5
  • Wall time is approximately linear in Nt, so Nt and Nb must be chosen as a trade-off under a given time budget.The constant cost is independent of Nt; the slope reflects matrix-element evaluation, which scales linearly with trial count.Source: A. Convergence of the SVM method · pp. 5–6

B. Code Performance

The tuned five-body SVM benchmark compares CPU and GPU wall times and breaks them down by computation stage. Results highlight the matrix-element workload, platform differences, and the eigensolver’s refresh schedule.

B. Code Performance

  • The GH200 port reaches 168× speedup, while CPU tuning achieves up to 10× over the baseline.Fig. 3 breaks down six wall-time components across platforms; its CPU panel uses a different vertical scale.Source: B. Code Performance · p. 6, Figure 3 · p. 6
  • On Sirius, 23 OpenMP threads per APU give better performance and more stable wall times than 24, despite each APU having 24 CPU cores.The 24-thread configuration showed large synchronization delays among MPI processes.Source: B. Code Performance · p. 6
  • Matrix elements dominate wall time and grow rapidly with particle count; MI300A is 10% faster than GH200 despite an 80% higher VALU peak.Matrix-element FP64 throughput is 16.4 TFLOP/s (48% of GH200 VALU peak) and 18.0 TFLOP/s (29% of MI300A VALU peak).Source: B. Code Performance · p. 6
  • The full eigensolver runs every 100 basis states to avoid accumulated numerical errors, and the comparison indicates greater maturity for cuSOLVER than rocSOLVER.Source: B. Code Performance · p. 6

C. 6-body system

The six-body discussion compares a preliminary measured cost with a scaling-based estimate and considers the projected resource requirement for larger calculations.

C. 6-body system

  • The Sirius benchmark uses 23 OpenMP threads per APU, one fewer than its 24 CPU cores, to avoid OS noise.Source: C. 6-body system · p. 7
  • The preliminary six-body benchmark took 10.2× the five-body time, below the 12.4× cost projected from scaling at the same Nt and Nb.Source: C. 6-body system · pp. 6–7
  • If both Nt and Nb are doubled, the six-body calculation is estimated to require 237 node×hours, which the authors consider feasible with current resources.Source: C. 6-body system · p. 7

V. Summary

The paper reports an AI-assisted, highly tuned SVM code for five-body 4He systems and large speedups over CPU implementations. Its conclusions emphasize the secular-equation approach and the hardware knowledge needed to exploit current GPUs.

V. Summary

  • The five-body 4He SVM code achieves 168-fold acceleration on GH200 over baseline CPU code and 16-fold over tuned CPU code.The authors state that the code makes six-body calculations and beyond computationally feasible.Source: §V · p. 7
  • The secular-equation solver and Gu–Eisenstat prescription support efficient parallel computation while preserving eigenvector orthogonality.The authors note that implementing these methods is burdensome for physicists, while AI coding enabled implementation in a short time.Source: §V · p. 7
  • Gradient-based and Bayesian optimization are identified as possible future improvements to the SVM algorithm.Source: §V · p. 7
Loading 2610.02745v1…