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 57242047a58..1d2221c0255 100644 --- a/Docs/source/usage/parameters.rst +++ b/Docs/source/usage/parameters.rst @@ -450,6 +450,33 @@ 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 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``, 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. + + 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 + 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..3d890b65993 --- /dev/null +++ b/Examples/Tests/electrostatic_sphere/inputs_test_3d_electrostatic_sphere_lab_frame_petsc @@ -0,0 +1,10 @@ +# 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 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..8834f1e6cca --- /dev/null +++ b/Examples/Tests/electrostatic_sphere_eb/inputs_test_3d_electrostatic_sphere_eb_petsc @@ -0,0 +1,7 @@ +# 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 + +# 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 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..dea05c86ec8 100755 --- a/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp +++ b/Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp @@ -47,6 +47,20 @@ 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("use_mlmg_preconditioner", m_petsc_options.use_mlmg_preconditioner); + 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 +235,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.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/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 + + +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`). + */ +struct PETScPoissonOptions +{ + /** Solve the Poisson equation with PETSc's GMRES instead of MLMG */ + bool use_petsc_ksp = false; + + /** Use MLMG V-cycles as the preconditioner (through PETSc's `PCShell`). + * When false, GMRES 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 + +/** \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 `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 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. + * + * 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; 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 + * \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..98df74b6c13 --- /dev/null +++ b/Source/ablastr/fields/PETScPoissonSolver.cpp @@ -0,0 +1,447 @@ +/* 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 + +// 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 { + +//! RAII 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; +}; + +//! RAII 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; +}; + +//! RAII 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; +}; + +/** 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` + * + * 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()"); + + 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; + + 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)); + + 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)); + + PetscFunctionReturn(PETSC_SUCCESS); +} + +//! Apply the multigrid preconditioner, called back by PETSc +PetscErrorCode applyPreconditioner (PC a_pc, Vec a_in, Vec a_out) +{ + PetscFunctionBeginUser; + + 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)); + + 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)); + + 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); +} + +} // anonymous namespace + +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("petscPoissonSolve()"); + + 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(); + auto & linop = mlmg.getLinOp(); + + PoissonCtx ctx; + ctx.mlmg = &mlmg; + ctx.geom = geom; + ctx.options = options; + ctx.buildDOFMap(phi); + // 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; + 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); +#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); + 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(ksp.obj, &pc); + if (options.use_mlmg_preconditioner) { + PCSetType(pc, PCSHELL); + PCShellSetApply(pc, applyPreconditioner); + PCShellSetContext(pc, &ctx); + PCShellSetName(pc, "AMReX MLMG"); + } else { + PCSetType(pc, PCNONE); + } + 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); + } + // 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"; + } + + // 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 = 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) + { + mlmg.setBottomSolver(amrex::BottomSolver::smoother); + } + 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. `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 = ρ + 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(res, -1._rt, *rhs, 0, 0, 1, amrex::IntVect(0)); + linop.setDirichletNodesToZero(0, 0, res); + + // Solve L(cor) = res for the correction + { + PetscScalar * b_arr = nullptr; + VecGetArrayWrite(b.obj, &b_arr); + ctx.copyToArray(res, static_cast(b_arr)); + VecRestoreArrayWrite(b.obj, &b_arr); + } + 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(x.obj, &x_arr); + ctx.copyFromArray(cor, static_cast(x_arr)); + VecRestoreArrayRead(x.obj, &x_arr); + } + 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 + linop.postSolve({&phi}); + phi.FillBoundary(geom.periodicity()); + + // Report on the solve, and abort if it failed (as MLMG does) + PetscInt niters = -1; + KSPGetIterationNumber(ksp.obj, &niters); + PetscReal norm = -1; + KSPGetResidualNorm(ksp.obj, &norm); + KSPConvergedReason reason; + KSPGetConvergedReason(ksp.obj, &reason); + char const * reason_string = nullptr; + 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 + mlmg.setBottomSolver(bottom_solver); + mlmg.setVerbose(mlmg_verbose); + mlmg.setBottomVerbose(mlmg_bottom_verbose); +} + +} // namespace ablastr::fields + +#endif // AMREX_USE_PETSC diff --git a/Source/ablastr/fields/PoissonSolver.H b/Source/ablastr/fields/PoissonSolver.H index db9687e582b..e54cd4cb0ac 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,18 @@ 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 + 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 ); + } 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(); @@ -446,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