Feature(pw): batched FFTs, small ecutexx grid and band chunking for EXX on CPU and GPU - #8018
Open
Flying-dragon-boxing wants to merge 6 commits into
Open
Flying-dragon-boxing wants to merge 6 commits into
Flying-dragon-boxing wants to merge 6 commits into
Conversation
Unify all EXX PW entry points (act_op, act_op_kpar, cal_exx_energy_op)
on one code path built on shared primitives; the batched kernels are a
device specialization of the per-band operations, selected inside the
primitives (batch_active).
- exx_batch kernels templated on Device: host loops + FFTW plan_many on
CPU (single-precision plans compiled only with ENABLE_FLOAT_FFTW, stubs
otherwise), CUDA kernels + cuFFT as DEVICE_GPU specializations.
- QE ecutfock-style small FFT grid from ecutexx when every |k+G|^2 fits,
on CPU and GPU; falls back to the full grid with a warning otherwise.
The full-grid batched path runs on both devices when the box is local.
- Physics fix: the Fock operator weight now uses the source-state
occupation f_{mq} and the source k-point weight (was the target-k wg in
act_op and target-k wk in act_op_kpar). With k-dependent occupations
(smearing) the old operator was inconsistent with the energy and the
EXX outer loop never converged.
- The per-(q,m) scalar MPI_Bcast of wg is replaced by one broadcast of
the occupation row + wk per source k-point via Parallel_Common wrappers.
- stress_exx G-sum truncated to the ecutexx sphere (CPU), consistent with
the operator and energy.
- Docs: ecutexx describes the small-grid behavior and fallback.
Governance exception: the PARAM/GlobalV budget flags are migration-
neutral moves - the refactor rewrites existing blocks (act_op_kpar,
cal_exx_energy_op, setup) that already read PARAM.inp.nspin/ecutexx and
GlobalV::MY_POOL in the same style as the surrounding module.
Verified against main (4f2a397): build/rel (g++ MPI) and
build_abacus_gnu (CUDA 13.1); 19-case regression matrix
(097_PW_PBE0{,_FM,_COND} x ACE/noACE x full/small grid x CPU/GPU) -
full-grid cases match the pre-change code at 1e-14, CPU small grid
matches GPU small grid to 1e-13, metallic noACE case converges in 3 EXX
outer iterations (did not converge before the wg fix).
The small ecut_exx grid and band batching were bound in a single
predicate; they are orthogonal concerns (grid = which FFT box, chunk =
how many bands per batched round). Split them:
- exx_grid_active() answers the grid question (small grid usable, or
the full box local); exx_band_chunk() answers the band question via
the new exx_batch_size INPUT (default 0 = all bands, identical to the
previous behavior; a positive value processes bands in chunks of that
width with a proportionally smaller work-buffer footprint).
- cache_psi_nk_real, apply_exx_nbatched and the energy pair-density
loop process bands in chunks; psi_nk_real_cache still holds all bands
(it is the reuse floor across (iq, m)), only the work buffers shrink.
- Docs: exx_batch_size in parameters.yaml and input-main.md.
Governance exception: exx_band_chunk() reads PARAM.inp.exx_batch_size,
the module's established style for INPUT values (same as ecutexx).
Verified: 24-case regression (097_PW_PBE0{,_FM,_COND} x ACE/noACE x
full/small grid x CPU/GPU, plus exx_batch_size 1/3/5 variants) - the
default keeps every previous value bitwise, chunked runs are bitwise
identical to unchunked (band blocks are disjoint and the (q,m)
accumulation order is unchanged).
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Unresolved critical and moderate findings remain in CUDA FFT distances, small-grid bounds validation, memory allocation, global dependencies, and documentation accuracy.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 2
Open (4)
What changed in this PR
Adds batched CPU/GPU EXX processing, optional smaller ecutexx FFT grids, band chunking, and corrected occupation weighting.
Changes:
- Adds shared CPU, CUDA, and cuFFT batching kernels.
- Introduces and documents
exx_batch_sizeand small-grid behavior. - Updates EXX operator, energy, stress, and occupation-weighting logic.
| File | Description |
|---|---|
source/source_pw/module_pwdft/stress_exx.cpp |
Applies the ecutexx cutoff to EXX stress sums. |
source/source_pw/module_pwdft/op_pw_exx.h |
Adds batching and grid-management state. |
source/source_pw/module_pwdft/op_pw_exx.cpp |
Implements unified EXX paths, chunking, grid handling, and weighting. |
source/source_pw/module_pwdft/kernels/exx_batch_op.h |
Declares batched EXX kernel APIs. |
source/source_pw/module_pwdft/kernels/exx_batch_op.cpp |
Implements CPU double-precision FFTW kernels. |
source/source_pw/module_pwdft/kernels/exx_batch_op_float.cpp |
Implements CPU float FFTW plans. |
source/source_pw/module_pwdft/kernels/exx_batch_op_float_stub.cpp |
Provides float-FFTW fallback stubs. |
source/source_pw/module_pwdft/kernels/cuda/exx_batch_op.cu |
Implements CUDA and cuFFT batching. |
source/source_pw/module_pwdft/CMakeLists.txt |
Registers EXX kernel sources. |
source/source_io/module_parameter/read_inp_other.cpp |
Parses and validates exx_batch_size. |
source/source_io/module_parameter/input_parameter.h |
Adds the input parameter field. |
source/CMakeLists.txt |
Links the CUDA EXX source. |
docs/parameters.yaml |
Documents EXX parameters and grid behavior. |
docs/advanced/input_files/input-main.md |
Updates user-facing input documentation. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+41
to
+42
| nullptr, 1, 0, // input: contiguous, distance = nx*ny*nz | ||
| nullptr, 1, 0, // output: same (in-place) |
Comment on lines
+587
to
+594
| if (gridt.nx == wfcpw->nx && gridt.ny == wfcpw->ny && gridt.nz == wfcpw->nz) | ||
| { | ||
| if (user_set && GlobalV::MY_RANK == 0) | ||
| { | ||
| ModuleBase::WARNING("OperatorEXXPW", | ||
| "ecutexx gives no smaller FFT grid than ecutrho; EXX stays on the full grid"); | ||
| } | ||
| return; // no gain over the current grid |
…ead of globals The EXX PW operator and stress read their configuration (nspin, ecutexx, exx_batch_size, exxace, exx_gamma_extrapolation) and the MPI layout (KPAR, MY_RANK, MY_POOL) directly from PARAM/GlobalV. Inject them explicitly instead: - General_Exx_Info carries the PW EXX INPUT values, resolved once in init_general_exx_info: exxace, gamma_extrapolation, exx_batch_size, and ecut_exx (ecutexx when set, else ecutrho) with the user-set flag. - OperatorEXXPW takes the General_Exx_Info plus the runtime values (nspin, kpar, my_rank, my_pool) at construction and snapshots them as members; all internal PARAM/GlobalV reads are gone. - Stress_PW::stress_exx takes the General_Exx_Info in place of the separate hybrid_alpha/coulomb_param parameters. Governance: the global-dependency budget of the PR goes from +12 to -28 (added=4, removed=32); the remaining 4 added references are the single snapshot point in HamiltPW, the orchestration layer that owns these values. Verified: build/rel (g++ MPI) rebuilds; tests/01_PW/097_PW_PBE0 matches result.ref (etot to 1.8e-14 eV, stress sum exact at 1e-6 kbar); exx_batch_size=1 is bitwise identical to unchunked; ecutexx=20 engages the small FFT grid at 1 rank and prints the distributed-box fallback warning once at 2 ranks, both converging to the same energy (2.4e-14 eV).
Match the --generate-parameters-yaml output exactly; the documentation consistency CI check requires docs/parameters.yaml to be byte-identical to the generated file.
The batched-EXX commits added kernels/exx_batch_op.cpp, kernels/exx_batch_op_float.cpp (ENABLE_FLOAT_FFTW) and its stub, plus kernels/cuda/exx_batch_op.cu, and wired them into CMake only. The Makefile build then failed at link time with undefined references to hamilt::exx_batch_* from the rewritten op_pw_exx.cpp. Add exx_batch_op.o and exx_batch_op_float_stub.o to OBJS_HAMILT; the legacy Makefile has no float-FFTW switch, matching the CMake ENABLE_FLOAT_FFTW=OFF branch. VPATH already covers the kernels dir.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Reminder
AGENTS.mdanddocs/developers_guide/agent_governance.md.source/changes.Linked Issue
No linked issue: this is a performance/feature PR (EXX PW acceleration) plus a physics fix for the Fock operator occupation weights, self-contained in
source_pw/module_pwdft.Unit Tests and/or Case Tests for my changes
097_PW_PBE0{,_FM,_COND}x {ACE, noACE} x {full grid, small ecutexx grid} x {CPU, GPU}, plusexx_batch_size= 1/3/5 variants of the chunked path, usingbuild/rel(g++ MPI) andbuild_abacus_gnu(CUDA 13.1) executables.exx_batch_size> 0) are bitwise identical to unchunked ones (band blocks are disjoint and the (q,m) accumulation order is unchanged); full-grid cases match the pre-change code at 1e-14; CPU small grid matches GPU small grid to 1e-13; the metallic noACE case (k-dependent occupations) converges in 3 EXX outer iterations where the old operator never converged.tests/in this PR; the existing097_PW_PBE0*cases already exercise the touched code paths, and the small-grid/chunked variants were verified locally as described above. Happy to add anecutexxsmall-grid case totests/01_PWif reviewers want one.What's changed?
act_op,act_op_kpar,cal_exx_energy_op) on one code path built on shared primitives; the batched kernels are a device specialization of the per-band operations, selected inside the primitives.kernels/exx_batch_op*: host loops + FFTWplan_manyon CPU (single-precision plans only withENABLE_FLOAT_FFTW, stubs otherwise), CUDA kernels + cuFFT asDEVICE_GPUspecializations.ecutfock-style small FFT grid fromecutexxwhen every |k+G|^2 fits, on CPU and GPU; falls back to the full grid with a warning otherwise.exx_batch_size(Integer, default 0 = all bands, identical to previous behavior): processes bands in chunks of that width with a proportionally smaller work-buffer footprint; results are independent of the chunking.wginact_opand target-kwkinact_op_kpar). With k-dependent occupations (smearing) the old operator was inconsistent with the energy and the EXX outer loop never converged.MPI_Bcastofwgis replaced by one broadcast of the occupation row +wkper source k-point viaParallel_Commonwrappers.stress_exxG-sum truncated to theecutexxsphere (CPU), consistent with the operator and energy.General_Exx_Infocarries the PW EXX INPUT values (resolved once ininit_general_exx_info),OperatorEXXPWsnapshots it plus the runtime layout (nspin,kpar,my_rank,my_pool) at construction, andStress_PW::stress_exxtakes theGeneral_Exx_Infoin place of the separatehybrid_alpha/coulomb_paramparameters.Governance Notes
exx_batch_sizeregistered inread_inp_other.cpp(with a>= 0value check) and documented indocs/parameters.yamlanddocs/advanced/input_files/input-main.md; theecutexxdescription is updated to document the small-grid behavior and fallback. No default behavior changes.source_pw/module_pwdft(EXX operator + stress), the INPUT registration, and theGeneral_Exx_Infoextension insource_hamilt/module_xc. All EXX PW paths (SCF operator action and energy) go through the new unified code path; default settings are bitwise identical to the previous results.agent_governance_check.py --base 4f2a39783 --head HEAD--op_pw_exx.cpp/stress_exx.cppno longer readPARAM/GlobalVat all; the 4 remaining added references are the single snapshot point inHamiltPW(PARAM.inp.nspin,GlobalV::KPAR/MY_RANK/MY_POOL), the orchestration layer that owns these values.