[WIP] Expose PETSc as a solver for the Poisson equation - #7214
Closed
RemiLehe wants to merge 5 commits into
Closed
Conversation
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.