Skip to content

feat(lcao): add optional LibRI GPU contractions through DDLA - #8169

Open
Stardust0831 wants to merge 15 commits into
deepmodeling:developfrom
Stardust0831:experiment/lcao-hse-gpu
Open

Stardust0831 wants to merge 15 commits into
deepmodeling:developfrom
Stardust0831:experiment/lcao-hse-gpu

Conversation

@Stardust0831

@Stardust0831 Stardust0831 commented Oct 18, 2026 •

Copy link
Copy Markdown
Collaborator

What's changed?

Enable optional GPU acceleration of LCAO exact-exchange RI tensor contractions through LibRI's DDLA CUDA backend, with a build option and a runtime INPUT keyword:

  • ENABLE_LIBRI_GPU: new CMake option, default OFF. Enable it together with ENABLE_LIBRI=ON and USE_CUDA=ON to build the DDLA GPU backend.
  • device_exx: new INPUT keyword, default cpu. Set device_exx gpu together with device gpu to run RI contractions on the GPU. device gpu alone continues to use CPU RI, so the solver and RI backend can be selected separately.

Build with compatible LibRI and DDLA installations:

cmake -B build -DENABLE_LIBRI=ON -DUSE_CUDA=ON -DENABLE_LIBRI_GPU=ON \
  -DLIBRI_DIR=/path/to/patched/LibRI -DDDLA_ROOT=/path/to/DDLA

Then select the GPU RI backend in INPUT:

device      gpu
device_exx  gpu

The backend is experimental. Installation, dependency versions and runtime requirements are documented in docs/advanced/install.md.

Dependencies and supported scope
  • Requires LibRI 6ce3c84 with tests/lcao_hse_gpu/libri-compat.patch and CUDA DDLA v0.0.4 (9c404ac). The patch fixes API compatibility and accumulation needed for forces and stress. CMake locates DDLA through DDLA_ROOT or CMAKE_PREFIX_PATH and checks GPU API compilation/linking.
  • Applies to LCAO Kohn-Sham exchange. Independent LR/RDMFT exchange engines retain CPU mode; LR inherits the selected backend when reusing the Kohn-Sham exchange object.
  • GPU RI does not support real-space EXX symmetry reduction. With symmetry 1, use device_exx cpu or set exx_symmetry_realspace 0; unsupported combinations are rejected.
  • All MPI ranks on a node must see the same allocated GPU list, with at least as many visible GPUs as local ranks. Per-rank single-GPU masking is unsupported by the pinned LibRI backend.

Unit Tests and/or Case Tests for my changes

Functional checks used SAI V100 GPUs with OMP_NUM_THREADS=1; the performance runs below tune MPI/OpenMP:

  • At 3fa1f972b, after merging current develop, a fresh GPU build, INPUT unit tests with GPU RI enabled and disabled (10 tests in each configuration), CPU/GPU dependency consumers and the optional numerical GPU test passed. CLI help, generated metadata and --check-input use device_exx; the former exx_device keyword is rejected.
  • On the same revision, HSE Gamma and two-rank complex multi-k passed all eight solver/RI combinations against unchanged references. Gamma checks energy, forces and stress; multi-k checks energy. Both valid backend selections and three invalid-input diagnostics passed.
  • Earlier validation also covered full CPU builds, static/shared MKL dependency consumers, spin-polarized and symmetry cases. The numerical GPU test covers real and complex signed accumulation across all 18 contraction branches with nonempty outputs; reverting the accumulation fix previously produced the expected failure.

Si64 force benchmark

A displaced Si64 HSE/DZP Gamma-point case (832 orbitals) was tested at 63ca5ef on SAI V100-SXM2 16 GB GPUs and an AMD Ryzen Threadripper PRO 5995WX CPU. These performance measurements predate the keyword rename; archived inputs use exx_device, while examples below use device_exx. The performance benchmark was not rerun for this rename. Each GPU is allocated 4 physical CPU cores / 8 SMT hardware threads. CPU-only and single-GPU runs use the same four physical cores; two-GPU runs use eight. Runs execute sequentially on the same node, with explicit CPU affinity.

CPU-only uses ScaLAPACK and CPU RI; GPU runs use cuSOLVER and GPU RI. Seven CPU MPI/OpenMP ratios and single-GPU 1×4/1×8 were tested (details below). The fastest CPU configurations are 1 MPI × 8 OpenMP for forces and 4 MPI × 2 OpenMP for the complete calculation; each row uses its corresponding CPU baseline. The selected GPU configuration uses 4 OpenMP threads per rank, with one rank per GPU. Selected configurations are run twice; the table shows their mean times.

