Skip to content

perf(lcao): batch RI radial table construction on CUDA - #8257

Merged
mohanchen merged 5 commits into
deepmodeling:developfrom
Stardust0831:perf/ri-radial-gpu
Nov 3, 2026
Merged

mohanchen merged 5 commits into
deepmodeling:developfrom
Stardust0831:perf/ri-radial-gpu

Conversation

@Stardust0831

Copy link
Copy Markdown
Collaborator

What's changed?

LCAO EXX initialization builds radial tables with many small quadrature operations and repeated product-orbital transforms. For CUDA calculations, batch the existing Simpson quadrature into FP64 GEMM using ABACUS's tensor/BLAS interfaces. The quadrature grids, weights, derivative recurrence and analytic origin values are preserved.

The existing resolved device selection is passed explicitly through initialization and Coulomb refresh. CPU full-table construction retains individual quadrature. Product orbitals are shared during batched and selective construction, and their unused origin derivative is omitted. Required integral gradients remain calculated.

GPU batching uses device gpu in a CUDA build. It works with the default LibRI revision b0eff7a and requires no DDLA, LibRI GPU backend, or new build/INPUT option. Other RI consumers retain CPU execution.

Performance

System Stage Baseline (s) This PR (s) Paired speedup
Si32 RI initialization 6.52–7.15 1.34–1.41 4.86–5.09x
Si32 Complete SCF + forces + stress 182.22–203.71 181.38–187.34 1.00–1.09x
CsPbI3 RI initialization 343.48–344.15 14.73–15.38 22.33–23.37x
CsPbI3 Complete SCF + forces + stress 424.48–426.14 95.52–96.23 4.43–4.44x

LCAO HSE calculations on displaced Si32 and CsPbI3 (5 atoms), using one V100, one MPI rank, four physical CPU cores/OpenMP threads and serial BLAS on SAI. Both versions use cuSOLVER diagonalization and CPU LibRI contractions. The baseline is unmodified d20090cc137f.

Ranges cover two paired runs in separate allocations on shared nodes, with reversed execution order. Ratios are calculated within each pair; SCF iteration counts match. Complete time includes process startup. The Exx_LRI::init speedup includes GPU batching and product preparation reuse.

Explicit GPU workspace and corresponding host staging are each bounded at 160 MiB per MPI rank and released before SCF. CUDA context/library allocations and existing host tables are outside this bound. Sampled process GPU peaks were 442–450 / 460–462 MiB for Si32 and 438 / 530 MiB for CsPbI3 (baseline / this PR; 0.5 s sampling).

Unit Tests and/or Case Tests for my changes

The full CUDA executable build and CPU/CUDA radial-table tests passed with one and four OpenMP threads, including execution of the GPU test. Tests cover individual versus batched tables, product sharing, high angular momenta, origin/cutoff behavior, radial tiling and gradients.

cmake --build build --target abacus_std_gpu MODULE_LCAO_radial_tables MODULE_LCAO_radial_tables_gpu --parallel 24
OMP_NUM_THREADS=1 ctest --test-dir build -V -R '^(MODULE_LCAO_radial_tables|MODULE_LCAO_radial_tables_gpu)$'
OMP_NUM_THREADS=4 ctest --test-dir build -V -R '^(MODULE_LCAO_radial_tables|MODULE_LCAO_radial_tables_gpu)$'

The benchmark cases, a CPU execution control, and complex-tensor Si32 on a regular 2x1x1 k mesh with two MPI ranks/GPUs all converged with matching iteration counts. Maximum differences were 2.2e-11 eV/atom, 2.3e-9 eV/Angstrom and 9.6e-9 kbar. These comparisons establish agreement with the baseline. The combined changes also passed a build and regression checks with CUDA disabled.

Inputs, timings, numerical comparisons and test logs. Benchmark executables enable identical existing timer JSON output; this instrumentation is excluded from the change.

Linked Issue

Supersedes the radial-table preparation part of #8193 and can be merged independently of the integral-evaluation and Libxc optimizations. LibRI/DDLA GPU contractions remain separate in #8169.

Governance Notes

Core impact is limited to RI table construction and explicit device propagation. No documentation update required: the existing device option and its metadata are unchanged. <memory> is required for the owning std::shared_ptr<const ProductOrbitals> member.

Comment thread source/source_lcao/test/CMakeLists.txt Resolved

@ABACUS-PR-BOT ABACUS-PR-BOT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No confirmed correctness defect found in the current in-tree paths. The inline comments ask to clarify selective-table reuse and name the workspace-limit calculations.

Static review of dfc71eb1 against d20090cc: whitespace passed and governance checks completed with no blockers. No builds or runtime tests were performed; the reported numerical and performance results remain author-reported.

Comment thread source/source_lcao/center2orb_orb21.cpp Resolved
Comment thread source/source_lcao/center2orb_orb11.cpp Resolved
Keep existing product orbitals and radial entries when selective
initialization is called again with additional radii. Prepare only new
objects in selective batches so existing Orb11 references remain valid.

Add a regression for individual and mixed-batch initialization, checking
overlaps and gradients against complete tables.
Retain the CPU and CUDA radial-table tests with the upstream CMake format.

Validation: clean CUDA build and both radial-table test targets on SAI,
with OMP_NUM_THREADS=1 and 4 (14 test executions passed, no skips).

@ABACUS-PR-BOT ABACUS-PR-BOT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The selective-initialization update preserves existing entries and covers repeated individual and mixed-batch calls. One additional ownership-lifetime regression would strengthen coverage. Static review only; no builds or runtime tests were run.

Exercise destruction and full reinitialization of the first Orb21 in a
shared batch. Check the remaining object after a selective update against
independent full tables, including retained and new overlaps and gradients.

Validation: CPU/CUDA radial-table tests with one and four OpenMP threads
(18 test executions passed), plus the ownership regression under
AddressSanitizer with CUDA disabled.

@AsTonyshment AsTonyshment left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@mohanchen mohanchen added Refactor Refactor ABACUS codes GPU & DCU & HPC GPU and DCU and HPC related any issues EXX/lr-TDDFT/RPA/BSE/GW Related to beyond DFT methods Performance Issues related to fail running ABACUS labels Nov 3, 2026
@mohanchen
mohanchen merged commit fb2e70c into deepmodeling:develop Nov 3, 2026
22 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

EXX/lr-TDDFT/RPA/BSE/GW Related to beyond DFT methods GPU & DCU & HPC GPU and DCU and HPC related any issues Performance Issues related to fail running ABACUS Refactor Refactor ABACUS codes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants