Multiple RHS support in PCPFLAREINV via PCMatApply, with true block matrix-free applies - #265
Merged
Conversation
PCPFLAREINV now registers pc->ops->matapply, so a block of right-hand sides is applied with a single sparse matrix by dense matrix product rather than PETSc's default loop over columns. On the GPU backends this keeps the whole multiple-rhs apply on the device. The assembled inverse goes through MatProduct with MATPRODUCT_AB; the matrix-free inverse is a MatShell with only MATOP_MULT registered, so that case still loops over the columns. The new tests/adv_1d_multi_rhs.c driver builds the same 1D upwind finite difference advection operator as adv_1d.c, creates the block of right-hand sides with MatCreateDenseFromVecType so its backend follows the operator's, and solves with KSPMatSolve. It defaults to preonly and pflareinv because KSPMatSolve only reaches PCMatApply for KSPPREONLY and KSPHPDDM - any other KSP type falls back to solving column by column. The block solution is checked against a column-by-column reference solve, and the solve is split into log stages so it can be timed on a GPU machine. PCAIR does not support multiple right-hand sides yet, so it gets no run line in tests/Makefile. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The matrix-free branch of PCMatApply_PFLAREINV_c previously applied the polynomial matshell one dense column at a time. The power/arnoldi polynomials are now applied blockwise: a new petsc_horner_block runs the Horner iteration with sparse matrix-dense matrix products on the underlying operator, with dense scratch blocks cached in the matshell context and sized lazily to the number of rhs. The right diagonally scaled variant q(D^-1 A) D^-1 (used by PCAIR's smoothers) is supported by bypassing the inner D^-1 A matshell and scaling the products directly. Anything unsupported falls back to the column-by-column loop. The newton family dispatch is stubbed to fall back for now. Adds tests/shell_block_apply.c to drive the block apply directly, covering the diagonally scaled mode that PCPFLAREINV cannot reach. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Adds petsc_newton_block, the multiple rhs version of petsc_newton, which iterates over the Newton basis roots with sparse matrix-dense matrix products in place of the matvecs, handling the complex conjugate root pairs with the same real arithmetic as the scalar kernel. The diagonally scaled variant scales every product by D^-1 immediately, replicating the inner D^-1 A matshell blockwise. The newton branch of the block dispatch is filled in, so matrix-free newton/newton_no_extra PCPFLAREINV now take the block path in PCMatApply rather than falling back to a column by column apply. Adds serial and parallel run lines for the matrix-free Newton block solves and the direct Newton matshell block apply. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The matrix-free Neumann polynomial is q(I - D^-1 A) D^-1 applied through the same right scaled Horner machinery as the gmres polynomials, so its block apply is petsc_horner_block with a different inner operator: after each product and diagonal scale, y = temp - D^-1 A temp finishes the I - D^-1 A arithmetic blockwise. A new neumann_inner flag in the matshell context records which inner operator the mat_scaled shell applies, which replaces caller gating as the discriminator between the Neumann and diagonally scaled gmres shells - previously nothing in the context could tell them apart and misrouting would have been silently wrong. PCMatApply now dispatches neumann to the block path too. Adds serial and parallel neumann run lines for both drivers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…cation fix The line truncation fix hoisted the context fields into local handle copies before the PetscObjectIsNull checks, but PETSc's create and destroy routines return the new handle through the argument, so two of those copies needed writing back to the context: - shell_poly_block_apply duplicated the reciprocal of the diagonal into the local copy and never stored it, so every diagonally scaled block apply created (and leaked) a new vec and the teardown had nothing to destroy - this is what tripped the malloc dump and valgrind CI jobs on the matrix-free Neumann test. - ensure_block_temp_mats destroyed the cached dense temporaries through a local copy on a rebuild, leaving dangling handles in the context that the creation loop then skipped - a latent use after free if the number of rhs ever changed between applies. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Same pattern as the previous fix, applied to every destroy that goes through a local handle copy: destroy nulls the local, not the context field, so copy it back. In reset_inverse_mat this is only defensive - the context is deallocated straight after - but it keeps the idiom uniform so no destroy-through-a-local ever leaves a dangling handle. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
calculate_and_build_approximate_inverse heap-allocates a coefficient array for every matrix-free inverse on the assumption the matshell takes ownership of it, but the Neumann builder allocated its own private all-ones array instead, orphaning the wrapper's allocation on every fresh setup - valgrind caught this in CI once the new matrix-free Neumann test line exercised the path (it predates this branch, but was never run under valgrind before). The Neumann builder now takes the coefficients through its argument list exactly like the gmres builders (finish_approximate_inverse already sets them all to one), points the matshell context at that storage, and frees any previously owned coefficients before reassigning on reuse - so the ownership rules in calculate_and_build_approximate_inverse and the C binding now hold uniformly for all the polynomial types. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The previous fix set own_coefficients inside calculate_and_build_neumann_polynomial_inverse, but PCAIR calls finish_approximate_inverse directly with its own coefficient storage - claiming ownership there made reset_inverse_mat deallocate PCAIR's memory and corrupt the heap (caught by the PCAIR matrix-free Neumann smoothing test in tests_short). Ownership is the caller's decision, exactly as with the gmres builders: calculate_and_build_approximate_inverse sets own_coefficients for the PCPFLAREINV paths after the builder returns, and for PCAIR it stays false. The leak fix is unaffected. Co-Authored-By: Claude Fable 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.
Summary
Adds multiple right-hand side support to PCPFLAREINV through a new
PCMatApplycallback, soKSPMatSolveapplies a whole dense block of right-hand sides at once instead of PETSc's default loop over columns.MatProduct, which keeps the whole multiple-RHS apply on the device for the GPU backends.power,arnoldi,newton,newton_no_extra,neumann): true block kernels rather than loopingMatMulton the matshell.petsc_horner_blockandpetsc_newton_blockrun the Horner / Newton-basis root iterations with SpMMs on the underlying operator, with dense scratch blocks cached in the matshell context and lazily sized to the number of RHS (unknown atPCSetUp).q(D^-1 A) D^-1(what PCAIR builds for its smoothers) and the Neumann polynomial's intrinsically scaledq(I - D^-1 A) D^-1are both supported by bypassing the inner matshell algebraically: real products with the unscaled operator followed byMatDiagonalScale, plus the extraI - D^-1 Astep for Neumann. A newneumann_innerflag in the matshell context discriminates the two inner operators, since nothing else in the context can tell them apart. This makes the block kernels ready for future PCAIR multiple-RHS work.-inforeports which path ran.Note
KSPMatSolveonly reachesPCMatApplyfor-ksp_type preonly(or HPDDM); other KSP types fall back to solving column by column inside PETSc.Test plan
tests/adv_1d_multi_rhs.c: builds the RHS block withMatCreateDenseFromVecType, solves withKSPMatSolve, and checks the block result against a column-by-columnKSPSolvereference at 1e-10 relative Frobenius. Run lines cover default/arnoldi/newton/newton_no_extra/neumann, assembled and matrix-free, serial and 2 ranks.tests/shell_block_apply.c: drives the block apply of the matshells directly (the routines exposed here are internal, surfaced for unit testing), covering the diagonally scaled mode that PCPFLAREINV itself never builds, with a second apply to exercise the cached temporaries and the per-apply diagonal-reciprocal refresh.make checkandmake tests_shortpass;-info :pcconfirms the block kernel runs for all five polynomial families; poly_order edge cases (0, 1, 3, 10) verified; Kokkos backend smoke tested (-mat_type aijkokkos -vec_type kokkos); complex conjugate root pairs in the Newton kernel confirmed exercised.🤖 Generated with Claude Code