Stage Tuned CPU only (s) 1 GPU + 4 CPU cores (s) 2 GPUs + 8 CPU cores (s) CPU / 1 GPU
EXX force 72.75 21.45 19.42 3.39x
All force work 75.44 22.91 20.37 3.29x
SCF + forces 352.58 238.50 192.07 1.48x
  • Scaling: one to two GPUs, with proportional CPU resources, gives 1.12x for all force work (56.2% efficiency) and 1.24x for the complete calculation (62.1%).
  • RI-only comparison: holding cuSOLVER, one MPI rank, OpenMP threads and CPU affinity fixed, switching device_exx cpu to gpu gives 3.87x for all force work and 1.49x overall. This uses one CPU-RI run and the mean of two GPU-RI runs; the CPU-only table also includes the solver change.
  • Correctness: All 15 runs converge in 7 outer and 96 inner iterations. Across all 192 force components, the maximum difference from the pure-CPU reference is 1.5e-09 eV/Angstrom; maximum energy difference is 1.1e-10 eV/atom. The largest reference force component is 4.89 eV/Angstrom.

EXX and all-force timings are nested ABACUS timers. Complete-calculation times include launch, initialization and SCF. These measurements cover one case on one node with the supplied ScaLAPACK/OpenBLAS build.

Verification commands and validation limits

Fresh validation of 3fa1f972b used the commands below (SAI job 1868744). GPU_BUILD_DIR has CUDA/GPU RI enabled; CPU_BUILD_DIR has both disabled:

cmake --build "$GPU_BUILD_DIR" --target abacus_std_gpu --parallel 12
cmake --build "$GPU_BUILD_DIR" --target MODULE_IO_read_item_serial --parallel 12
ctest --test-dir "$GPU_BUILD_DIR" -V -R '^MODULE_IO_read_item_serial$'
cmake --build "$CPU_BUILD_DIR" --target MODULE_IO_read_item_serial --parallel 12
ctest --test-dir "$CPU_BUILD_DIR" -V -R '^MODULE_IO_read_item_serial$'
"$ABACUS_BIN" --version
"$ABACUS_BIN" --info
"$ABACUS_BIN" -h device_exx
"$ABACUS_BIN" --generate-parameters-yaml
"$ABACUS_BIN" --check-input

From copies of tests/08_RI, using the existing references:

bash ../integrate/Autotest.sh -a "$ABACUS_BIN" -n 1 -o 1 -f CASES.task.txt
bash ../integrate/Autotest.sh -a "$ABACUS_BIN" -n 2 -o 1 -f CASES.task.txt

Variants cover CPU and GPU solvers with the default and explicit RI backend selections. Symmetry cases cover CPU RI with real-space symmetry enabled and CPU/GPU RI with it disabled.

Standalone dependency checks:

cmake -S tests/lcao_hse_gpu/cmake -B build-consumer \
  -DLIBRI_DIR="$LIBRI_PREFIX" -DLIBCOMM_DIR="$LIBCOMM_PREFIX" \
  -Dcereal_DIR="$CEREAL_CONFIG" \
  -DENABLE_LIBRI_GPU=ON -DENABLE_LIBRI_GPU_TESTS=ON -DDDLA_ROOT="$DDLA_PREFIX"
cmake --build build-consumer --parallel 4
ctest --test-dir build-consumer -V

The numerical test requires one GPU and MPI and is off by default. The CPU consumer uses ENABLE_LIBRI_GPU=OFF and ENABLE_LIBRI_GPU_TESTS=OFF. Earlier MKL consumers additionally set MKL_ROOT, MKL_MPI=openmpi and MKL_LINK=static or dynamic.

The Si64 hybrid sweep was launched with sbatch control/run.sbatch (SAI job 1862952, exit 0:0). Each calculation runs:

OMP_NUM_THREADS=<threads> OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
  mpirun --bind-to none --map-by ppr:<ranks>:node -np <ranks> \
  python3 control/rank_launcher.py <ABACUS_BIN>

