Feature(pw): batched FFTs, small ecutexx grid and band chunking for EXX on CPU and GPU - #8018
Open
Flying-dragon-boxing wants to merge 2 commits into
Open
Flying-dragon-boxing wants to merge 2 commits into
Flying-dragon-boxing wants to merge 2 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 |
| item.category = "Exact Exchange (PW)"; | ||
| item.type = "Real"; | ||
| item.description = "The energy cutoff for EXX (Fock) exchange operator in plane wave basis calculations. Reducing ecutexx below ecutrho may significantly accelerate EXX computations. This speed improvement comes with a reduced numerical accuracy in the exchange energy calculation."; | ||
| item.description = "The energy cutoff for EXX (Fock) exchange operator in plane wave basis calculations. The pair-density G-sphere of the exchange operator, the EXX energy, and the EXX stress are all truncated at this value. If ecutexx yields a smaller FFT box and every |k+G|^2 of the wavefunctions fits inside it (i.e. ecutexx should not be smaller than ecutwfc), all EXX FFTs run on that smaller grid (QE ecutfock-style), which can significantly accelerate EXX computations. If the small grid is not usable (box not smaller, wavefunctions do not fit, or the FFT box is distributed over MPI), a warning is printed and the full grid is used. Reducing ecutexx below ecutrho reduces the numerical accuracy of the exchange contribution."; |
| void setup_full_grid_maps() const; | ||
| mutable bool exx_sg_init = false; | ||
| mutable bool exx_sg_ok = false; | ||
| mutable int sg_nx = 0, sg_ny = 0, sg_nz = 0, sg_nxyz = 0; |
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.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_pwdftonly (EXX operator + stress), plus the INPUT registration. 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.exx_band_chunk()readsPARAM.inp.exx_batch_size, the module's established style for INPUT values (same asecutexx); thePARAM/GlobalVuses inop_pw_exx.cpp/stress_exx.cppare migration-neutral moves — the refactor rewrites existing blocks (act_op_kpar,cal_exx_energy_op, setup) that already readPARAM.inp.nspin/ecutexxandGlobalV::MY_POOLin the same style, so no new global-dependency direction is introduced.