From d9397dc661bc255f7a184005d2c3f7b006b389a6 Mon Sep 17 00:00:00 2001 From: Remi Lehe Date: Thu, 27 Aug 2026 15:11:57 -0700 Subject: [PATCH 1/5] Expose PETSc as a solver for the Poisson equation WarpX already exposes PETSc as an interface to solve the curl-curl equation of the implicit electromagnetic solvers (`newton.linear_solver = petsc_ksp`, `jacobian.pc_type = pc_petsc`). This adds the electrostatic counterpart: with `warpx.poisson_solver = petsc`, the Poisson equation of `computePhi` is solved by one of PETSc's Krylov solvers, preconditioned by the AMReX geometric multigrid. As in `amrex::GMRESMLMG`, the linear system is not assembled into a sparse matrix. It is handed to PETSc as a `MatShell` whose action is computed by the very same `amrex::MLNodeLinOp` that MLMG would use, and the multigrid V-cycles are handed to PETSc as a `PCShell`. The Krylov solver iterates on a correction, whose residual is computed with the inhomogeneous operator, so that non-zero Dirichlet boundary values are accounted for. The degrees of freedom of the PETSc vectors are the nodes that a box owns (nodes shared with a neighboring box, or with a periodic image, appear only once) and that are not Dirichlet nodes. Since `petsc` discretizes Poisson's equation exactly like `multigrid` does and accepts the same boundary conditions, the sites that used to test `poisson_solver_id == Multigrid` (the Poisson boundary handler and the initial div(B) cleaning) now test `!= IntegratedGreenFunction`, which is what they actually mean. The tuning knobs of the Krylov solver live in the new `petsc_poisson` namespace. All of the new code is guarded by `AMREX_USE_PETSC`. Embedded boundaries are not supported yet, since the calculation of the E field that follows the solve needs the internal state of a completed MLMG solve; the effective-potential solver is not supported either, since it does not go through the same Poisson solve. Co-Authored-By: Claude Opus 5 --- Docs/source/usage/parameters.rst | 31 ++ .../Tests/electrostatic_sphere/CMakeLists.txt | 15 + ...st_3d_electrostatic_sphere_lab_frame_petsc | 12 + Python/pywarpx/PETScPoisson.py | 11 + Python/pywarpx/WarpX.py | 3 + Python/pywarpx/__init__.py | 1 + .../ElectrostaticSolver.H | 4 + .../ElectrostaticSolver.cpp | 18 + .../PoissonBoundaryHandler.cpp | 8 +- .../DivCleaner/ProjectionDivCleaner.cpp | 2 +- Source/Initialization/WarpXInitData.cpp | 3 + Source/Utils/WarpXAlgorithmSelection.H | 1 + Source/WarpX.cpp | 20 +- Source/ablastr/fields/CMakeLists.txt | 1 + Source/ablastr/fields/Make.package | 1 + Source/ablastr/fields/PETScPoissonSolver.H | 229 ++++++++ Source/ablastr/fields/PETScPoissonSolver.cpp | 493 ++++++++++++++++++ Source/ablastr/fields/PoissonSolver.H | 21 +- 18 files changed, 867 insertions(+), 7 deletions(-) create mode 100644 Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc create mode 100644 Python/pywarpx/PETScPoisson.py create mode 100644 Source/ablastr/fields/PETScPoissonSolver.H create mode 100644 Source/ablastr/fields/PETScPoissonSolver.cpp diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst index 57242047a58..c4ea9c305e9 100644 --- a/Docs/source/usage/parameters.rst +++ b/Docs/source/usage/parameters.rst @@ -450,6 +450,37 @@ Overall simulation parameters The extended simulation box size in real space is :math:`2n_x-1, 2n_y-1, 2n_z-1` with the 3D solver, :math:`2n_x-1, 2n_y -1, n_z` with the 2D solver. The extended simulation box size in spectral space is :math:`n_x, 2n_y-1, 2n_z-1` with the 3D solver, :math:`n_x, 2n_y-1, n_z` with the 2D solver. + * ``petsc``: Poisson's equation is solved by one of PETSc's Krylov solvers, preconditioned by the AMReX multigrid. + This is the electrostatic counterpart of the PETSc interface that the implicit electromagnetic solvers + offer for the curl-curl equation (``newton.linear_solver = petsc_ksp``): the linear system is + handed to PETSc as a matrix-free operator, whose action, as well as that of the multigrid preconditioner, + is computed by AMReX. It therefore discretizes Poisson's equation exactly like ``multigrid`` does, and + accepts the same boundary conditions; only the outer iteration differs. + It requires the compilation flag ``-DWarpX_PETSC=ON``. + It is not supported with embedded boundaries, nor in ``labframe-effective-potential`` mode. + Note that in 1D with ``warpx.do_electrostatic = labframe``, and with the ``poissonsolver`` Python callback, + Poisson's equation is solved by a dedicated solver and this option has no effect. + + Any PETSc option that is set on the command line or in a PETSc option file (e.g. ``-ksp_type``) + takes precedence over the parameters below. + + * ``petsc_poisson.ksp_type`` (``string``) optional (default: ``gmres``): The type of the PETSc ``KSP`` + object, e.g. ``gmres``, ``fgmres``, ``bcgs`` or ``cg``. + + * ``petsc_poisson.restart_length`` (``int``) optional (default: 30): The restart length, only used by + the GMRES variants. + + * ``petsc_poisson.use_mlmg_preconditioner`` (``bool``) optional (default: 1): Whether to use the AMReX + multigrid V-cycles as a preconditioner, through PETSc's ``PCShell`` interface. Setting this to 0 runs + the Krylov solver unpreconditioned, which is much slower and mostly useful for debugging. + + * ``petsc_poisson.precond_num_iters`` (``int``) optional (default: 1): The number of multigrid V-cycles + per application of the preconditioner. + + * ``petsc_poisson.verbose`` (``int``) optional (default: value of + :pp:param:`warpx.self_fields_verbosity`): 0 is silent, 1 prints the exit status of every solve, + 2 additionally prints the residual at every Krylov iteration. + .. pp:param:: warpx.self_fields_required_precision :type: ``float`` :default: 1.e-11 diff --git a/Examples/Tests/electrostatic_sphere/CMakeLists.txt b/Examples/Tests/electrostatic_sphere/CMakeLists.txt index 85ec6ceb74e..5fede12b68a 100644 --- a/Examples/Tests/electrostatic_sphere/CMakeLists.txt +++ b/Examples/Tests/electrostatic_sphere/CMakeLists.txt @@ -21,6 +21,21 @@ add_warpx_test( OFF # dependency ) +if(WarpX_PETSC) + # No checksum: the reference values still have to be generated on a build + # with PETSc. The physics analysis below already checks the solution against + # the analytical solution of the expanding sphere. + add_warpx_test( + test_3d_electrostatic_sphere_lab_frame_petsc # name + 3 # dims + 2 # nprocs + inputs_test_3d_electrostatic_sphere_lab_frame_petsc # inputs + "analysis_electrostatic_sphere.py diags/diag1000030" # analysis + OFF # checksum + OFF # dependency + ) +endif() + add_warpx_test( test_3d_electrostatic_sphere_lab_frame_mr_emass_10 # name 3 # dims diff --git a/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc b/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc new file mode 100644 index 00000000000..9f1514fec4a --- /dev/null +++ b/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc @@ -0,0 +1,12 @@ +# base input parameters +FILE = inputs_base_3d + +# test input parameters +diag2.electron.variables = x y z ux uy uz w phi +warpx.do_electrostatic = labframe + +# solve the Poisson equation with PETSc's GMRES, preconditioned by the AMReX +# multigrid, instead of using the multigrid directly +warpx.poisson_solver = petsc +petsc_poisson.ksp_type = gmres +petsc_poisson.precond_num_iters = 1 diff --git a/Python/pywarpx/PETScPoisson.py b/Python/pywarpx/PETScPoisson.py new file mode 100644 index 00000000000..e2fe22d49d4 --- /dev/null +++ b/Python/pywarpx/PETScPoisson.py @@ -0,0 +1,11 @@ +# Copyright 2026 The WarpX Community +# +# This file is part of WarpX. +# +# License: BSD-3-Clause-LBNL + +from .Bucket import Bucket + +# Options of the PETSc Krylov solver that is used for the Poisson equation +# when warpx.poisson_solver = petsc (requires compiling with -DWarpX_PETSC=ON) +petsc_poisson = Bucket("petsc_poisson") diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py index 8b5f07762cf..d9d867084ac 100644 --- a/Python/pywarpx/WarpX.py +++ b/Python/pywarpx/WarpX.py @@ -23,6 +23,7 @@ from .Interpolation import interpolation from .Lasers import lasers, lasers_list from .Particles import particles, particles_list +from .PETScPoisson import petsc_poisson from .PSATD import psatd @@ -55,6 +56,7 @@ def create_argv_list(self, **kw): argv += interpolation.attrlist() argv += psatd.attrlist() argv += eb2.attrlist() + argv += petsc_poisson.attrlist() argv += particles.attrlist() for particle in particles_list: @@ -223,6 +225,7 @@ def finalize(self, finalize_mpi=1): lasers, my_constants, particles, + petsc_poisson, psatd, reduced_diagnostics, self, diff --git a/Python/pywarpx/__init__.py b/Python/pywarpx/__init__.py index 8376a9beb28..f9abd9858c4 100644 --- a/Python/pywarpx/__init__.py +++ b/Python/pywarpx/__init__.py @@ -38,6 +38,7 @@ from .Lasers import lasers # noqa from .LoadThirdParty import load_cupy # noqa from .Particles import new_species, particles # noqa +from .PETScPoisson import petsc_poisson # noqa from .PSATD import psatd # noqa from .WarpX import warpx # noqa diff --git a/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.H b/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.H index c0fe04461d6..a7eecbdcb71 100755 --- a/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.H +++ b/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.H @@ -13,6 +13,7 @@ #include "Fluids/MultiFluidContainer.H" #include "Particles/MultiParticleContainer.H" +#include #include #include @@ -159,6 +160,9 @@ public: /** MLGM number of smoothing sweeps */ int self_fields_num_final_sweeps = 8; + /** Settings of the optional PETSc Krylov solver for the Poisson equation */ + ablastr::fields::PETScPoissonOptions m_petsc_options; + /** Parameters for FFT Poisson solver aka IGF */ // 0: full 3D, 1: many 2D z-slices (quasi-3D) bool is_igf_2d_slices = false; diff --git a/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp b/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp index 587382c12d4..ddcd5f8cc81 100755 --- a/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp +++ b/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp @@ -47,6 +47,23 @@ void ElectrostaticSolver::ReadParameters () { self_fields_num_final_sweeps > 0, "warpx.self_fields_num_final_sweeps must be > 0"); } + + // With warpx.poisson_solver = petsc, the Poisson equation is solved by one + // of PETSc's Krylov solvers, preconditioned by MLMG + m_petsc_options.use_petsc_ksp = + (WarpX::poisson_solver_id == PoissonSolverAlgo::PETSc); + if (m_petsc_options.use_petsc_ksp) { + ParmParse const pp_petsc("petsc_poisson"); + pp_petsc.query("ksp_type", m_petsc_options.ksp_type); + pp_petsc.query("use_mlmg_preconditioner", m_petsc_options.use_mlmg_preconditioner); + utils::parser::queryWithParser( + pp_petsc, "restart_length", m_petsc_options.restart_length); + utils::parser::queryWithParser( + pp_petsc, "precond_num_iters", m_petsc_options.precond_num_iters); + m_petsc_options.verbosity = self_fields_verbosity; + utils::parser::queryWithParser(pp_petsc, "verbose", m_petsc_options.verbosity); + } + // FFT solver flags utils::parser::queryWithParser( pp_warpx, "use_2d_slices_fft_solver", is_igf_2d_slices); @@ -221,6 +238,7 @@ ElectrostaticSolver::computePhi ( WarpX::do_single_precision_comms, warpx.refRatio(), self_fields_num_final_sweeps, + m_petsc_options, post_phi_calculation, *m_poisson_boundary_handler, warpx.gett_new(0), diff --git a/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.cpp b/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.cpp index 1e3697a12ee..6e7291eef86 100644 --- a/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.cpp +++ b/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.cpp @@ -81,7 +81,9 @@ void PoissonBoundaryHandler::DefinePhiBCs (const amrex::Geometry& geom) amrex::ignore_unused(geom); #endif for (int idim=dim_start; idim +#include +#include +#include +#include +#include +#include + +#include +#include + + +namespace ablastr::fields { + +/** Settings of the PETSc Krylov solver that can be used for the Poisson equation + * + * This is the electrostatic counterpart of the PETSc interface that the implicit + * electromagnetic solvers use for the curl-curl equation (@see PETScKSP in + * `Source/NonlinearSolvers/PETScKSP_Wrapper.H`): the linear system is handed to a + * PETSc `KSP` object as a matrix-free operator, and the AMReX geometric multigrid + * is used as a preconditioner through PETSc's `PCShell` interface. + */ +struct PETScPoissonOptions +{ + /** Solve the Poisson equation with a PETSc Krylov solver instead of MLMG */ + bool use_petsc_ksp = false; + + /** PETSc `KSP` type, e.g. "gmres", "fgmres", "bcgs", "cg" */ + std::string ksp_type = "gmres"; + + /** Restart length, only used by the GMRES variants (PETSc default: 30) */ + int restart_length = 30; + + /** Use MLMG V-cycles as the preconditioner (through PETSc's `PCShell`). + * When false, the Krylov solver runs unpreconditioned, which is usually + * much slower and is mostly useful for debugging. + */ + bool use_mlmg_preconditioner = true; + + /** Number of MLMG V-cycles per application of the preconditioner */ + int precond_num_iters = 1; + + /** Verbosity: 0 is silent, 1 prints the exit status of every solve, + * 2 additionally prints the residual at every Krylov iteration. + */ + int verbosity = 0; +}; + +#ifdef AMREX_USE_PETSC + +// The PETSc objects are wrapped in these structs, which are only defined in +// PETScPoissonSolver.cpp, so that the PETSc headers do not have to be included +// here (they conflict with some of the WarpX and AMReX names). +namespace petsc_poisson +{ + struct KSPObj; + struct MatObj; + struct VecObj; +} + +/** \brief Solve the Poisson equation of one MR level with a PETSc Krylov solver + * + * The linear system that is solved is the very same one that `amrex::MLMG` would + * solve, i.e. it is defined by the `amrex::MLNodeLinOp` that the `amrex::MLMG` + * object passed to the constructor was built from. The operator is never + * assembled into a sparse matrix: it is passed to PETSc as a `MatShell` whose + * action is computed by the AMReX linear operator, and the AMReX multigrid + * V-cycles are passed to PETSc as a `PCShell`. This mirrors `amrex::GMRESMLMG`, + * with PETSc's Krylov solvers in place of the AMReX GMRES implementation. + * + * As in `amrex::GMRESMLMG`, the Krylov solver does not iterate on the potential + * itself but on a correction: the residual of the initial guess is computed with + * the inhomogeneous operator, so that non-zero Dirichlet boundary values stored + * in `phi` are accounted for, and the correction that the Krylov solver computes + * vanishes on the Dirichlet nodes. + * + * The degrees of freedom of the PETSc vectors are the nodes that are owned by a + * box (nodes shared between boxes, or with a periodic image, appear only once) + * and that are not Dirichlet nodes. + */ +class PETScPoissonSolver +{ +public: + + /** Construct the solver and create the PETSc objects + * + * \param[in] mlmg the multigrid solver that defines the linear operator; it + * is used both for the operator and for the preconditioner, + * and must outlive this object + * \param[in] phi_prototype a nodal MultiFab with the layout of the unknowns + * \param[in] geom the geometry of the MR level that is solved + * \param[in] options the settings of the PETSc solver + */ + PETScPoissonSolver (amrex::MLMG & mlmg, + amrex::MultiFab const & phi_prototype, + amrex::Geometry const & geom, + PETScPoissonOptions const & options); + + ~PETScPoissonSolver (); + + // Prohibit Move and Copy operations + PETScPoissonSolver (PETScPoissonSolver const &) = delete; + PETScPoissonSolver & operator= (PETScPoissonSolver const &) = delete; + PETScPoissonSolver (PETScPoissonSolver &&) noexcept = delete; + PETScPoissonSolver & operator= (PETScPoissonSolver &&) noexcept = delete; + + /** Solve the Poisson equation + * + * \param[inout] phi on input the initial guess, which must hold the Dirichlet + * boundary values; on output the computed potential + * \param[in] rho the (already scaled) right-hand side + * \param[in] relative_tolerance the relative convergence threshold + * \param[in] absolute_tolerance the absolute convergence threshold + * \param[in] max_iters the maximum number of Krylov iterations + */ + void solve (amrex::MultiFab & phi, + amrex::MultiFab const & rho, + amrex::Real relative_tolerance, + amrex::Real absolute_tolerance, + int max_iters); + + //! Number of Krylov iterations of the most recent solve + [[nodiscard]] int getNumIters () const { return m_num_iters; } + + //! Residual norm of the most recent solve + [[nodiscard]] amrex::Real getResidualNorm () const { return m_residual_norm; } + + /** Apply the linear operator: `out = L(in)` + * + * This is called back by PETSc, with `in` and `out` pointing to the local + * part of the PETSc vectors. + */ + void applyOperator (amrex::Real * out, amrex::Real const * in); + + /** Apply the multigrid preconditioner: `out = P^{-1}(in)` + * + * This is called back by PETSc, with `in` and `out` pointing to the local + * part of the PETSc vectors. + */ + void applyPreconditioner (amrex::Real * out, amrex::Real const * in); + +private: + + //! Number the degrees of freedom that this MPI rank owns + void buildDOFMap (amrex::MultiFab const & phi_prototype); + + //! Gather the degrees of freedom of `mf` into the PETSc array `arr` + void copyToArray (amrex::MultiFab const & mf, amrex::Real * arr) const; + + //! Scatter the PETSc array `arr` into `mf`, and make `mf` consistent + void copyFromArray (amrex::MultiFab & mf, amrex::Real const * arr) const; + + //! Multigrid solver that provides the operator and the preconditioner + amrex::MLMG * m_mlmg = nullptr; + //! Geometry of the MR level that is solved + amrex::Geometry m_geom; + //! Settings of the PETSc solver + PETScPoissonOptions m_options; + + //! Local index of the degree of freedom of each node, -1 if it is not one + std::unique_ptr m_dof; + //! Number of degrees of freedom owned by this MPI rank + amrex::Long m_ndofs_local = 0; + //! Total number of degrees of freedom + amrex::Long m_ndofs_global = 0; + + //! Work arrays for the operator (`_in` also holds the ghost nodes) + amrex::MultiFab m_op_in, m_op_out; + //! Work arrays for the preconditioner + amrex::MultiFab m_pc_in, m_pc_out; + //! Residual of the initial guess, and the correction that is solved for + amrex::MultiFab m_res, m_cor; + + //! Matrix-free linear operator + std::unique_ptr m_A; + //! Solution vector + std::unique_ptr m_x; + //! Right-hand-side vector + std::unique_ptr m_b; + //! Krylov solver + std::unique_ptr m_ksp; + + //! Number of iterations of the most recent solve + int m_num_iters = -1; + //! Residual norm of the most recent solve + amrex::Real m_residual_norm = amrex::Real(-1.0); +}; + +/** Solve the Poisson equation of one MR level with a PETSc Krylov solver + * + * This is a convenience wrapper that creates a @see PETScPoissonSolver, solves, + * and destroys it again. The PETSc objects are cheap to create compared to the + * `amrex::MLMG` object that `mlmg` refers to, which `computePhi` also recreates + * at every call. + * + * \param[in] mlmg the multigrid solver that defines the linear operator + * \param[inout] phi on input the initial guess (holding the Dirichlet boundary + * values), on output the computed potential + * \param[in] rho the (already scaled) right-hand side + * \param[in] geom the geometry of the MR level that is solved + * \param[in] relative_tolerance the relative convergence threshold + * \param[in] absolute_tolerance the absolute convergence threshold + * \param[in] max_iters the maximum number of Krylov iterations + * \param[in] options the settings of the PETSc solver + */ +void +petscPoissonSolve (amrex::MLMG & mlmg, + amrex::MultiFab & phi, + amrex::MultiFab const & rho, + amrex::Geometry const & geom, + amrex::Real relative_tolerance, + amrex::Real absolute_tolerance, + int max_iters, + PETScPoissonOptions const & options); + +#endif // AMREX_USE_PETSC + +} // namespace ablastr::fields + +#endif // ABLASTR_FIELDS_PETSC_POISSON_SOLVER_H diff --git a/Source/ablastr/fields/PETScPoissonSolver.cpp b/Source/ablastr/fields/PETScPoissonSolver.cpp new file mode 100644 index 00000000000..4570074a875 --- /dev/null +++ b/Source/ablastr/fields/PETScPoissonSolver.cpp @@ -0,0 +1,493 @@ +/* Copyright 2026 The WarpX Community + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include + +#ifdef AMREX_USE_PETSC + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +// The PETSc headers must be included before PETScPoissonSolver.H, see the +// comment in Source/NonlinearSolvers/WarpX_PETSc.cpp +#include +#include +#include +#include + +#include + + +namespace ablastr::fields { + +namespace petsc_poisson { + +//! Wrapper for a PETSc KSP object +struct KSPObj +{ + KSPObj () = default; + ~KSPObj () { if (obj) { KSPDestroy(&obj); } } + KSPObj (KSPObj const &) = delete; + KSPObj (KSPObj &&) = delete; + KSPObj & operator= (KSPObj const &) = delete; + KSPObj & operator= (KSPObj &&) = delete; + KSP obj = nullptr; +}; + +//! Wrapper for a PETSc Mat object +struct MatObj +{ + MatObj () = default; + ~MatObj () { if (obj) { MatDestroy(&obj); } } + MatObj (MatObj const &) = delete; + MatObj (MatObj &&) = delete; + MatObj & operator= (MatObj const &) = delete; + MatObj & operator= (MatObj &&) = delete; + Mat obj = nullptr; +}; + +//! Wrapper for a PETSc Vec object +struct VecObj +{ + VecObj () = default; + ~VecObj () { if (obj) { VecDestroy(&obj); } } + VecObj (VecObj const &) = delete; + VecObj (VecObj &&) = delete; + VecObj & operator= (VecObj const &) = delete; + VecObj & operator= (VecObj &&) = delete; + Vec obj = nullptr; +}; + +//! Apply the matrix-free linear operator, called back by PETSc +PetscErrorCode applyOperator (Mat a_A, Vec a_in, Vec a_out) +{ + PetscFunctionBeginUser; + + PETScPoissonSolver * solver = nullptr; + PetscCall(MatShellGetContext(a_A, &solver)); + + PetscScalar const * in_arr = nullptr; + PetscScalar * out_arr = nullptr; + PetscCall(VecGetArrayRead(a_in, &in_arr)); + PetscCall(VecGetArrayWrite(a_out, &out_arr)); + + solver->applyOperator( static_cast(out_arr), + static_cast(in_arr) ); + + PetscCall(VecRestoreArrayWrite(a_out, &out_arr)); + PetscCall(VecRestoreArrayRead(a_in, &in_arr)); + + PetscFunctionReturn(PETSC_SUCCESS); +} + +//! Apply the multigrid preconditioner, called back by PETSc +PetscErrorCode applyPreconditioner (PC a_pc, Vec a_in, Vec a_out) +{ + PetscFunctionBeginUser; + + PETScPoissonSolver * solver = nullptr; + PetscCall(PCShellGetContext(a_pc, &solver)); + + PetscScalar const * in_arr = nullptr; + PetscScalar * out_arr = nullptr; + PetscCall(VecGetArrayRead(a_in, &in_arr)); + PetscCall(VecGetArrayWrite(a_out, &out_arr)); + + solver->applyPreconditioner( static_cast(out_arr), + static_cast(in_arr) ); + + PetscCall(VecRestoreArrayWrite(a_out, &out_arr)); + PetscCall(VecRestoreArrayRead(a_in, &in_arr)); + + PetscFunctionReturn(PETSC_SUCCESS); +} + +//! Print the residual of every Krylov iteration +PetscErrorCode printResidual (KSP a_ksp, PetscInt a_n, PetscReal a_rnorm, void * a_ctxt) +{ + PetscFunctionBeginUser; + amrex::ignore_unused(a_ksp, a_ctxt); + amrex::Print() << "Poisson (PETSc KSP): iter = " << a_n + << ", residual = " << a_rnorm << "\n"; + PetscFunctionReturn(PETSC_SUCCESS); +} + +//! Is `a_type` one of the GMRES variants of PETSc? +bool isGMRES (std::string const & a_type) +{ + return (a_type == "gmres") || (a_type == "fgmres") + || (a_type == "lgmres") || (a_type == "dgmres") + || (a_type == "pgmres") || (a_type == "pipefgmres"); +} + +} // namespace petsc_poisson + + +PETScPoissonSolver::PETScPoissonSolver (amrex::MLMG & mlmg, + amrex::MultiFab const & phi_prototype, + amrex::Geometry const & geom, + PETScPoissonOptions const & options) + : m_mlmg(&mlmg), m_geom(geom), m_options(options) +{ + ABLASTR_PROFILE("PETScPoissonSolver::PETScPoissonSolver()"); + + // This builds the multigrid hierarchy and the masks of the linear operator, + // which the operator, the preconditioner and buildDOFMap() below all need. + m_mlmg->preparePrecond(); + + buildDOFMap(phi_prototype); + + // The work arrays are created by the linear operator itself, so that they + // have the right layout and factory. The inputs of the operator and of the + // preconditioner need one layer of ghost nodes, as in + // amrex::GMRESMLMG::makeVecLHS(). + auto & linop = m_mlmg->getLinOp(); + m_op_in = linop.make(0, 0, amrex::IntVect(1)); + m_op_out = linop.make(0, 0, amrex::IntVect(0)); + m_pc_in = linop.make(0, 0, amrex::IntVect(1)); + m_pc_out = linop.make(0, 0, amrex::IntVect(1)); + m_res = linop.make(0, 0, amrex::IntVect(1)); + m_cor = linop.make(0, 0, amrex::IntVect(1)); + + m_A = std::make_unique(); + m_x = std::make_unique(); + m_b = std::make_unique(); + m_ksp = std::make_unique(); + + // Vectors + VecCreate(PETSC_COMM_WORLD, &m_x->obj); +#ifdef AMREX_USE_GPU +# if defined(AMREX_USE_CUDA) + VecSetType(m_x->obj, VECCUDA); +# elif defined(AMREX_USE_HIP) + VecSetType(m_x->obj, VECHIP); +# else + ABLASTR_ABORT_WITH_MESSAGE( + "The PETSc Poisson solver is not yet implemented for non-CUDA/HIP GPUs"); +# endif +#else + VecSetType(m_x->obj, VECSTANDARD); +#endif + auto const ndofs_local = static_cast(m_ndofs_local); + auto const ndofs_global = static_cast(m_ndofs_global); + VecSetSizes(m_x->obj, ndofs_local, ndofs_global); + VecSetFromOptions(m_x->obj); + VecDuplicate(m_x->obj, &m_b->obj); + + // Matrix-free linear operator + MatCreateShell( PETSC_COMM_WORLD, + ndofs_local, ndofs_local, + ndofs_global, ndofs_global, + this, &m_A->obj ); + MatShellSetOperation( m_A->obj, MATOP_MULT, + (void(*)())petsc_poisson::applyOperator ); // NOLINT + MatSetUp(m_A->obj); + + // Krylov solver + KSPCreate(PETSC_COMM_WORLD, &m_ksp->obj); + KSPSetType(m_ksp->obj, m_options.ksp_type.c_str()); + KSPSetOperators(m_ksp->obj, m_A->obj, m_A->obj); + if (petsc_poisson::isGMRES(m_options.ksp_type)) { + KSPGMRESSetRestart(m_ksp->obj, m_options.restart_length); + // Right preconditioning, so that the residual that PETSc monitors and + // uses for its convergence test is the residual of the actual system + KSPSetPCSide(m_ksp->obj, PC_RIGHT); + KSPSetNormType(m_ksp->obj, KSP_NORM_UNPRECONDITIONED); + } + + PC pc = nullptr; + KSPGetPC(m_ksp->obj, &pc); + if (m_options.use_mlmg_preconditioner) { + PCSetType(pc, PCSHELL); + PCShellSetApply(pc, petsc_poisson::applyPreconditioner); + PCShellSetContext(pc, this); + PCShellSetName(pc, "AMReX MLMG"); + } else { + PCSetType(pc, PCNONE); + } + + if (m_options.verbosity > 1) { + KSPMonitorSet(m_ksp->obj, petsc_poisson::printResidual, nullptr, nullptr); + } + // Command-line and input-file PETSc options take precedence over the above + KSPSetFromOptions(m_ksp->obj); + + if (m_options.verbosity > 0) { + amrex::Print() << "PETScPoissonSolver: using PETSc's KSP (" << m_options.ksp_type + << ") with " + << (m_options.use_mlmg_preconditioner ? "the AMReX MLMG" : "no") + << " preconditioner (total DOFs = " << m_ndofs_global << ").\n"; + } +} + +PETScPoissonSolver::~PETScPoissonSolver () = default; + +void PETScPoissonSolver::buildDOFMap (amrex::MultiFab const & phi_prototype) +{ + ABLASTR_PROFILE("PETScPoissonSolver::buildDOFMap()"); + + using namespace amrex::literals; + + // The nodes that sit on the boundary between two boxes (or on the boundary + // between a box and the periodic image of another one) belong to the valid + // region of both boxes, but they are a single unknown of the linear system: + // only the node of the "owner" box is a degree of freedom. + auto const owner_mask = amrex::OwnerMask(phi_prototype, m_geom.periodicity()); + + // The nodes on which a Dirichlet boundary condition is applied are not + // unknowns of the linear system either. `setDirichletNodesToZero` is the + // public interface through which the AMReX linear operator exposes them. + auto & linop = m_mlmg->getLinOp(); + amrex::MultiFab dirichlet_indicator = linop.make(0, 0, amrex::IntVect(0)); + dirichlet_indicator.setVal(1._rt); + linop.setDirichletNodesToZero(0, 0, dirichlet_indicator); + + m_dof = std::make_unique(phi_prototype.boxArray(), + phi_prototype.DistributionMap(), 1, 0); + m_dof->setVal(-1); + + m_ndofs_local = 0; + for (amrex::MFIter mfi(*m_dof); mfi.isValid(); ++mfi) + { + amrex::Box const & bx = mfi.validbox(); + auto const npts = static_cast(bx.numPts()); + amrex::BoxIndexer const box_indexer(bx); + + auto const & owner_arr = owner_mask->const_array(mfi); + auto const & dirichlet_arr = dirichlet_indicator.const_array(mfi); + auto const & dof_arr = m_dof->array(mfi); + auto const first_dof = static_cast(m_ndofs_local); + + auto const ndofs = amrex::Scan::PrefixSum( + npts, + [=] AMREX_GPU_DEVICE (int offset) -> int + { + auto const [i,j,k] = box_indexer(offset); + return (owner_arr(i,j,k) && (dirichlet_arr(i,j,k) > 0.5_rt)) ? 1 : 0; + }, + [=] AMREX_GPU_DEVICE (int offset, int ps) + { + auto const [i,j,k] = box_indexer(offset); + if (owner_arr(i,j,k) && (dirichlet_arr(i,j,k) > 0.5_rt)) { + dof_arr(i,j,k) = ps + first_dof; + } + }, + amrex::Scan::Type::exclusive, amrex::Scan::retSum); + + m_ndofs_local += ndofs; + } + + m_ndofs_global = m_ndofs_local; + amrex::ParallelDescriptor::ReduceLongSum(m_ndofs_global); + + ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(m_ndofs_global > 0, + "PETScPoissonSolver: the linear system has no degree of freedom"); +} + +void PETScPoissonSolver::copyToArray (amrex::MultiFab const & mf, amrex::Real * arr) const +{ + ABLASTR_PROFILE("PETScPoissonSolver::copyToArray()"); + + for (amrex::MFIter mfi(*m_dof); mfi.isValid(); ++mfi) + { + amrex::Box const & bx = mfi.validbox(); + auto const & mf_arr = mf.const_array(mfi); + auto const & dof_arr = m_dof->const_array(mfi); + amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + int const dof = dof_arr(i,j,k); + if (dof >= 0) { arr[dof] = mf_arr(i,j,k); } + }); + } + amrex::Gpu::streamSynchronize(); +} + +void PETScPoissonSolver::copyFromArray (amrex::MultiFab & mf, amrex::Real const * arr) const +{ + ABLASTR_PROFILE("PETScPoissonSolver::copyFromArray()"); + + using namespace amrex::literals; + + // The nodes that are not degrees of freedom (Dirichlet nodes, and the nodes + // that another box owns) are set to zero here, and the ones that another box + // owns are then filled from their owner by `OverrideSync` below. + mf.setVal(0._rt); + + for (amrex::MFIter mfi(*m_dof); mfi.isValid(); ++mfi) + { + amrex::Box const & bx = mfi.validbox(); + auto const & mf_arr = mf.array(mfi); + auto const & dof_arr = m_dof->const_array(mfi); + amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + int const dof = dof_arr(i,j,k); + if (dof >= 0) { mf_arr(i,j,k) = arr[dof]; } + }); + } + amrex::Gpu::streamSynchronize(); + + // `OverrideSync` uses the same `amrex::OwnerMask` as `buildDOFMap` above, + // so the nodes that are shared between boxes are filled from the very box + // whose node was numbered as a degree of freedom. + mf.OverrideSync(m_geom.periodicity()); + mf.FillBoundary(m_geom.periodicity()); +} + +void PETScPoissonSolver::applyOperator (amrex::Real * out, amrex::Real const * in) +{ + ABLASTR_PROFILE("PETScPoissonSolver::applyOperator()"); + + copyFromArray(m_op_in, in); + // `applyPrecond` applies the operator with homogeneous boundary conditions, + // which is the operator that the correction equation solved here uses + m_mlmg->applyPrecond({&m_op_out}, {&m_op_in}); + m_mlmg->getLinOp().setDirichletNodesToZero(0, 0, m_op_out); + copyToArray(m_op_out, out); +} + +void PETScPoissonSolver::applyPreconditioner (amrex::Real * out, amrex::Real const * in) +{ + ABLASTR_PROFILE("PETScPoissonSolver::applyPreconditioner()"); + + using namespace amrex::literals; + + copyFromArray(m_pc_in, in); + m_mlmg->setPrecondIter(m_options.precond_num_iters); + m_pc_out.setVal(0._rt); + m_mlmg->precond({&m_pc_out}, {&m_pc_in}, 0._rt, 0._rt); + copyToArray(m_pc_out, out); +} + +void PETScPoissonSolver::solve (amrex::MultiFab & phi, + amrex::MultiFab const & rho, + amrex::Real relative_tolerance, + amrex::Real absolute_tolerance, + int max_iters) +{ + ABLASTR_PROFILE("PETScPoissonSolver::solve()"); + + using namespace amrex::literals; + + auto & linop = m_mlmg->getLinOp(); + + // MLMG is only used as a preconditioner here, so its bottom solve must be + // cheap and linear; this mirrors what amrex::GMRESMLMG does. + auto const bottom_solver = m_mlmg->getBottomSolver(); + auto const mlmg_verbose = m_mlmg->getVerbose(); + auto const mlmg_bottom_verbose = m_mlmg->getBottomVerbose(); + if (bottom_solver != amrex::BottomSolver::smoother && + bottom_solver != amrex::BottomSolver::hypre && + bottom_solver != amrex::BottomSolver::petsc) + { + m_mlmg->setBottomSolver(amrex::BottomSolver::smoother); + } + m_mlmg->setVerbose(0); + m_mlmg->setBottomVerbose(0); + + // Residual of the initial guess: res = L(phi) - rho. Note that `apply` uses + // the inhomogeneous operator, so that the Dirichlet values that `phi` holds + // contribute to the residual. + m_res.setVal(0._rt); + m_mlmg->apply({&m_res}, {&phi}); + + amrex::MultiFab scaled_rho; + amrex::MultiFab const * rhs = ρ + if (linop.scaleRHS(0, nullptr)) { + scaled_rho.define(rho.boxArray(), rho.DistributionMap(), 1, 0); + amrex::MultiFab::Copy(scaled_rho, rho, 0, 0, 1, 0); + auto const scaled = linop.scaleRHS(0, &scaled_rho); + amrex::ignore_unused(scaled); + rhs = &scaled_rho; + } + amrex::MultiFab::Saxpy(m_res, -1._rt, *rhs, 0, 0, 1, amrex::IntVect(0)); + linop.setDirichletNodesToZero(0, 0, m_res); + + // Solve L(cor) = res for the correction, with PETSc's Krylov solver + { + PetscScalar * b_arr = nullptr; + VecGetArrayWrite(m_b->obj, &b_arr); + copyToArray(m_res, static_cast(b_arr)); + VecRestoreArrayWrite(m_b->obj, &b_arr); + } + VecZeroEntries(m_x->obj); + + KSPSetTolerances( m_ksp->obj, + relative_tolerance, + absolute_tolerance, + PETSC_CURRENT, + (max_iters > 0 ? max_iters : PETSC_CURRENT) ); + KSPSolve(m_ksp->obj, m_b->obj, m_x->obj); + + { + PetscScalar const * x_arr = nullptr; + VecGetArrayRead(m_x->obj, &x_arr); + copyFromArray(m_cor, static_cast(x_arr)); + VecRestoreArrayRead(m_x->obj, &x_arr); + } + + // phi = phi - cor + amrex::MultiFab::Saxpy(phi, -1._rt, m_cor, 0, 0, 1, amrex::IntVect(0)); + phi.FillBoundary(m_geom.periodicity()); + + // Report on the solve + PetscInt niters = -1; + KSPGetIterationNumber(m_ksp->obj, &niters); + m_num_iters = static_cast(niters); + PetscReal norm = -1; + KSPGetResidualNorm(m_ksp->obj, &norm); + m_residual_norm = static_cast(norm); + + KSPConvergedReason reason; + KSPGetConvergedReason(m_ksp->obj, &reason); + char const * reason_string = nullptr; + KSPGetConvergedReasonString(m_ksp->obj, &reason_string); + + if (m_options.verbosity > 0) { + amrex::Print() << "Poisson (PETSc KSP): " << m_num_iters << " iterations, exited due to \"" + << reason_string << "\" (abs. norm = " << m_residual_norm << ").\n"; + } + ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(reason > 0, + std::string("The PETSc Poisson solver failed to converge: ") + reason_string); + + // Restore the settings of the multigrid solver + m_mlmg->setBottomSolver(bottom_solver); + m_mlmg->setVerbose(mlmg_verbose); + m_mlmg->setBottomVerbose(mlmg_bottom_verbose); +} + +void +petscPoissonSolve (amrex::MLMG & mlmg, + amrex::MultiFab & phi, + amrex::MultiFab const & rho, + amrex::Geometry const & geom, + amrex::Real relative_tolerance, + amrex::Real absolute_tolerance, + int max_iters, + PETScPoissonOptions const & options) +{ + PETScPoissonSolver solver(mlmg, phi, geom, options); + solver.solve(phi, rho, relative_tolerance, absolute_tolerance, max_iters); +} + +} // namespace ablastr::fields + +#endif // AMREX_USE_PETSC diff --git a/Source/ablastr/fields/PoissonSolver.H b/Source/ablastr/fields/PoissonSolver.H index db9687e582b..fb3cff25afc 100755 --- a/Source/ablastr/fields/PoissonSolver.H +++ b/Source/ablastr/fields/PoissonSolver.H @@ -15,6 +15,7 @@ #include #include #include +#include #include #if defined(ABLASTR_USE_FFT) && defined(WARPX_DIM_3D) @@ -185,6 +186,7 @@ inline void interpolatePhiBetweenLevels ( * \param[in] rel_ref_ratio mesh refinement ratio between levels (default: 1) * \param[in] num_final_sweeps Optional MLMG final smoothing count. If set, it is used for final smoothing. * Otherwise, the default AMReX MLMG value (8) is used. + * \param[in] petsc_options Settings of the optional PETSc Krylov solver (default: use MLMG) * \param[in] post_phi_calculation perform a calculation per level directly after phi was calculated; required for embedded boundaries (default: none) * \param[in] boundary_handler a handler for boundary conditions, for example @see ElectrostaticSolver::PoissonBoundaryHandler * \param[in] current_time the current time; required for embedded boundaries (default: none) @@ -214,6 +216,7 @@ computePhi ( bool do_single_precision_comms = false, std::optional > rel_ref_ratio = std::nullopt, std::optional num_final_sweeps = std::nullopt, + [[maybe_unused]] PETScPoissonOptions const & petsc_options = PETScPoissonOptions{}, [[maybe_unused]] T_PostPhiCalculationFunctor post_phi_calculation = std::nullopt, [[maybe_unused]] T_BoundaryHandler const& boundary_handler = std::nullopt, [[maybe_unused]] std::optional current_time = std::nullopt, // only used for EB @@ -422,8 +425,22 @@ computePhi ( rho[lev]->OverrideSync(geom[lev].periodicity()); // Solve Poisson equation at lev - mlmg.solve( {phi[lev]}, {rho[lev]}, - relative_tolerance, absolute_tolerance ); +#ifdef AMREX_USE_PETSC + if (petsc_options.use_petsc_ksp) { + // The PETSc solver only computes `phi`, whereas the operations below + // (currently only the calculation of the E field with embedded + // boundaries) need the internal state of a completed MLMG solve. + ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(!eb_enabled, + "The PETSc Poisson solver does not support embedded boundaries yet"); + petscPoissonSolve( mlmg, *phi[lev], *rho[lev], geom[lev], + relative_tolerance, absolute_tolerance, + max_iters, petsc_options ); + } else +#endif + { + mlmg.solve( {phi[lev]}, {rho[lev]}, + relative_tolerance, absolute_tolerance ); + } const amrex::IntVect& refratio = rel_ref_ratio.value()[lev]; const int ncomp = linop->getNComp(); From c6169c693cd15cd1ff49b54f1c9b854e5c4f7177 Mon Sep 17 00:00:00 2001 From: Remi Lehe Date: Thu, 27 Aug 2026 15:23:26 -0700 Subject: [PATCH 2/5] PETSc Poisson solver: support embedded boundaries With embedded boundaries, WarpX does not compute E from phi itself but asks AMReX to do it, through the `EBCalcEfromPhiPerLevel` functor that `computePhi` calls after the solve. That functor used `amrex::MLMG::getGradSolution`, which reads the potential from the state that `amrex::MLMG::solve` leaves behind, so it could not be used after a solve that bypasses MLMG. `EBCalcEfromPhiPerLevel` gains an overload that computes the gradient directly from the linear operator and `phi`, via the public `MLNodeLinOp::applyBC` and `MLLinOp::compGrad`. The MLMG path keeps calling the original overload, so its behaviour is unchanged. The PETSc solver additionally calls `MLLinOp::postSolve`, which is what writes the prescribed potential into the nodes covered by the embedded boundary at the end of an MLMG solve. No change was needed for the degrees of freedom: `MLEBNodeFDLaplacian` marks the EB-covered nodes as Dirichlet (with a negative value) in its Dirichlet mask, so the `setDirichletNodesToZero` probe that builds the DOF map already excludes them. Co-Authored-By: Claude Opus 5 --- Docs/source/usage/parameters.rst | 2 +- .../electrostatic_sphere_eb/CMakeLists.txt | 15 +++++++++ ...puts_test_3d_electrostatic_sphere_eb_petsc | 32 +++++++++++++++++++ .../PoissonBoundaryHandler.H | 27 ++++++++++++++++ Source/ablastr/fields/PETScPoissonSolver.cpp | 5 +++ Source/ablastr/fields/PoissonSolver.H | 19 +++++++---- 6 files changed, 92 insertions(+), 8 deletions(-) create mode 100644 Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst index c4ea9c305e9..3ef7e690f46 100644 --- a/Docs/source/usage/parameters.rst +++ b/Docs/source/usage/parameters.rst @@ -457,7 +457,7 @@ Overall simulation parameters is computed by AMReX. It therefore discretizes Poisson's equation exactly like ``multigrid`` does, and accepts the same boundary conditions; only the outer iteration differs. It requires the compilation flag ``-DWarpX_PETSC=ON``. - It is not supported with embedded boundaries, nor in ``labframe-effective-potential`` mode. + It is not supported in ``labframe-effective-potential`` mode. Note that in 1D with ``warpx.do_electrostatic = labframe``, and with the ``poissonsolver`` Python callback, Poisson's equation is solved by a dedicated solver and this option has no effect. diff --git a/Examples/Tests/electrostatic_sphere_eb/CMakeLists.txt b/Examples/Tests/electrostatic_sphere_eb/CMakeLists.txt index 2be46f2e0e6..e9ef1f6ac87 100644 --- a/Examples/Tests/electrostatic_sphere_eb/CMakeLists.txt +++ b/Examples/Tests/electrostatic_sphere_eb/CMakeLists.txt @@ -25,6 +25,21 @@ if(WarpX_EB) ) endif() +if(WarpX_EB AND WarpX_PETSC) + # No checksum: the reference values still have to be generated on a build + # with PETSc. Same setup as test_3d_electrostatic_sphere_eb_mixed_bc, but + # solved with PETSc, so the two can be compared against each other. + add_warpx_test( + test_3d_electrostatic_sphere_eb_petsc # name + 3 # dims + 2 # nprocs + inputs_test_3d_electrostatic_sphere_eb_petsc # inputs + OFF # analysis + OFF # checksum + OFF # dependency + ) +endif() + if(WarpX_EB) add_warpx_test( test_3d_electrostatic_sphere_eb_picmi # name diff --git a/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc b/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc new file mode 100644 index 00000000000..f79ea54b525 --- /dev/null +++ b/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc @@ -0,0 +1,32 @@ +max_step = 1 +amr.n_cell = 64 64 64 +amr.max_level = 0 +amr.blocking_factor = 8 +amr.max_grid_size = 128 +geometry.dims = 3 +boundary.field_lo = pec pec neumann +boundary.field_hi = pec neumann neumann +boundary.potential_lo_x = 0 +boundary.potential_hi_x = 0 +boundary.potential_lo_y = 0 +geometry.prob_lo = -0.5 -0.5 -0.5 +geometry.prob_hi = 0.5 0.5 0.5 +warpx.const_dt = 1e-6 + +warpx.do_electrostatic = labframe +warpx.eb_implicit_function = "-(x**2+y**2+z**2-0.3**2)" +warpx.eb_potential(x,y,z,t) = "1." +warpx.self_fields_required_precision = 1.e-7 +warpx.abort_on_warning_threshold = medium + +# solve the Poisson equation with PETSc's GMRES, preconditioned by the AMReX +# multigrid. This exercises the embedded boundary (with a non-zero potential) +# and mixed PEC/Neumann domain boundaries. +warpx.poisson_solver = petsc + +algo.field_gathering = momentum-conserving + +diagnostics.diags_names = diag1 +diag1.intervals = 1 +diag1.diag_type = Full +diag1.fields_to_plot = Ex Ey Ez rho phi diff --git a/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.H b/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.H index 5d608354097..119501c5008 100644 --- a/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.H +++ b/Source/FieldSolver/ElectrostaticSolvers/PoissonBoundaryHandler.H @@ -14,7 +14,9 @@ #include #include #include +#include #include +#include #include #include @@ -136,6 +138,31 @@ class EBCalcEfromPhiPerLevel { field->mult(-1._rt); } } + + /** \brief Same, but computed directly from the linear operator and `phi` + * + * `amrex::MLMG::getGradSolution` takes the potential from the internal state + * that `amrex::MLMG::solve` leaves behind, so it can only be used right after + * an MLMG solve. This overload is for the solvers that compute `phi` without + * going through `amrex::MLMG::solve`, such as the PETSc Poisson solver. + * + * \param[in] linop the linear operator of the Poisson solve + * \param[inout] phi the potential; only its ghost nodes are modified + * \param[in] lev the mesh refinement level + */ + void operator()(amrex::MLNodeLinOp & linop, amrex::MultiFab & phi, int const lev) { + using namespace amrex::literals; + + // Fill the ghost nodes of `phi`, in particular the ones that lie outside + // of a Neumann boundary, which `compGrad` reads. `amrex::MLMG::solve` + // does this internally on the potential that it hands to `compGrad`. + linop.applyBC(0, 0, phi, amrex::MLNodeLinOp::BCMode::Inhomogeneous, + amrex::MLNodeLinOp::StateMode::Solution); + linop.compGrad(0, m_e_field[lev], phi, amrex::MLNodeLinOp::Location::FaceCenter); + for (auto &field: m_e_field[lev]) { + field->mult(-1._rt); + } + } }; #endif // WARPX_BOUNDARYHANDLER_H_ diff --git a/Source/ablastr/fields/PETScPoissonSolver.cpp b/Source/ablastr/fields/PETScPoissonSolver.cpp index 4570074a875..9b27b2bc88d 100644 --- a/Source/ablastr/fields/PETScPoissonSolver.cpp +++ b/Source/ablastr/fields/PETScPoissonSolver.cpp @@ -446,6 +446,11 @@ void PETScPoissonSolver::solve (amrex::MultiFab & phi, // phi = phi - cor amrex::MultiFab::Saxpy(phi, -1._rt, m_cor, 0, 0, 1, amrex::IntVect(0)); + + // `amrex::MLMG::solve` ends with this; for the embedded-boundary operator it + // writes the prescribed potential into the nodes that the EB covers. + linop.postSolve({&phi}); + phi.FillBoundary(m_geom.periodicity()); // Report on the solve diff --git a/Source/ablastr/fields/PoissonSolver.H b/Source/ablastr/fields/PoissonSolver.H index fb3cff25afc..e54cd4cb0ac 100755 --- a/Source/ablastr/fields/PoissonSolver.H +++ b/Source/ablastr/fields/PoissonSolver.H @@ -426,12 +426,8 @@ computePhi ( // Solve Poisson equation at lev #ifdef AMREX_USE_PETSC - if (petsc_options.use_petsc_ksp) { - // The PETSc solver only computes `phi`, whereas the operations below - // (currently only the calculation of the E field with embedded - // boundaries) need the internal state of a completed MLMG solve. - ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(!eb_enabled, - "The PETSc Poisson solver does not support embedded boundaries yet"); + bool const use_petsc_ksp = petsc_options.use_petsc_ksp; + if (use_petsc_ksp) { petscPoissonSolve( mlmg, *phi[lev], *rho[lev], geom[lev], relative_tolerance, absolute_tolerance, max_iters, petsc_options ); @@ -463,7 +459,16 @@ computePhi ( // Run additional operations, such as calculation of the E field for embedded boundaries if constexpr (!std::is_same_v) { if (post_phi_calculation.has_value()) { - post_phi_calculation.value()(mlmg, lev); +#ifdef AMREX_USE_PETSC + if (use_petsc_ksp) { + // `mlmg` did not solve, so it holds no solution to work from: + // hand over the linear operator and `phi` instead + post_phi_calculation.value()(*linop, *phi[lev], lev); + } else +#endif + { + post_phi_calculation.value()(mlmg, lev); + } } } rho[lev]->mult(-ablastr::constant::SI::epsilon_0); // Multiply rho by epsilon again From b369f7abf690fc6e39b74af4352a35751444892c Mon Sep 17 00:00:00 2001 From: Remi Lehe Date: Thu, 27 Aug 2026 15:46:30 -0700 Subject: [PATCH 3/5] Simplify the PETSc Poisson solver Review pass to reduce the amount of code and make it easier to digest: - Collapse the PETScPoissonSolver class into the petscPoissonSolve free function, with a file-local context struct for the PETSc callbacks. The class was created and destroyed inside petscPoissonSolve at every call anyway, so nothing persisted; the public header shrinks to the options struct and one function declaration. - Keep Dirichlet nodes as degrees of freedom instead of excluding them from the DOF map, which is also how amrex::GMRESMLMG treats them. Both the right-hand side and the operator output are zeroed on those nodes, and GMRES only forms linear combinations, so they remain exactly zero throughout the solve. This removes the probe that extracted the Dirichlet mask through setDirichletNodesToZero, the subtlest part of the previous version. - Hardcode GMRES with right preconditioning, and drop the petsc_poisson.ksp_type and petsc_poisson.restart_length parameters (and the fragile list of GMRES variants that supported them). These duplicated PETSc's own runtime options (-ksp_type, -ksp_gmres_restart), which already take precedence since KSPSetFromOptions is called last. - Share two work MultiFabs between the operator and preconditioner callbacks and the residual/correction computations, instead of allocating six; the phases never overlap. - Replace the untested GPU code paths (VECCUDA/VECHIP) with an abort; no CI compiles PETSc with CUDA/HIP, and the VecGetArray semantics on device vectors deserve a dedicated, tested implementation. Also enable PETSc in the 3D CPU job of the Azure pipeline (the 1D, 2D and RZ jobs already build with it), so that the two new tests actually run in CI. Co-Authored-By: Claude Opus 5 --- .azure-pipelines.yml | 2 +- Docs/source/usage/parameters.rst | 15 +- ...st_3d_electrostatic_sphere_lab_frame_petsc | 2 - ...puts_test_3d_electrostatic_sphere_eb_petsc | 35 +- .../ElectrostaticSolver.cpp | 3 - Source/ablastr/fields/PETScPoissonSolver.H | 180 +----- Source/ablastr/fields/PETScPoissonSolver.cpp | 566 ++++++++---------- 7 files changed, 277 insertions(+), 526 deletions(-) diff --git a/.azure-pipelines.yml b/.azure-pipelines.yml index 39b0fa7a171..bd7123144cc 100644 --- a/.azure-pipelines.yml +++ b/.azure-pipelines.yml @@ -54,7 +54,7 @@ jobs: # Cartesian 3D cartesian_3d: CDASH_BUILD_NAME: CPU-3D - WARPX_CMAKE_FLAGS: -DWarpX_DIMS=3 -DWarpX_FFT=ON -DWarpX_PYTHON=ON + WARPX_CMAKE_FLAGS: -DWarpX_DIMS=3 -DWarpX_FFT=ON -DWarpX_PYTHON=ON -DWarpX_PETSC=ON # Cylindrical RZ cylindrical_rz: CDASH_BUILD_NAME: CPU-RZ diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst index 3ef7e690f46..06de15998d2 100644 --- a/Docs/source/usage/parameters.rst +++ b/Docs/source/usage/parameters.rst @@ -450,25 +450,20 @@ Overall simulation parameters The extended simulation box size in real space is :math:`2n_x-1, 2n_y-1, 2n_z-1` with the 3D solver, :math:`2n_x-1, 2n_y -1, n_z` with the 2D solver. The extended simulation box size in spectral space is :math:`n_x, 2n_y-1, 2n_z-1` with the 3D solver, :math:`n_x, 2n_y-1, n_z` with the 2D solver. - * ``petsc``: Poisson's equation is solved by one of PETSc's Krylov solvers, preconditioned by the AMReX multigrid. + * ``petsc``: Poisson's equation is solved by PETSc's GMRES, right-preconditioned by the AMReX multigrid. This is the electrostatic counterpart of the PETSc interface that the implicit electromagnetic solvers offer for the curl-curl equation (``newton.linear_solver = petsc_ksp``): the linear system is handed to PETSc as a matrix-free operator, whose action, as well as that of the multigrid preconditioner, is computed by AMReX. It therefore discretizes Poisson's equation exactly like ``multigrid`` does, and accepts the same boundary conditions; only the outer iteration differs. - It requires the compilation flag ``-DWarpX_PETSC=ON``. + It requires the compilation flag ``-DWarpX_PETSC=ON``, and is not yet implemented on GPUs. It is not supported in ``labframe-effective-potential`` mode. Note that in 1D with ``warpx.do_electrostatic = labframe``, and with the ``poissonsolver`` Python callback, Poisson's equation is solved by a dedicated solver and this option has no effect. - Any PETSc option that is set on the command line or in a PETSc option file (e.g. ``-ksp_type``) - takes precedence over the parameters below. - - * ``petsc_poisson.ksp_type`` (``string``) optional (default: ``gmres``): The type of the PETSc ``KSP`` - object, e.g. ``gmres``, ``fgmres``, ``bcgs`` or ``cg``. - - * ``petsc_poisson.restart_length`` (``int``) optional (default: 30): The restart length, only used by - the GMRES variants. + Further customization of the Krylov solver (e.g. ``-ksp_type``, ``-ksp_gmres_restart``) is available + through PETSc's own runtime options (command line or option file), which take precedence over the + parameters below. * ``petsc_poisson.use_mlmg_preconditioner`` (``bool``) optional (default: 1): Whether to use the AMReX multigrid V-cycles as a preconditioner, through PETSc's ``PCShell`` interface. Setting this to 0 runs diff --git a/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc b/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc index 9f1514fec4a..3d890b65993 100644 --- a/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc +++ b/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc @@ -8,5 +8,3 @@ warpx.do_electrostatic = labframe # solve the Poisson equation with PETSc's GMRES, preconditioned by the AMReX # multigrid, instead of using the multigrid directly warpx.poisson_solver = petsc -petsc_poisson.ksp_type = gmres -petsc_poisson.precond_num_iters = 1 diff --git a/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc b/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc index f79ea54b525..8834f1e6cca 100644 --- a/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc +++ b/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc @@ -1,32 +1,7 @@ -max_step = 1 -amr.n_cell = 64 64 64 -amr.max_level = 0 -amr.blocking_factor = 8 -amr.max_grid_size = 128 -geometry.dims = 3 -boundary.field_lo = pec pec neumann -boundary.field_hi = pec neumann neumann -boundary.potential_lo_x = 0 -boundary.potential_hi_x = 0 -boundary.potential_lo_y = 0 -geometry.prob_lo = -0.5 -0.5 -0.5 -geometry.prob_hi = 0.5 0.5 0.5 -warpx.const_dt = 1e-6 +# base input parameters: same physics as the mixed-BC test (embedded boundary +# with a non-zero potential, mixed PEC/Neumann domain boundaries) +FILE = inputs_test_3d_electrostatic_sphere_eb_mixed_bc -warpx.do_electrostatic = labframe -warpx.eb_implicit_function = "-(x**2+y**2+z**2-0.3**2)" -warpx.eb_potential(x,y,z,t) = "1." -warpx.self_fields_required_precision = 1.e-7 -warpx.abort_on_warning_threshold = medium - -# solve the Poisson equation with PETSc's GMRES, preconditioned by the AMReX -# multigrid. This exercises the embedded boundary (with a non-zero potential) -# and mixed PEC/Neumann domain boundaries. +# test input parameters: solve the Poisson equation with PETSc's GMRES, +# preconditioned by the AMReX multigrid, instead of using the multigrid directly warpx.poisson_solver = petsc - -algo.field_gathering = momentum-conserving - -diagnostics.diags_names = diag1 -diag1.intervals = 1 -diag1.diag_type = Full -diag1.fields_to_plot = Ex Ey Ez rho phi diff --git a/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp b/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp index ddcd5f8cc81..dea05c86ec8 100755 --- a/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp +++ b/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp @@ -54,10 +54,7 @@ void ElectrostaticSolver::ReadParameters () { (WarpX::poisson_solver_id == PoissonSolverAlgo::PETSc); if (m_petsc_options.use_petsc_ksp) { ParmParse const pp_petsc("petsc_poisson"); - pp_petsc.query("ksp_type", m_petsc_options.ksp_type); pp_petsc.query("use_mlmg_preconditioner", m_petsc_options.use_mlmg_preconditioner); - utils::parser::queryWithParser( - pp_petsc, "restart_length", m_petsc_options.restart_length); utils::parser::queryWithParser( pp_petsc, "precond_num_iters", m_petsc_options.precond_num_iters); m_petsc_options.verbosity = self_fields_verbosity; diff --git a/Source/ablastr/fields/PETScPoissonSolver.H b/Source/ablastr/fields/PETScPoissonSolver.H index ee7bc86aa90..6a2b9d6d471 100644 --- a/Source/ablastr/fields/PETScPoissonSolver.H +++ b/Source/ablastr/fields/PETScPoissonSolver.H @@ -9,14 +9,9 @@ #include #include -#include #include #include #include -#include - -#include -#include namespace ablastr::fields { @@ -25,24 +20,16 @@ namespace ablastr::fields { * * This is the electrostatic counterpart of the PETSc interface that the implicit * electromagnetic solvers use for the curl-curl equation (@see PETScKSP in - * `Source/NonlinearSolvers/PETScKSP_Wrapper.H`): the linear system is handed to a - * PETSc `KSP` object as a matrix-free operator, and the AMReX geometric multigrid - * is used as a preconditioner through PETSc's `PCShell` interface. + * `Source/NonlinearSolvers/PETScKSP_Wrapper.H`). */ struct PETScPoissonOptions { - /** Solve the Poisson equation with a PETSc Krylov solver instead of MLMG */ + /** Solve the Poisson equation with PETSc's GMRES instead of MLMG */ bool use_petsc_ksp = false; - /** PETSc `KSP` type, e.g. "gmres", "fgmres", "bcgs", "cg" */ - std::string ksp_type = "gmres"; - - /** Restart length, only used by the GMRES variants (PETSc default: 30) */ - int restart_length = 30; - /** Use MLMG V-cycles as the preconditioner (through PETSc's `PCShell`). - * When false, the Krylov solver runs unpreconditioned, which is usually - * much slower and is mostly useful for debugging. + * When false, GMRES runs unpreconditioned, which is usually much slower + * and is mostly useful for debugging. */ bool use_mlmg_preconditioner = true; @@ -57,154 +44,25 @@ struct PETScPoissonOptions #ifdef AMREX_USE_PETSC -// The PETSc objects are wrapped in these structs, which are only defined in -// PETScPoissonSolver.cpp, so that the PETSc headers do not have to be included -// here (they conflict with some of the WarpX and AMReX names). -namespace petsc_poisson -{ - struct KSPObj; - struct MatObj; - struct VecObj; -} - -/** \brief Solve the Poisson equation of one MR level with a PETSc Krylov solver +/** \brief Solve the Poisson equation of one MR level with PETSc's GMRES * - * The linear system that is solved is the very same one that `amrex::MLMG` would - * solve, i.e. it is defined by the `amrex::MLNodeLinOp` that the `amrex::MLMG` - * object passed to the constructor was built from. The operator is never - * assembled into a sparse matrix: it is passed to PETSc as a `MatShell` whose + * The linear system that is solved is the very same one that `mlmg` would + * solve: the operator is handed to PETSc as a matrix-free `MatShell` whose * action is computed by the AMReX linear operator, and the AMReX multigrid - * V-cycles are passed to PETSc as a `PCShell`. This mirrors `amrex::GMRESMLMG`, - * with PETSc's Krylov solvers in place of the AMReX GMRES implementation. - * - * As in `amrex::GMRESMLMG`, the Krylov solver does not iterate on the potential - * itself but on a correction: the residual of the initial guess is computed with - * the inhomogeneous operator, so that non-zero Dirichlet boundary values stored - * in `phi` are accounted for, and the correction that the Krylov solver computes - * vanishes on the Dirichlet nodes. - * - * The degrees of freedom of the PETSc vectors are the nodes that are owned by a - * box (nodes shared between boxes, or with a periodic image, appear only once) - * and that are not Dirichlet nodes. - */ -class PETScPoissonSolver -{ -public: - - /** Construct the solver and create the PETSc objects - * - * \param[in] mlmg the multigrid solver that defines the linear operator; it - * is used both for the operator and for the preconditioner, - * and must outlive this object - * \param[in] phi_prototype a nodal MultiFab with the layout of the unknowns - * \param[in] geom the geometry of the MR level that is solved - * \param[in] options the settings of the PETSc solver - */ - PETScPoissonSolver (amrex::MLMG & mlmg, - amrex::MultiFab const & phi_prototype, - amrex::Geometry const & geom, - PETScPoissonOptions const & options); - - ~PETScPoissonSolver (); - - // Prohibit Move and Copy operations - PETScPoissonSolver (PETScPoissonSolver const &) = delete; - PETScPoissonSolver & operator= (PETScPoissonSolver const &) = delete; - PETScPoissonSolver (PETScPoissonSolver &&) noexcept = delete; - PETScPoissonSolver & operator= (PETScPoissonSolver &&) noexcept = delete; - - /** Solve the Poisson equation - * - * \param[inout] phi on input the initial guess, which must hold the Dirichlet - * boundary values; on output the computed potential - * \param[in] rho the (already scaled) right-hand side - * \param[in] relative_tolerance the relative convergence threshold - * \param[in] absolute_tolerance the absolute convergence threshold - * \param[in] max_iters the maximum number of Krylov iterations - */ - void solve (amrex::MultiFab & phi, - amrex::MultiFab const & rho, - amrex::Real relative_tolerance, - amrex::Real absolute_tolerance, - int max_iters); - - //! Number of Krylov iterations of the most recent solve - [[nodiscard]] int getNumIters () const { return m_num_iters; } - - //! Residual norm of the most recent solve - [[nodiscard]] amrex::Real getResidualNorm () const { return m_residual_norm; } - - /** Apply the linear operator: `out = L(in)` - * - * This is called back by PETSc, with `in` and `out` pointing to the local - * part of the PETSc vectors. - */ - void applyOperator (amrex::Real * out, amrex::Real const * in); - - /** Apply the multigrid preconditioner: `out = P^{-1}(in)` - * - * This is called back by PETSc, with `in` and `out` pointing to the local - * part of the PETSc vectors. - */ - void applyPreconditioner (amrex::Real * out, amrex::Real const * in); - -private: - - //! Number the degrees of freedom that this MPI rank owns - void buildDOFMap (amrex::MultiFab const & phi_prototype); - - //! Gather the degrees of freedom of `mf` into the PETSc array `arr` - void copyToArray (amrex::MultiFab const & mf, amrex::Real * arr) const; - - //! Scatter the PETSc array `arr` into `mf`, and make `mf` consistent - void copyFromArray (amrex::MultiFab & mf, amrex::Real const * arr) const; - - //! Multigrid solver that provides the operator and the preconditioner - amrex::MLMG * m_mlmg = nullptr; - //! Geometry of the MR level that is solved - amrex::Geometry m_geom; - //! Settings of the PETSc solver - PETScPoissonOptions m_options; - - //! Local index of the degree of freedom of each node, -1 if it is not one - std::unique_ptr m_dof; - //! Number of degrees of freedom owned by this MPI rank - amrex::Long m_ndofs_local = 0; - //! Total number of degrees of freedom - amrex::Long m_ndofs_global = 0; - - //! Work arrays for the operator (`_in` also holds the ghost nodes) - amrex::MultiFab m_op_in, m_op_out; - //! Work arrays for the preconditioner - amrex::MultiFab m_pc_in, m_pc_out; - //! Residual of the initial guess, and the correction that is solved for - amrex::MultiFab m_res, m_cor; - - //! Matrix-free linear operator - std::unique_ptr m_A; - //! Solution vector - std::unique_ptr m_x; - //! Right-hand-side vector - std::unique_ptr m_b; - //! Krylov solver - std::unique_ptr m_ksp; - - //! Number of iterations of the most recent solve - int m_num_iters = -1; - //! Residual norm of the most recent solve - amrex::Real m_residual_norm = amrex::Real(-1.0); -}; - -/** Solve the Poisson equation of one MR level with a PETSc Krylov solver + * V-cycles are handed to PETSc as a right `PCShell` preconditioner. This + * mirrors `amrex::GMRESMLMG`, with PETSc's GMRES in place of the AMReX GMRES + * implementation. Further PETSc customization (e.g. `-ksp_type`) is available + * through PETSc's own runtime options, which take precedence. * - * This is a convenience wrapper that creates a @see PETScPoissonSolver, solves, - * and destroys it again. The PETSc objects are cheap to create compared to the - * `amrex::MLMG` object that `mlmg` refers to, which `computePhi` also recreates - * at every call. + * As in `amrex::GMRESMLMG`, GMRES does not iterate on the potential itself but + * on a correction: the residual of the initial guess is computed with the + * inhomogeneous operator, so that non-zero Dirichlet boundary values stored in + * `phi` are accounted for, and the correction vanishes on the Dirichlet nodes. * - * \param[in] mlmg the multigrid solver that defines the linear operator - * \param[inout] phi on input the initial guess (holding the Dirichlet boundary - * values), on output the computed potential + * \param[in] mlmg the multigrid solver that defines the linear operator; it is + * used both for the operator and for the preconditioner + * \param[inout] phi on input the initial guess, which must hold the Dirichlet + * boundary values; on output the computed potential * \param[in] rho the (already scaled) right-hand side * \param[in] geom the geometry of the MR level that is solved * \param[in] relative_tolerance the relative convergence threshold diff --git a/Source/ablastr/fields/PETScPoissonSolver.cpp b/Source/ablastr/fields/PETScPoissonSolver.cpp index 9b27b2bc88d..2d7906d0430 100644 --- a/Source/ablastr/fields/PETScPoissonSolver.cpp +++ b/Source/ablastr/fields/PETScPoissonSolver.cpp @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include #include @@ -23,9 +21,10 @@ #include #include #include +#include +#include #include -#include // The PETSc headers must be included before PETScPoissonSolver.H, see the // comment in Source/NonlinearSolvers/WarpX_PETSc.cpp @@ -39,9 +38,9 @@ namespace ablastr::fields { -namespace petsc_poisson { +namespace { -//! Wrapper for a PETSc KSP object +//! RAII wrapper for a PETSc KSP object struct KSPObj { KSPObj () = default; @@ -53,7 +52,7 @@ struct KSPObj KSP obj = nullptr; }; -//! Wrapper for a PETSc Mat object +//! RAII wrapper for a PETSc Mat object struct MatObj { MatObj () = default; @@ -65,7 +64,7 @@ struct MatObj Mat obj = nullptr; }; -//! Wrapper for a PETSc Vec object +//! RAII wrapper for a PETSc Vec object struct VecObj { VecObj () = default; @@ -77,21 +76,143 @@ struct VecObj Vec obj = nullptr; }; +/** Data shared between petscPoissonSolve() and the PETSc callbacks + * + * The degrees of freedom of the PETSc vectors are the nodes that this MPI rank + * owns: nodes shared between boxes, or with a periodic image, appear only once. + * Dirichlet nodes are kept as degrees of freedom; they simply remain zero + * throughout the Krylov solve, since both the right-hand side and the operator + * output are zeroed on them (this is also how `amrex::GMRESMLMG` treats them). + */ +struct PoissonCtx +{ + amrex::MLMG * mlmg = nullptr; + amrex::Geometry geom; + PETScPoissonOptions options; + + //! Local index of the degree of freedom of each node, -1 if it is not one + std::unique_ptr dof; + //! Number of degrees of freedom owned by this MPI rank / in total + amrex::Long ndofs_local = 0; + amrex::Long ndofs_global = 0; + + //! Work arrays (one ghost layer), reused by the operator and the + //! preconditioner callbacks, and by petscPoissonSolve() itself + amrex::MultiFab work_in; + amrex::MultiFab work_out; + + //! Number the degrees of freedom that this MPI rank owns + void buildDOFMap (amrex::MultiFab const & phi) + { + ABLASTR_PROFILE("petsc_poisson::buildDOFMap()"); + + // Owner is the box with the lowest index containing the node; the same + // convention is used by OverrideSync in copyFromArray() below. + auto const owner_mask = amrex::OwnerMask(phi, geom.periodicity()); + + dof = std::make_unique(phi.boxArray(), + phi.DistributionMap(), 1, 0); + dof->setVal(-1); + + for (amrex::MFIter mfi(*dof); mfi.isValid(); ++mfi) + { + amrex::Box const & bx = mfi.validbox(); + auto const npts = static_cast(bx.numPts()); + amrex::BoxIndexer const box_indexer(bx); + + auto const & owner_arr = owner_mask->const_array(mfi); + auto const & dof_arr = dof->array(mfi); + auto const first_dof = static_cast(ndofs_local); + + auto const ndofs = amrex::Scan::PrefixSum( + npts, + [=] AMREX_GPU_DEVICE (int offset) -> int + { + auto const [i,j,k] = box_indexer(offset); + return owner_arr(i,j,k) ? 1 : 0; + }, + [=] AMREX_GPU_DEVICE (int offset, int ps) + { + auto const [i,j,k] = box_indexer(offset); + if (owner_arr(i,j,k)) { + dof_arr(i,j,k) = ps + first_dof; + } + }, + amrex::Scan::Type::exclusive, amrex::Scan::retSum); + + ndofs_local += ndofs; + } + + ndofs_global = ndofs_local; + amrex::ParallelDescriptor::ReduceLongSum(ndofs_global); + } + + //! Gather the degrees of freedom of `mf` into the PETSc array `arr` + void copyToArray (amrex::MultiFab const & mf, amrex::Real * arr) const + { + ABLASTR_PROFILE("petsc_poisson::copyToArray()"); + + for (amrex::MFIter mfi(*dof); mfi.isValid(); ++mfi) + { + auto const & mf_arr = mf.const_array(mfi); + auto const & dof_arr = dof->const_array(mfi); + amrex::ParallelFor(mfi.validbox(), + [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + int const idx = dof_arr(i,j,k); + if (idx >= 0) { arr[idx] = mf_arr(i,j,k); } + }); + } + amrex::Gpu::streamSynchronize(); + } + + //! Scatter the PETSc array `arr` into `mf`, and make `mf` consistent + void copyFromArray (amrex::MultiFab & mf, amrex::Real const * arr) const + { + ABLASTR_PROFILE("petsc_poisson::copyFromArray()"); + + using namespace amrex::literals; + + mf.setVal(0._rt); + for (amrex::MFIter mfi(*dof); mfi.isValid(); ++mfi) + { + auto const & mf_arr = mf.array(mfi); + auto const & dof_arr = dof->const_array(mfi); + amrex::ParallelFor(mfi.validbox(), + [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + int const idx = dof_arr(i,j,k); + if (idx >= 0) { mf_arr(i,j,k) = arr[idx]; } + }); + } + amrex::Gpu::streamSynchronize(); + + // Fill the nodes owned by another box from their owner (OverrideSync + // uses the same OwnerMask convention as buildDOFMap), then the ghosts + mf.OverrideSync(geom.periodicity()); + mf.FillBoundary(geom.periodicity()); + } +}; + //! Apply the matrix-free linear operator, called back by PETSc PetscErrorCode applyOperator (Mat a_A, Vec a_in, Vec a_out) { PetscFunctionBeginUser; - PETScPoissonSolver * solver = nullptr; - PetscCall(MatShellGetContext(a_A, &solver)); + PoissonCtx * ctx = nullptr; + PetscCall(MatShellGetContext(a_A, &ctx)); PetscScalar const * in_arr = nullptr; PetscScalar * out_arr = nullptr; PetscCall(VecGetArrayRead(a_in, &in_arr)); PetscCall(VecGetArrayWrite(a_out, &out_arr)); - solver->applyOperator( static_cast(out_arr), - static_cast(in_arr) ); + ctx->copyFromArray(ctx->work_in, static_cast(in_arr)); + // `applyPrecond` applies the operator with homogeneous boundary conditions, + // which is the operator that the correction equation uses + ctx->mlmg->applyPrecond({&ctx->work_out}, {&ctx->work_in}); + ctx->mlmg->getLinOp().setDirichletNodesToZero(0, 0, ctx->work_out); + ctx->copyToArray(ctx->work_out, static_cast(out_arr)); PetscCall(VecRestoreArrayWrite(a_out, &out_arr)); PetscCall(VecRestoreArrayRead(a_in, &in_arr)); @@ -104,16 +225,21 @@ PetscErrorCode applyPreconditioner (PC a_pc, Vec a_in, Vec a_out) { PetscFunctionBeginUser; - PETScPoissonSolver * solver = nullptr; - PetscCall(PCShellGetContext(a_pc, &solver)); + using namespace amrex::literals; + + PoissonCtx * ctx = nullptr; + PetscCall(PCShellGetContext(a_pc, &ctx)); PetscScalar const * in_arr = nullptr; PetscScalar * out_arr = nullptr; PetscCall(VecGetArrayRead(a_in, &in_arr)); PetscCall(VecGetArrayWrite(a_out, &out_arr)); - solver->applyPreconditioner( static_cast(out_arr), - static_cast(in_arr) ); + ctx->copyFromArray(ctx->work_in, static_cast(in_arr)); + ctx->mlmg->setPrecondIter(ctx->options.precond_num_iters); + ctx->work_out.setVal(0._rt); + ctx->mlmg->precond({&ctx->work_out}, {&ctx->work_in}, 0._rt, 0._rt); + ctx->copyToArray(ctx->work_out, static_cast(out_arr)); PetscCall(VecRestoreArrayWrite(a_out, &out_arr)); PetscCall(VecRestoreArrayRead(a_in, &in_arr)); @@ -131,283 +257,110 @@ PetscErrorCode printResidual (KSP a_ksp, PetscInt a_n, PetscReal a_rnorm, void * PetscFunctionReturn(PETSC_SUCCESS); } -//! Is `a_type` one of the GMRES variants of PETSc? -bool isGMRES (std::string const & a_type) -{ - return (a_type == "gmres") || (a_type == "fgmres") - || (a_type == "lgmres") || (a_type == "dgmres") - || (a_type == "pgmres") || (a_type == "pipefgmres"); -} - -} // namespace petsc_poisson - +} // anonymous namespace -PETScPoissonSolver::PETScPoissonSolver (amrex::MLMG & mlmg, - amrex::MultiFab const & phi_prototype, - amrex::Geometry const & geom, - PETScPoissonOptions const & options) - : m_mlmg(&mlmg), m_geom(geom), m_options(options) +void +petscPoissonSolve (amrex::MLMG & mlmg, + amrex::MultiFab & phi, + amrex::MultiFab const & rho, + amrex::Geometry const & geom, + amrex::Real relative_tolerance, + amrex::Real absolute_tolerance, + int max_iters, + PETScPoissonOptions const & options) { - ABLASTR_PROFILE("PETScPoissonSolver::PETScPoissonSolver()"); + ABLASTR_PROFILE("petscPoissonSolve()"); - // This builds the multigrid hierarchy and the masks of the linear operator, - // which the operator, the preconditioner and buildDOFMap() below all need. - m_mlmg->preparePrecond(); - - buildDOFMap(phi_prototype); - - // The work arrays are created by the linear operator itself, so that they - // have the right layout and factory. The inputs of the operator and of the - // preconditioner need one layer of ghost nodes, as in - // amrex::GMRESMLMG::makeVecLHS(). - auto & linop = m_mlmg->getLinOp(); - m_op_in = linop.make(0, 0, amrex::IntVect(1)); - m_op_out = linop.make(0, 0, amrex::IntVect(0)); - m_pc_in = linop.make(0, 0, amrex::IntVect(1)); - m_pc_out = linop.make(0, 0, amrex::IntVect(1)); - m_res = linop.make(0, 0, amrex::IntVect(1)); - m_cor = linop.make(0, 0, amrex::IntVect(1)); - - m_A = std::make_unique(); - m_x = std::make_unique(); - m_b = std::make_unique(); - m_ksp = std::make_unique(); - - // Vectors - VecCreate(PETSC_COMM_WORLD, &m_x->obj); #ifdef AMREX_USE_GPU -# if defined(AMREX_USE_CUDA) - VecSetType(m_x->obj, VECCUDA); -# elif defined(AMREX_USE_HIP) - VecSetType(m_x->obj, VECHIP); -# else ABLASTR_ABORT_WITH_MESSAGE( - "The PETSc Poisson solver is not yet implemented for non-CUDA/HIP GPUs"); -# endif -#else - VecSetType(m_x->obj, VECSTANDARD); + "The PETSc Poisson solver is not yet implemented on GPUs"); #endif - auto const ndofs_local = static_cast(m_ndofs_local); - auto const ndofs_global = static_cast(m_ndofs_global); - VecSetSizes(m_x->obj, ndofs_local, ndofs_global); - VecSetFromOptions(m_x->obj); - VecDuplicate(m_x->obj, &m_b->obj); - - // Matrix-free linear operator - MatCreateShell( PETSC_COMM_WORLD, - ndofs_local, ndofs_local, - ndofs_global, ndofs_global, - this, &m_A->obj ); - MatShellSetOperation( m_A->obj, MATOP_MULT, - (void(*)())petsc_poisson::applyOperator ); // NOLINT - MatSetUp(m_A->obj); - - // Krylov solver - KSPCreate(PETSC_COMM_WORLD, &m_ksp->obj); - KSPSetType(m_ksp->obj, m_options.ksp_type.c_str()); - KSPSetOperators(m_ksp->obj, m_A->obj, m_A->obj); - if (petsc_poisson::isGMRES(m_options.ksp_type)) { - KSPGMRESSetRestart(m_ksp->obj, m_options.restart_length); - // Right preconditioning, so that the residual that PETSc monitors and - // uses for its convergence test is the residual of the actual system - KSPSetPCSide(m_ksp->obj, PC_RIGHT); - KSPSetNormType(m_ksp->obj, KSP_NORM_UNPRECONDITIONED); - } + using namespace amrex::literals; + + // This builds the multigrid hierarchy and the masks of the linear operator, + // which the operator, the preconditioner and the DOF map all need + mlmg.preparePrecond(); + auto & linop = mlmg.getLinOp(); + + PoissonCtx ctx; + ctx.mlmg = &mlmg; + ctx.geom = geom; + ctx.options = options; + ctx.buildDOFMap(phi); + // The work arrays are created by the linear operator itself, so that they + // have the right layout; their ghost layer is needed by the AMReX operators + // (as in amrex::GMRESMLMG::makeVecLHS) + ctx.work_in = linop.make(0, 0, amrex::IntVect(1)); + ctx.work_out = linop.make(0, 0, amrex::IntVect(1)); + + // PETSc vectors and matrix-free operator + VecObj x, b; + MatObj A; + KSPObj ksp; + auto const ndofs_l = static_cast(ctx.ndofs_local); + auto const ndofs_g = static_cast(ctx.ndofs_global); + VecCreate(PETSC_COMM_WORLD, &x.obj); + VecSetType(x.obj, VECSTANDARD); + VecSetSizes(x.obj, ndofs_l, ndofs_g); + VecDuplicate(x.obj, &b.obj); + MatCreateShell(PETSC_COMM_WORLD, ndofs_l, ndofs_l, ndofs_g, ndofs_g, + &ctx, &A.obj); + MatShellSetOperation(A.obj, MATOP_MULT, (void (*)(void))applyOperator); + MatSetUp(A.obj); + + // GMRES, right-preconditioned so that the monitored residual is the + // residual of the actual system + KSPCreate(PETSC_COMM_WORLD, &ksp.obj); + KSPSetType(ksp.obj, KSPGMRES); + KSPSetOperators(ksp.obj, A.obj, A.obj); + KSPSetPCSide(ksp.obj, PC_RIGHT); + KSPSetNormType(ksp.obj, KSP_NORM_UNPRECONDITIONED); PC pc = nullptr; - KSPGetPC(m_ksp->obj, &pc); - if (m_options.use_mlmg_preconditioner) { + KSPGetPC(ksp.obj, &pc); + if (options.use_mlmg_preconditioner) { PCSetType(pc, PCSHELL); - PCShellSetApply(pc, petsc_poisson::applyPreconditioner); - PCShellSetContext(pc, this); + PCShellSetApply(pc, applyPreconditioner); + PCShellSetContext(pc, &ctx); PCShellSetName(pc, "AMReX MLMG"); } else { PCSetType(pc, PCNONE); } - - if (m_options.verbosity > 1) { - KSPMonitorSet(m_ksp->obj, petsc_poisson::printResidual, nullptr, nullptr); + KSPSetTolerances(ksp.obj, relative_tolerance, absolute_tolerance, + PETSC_CURRENT, (max_iters > 0 ? max_iters : PETSC_CURRENT)); + if (options.verbosity > 1) { + KSPMonitorSet(ksp.obj, printResidual, nullptr, nullptr); } - // Command-line and input-file PETSc options take precedence over the above - KSPSetFromOptions(m_ksp->obj); - - if (m_options.verbosity > 0) { - amrex::Print() << "PETScPoissonSolver: using PETSc's KSP (" << m_options.ksp_type - << ") with " - << (m_options.use_mlmg_preconditioner ? "the AMReX MLMG" : "no") - << " preconditioner (total DOFs = " << m_ndofs_global << ").\n"; + // PETSc runtime options (e.g. -ksp_type) take precedence over the above + KSPSetFromOptions(ksp.obj); + + if (options.verbosity > 0) { + amrex::Print() << "Poisson (PETSc KSP): " + << (options.use_mlmg_preconditioner ? "MLMG-preconditioned" + : "unpreconditioned") + << " solve, total DOFs = " << ctx.ndofs_global << ".\n"; } -} - -PETScPoissonSolver::~PETScPoissonSolver () = default; - -void PETScPoissonSolver::buildDOFMap (amrex::MultiFab const & phi_prototype) -{ - ABLASTR_PROFILE("PETScPoissonSolver::buildDOFMap()"); - - using namespace amrex::literals; - - // The nodes that sit on the boundary between two boxes (or on the boundary - // between a box and the periodic image of another one) belong to the valid - // region of both boxes, but they are a single unknown of the linear system: - // only the node of the "owner" box is a degree of freedom. - auto const owner_mask = amrex::OwnerMask(phi_prototype, m_geom.periodicity()); - - // The nodes on which a Dirichlet boundary condition is applied are not - // unknowns of the linear system either. `setDirichletNodesToZero` is the - // public interface through which the AMReX linear operator exposes them. - auto & linop = m_mlmg->getLinOp(); - amrex::MultiFab dirichlet_indicator = linop.make(0, 0, amrex::IntVect(0)); - dirichlet_indicator.setVal(1._rt); - linop.setDirichletNodesToZero(0, 0, dirichlet_indicator); - - m_dof = std::make_unique(phi_prototype.boxArray(), - phi_prototype.DistributionMap(), 1, 0); - m_dof->setVal(-1); - - m_ndofs_local = 0; - for (amrex::MFIter mfi(*m_dof); mfi.isValid(); ++mfi) - { - amrex::Box const & bx = mfi.validbox(); - auto const npts = static_cast(bx.numPts()); - amrex::BoxIndexer const box_indexer(bx); - - auto const & owner_arr = owner_mask->const_array(mfi); - auto const & dirichlet_arr = dirichlet_indicator.const_array(mfi); - auto const & dof_arr = m_dof->array(mfi); - auto const first_dof = static_cast(m_ndofs_local); - - auto const ndofs = amrex::Scan::PrefixSum( - npts, - [=] AMREX_GPU_DEVICE (int offset) -> int - { - auto const [i,j,k] = box_indexer(offset); - return (owner_arr(i,j,k) && (dirichlet_arr(i,j,k) > 0.5_rt)) ? 1 : 0; - }, - [=] AMREX_GPU_DEVICE (int offset, int ps) - { - auto const [i,j,k] = box_indexer(offset); - if (owner_arr(i,j,k) && (dirichlet_arr(i,j,k) > 0.5_rt)) { - dof_arr(i,j,k) = ps + first_dof; - } - }, - amrex::Scan::Type::exclusive, amrex::Scan::retSum); - - m_ndofs_local += ndofs; - } - - m_ndofs_global = m_ndofs_local; - amrex::ParallelDescriptor::ReduceLongSum(m_ndofs_global); - - ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(m_ndofs_global > 0, - "PETScPoissonSolver: the linear system has no degree of freedom"); -} - -void PETScPoissonSolver::copyToArray (amrex::MultiFab const & mf, amrex::Real * arr) const -{ - ABLASTR_PROFILE("PETScPoissonSolver::copyToArray()"); - - for (amrex::MFIter mfi(*m_dof); mfi.isValid(); ++mfi) - { - amrex::Box const & bx = mfi.validbox(); - auto const & mf_arr = mf.const_array(mfi); - auto const & dof_arr = m_dof->const_array(mfi); - amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) - { - int const dof = dof_arr(i,j,k); - if (dof >= 0) { arr[dof] = mf_arr(i,j,k); } - }); - } - amrex::Gpu::streamSynchronize(); -} - -void PETScPoissonSolver::copyFromArray (amrex::MultiFab & mf, amrex::Real const * arr) const -{ - ABLASTR_PROFILE("PETScPoissonSolver::copyFromArray()"); - - using namespace amrex::literals; - - // The nodes that are not degrees of freedom (Dirichlet nodes, and the nodes - // that another box owns) are set to zero here, and the ones that another box - // owns are then filled from their owner by `OverrideSync` below. - mf.setVal(0._rt); - - for (amrex::MFIter mfi(*m_dof); mfi.isValid(); ++mfi) - { - amrex::Box const & bx = mfi.validbox(); - auto const & mf_arr = mf.array(mfi); - auto const & dof_arr = m_dof->const_array(mfi); - amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) - { - int const dof = dof_arr(i,j,k); - if (dof >= 0) { mf_arr(i,j,k) = arr[dof]; } - }); - } - amrex::Gpu::streamSynchronize(); - - // `OverrideSync` uses the same `amrex::OwnerMask` as `buildDOFMap` above, - // so the nodes that are shared between boxes are filled from the very box - // whose node was numbered as a degree of freedom. - mf.OverrideSync(m_geom.periodicity()); - mf.FillBoundary(m_geom.periodicity()); -} - -void PETScPoissonSolver::applyOperator (amrex::Real * out, amrex::Real const * in) -{ - ABLASTR_PROFILE("PETScPoissonSolver::applyOperator()"); - - copyFromArray(m_op_in, in); - // `applyPrecond` applies the operator with homogeneous boundary conditions, - // which is the operator that the correction equation solved here uses - m_mlmg->applyPrecond({&m_op_out}, {&m_op_in}); - m_mlmg->getLinOp().setDirichletNodesToZero(0, 0, m_op_out); - copyToArray(m_op_out, out); -} - -void PETScPoissonSolver::applyPreconditioner (amrex::Real * out, amrex::Real const * in) -{ - ABLASTR_PROFILE("PETScPoissonSolver::applyPreconditioner()"); - - using namespace amrex::literals; - - copyFromArray(m_pc_in, in); - m_mlmg->setPrecondIter(m_options.precond_num_iters); - m_pc_out.setVal(0._rt); - m_mlmg->precond({&m_pc_out}, {&m_pc_in}, 0._rt, 0._rt); - copyToArray(m_pc_out, out); -} - -void PETScPoissonSolver::solve (amrex::MultiFab & phi, - amrex::MultiFab const & rho, - amrex::Real relative_tolerance, - amrex::Real absolute_tolerance, - int max_iters) -{ - ABLASTR_PROFILE("PETScPoissonSolver::solve()"); - - using namespace amrex::literals; - - auto & linop = m_mlmg->getLinOp(); // MLMG is only used as a preconditioner here, so its bottom solve must be - // cheap and linear; this mirrors what amrex::GMRESMLMG does. - auto const bottom_solver = m_mlmg->getBottomSolver(); - auto const mlmg_verbose = m_mlmg->getVerbose(); - auto const mlmg_bottom_verbose = m_mlmg->getBottomVerbose(); + // cheap and linear; this mirrors what amrex::GMRESMLMG does + auto const bottom_solver = mlmg.getBottomSolver(); + auto const mlmg_verbose = mlmg.getVerbose(); + auto const mlmg_bottom_verbose = mlmg.getBottomVerbose(); if (bottom_solver != amrex::BottomSolver::smoother && bottom_solver != amrex::BottomSolver::hypre && bottom_solver != amrex::BottomSolver::petsc) { - m_mlmg->setBottomSolver(amrex::BottomSolver::smoother); + mlmg.setBottomSolver(amrex::BottomSolver::smoother); } - m_mlmg->setVerbose(0); - m_mlmg->setBottomVerbose(0); + mlmg.setVerbose(0); + mlmg.setBottomVerbose(0); // Residual of the initial guess: res = L(phi) - rho. Note that `apply` uses // the inhomogeneous operator, so that the Dirichlet values that `phi` holds - // contribute to the residual. - m_res.setVal(0._rt); - m_mlmg->apply({&m_res}, {&phi}); + // contribute to the residual. `work_in` is free until KSPSolve starts. + amrex::MultiFab & res = ctx.work_in; + res.setVal(0._rt); + mlmg.apply({&res}, {&phi}); amrex::MultiFab scaled_rho; amrex::MultiFab const * rhs = ρ @@ -418,79 +371,54 @@ void PETScPoissonSolver::solve (amrex::MultiFab & phi, amrex::ignore_unused(scaled); rhs = &scaled_rho; } - amrex::MultiFab::Saxpy(m_res, -1._rt, *rhs, 0, 0, 1, amrex::IntVect(0)); - linop.setDirichletNodesToZero(0, 0, m_res); + amrex::MultiFab::Saxpy(res, -1._rt, *rhs, 0, 0, 1, amrex::IntVect(0)); + linop.setDirichletNodesToZero(0, 0, res); - // Solve L(cor) = res for the correction, with PETSc's Krylov solver + // Solve L(cor) = res for the correction { PetscScalar * b_arr = nullptr; - VecGetArrayWrite(m_b->obj, &b_arr); - copyToArray(m_res, static_cast(b_arr)); - VecRestoreArrayWrite(m_b->obj, &b_arr); + VecGetArrayWrite(b.obj, &b_arr); + ctx.copyToArray(res, static_cast(b_arr)); + VecRestoreArrayWrite(b.obj, &b_arr); } - VecZeroEntries(m_x->obj); - - KSPSetTolerances( m_ksp->obj, - relative_tolerance, - absolute_tolerance, - PETSC_CURRENT, - (max_iters > 0 ? max_iters : PETSC_CURRENT) ); - KSPSolve(m_ksp->obj, m_b->obj, m_x->obj); + KSPSolve(ksp.obj, b.obj, x.obj); + // phi = phi - cor. `work_in` is free again once KSPSolve has returned. + amrex::MultiFab & cor = ctx.work_in; { PetscScalar const * x_arr = nullptr; - VecGetArrayRead(m_x->obj, &x_arr); - copyFromArray(m_cor, static_cast(x_arr)); - VecRestoreArrayRead(m_x->obj, &x_arr); + VecGetArrayRead(x.obj, &x_arr); + ctx.copyFromArray(cor, static_cast(x_arr)); + VecRestoreArrayRead(x.obj, &x_arr); } - - // phi = phi - cor - amrex::MultiFab::Saxpy(phi, -1._rt, m_cor, 0, 0, 1, amrex::IntVect(0)); + amrex::MultiFab::Saxpy(phi, -1._rt, cor, 0, 0, 1, amrex::IntVect(0)); // `amrex::MLMG::solve` ends with this; for the embedded-boundary operator it - // writes the prescribed potential into the nodes that the EB covers. + // writes the prescribed potential into the nodes that the EB covers linop.postSolve({&phi}); + phi.FillBoundary(geom.periodicity()); - phi.FillBoundary(m_geom.periodicity()); - - // Report on the solve + // Report on the solve, and abort if it failed (as MLMG does) PetscInt niters = -1; - KSPGetIterationNumber(m_ksp->obj, &niters); - m_num_iters = static_cast(niters); + KSPGetIterationNumber(ksp.obj, &niters); PetscReal norm = -1; - KSPGetResidualNorm(m_ksp->obj, &norm); - m_residual_norm = static_cast(norm); - + KSPGetResidualNorm(ksp.obj, &norm); KSPConvergedReason reason; - KSPGetConvergedReason(m_ksp->obj, &reason); + KSPGetConvergedReason(ksp.obj, &reason); char const * reason_string = nullptr; - KSPGetConvergedReasonString(m_ksp->obj, &reason_string); - - if (m_options.verbosity > 0) { - amrex::Print() << "Poisson (PETSc KSP): " << m_num_iters << " iterations, exited due to \"" - << reason_string << "\" (abs. norm = " << m_residual_norm << ").\n"; + KSPGetConvergedReasonString(ksp.obj, &reason_string); + if (options.verbosity > 0) { + amrex::Print() << "Poisson (PETSc KSP): " << niters + << " iterations, exited due to \"" << reason_string + << "\" (abs. norm = " << norm << ").\n"; } ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(reason > 0, std::string("The PETSc Poisson solver failed to converge: ") + reason_string); // Restore the settings of the multigrid solver - m_mlmg->setBottomSolver(bottom_solver); - m_mlmg->setVerbose(mlmg_verbose); - m_mlmg->setBottomVerbose(mlmg_bottom_verbose); -} - -void -petscPoissonSolve (amrex::MLMG & mlmg, - amrex::MultiFab & phi, - amrex::MultiFab const & rho, - amrex::Geometry const & geom, - amrex::Real relative_tolerance, - amrex::Real absolute_tolerance, - int max_iters, - PETScPoissonOptions const & options) -{ - PETScPoissonSolver solver(mlmg, phi, geom, options); - solver.solve(phi, rho, relative_tolerance, absolute_tolerance, max_iters); + mlmg.setBottomSolver(bottom_solver); + mlmg.setVerbose(mlmg_verbose); + mlmg.setBottomVerbose(mlmg_bottom_verbose); } } // namespace ablastr::fields From d47becc1fad5ffd3913407e90d4673df4c45efeb Mon Sep 17 00:00:00 2001 From: Remi Lehe Date: Thu, 27 Aug 2026 16:31:30 -0700 Subject: [PATCH 4/5] Fix CI: MLLinOp::make is protected All four Azure jobs failed to compile PETScPoissonSolver.cpp: `amrex::MLLinOpT::make` is a protected method, which amrex::GMRESMLMG can only call because it is declared a friend of MLLinOpT. Create the work arrays directly from the layout (and factory) of `phi` instead, and assert that `phi` carries the ghost layer that the AMReX operators need. Co-Authored-By: Claude Opus 5 --- Source/ablastr/fields/PETScPoissonSolver.cpp | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/Source/ablastr/fields/PETScPoissonSolver.cpp b/Source/ablastr/fields/PETScPoissonSolver.cpp index 2d7906d0430..986dfe174e5 100644 --- a/Source/ablastr/fields/PETScPoissonSolver.cpp +++ b/Source/ablastr/fields/PETScPoissonSolver.cpp @@ -278,6 +278,9 @@ petscPoissonSolve (amrex::MLMG & mlmg, using namespace amrex::literals; + ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(phi.nGrowVect().allGE(amrex::IntVect(1)), + "petscPoissonSolve: phi must have (at least) one ghost layer"); + // This builds the multigrid hierarchy and the masks of the linear operator, // which the operator, the preconditioner and the DOF map all need mlmg.preparePrecond(); @@ -288,11 +291,14 @@ petscPoissonSolve (amrex::MLMG & mlmg, ctx.geom = geom; ctx.options = options; ctx.buildDOFMap(phi); - // The work arrays are created by the linear operator itself, so that they - // have the right layout; their ghost layer is needed by the AMReX operators - // (as in amrex::GMRESMLMG::makeVecLHS) - ctx.work_in = linop.make(0, 0, amrex::IntVect(1)); - ctx.work_out = linop.make(0, 0, amrex::IntVect(1)); + // The work arrays share the layout (and factory) of `phi`; their ghost + // layer is needed by the AMReX operators (as in amrex::GMRESMLMG::makeVecLHS). + // Note that `phi` must be nodal with one ghost layer, like the vectors that + // amrex::MLMG::solve would create internally. + ctx.work_in.define(phi.boxArray(), phi.DistributionMap(), 1, 1, + amrex::MFInfo(), phi.Factory()); + ctx.work_out.define(phi.boxArray(), phi.DistributionMap(), 1, 1, + amrex::MFInfo(), phi.Factory()); // PETSc vectors and matrix-free operator VecObj x, b; From 086e37227333cbc9a00421c41e3ea90991b3696f Mon Sep 17 00:00:00 2001 From: Remi Lehe Date: Fri, 28 Aug 2026 07:15:20 -0700 Subject: [PATCH 5/5] PETSc Poisson solver: enable GPU builds Select the PETSc vector type per compute backend (VECCUDA, VECHIP, or VECSTANDARD) and remove the abort that disabled the solver on GPUs. This mirrors what KSP_impl::createObjects does for the curl-curl solver, down to calling VecSetFromOptions after VecSetSizes so that -vec_type can still override the choice. The copy kernels are unchanged: as in WarpXSolverVec::copyTo/copyFrom, they use VecGetArray and write the resulting array from amrex::ParallelFor. Note that on a GPU this array is the host-side array of the PETSc vector, so this assumes the allocation is addressable from the device; a comment in copyToArray now records that assumption. Compile-tested for CUDA (sm_86) against PETSc 3.24.0 built with --with-cuda=1, in a 3D build configured with -DWarpX_COMPUTE=CUDA -DWarpX_PETSC=ON -DWarpX_EB=ON. Not run: the build machine has no GPU. Co-Authored-By: Claude Opus 5 --- Docs/source/usage/parameters.rst | 3 ++- Source/ablastr/fields/PETScPoissonSolver.cpp | 27 +++++++++++++++----- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst index 06de15998d2..1d2221c0255 100644 --- a/Docs/source/usage/parameters.rst +++ b/Docs/source/usage/parameters.rst @@ -456,7 +456,8 @@ Overall simulation parameters handed to PETSc as a matrix-free operator, whose action, as well as that of the multigrid preconditioner, is computed by AMReX. It therefore discretizes Poisson's equation exactly like ``multigrid`` does, and accepts the same boundary conditions; only the outer iteration differs. - It requires the compilation flag ``-DWarpX_PETSC=ON``, and is not yet implemented on GPUs. + It requires the compilation flag ``-DWarpX_PETSC=ON``, and PETSc itself must be built with + CUDA (or HIP) support in order to run on GPUs. It is not supported in ``labframe-effective-potential`` mode. Note that in 1D with ``warpx.do_electrostatic = labframe``, and with the ``poissonsolver`` Python callback, Poisson's equation is solved by a dedicated solver and this option has no effect. diff --git a/Source/ablastr/fields/PETScPoissonSolver.cpp b/Source/ablastr/fields/PETScPoissonSolver.cpp index 986dfe174e5..98df74b6c13 100644 --- a/Source/ablastr/fields/PETScPoissonSolver.cpp +++ b/Source/ablastr/fields/PETScPoissonSolver.cpp @@ -147,7 +147,15 @@ struct PoissonCtx amrex::ParallelDescriptor::ReduceLongSum(ndofs_global); } - //! Gather the degrees of freedom of `mf` into the PETSc array `arr` + /** Gather the degrees of freedom of `mf` into the PETSc array `arr` + * + * Note that on GPUs `arr` comes from `VecGetArray`, i.e. it is the host-side + * array of the PETSc vector (PETSc copies it back to the device before the + * next device operation). The kernel below writes it from device code, which + * assumes that this allocation is addressable from the device. This is the + * same assumption that `WarpXSolverVec::copyTo/copyFrom` makes for the + * curl-curl solver, see Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.cpp. + */ void copyToArray (amrex::MultiFab const & mf, amrex::Real * arr) const { ABLASTR_PROFILE("petsc_poisson::copyToArray()"); @@ -271,11 +279,6 @@ petscPoissonSolve (amrex::MLMG & mlmg, { ABLASTR_PROFILE("petscPoissonSolve()"); -#ifdef AMREX_USE_GPU - ABLASTR_ABORT_WITH_MESSAGE( - "The PETSc Poisson solver is not yet implemented on GPUs"); -#endif - using namespace amrex::literals; ABLASTR_ALWAYS_ASSERT_WITH_MESSAGE(phi.nGrowVect().allGE(amrex::IntVect(1)), @@ -307,8 +310,20 @@ petscPoissonSolve (amrex::MLMG & mlmg, auto const ndofs_l = static_cast(ctx.ndofs_local); auto const ndofs_g = static_cast(ctx.ndofs_global); VecCreate(PETSC_COMM_WORLD, &x.obj); +#ifdef AMREX_USE_GPU +# if defined(AMREX_USE_CUDA) + VecSetType(x.obj, VECCUDA); +# elif defined(AMREX_USE_HIP) + VecSetType(x.obj, VECHIP); +# else + ABLASTR_ABORT_WITH_MESSAGE( + "The PETSc Poisson solver is not yet implemented for non-CUDA/HIP GPUs"); +# endif +#else VecSetType(x.obj, VECSTANDARD); +#endif VecSetSizes(x.obj, ndofs_l, ndofs_g); + VecSetFromOptions(x.obj); VecDuplicate(x.obj, &b.obj); MatCreateShell(PETSC_COMM_WORLD, ndofs_l, ndofs_l, ndofs_g, ndofs_g, &ctx, &A.obj);