Repository navigation
perf(lcao): batch RI radial table construction on CUDA - #8257
Stardust0831 wants to merge 4 commits into
Conversation
| if(USE_CUDA) | ||
| addtest( | ||
| TARGET | ||
| MODULE_LCAO_radial_tables_gpu | ||
| KEEP_FEATURE_DEFINITIONS | ||
| __CUDA | ||
| LIBS | ||
| parameter | ||
| base | ||
| device | ||
| container | ||
| SOURCES | ||
| test_center2orb_orb11.cpp | ||
| ../center2orb.cpp | ||
| ../center2orb_orb11.cpp | ||
| ../center2orb_orb21.cpp | ||
| ${ABACUS_SOURCE_DIR}/source_basis/module_ao/orb_atomic_lm.cpp | ||
| ${ABACUS_SOURCE_DIR}/source_basis/module_ao/orb_gaunt_table.cpp) | ||
| set_tests_properties(MODULE_LCAO_radial_tables_gpu PROPERTIES LABELS gpu) |
There was a problem hiding this comment.
I haven't looked into the implementation in detail (as I'm not particularly familiar with the RI code), but it seems that the newly added GPU unit test hasn't been included in the CI workflow in .github/workflows/cuda.yml, has it?
There was a problem hiding this comment.
Yes, that was missing. Added MODULE_LCAO_radial_tables_gpu to the CUDA workflow's CTest selection in 026ea36. That commit's CI log confirms that BatchedGpuMatchesIndividualTables executed and passed.
ABACUS-PR-BOT
left a comment
There was a problem hiding this comment.
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.
|
|
||
| void Center2_Orb::Orb21::bind_product_orbitals() | ||
| { | ||
| this->orb11s.clear(); |
There was a problem hiding this comment.
Non-blocking: Please make the selective-initialization contract explicit and add a test that calls it twice with different radius sets. orb11s.clear() discards previously populated rows, while the prior implementation retained them; the current in-tree subset callers initialize once.
There was a problem hiding this comment.
Thanks for catching this behavior change. Restored and documented accumulation across selective calls in a853060, retaining existing entries and the product orbitals they reference. The new regression covers repeated individual calls and mixed batches of initialized and new objects, checking overlaps and gradients against full tables. It fails before the fix. After the fix, both CPU and CUDA radial-table test targets pass with one and four OpenMP threads on SAI.
| { | ||
| throw std::invalid_argument("Unsupported radial quadrature grid for batched tables"); | ||
| } | ||
| const std::size_t radial_limit = std::min(table_limit / (2 * kmesh), output_limit / 2); |
There was a problem hiding this comment.
Non-blocking: Please assign the per-buffer capacity calculations and inner batch cap to named locals before these std::min calls. This follows AGENTS.md rule 14 and makes each part of the workspace bound directly inspectable.
There was a problem hiding this comment.
Thanks for the suggestion. The three device-buffer budgets already have named constants. I would prefer to retain the remaining capacity expressions in this revision.
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).
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
deviceselection 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 gpuin a CUDA build. It works with the default LibRI revisionb0eff7aand requires no DDLA, LibRI GPU backend, or new build/INPUT option. Other RI consumers retain CPU execution.Performance
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::initspeedup 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.
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
deviceoption and its metadata are unchanged.<memory>is required for the owningstd::shared_ptr<const ProductOrbitals>member.