The launcher pins each rank to a recorded, disjoint subset of the scheduler cpuset and sets explicit OMP_PLACES; an OpenMP probe verifies every thread's CPU before each run. OMP_DYNAMIC=FALSE and OMP_MAX_ACTIVE_LEVELS=1; BLAS stays single-threaded to avoid nested oversubscription. The physical-core sweep uses one hardware thread per core; the SMT sweep also uses its sibling. CPU-only runs have no visible GPU and no observed GPU process. GPU RI uses one rank per GPU and the same visible GPU list on all ranks.

CPU-only INPUT: device cpu, ks_solver scalapack_gvx, device_exx cpu. GPU RI INPUT: device gpu, ks_solver cusolver, device_exx gpu. The cpu_ri_gpu_solver control keeps the GPU settings and changes only device_exx cpu.

Run MPI×OpenMP EXX force (s) All forces (s) SCF + forces (s) Outer/inner iterations
cpu_4x1 4×1 115.30 117.31 412.57 7/96
cpu_2x2 2×2 97.36 99.67 407.46 7/96
cpu_1x4 1×4 85.35 88.23 410.15 7/96
gpu_1x4 1×4 21.32 22.76 237.24 7/96
cpu_4x2 4×2 94.51 96.24 351.86 7/96
cpu_2x4 2×4 83.11 85.21 385.06 7/96
cpu_1x8 1×8 73.28 75.99 379.31 7/96
cpu_8x1 8×1 121.31 123.28 450.10 7/96
gpu_1x8 1×8 21.39 22.86 237.88 7/96
cpu_1x8_repeat 1×8 72.22 74.89 375.46 7/96
cpu_4x2_repeat 4×2 95.28 97.02 353.30 7/96
gpu_1x4_repeat 1×4 21.58 23.06 239.75 7/96
cpu_ri_gpu_solver 1×4 87.11 88.64 354.62 7/96
gpu_2rank 2×4 19.64 20.60 192.39 7/96
gpu_2rank_repeat 2×4 19.20 20.14 191.75 7/96

Settings: cal_force=1, cal_stress=0, symmetry=0, exx_symmetry_realspace=0, ecutwfc=60, nbands=160, scf_thr=1e-9, scf_ene_thr=1e-8, exx_pca_threshold=1e-4 and exx_ccp_rmesh_times=1. Each run starts from atomic density. Independent post-run checks of full force vectors, structures, numerical settings, asset hashes, executable hash and CPU/GPU placement passed at tolerances of 1e-4 eV/Angstrom for forces and 1e-6 eV/atom for energy. The one-pass sweep selects configurations for one additional repeat; this is a small tuning experiment, not an exhaustive CPU solver/library comparison.

GPU regression builds used ENABLE_NCCL_PARALLEL_DEVICE=OFF because the input-test target had unresolved NCCL symbols with it enabled. NCCL runtime behavior, full MKL application builds, multi-node execution, SOC and long MD runs remain unverified. GPU memory use increases; initialization remains largely on CPU. A Si32 comparison needed scf_ene_thr=1e-8 for consistent outer-loop stopping, so default-threshold convergence needs further validation.

Linked Issue

No linked issue; this adds optional support for an existing LibRI backend.

Governance Notes

  • INPUT/docs: device_exx includes parser validation, MPI broadcast, help metadata, installation documentation and focused tests. Generated parameter documentation is not committed.
  • Core impact: select the backend before LCAO EXX initialization and encapsulate dependencies in LibRI::LibRI and DDLA::DDLA CMake targets. No new global dependencies. GPU RI requires C++17; CPU RI retains C++14.
  • The standalone numerical test initializes/finalizes MPI directly because it links only the dependency target; communicator queries use the LibRI wrapper. Formatting, LF, whitespace and PR-diff governance checks pass.

Follow-up: add DDLA installation to the existing toolchain/scripts/ workflow.

Contribution checklist
  • I have read AGENTS.md and docs/developers_guide/agent_governance.md.
  • I have linked an issue or explained why this PR does not need one.
  • I have added adequate unit tests and/or case tests, or explained why not.
  • I have listed the exact verification commands run and their results.
  • I have described user-visible behavior changes, including INPUT parameter changes.
  • I have explained core-module impact for ESolver, HSolver, ElecState, Hamilt, Operator, Psi, or other source/ changes.
  • I have requested any needed governance exception below.

@userchba

Copy link
Copy Markdown

the LibDDLA version which is compatible with LibRI GPU is https://github.com/userchba/LibDDLA/releases/tag/v0.0.4, current LibRI GPU can't correctly compute symmetry situation which LibRI CPU can compute correctly

