Expose MLMG bottom solver and coarse level controls for the ES solvers - #7228
Open
roelof-groenewald wants to merge 6 commits into
Open
Expose MLMG bottom solver and coarse level controls for the ES solvers#7228roelof-groenewald wants to merge 6 commits into
roelof-groenewald wants to merge 6 commits into
Conversation
Add ablastr::fields::MLMGOptions, a small struct carrying the AMReX MLMG settings that were previously not reachable from a WarpX input deck, and plumb it through both electrostatic Poisson solvers (PoissonSolver.H and EffectivePotentialPoissonSolver.H). Members left negative keep the AMReX default, so behaviour is unchanged unless an input is set. Bottom solve controls: warpx.self_fields_bottom_solver warpx.self_fields_bottom_verbosity warpx.self_fields_bottom_max_iters warpx.self_fields_bottom_relative_tolerance warpx.self_fields_bottom_absolute_tolerance warpx.self_fields_max_coarsening_level Coarse level distribution controls (amrex::LPInfo): warpx.self_fields_agglomeration warpx.self_fields_agglomeration_grid_size warpx.self_fields_consolidation warpx.self_fields_consolidation_grid_size All are also available as warpx_self_fields_* keyword arguments of picmi.ElectrostaticSolver. Motivation: on GPUs the coarsest multigrid level is small enough that the bottom solve is dominated by kernel launch and reduction latency rather than arithmetic. On a 128x128x352 effective-potential run with an embedded boundary on A100s, the bottom solve was 55% of every Poisson solve and cost the same wall time on 1 and on 8 GPUs. Switching the bottom solver from BiCGStab to CG and capping its iterations cut it from 60 ms to 22 ms per solve for 0.9 extra MLMG iterations, a 16% reduction in total runtime on 8 GPUs. The remaining serialization is agglomeration: AMReX gathers the coarse levels onto a single box once the coarsened boxes fall below agg_grid_size (32 in 3D), so one rank performs all the work from that level down while the others wait in FabArray::ParallelCopy_finish. The LPInfo knobs make that tunable. Add test_3d_electrostatic_sphere_lab_frame_no_agglomeration, which runs the lab frame sphere on 2 ranks with several boxes per rank and the coarse levels left distributed. Its checksums agree with those of test_3d_electrostatic_sphere_lab_frame to 7e-16, confirming that the coarse level decomposition does not change the result. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Roelof Groenewald <rgroenewald@realtafusion.com>
The relative and absolute tolerance, the iteration limit, the verbosity and the final smoothing count of the electrostatic Poisson solve were loose members of ElectrostaticSolver that were threaded through four layers of call as positional arguments. Move them into ablastr::fields::MLMGOptions, next to the bottom solver and coarsening settings they belong with, so that a single struct describes an MLMG solve. applyMLMGOptions(amrex::MLMG&, ...) now also applies the verbosity, the iteration limit and the final smoothing count. The two tolerances remain at the call site because they are arguments of amrex::MLMG::solve. The defaults of the new members match the previous WarpX defaults and the input parameter names are unchanged, so existing input files behave identically. Two consequences are worth noting: - The relativistic solver takes the tolerances, iteration limit and verbosity per species. It now copies the shared options and overrides those four, which means a species solve also inherits the shared bottom solver and agglomeration settings, which it previously could not see. - The effective potential solver never called setFinalSmooth. It now does, through the shared applyMLMGOptions. Since both its default and AMReX's are 8, this only takes effect when warpx.self_fields_num_final_sweeps is set explicitly. This changes the signature of ablastr::fields::computePhi and computeEffectivePotentialPhi, which downstream users of ablastr need to follow. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The MLMG options are all handled uniformly now, so the existing electrostatic tests exercise this code path well enough on their own. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The 3D FFT build calls getMaxNormRho, which takes the absolute tolerance by non-const reference and raises it to 1e-6 when rho is zero everywhere, since a purely relative convergence criterion is meaningless there. That call site was missed when the tolerances moved into MLMGOptions, and it cannot read the value straight from the const options struct. Keep a local copy of the tolerance, initialized from the options and declared outside the level loop, so the behaviour of both the FFT and the non-FFT build is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
roelof-groenewald
added
cleaning
getMaxNormRho took the absolute tolerance by non-const reference and raised it when rho was zero everywhere, which forced its caller to keep a mutable local. That local is only ever written in the 3D FFT build, so clang-tidy asks for it to be const in every other configuration. Remove the out parameter: getMaxNormRho is now a pure query and the caller derives the tolerance from its result. The norm does not depend on the level, so it is also hoisted out of the level loop, where it was recomputed on every iteration while only being used on level 0. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
roelof-groenewald
changed the title
Member
Author
|
Note: This PR modifies public APIs and therefore likely requires updates to ImpactX. |
2 tasks
Member
|
@RemiLehe @roelof-groenewald looks good from the ImpactX side :) BLAST-ImpactX/impactx#1646 |
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
3 participants
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.
Add
ablastr::fields::MLMGOptions, a small struct carrying the AMReX MLMG settings that were previously not reachable from a WarpX input deck, and plumb it through both electrostatic Poisson solvers (PoissonSolver.HandEffectivePotentialPoissonSolver.H). Members left negative keep the AMReX default, so behaviour is unchanged unless an input is set.Bottom solve controls:
Coarse level distribution controls (amrex::LPInfo):
All are also available as
warpx_self_fields_*keyword arguments ofpicmi.ElectrostaticSolver.Motivation: on GPUs the coarsest multigrid level is small enough that the bottom solve is dominated by kernel launch and reduction latency rather than arithmetic. On a 128x128x352 effective-potential run with an embedded boundary on A100s, the bottom solve was 55% of every Poisson solve and cost the same wall time on 1 and on 8 GPUs. Switching the bottom solver from BiCGStab to CG and capping its iterations cut it from 60 ms to 22 ms per solve for 0.9 extra MLMG iterations, a 16% reduction in total runtime on 8 GPUs.
The remaining serialization is agglomeration: AMReX gathers the coarse levels onto a single box once the coarsened boxes fall below agg_grid_size (32 in 3D), so one rank performs all the work from that level down while the others wait in
FabArray::ParallelCopy_finish. The LPInfo knobs make that tunable.The new
ablastr::fields::MLMGOptionsstruct is now also used to collect outer solver parameters (tolerance, verbosity, etc.) making the Poisson solver API more stable (after this PR).