@mohanchen mohanchen added the GPU & DCU & HPC GPU and DCU and HPC related any issues label Oct 18, 2026

@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.

Static review of f07d8d14f39d26782490cc4bc15d0a58c5700e62 with independent source review and adjudication. No build/runtime/MPI/GPU tests were run; CI diagnostics were unavailable.

Comment thread docs/advanced/install.md Outdated
(default OFF), `USE_CUDA=ON`, and a CUDA build of DDLA. Like other dependent
options, GPU RI is disabled when its LibRI/CUDA prerequisites are disabled.
INPUT `exx_device` defaults to `cpu`, retaining CPU contractions even with
`device gpu`. To use GPU contractions, set both `device gpu` and `exx_device gpu`.

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.

[P2] State the GPU visibility requirement for MPI runs

For two ranks on one node, exposing one different GPU per rank through CUDA_VISIBLE_DEVICES passes ABACUS device setup but makes the pinned LibRI validate_local_gpu_count() throw at the first GPU contraction (local_size=2 > device_count=1). This backend therefore requires each process to see at least as many GPUs as there are node-local ranks, a constraint missing here. Please document it with a compatible launch layout, or extend the compatibility patch to respect the existing per-rank binding.

Comment thread source/source_lcao/setup_exx.cpp Outdated
Exx_Info& exx_info)
{
#ifdef __EXX
const bool use_gpu = inp.exx_device == "gpu";

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.

Non-blocking: Please clarify the selector's scope for independent LR/RDMFT exchange engines. Standalone LR, fresh LR-kernel setup after KS, and RDMFT construct their own Exx_LRI objects, which retain CPU mode despite accepted exx_device=gpu; matching-kernel KS pointer reuse retains GPU mode. If those engines are in scope, pass the selection to them; otherwise document or reject the unsupported combinations. A follow-up after merge is fine for these additional workflows.

{
return 1;
}
lri.cal_mode = RI::LRI_Cal_Mode::GPU;

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.

Non-blocking: Consider an opt-in numerical regression for the patch's signed accumulation in a follow-up after merge. This consumer switches cal_mode but then executes only CPU nrm2, and the configure probe only compiles/links. Alongside the reported force/stress runs, a tiny real/complex CPU-versus-GPU case with nonempty output and repeated +1/-1 contractions would preserve the specific accumulation fix across dependency updates.

@Stardust0831
Stardust0831 marked this pull request as ready for review October 19, 2026 14:30
@mohanchen mohanchen added the Compile & CICD & Docs & Dependencies Issues related to compiling ABACUS label Oct 20, 2026
@mohanchen

Copy link
Copy Markdown
Collaborator

To be consistent, I think device_exx maybe better than exx_device

@mohanchen mohanchen added Input&Output Suitable for coders without knowing too many DFT details EXX/lr-TDDFT/RPA/BSE/GW Related to beyond DFT methods labels Oct 20, 2026
@Stardust0831

Copy link
Copy Markdown
Collaborator Author

To be consistent, I think device_exx maybe better than exx_device

Thank you for the review and suggestions; I’ve renamed exx_device to device_exx.

@QuantumMisaka

QuantumMisaka commented Oct 20, 2026 •

Copy link
Copy Markdown
Collaborator

@PeizeLin @maki49 May you have a look? Thanks !

@mohanchen

Copy link
Copy Markdown
Collaborator

@Growl1234 Could you check the cmake changes for this PR?

if(ENABLE_LIBRI_GPU)
message(
FATAL_ERROR
"ENABLE_LIBRI_GPU requires an explicitly supplied, patched LIBRI_DIR; "

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.

Does it mean that LibRI master branch hasn't yet accepted the LibRI changes? If this is true then this PR should not be merged on abacus side first.

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.

LibDDLA has CMake config export, why is FindDDLA module needed here?

@maki49

maki49 commented Oct 21, 2026 •

Copy link
Copy Markdown
Collaborator

I have no further comment within this PR's scope.
LR and symmetry with LibRI on GPU are already supported in my local branch but they need to update LibRI. I will make a new PR in the future when everything is ready.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Compile & CICD & Docs & Dependencies Issues related to compiling ABACUS EXX/lr-TDDFT/RPA/BSE/GW Related to beyond DFT methods GPU & DCU & HPC GPU and DCU and HPC related any issues Input&Output Suitable for coders without knowing too many DFT details

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants