From 1f5fe16a9e3c8df01af21abc77cc7f42fbfa6b16 Mon Sep 17 00:00:00 2001 From: Timofey Mukha Date: Wed, 2 Sep 2026 10:44:27 +0200 Subject: [PATCH 1/3] Implement --- CHANGELOG.md | 3 + src/.depends | 6 +- src/krylov/bcknd/cpu/cg.f90 | 87 ++++++++----- src/krylov/bcknd/cpu/cg_coupled.f90 | 191 ++++++++++++++++------------ src/krylov/bcknd/cpu/gmres.f90 | 81 +++++++----- 5 files changed, 222 insertions(+), 146 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index aae6ee0382d4..7f73fb780ae8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,9 @@ ## Develop +- Changed the CPU CG, coupled CG and GMRES solvers to obtain their large work + arrays from the scratch registry for each solve instead of retaining + dedicated storage. - Added crystal router gather-scatter communication backends, `NEKO_GS_COMM=CRYSTAL` on the host and `CRYSTALGPU` on the device. They route the halo in recursive-bisection stages instead of sending one message diff --git a/src/.depends b/src/.depends index 24266055ad06..ababc0a7b8e2 100644 --- a/src/.depends +++ b/src/.depends @@ -227,14 +227,14 @@ krylov/krylov.lo : krylov/krylov.f90 common/log.lo config/neko_config.lo krylov/ krylov/pc_identity.lo : krylov/pc_identity.f90 config/num_types.lo krylov/precon.lo math/math.lo krylov/precon_fctry.lo : krylov/precon_fctry.f90 config/neko_config.lo common/utils.lo multigrid/phmg.lo krylov/pc_hsmg.lo krylov/bcknd/device/pc_jacobi_device.lo krylov/bcknd/sx/pc_jacobi_sx.lo krylov/bcknd/cpu/pc_jacobi.lo krylov/bcknd/device/pc_identity_device.lo krylov/pc_identity.lo krylov/precon.lo krylov/krylov_fctry.lo : krylov/krylov_fctry.f90 config/neko_config.lo common/utils.lo krylov/precon.lo config/num_types.lo krylov/bcknd/device/gmres_device.lo krylov/bcknd/sx/gmres_sx.lo krylov/bcknd/device/cheby_device.lo krylov/bcknd/cpu/cheby.lo krylov/bcknd/cpu/gmres.lo krylov/bcknd/cpu/bicgstab.lo krylov/bcknd/device/fusedcg_cpld_device.lo krylov/bcknd/device/fusedcg_device.lo krylov/bcknd/device/pipecg_device.lo krylov/bcknd/sx/pipecg_sx.lo krylov/bcknd/cpu/pipecg.lo krylov/bcknd/cpu/cacg.lo krylov/bcknd/device/cg_cpld_device.lo krylov/bcknd/device/cg_device.lo krylov/bcknd/cpu/cg_coupled.lo krylov/bcknd/sx/cg_sx.lo krylov/bcknd/cpu/cg.lo krylov/krylov.lo -krylov/bcknd/cpu/cg.lo : krylov/bcknd/cpu/cg.f90 comm/comm.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/num_types.lo config/neko_config.lo +krylov/bcknd/cpu/cg.lo : krylov/bcknd/cpu/cg.f90 comm/comm.lo math/math.lo registries/scratch_registry.lo data_types/host_array.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/num_types.lo config/neko_config.lo krylov/bcknd/cpu/cacg.lo : krylov/bcknd/cpu/cacg.f90 math/mxm_wrapper.lo comm/comm.lo common/utils.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/neko_config.lo config/num_types.lo krylov/bcknd/cpu/cheby.lo : krylov/bcknd/cpu/cheby.f90 math/math.lo math/schwarz.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/space.lo mesh/mesh.lo sem/coef.lo field/field.lo common/profiler.lo config/num_types.lo math/ax.lo krylov/precon.lo krylov/krylov.lo krylov/bcknd/cpu/pipecg.lo : krylov/bcknd/cpu/pipecg.f90 comm/comm.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo config/num_types.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/neko_config.lo krylov/bcknd/cpu/bicgstab.lo : krylov/bcknd/cpu/bicgstab.f90 comm/comm.lo common/utils.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/num_types.lo -krylov/bcknd/cpu/gmres.lo : krylov/bcknd/cpu/gmres.f90 comm/comm.lo config/neko_config.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo config/num_types.lo math/ax.lo krylov/precon.lo krylov/krylov.lo +krylov/bcknd/cpu/gmres.lo : krylov/bcknd/cpu/gmres.f90 comm/comm.lo config/neko_config.lo math/math.lo registries/scratch_registry.lo data_types/host_array.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo config/num_types.lo math/ax.lo krylov/precon.lo krylov/krylov.lo krylov/bcknd/cpu/pc_jacobi.lo : krylov/bcknd/cpu/pc_jacobi.f90 gs/gather_scatter.lo sem/dofmap.lo config/num_types.lo sem/coef.lo krylov/precon.lo math/math.lo -krylov/bcknd/cpu/cg_coupled.lo : krylov/bcknd/cpu/cg_coupled.f90 math/operators.lo common/utils.lo comm/comm.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/num_types.lo +krylov/bcknd/cpu/cg_coupled.lo : krylov/bcknd/cpu/cg_coupled.f90 math/operators.lo common/utils.lo comm/comm.lo math/math.lo registries/scratch_registry.lo data_types/host_array.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/num_types.lo krylov/bcknd/sx/cg_sx.lo : krylov/bcknd/sx/cg_sx.f90 math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo math/ax.lo krylov/precon.lo krylov/krylov.lo config/num_types.lo krylov/bcknd/sx/pipecg_sx.lo : krylov/bcknd/sx/pipecg_sx.f90 comm/comm.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo config/num_types.lo math/ax.lo krylov/precon.lo krylov/krylov.lo krylov/bcknd/sx/gmres_sx.lo : krylov/bcknd/sx/gmres_sx.f90 comm/comm.lo math/math.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo gs/gather_scatter.lo sem/coef.lo field/field.lo config/num_types.lo math/ax.lo krylov/precon.lo krylov/krylov.lo diff --git a/src/krylov/bcknd/cpu/cg.f90 b/src/krylov/bcknd/cpu/cg.f90 index 65e2bbaff7eb..0e6f1776fe85 100644 --- a/src/krylov/bcknd/cpu/cg.f90 +++ b/src/krylov/bcknd/cpu/cg.f90 @@ -43,6 +43,8 @@ module cg use scalar_bc_projector, only : scalar_bc_projector_t use vector_bc_projector, only : vector_bc_projector_t, & vector_bc_projector_components + use host_array, only : host_array_t + use scratch_registry, only : neko_scratch_registry use math, only : glsc3, abscmp use comm, only : MPI_EXTRA_PRECISION, MPI_REAL_PRECISION, NEKO_COMM use mpi_f08, only : MPI_Allreduce, MPI_IN_PLACE, MPI_SUM @@ -51,23 +53,35 @@ module cg integer, parameter :: CG_P_SPACE = 7 - !> Standard preconditioned conjugate gradient method + !> CPU implementation of the preconditioned conjugate gradient method. + !! + !! Workspace pointers are associated with host arrays from the scratch + !! registry only for the duration of a solve. type, public, extends(ksp_t) :: cg_t - real(kind=rp), allocatable :: w(:) - real(kind=rp), allocatable :: r(:) - real(kind=rp), allocatable :: p(:,:) - real(kind=rp), allocatable :: z(:) - real(kind=rp), allocatable :: alpha(:) + !> Operator action \f$w = A p\f$. + real(kind=rp), pointer :: w(:) => null() + !> Residual \f$r = f - A x\f$. + real(kind=rp), pointer :: r(:) => null() + !> Rolling space of search directions \f$p\f$. + real(kind=rp), pointer :: p(:,:) => null() + !> Preconditioned residual \f$z = M^{-1} r\f$. + real(kind=rp), pointer :: z(:) => null() + !> Step lengths associated with the stored search directions. + real(kind=rp), pointer :: alpha(:) => null() contains + !> Initialise a CPU PCG solver. procedure, pass(this) :: init => cg_init + !> Free a CPU PCG solver. procedure, pass(this) :: free => cg_free + !> Solve a linear system with the CPU PCG method. procedure, pass(this) :: solve => cg_solve + !> Solve three independent systems with the CPU PCG method. procedure, pass(this) :: solve_coupled => cg_solve_coupled end type cg_t contains - !> Initialise a standard PCG solver + !> Initialise a CPU PCG solver. subroutine cg_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) class(cg_t), intent(inout), target :: this integer, intent(in) :: max_iter @@ -79,12 +93,6 @@ subroutine cg_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) call this%free() - allocate(this%w(n)) - allocate(this%r(n)) - allocate(this%p(n, CG_P_SPACE)) - allocate(this%z(n)) - allocate(this%alpha(CG_P_SPACE)) - if (present(M)) then this%M => M end if @@ -108,37 +116,23 @@ subroutine cg_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) end subroutine cg_init - !> Deallocate a standard PCG solver + !> Free a CPU PCG solver. subroutine cg_free(this) class(cg_t), intent(inout) :: this call this%ksp_free() - if (allocated(this%w)) then - deallocate(this%w) - end if - - if (allocated(this%r)) then - deallocate(this%r) - end if - - if (allocated(this%p)) then - deallocate(this%p) - end if - - if (allocated(this%z)) then - deallocate(this%z) - end if - - if (allocated(this%alpha)) then - deallocate(this%alpha) - end if + nullify(this%w) + nullify(this%r) + nullify(this%p) + nullify(this%z) + nullify(this%alpha) nullify(this%M) end subroutine cg_free - !> Standard PCG solve + !> Solve a linear system with the CPU PCG method. function cg_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & result(ksp_results) class(cg_t), intent(inout) :: this @@ -154,6 +148,8 @@ function cg_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & integer :: iter, max_iter, i, j, k, p_cur, p_prev, ierr real(kind=rp) :: rnorm, rtr, rtz2, rtz1, x_plus(NEKO_BLK_SIZE) real(kind=rp) :: beta, pap, norm_fac, tmp + type(host_array_t), pointer :: w_tmp, r_tmp, p_tmp, z_tmp, alpha_tmp + integer :: temp_indices(5) if (present(niter)) then max_iter = niter @@ -162,6 +158,23 @@ function cg_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & end if norm_fac = 1.0_rp / sqrt(coef%volume) + call neko_scratch_registry%request_host_array(w_tmp, temp_indices(1), & + n, .false.) + call neko_scratch_registry%request_host_array(r_tmp, temp_indices(2), & + n, .false.) + call neko_scratch_registry%request_host_array(p_tmp, temp_indices(3), & + n * CG_P_SPACE, .false.) + call neko_scratch_registry%request_host_array(z_tmp, temp_indices(4), & + n, .false.) + call neko_scratch_registry%request_host_array(alpha_tmp, & + temp_indices(5), CG_P_SPACE, .false.) + + this%w => w_tmp%x + this%r => r_tmp%x + this%p(1:n, 1:CG_P_SPACE) => p_tmp%x + this%z => z_tmp%x + this%alpha => alpha_tmp%x + associate(w => this%w, r => this%r, p => this%p, & z => this%z, alpha => this%alpha) @@ -185,6 +198,8 @@ function cg_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & ksp_results%iter = 0 if (abscmp(rnorm, 0.0_rp)) then ksp_results%converged = .true. + nullify(this%w, this%r, this%p, this%z, this%alpha) + call neko_scratch_registry%relinquish_host_array(temp_indices) return end if @@ -254,6 +269,8 @@ function cg_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & end if end do end associate + nullify(this%w, this%r, this%p, this%z, this%alpha) + call neko_scratch_registry%relinquish_host_array(temp_indices) call this%monitor_stop() ksp_results%res_final = rnorm ksp_results%iter = iter @@ -280,7 +297,7 @@ subroutine second_cg_part(rtr, r, mult, w, alpha, n) end subroutine second_cg_part - !> Standard PCG coupled solve + !> Solve three independent systems with the CPU PCG method. function cg_solve_coupled(this, Ax, x, y, z, fx, fy, fz, & n, coef, bc_projector, gs_h, niter) result(ksp_results) class(cg_t), intent(inout) :: this diff --git a/src/krylov/bcknd/cpu/cg_coupled.f90 b/src/krylov/bcknd/cpu/cg_coupled.f90 index 5c50848de389..110e5625a480 100644 --- a/src/krylov/bcknd/cpu/cg_coupled.f90 +++ b/src/krylov/bcknd/cpu/cg_coupled.f90 @@ -30,7 +30,7 @@ ! ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE ! POSSIBILITY OF SUCH DAMAGE. ! -!> Defines a coupled Conjugate Gradient methods +!> Defines a coupled Conjugate Gradient method. module cg_cpld use num_types, only : rp, xp use krylov, only : ksp_t, ksp_monitor_t, KSP_MAX_ITER @@ -41,6 +41,8 @@ module cg_cpld use gather_scatter, only : gs_t, GS_OP_ADD use scalar_bc_projector, only : scalar_bc_projector_t use vector_bc_projector, only : vector_bc_projector_t + use host_array, only : host_array_t + use scratch_registry, only : neko_scratch_registry use math, only : abscmp use comm, only : MPI_EXTRA_PRECISION, NEKO_COMM use mpi_f08, only : MPI_Allreduce, MPI_IN_PLACE, MPI_SUM @@ -49,30 +51,58 @@ module cg_cpld implicit none private - !> Coupled preconditioned conjugate gradient method + !> CPU implementation of the coupled preconditioned conjugate gradient + !! method. + !! + !! Workspace pointers are associated with host arrays from the scratch + !! registry only for the duration of a solve. The method uses a single + !! Krylov recurrence and combined inner products over all three components. type, public, extends(ksp_t) :: cg_cpld_t - real(kind=rp), allocatable :: w1(:) - real(kind=rp), allocatable :: w2(:) - real(kind=rp), allocatable :: w3(:) - real(kind=rp), allocatable :: r1(:) - real(kind=rp), allocatable :: r2(:) - real(kind=rp), allocatable :: r3(:) - real(kind=rp), allocatable :: p1(:) - real(kind=rp), allocatable :: p2(:) - real(kind=rp), allocatable :: p3(:) - real(kind=rp), allocatable :: z1(:) - real(kind=rp), allocatable :: z2(:) - real(kind=rp), allocatable :: z3(:) + !> First component of the operator action \f$w = A p\f$. + real(kind=rp), pointer :: w1(:) => null() + !> Second component of the operator action \f$w = A p\f$. + real(kind=rp), pointer :: w2(:) => null() + !> Third component of the operator action \f$w = A p\f$. + real(kind=rp), pointer :: w3(:) => null() + !> First component of the residual \f$r = f - A x\f$. + real(kind=rp), pointer :: r1(:) => null() + !> Second component of the residual \f$r = f - A x\f$. + real(kind=rp), pointer :: r2(:) => null() + !> Third component of the residual \f$r = f - A x\f$. + real(kind=rp), pointer :: r3(:) => null() + !> First component of the search direction \f$p\f$. + real(kind=rp), pointer :: p1(:) => null() + !> Second component of the search direction \f$p\f$. + real(kind=rp), pointer :: p2(:) => null() + !> Third component of the search direction \f$p\f$. + real(kind=rp), pointer :: p3(:) => null() + !> First component of the preconditioned residual \f$z = M^{-1}r\f$. + real(kind=rp), pointer :: z1(:) => null() + !> Second component of the preconditioned residual \f$z = M^{-1}r\f$. + real(kind=rp), pointer :: z2(:) => null() + !> Third component of the preconditioned residual \f$z = M^{-1}r\f$. + real(kind=rp), pointer :: z3(:) => null() contains + !> Initialise a coupled CPU PCG solver. procedure, pass(this) :: init => cg_cpld_init + !> Free a coupled CPU PCG solver. procedure, pass(this) :: free => cg_cpld_free + !> Reject scalar solves, which are not supported by this type. procedure, pass(this) :: solve => cg_cpld_nop + !> Solve a three-component coupled system with the CPU PCG method. procedure, pass(this) :: solve_coupled => cg_cpld_solve end type cg_cpld_t contains - !> Initialise a coupled PCG solver + !> Initialise a coupled CPU PCG solver. + !! @param n Number of degrees of freedom per component. + !! @param max_iter Maximum number of iterations. + !! @param M Optional preconditioner. An identity preconditioner is used if + !! absent. + !! @param rel_tol Optional relative convergence tolerance. + !! @param abs_tol Optional absolute convergence tolerance. + !! @param monitor Optional switch for logging the residual at each iteration. subroutine cg_cpld_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) class(cg_cpld_t), target, intent(inout) :: this integer, intent(in) :: max_iter @@ -84,19 +114,6 @@ subroutine cg_cpld_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) call this%free() - allocate(this%w1(n)) - allocate(this%w2(n)) - allocate(this%w3(n)) - allocate(this%r1(n)) - allocate(this%r2(n)) - allocate(this%r3(n)) - allocate(this%p1(n)) - allocate(this%p2(n)) - allocate(this%p3(n)) - allocate(this%z1(n)) - allocate(this%z2(n)) - allocate(this%z3(n)) - if (present(M)) then this%M => M end if @@ -121,64 +138,31 @@ subroutine cg_cpld_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) end subroutine cg_cpld_init - !> Deallocate a coupled PCG solver + !> Free a coupled CPU PCG solver. subroutine cg_cpld_free(this) class(cg_cpld_t), intent(inout) :: this call this%ksp_free() - if (allocated(this%w1)) then - deallocate(this%w1) - end if - - if (allocated(this%w2)) then - deallocate(this%w2) - end if - - if (allocated(this%w3)) then - deallocate(this%w3) - end if - - if (allocated(this%r1)) then - deallocate(this%r1) - end if - - if (allocated(this%r2)) then - deallocate(this%r2) - end if - - if (allocated(this%r3)) then - deallocate(this%r3) - end if - - if (allocated(this%p1)) then - deallocate(this%p1) - end if - - if (allocated(this%p2)) then - deallocate(this%p2) - end if - - if (allocated(this%p3)) then - deallocate(this%p3) - end if - - if (allocated(this%z1)) then - deallocate(this%z1) - end if - - if (allocated(this%z2)) then - deallocate(this%z2) - end if - - if (allocated(this%z3)) then - deallocate(this%z3) - end if + nullify(this%w1, this%w2, this%w3) + nullify(this%r1, this%r2, this%r3) + nullify(this%p1, this%p2, this%p3) + nullify(this%z1, this%z2, this%z3) nullify(this%M) end subroutine cg_cpld_free + !> Reject scalar solves, which are not supported by this type. + !! @param Ax Linear operator. + !! @param x Solution field. + !! @param f Right-hand side. + !! @param n Number of degrees of freedom. + !! @param coef Spectral element coefficients and multiplicity weights. + !! @param bc_projector Projector for Dirichlet boundary nodes. + !! @param gs_h Gather-scatter handle used to assemble the operator result. + !! @param niter Optional maximum number of iterations. + !! @return Unused convergence information. function cg_cpld_nop(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & result(ksp_results) class(cg_cpld_t), intent(inout) :: this @@ -192,14 +176,31 @@ function cg_cpld_nop(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & type(ksp_monitor_t) :: ksp_results integer, optional, intent(in) :: niter - ! Throw and error call neko_error('The cpldcg solver is only defined for coupled solves') ksp_results%res_final = 0.0 ksp_results%iter = 0 end function cg_cpld_nop - !> Coupled PCG solve + !> Solve a three-component coupled system with the CPU PCG method. + !! + !! The initial guesses are discarded. All inner products combine the three + !! components, so the method advances one Krylov recurrence for the complete + !! coupled system. + !! @param Ax Coupled linear operator. + !! @param x Solution field for the first component. + !! @param y Solution field for the second component. + !! @param z Solution field for the third component. + !! @param fx Right-hand side for the first component. + !! @param fy Right-hand side for the second component. + !! @param fz Right-hand side for the third component. + !! @param n Number of degrees of freedom per component. + !! @param coef Spectral element coefficients and multiplicity weights. + !! @param bc_projector Projector for vector boundary nodes. + !! @param gs_h Gather-scatter handle used to assemble operator results. + !! @param niter Optional maximum number of iterations, overriding the + !! configured value. + !! @return Identical combined convergence information for all components. function cg_cpld_solve(this, Ax, x, y, z, fx, fy, fz, & n, coef, bc_projector, gs_h, niter) result(ksp_results) class(cg_cpld_t), intent(inout) :: this @@ -220,6 +221,8 @@ function cg_cpld_solve(this, Ax, x, y, z, fx, fy, fz, & real(kind=rp) :: rnorm, rtr, rtr0, rtz2, rtz1 real(kind=rp) :: beta, pap, alpha, norm_fac real(kind=xp) :: tmp_xp, r1_xp, r2_xp, r3_xp, mult_xp + type(host_array_t), pointer :: w_tmp, r_tmp, p_tmp, z_tmp + integer :: temp_indices(4) if (present(niter)) then max_iter = niter @@ -228,6 +231,28 @@ function cg_cpld_solve(this, Ax, x, y, z, fx, fy, fz, & end if norm_fac = 1.0_rp / sqrt(coef%volume) + call neko_scratch_registry%request_host_array(w_tmp, temp_indices(1), & + 3 * n, .false.) + call neko_scratch_registry%request_host_array(r_tmp, temp_indices(2), & + 3 * n, .false.) + call neko_scratch_registry%request_host_array(p_tmp, temp_indices(3), & + 3 * n, .false.) + call neko_scratch_registry%request_host_array(z_tmp, temp_indices(4), & + 3 * n, .false.) + + this%w1 => w_tmp%x(1:n) + this%w2 => w_tmp%x(n+1:2*n) + this%w3 => w_tmp%x(2*n+1:3*n) + this%r1 => r_tmp%x(1:n) + this%r2 => r_tmp%x(n+1:2*n) + this%r3 => r_tmp%x(2*n+1:3*n) + this%p1 => p_tmp%x(1:n) + this%p2 => p_tmp%x(n+1:2*n) + this%p3 => p_tmp%x(2*n+1:3*n) + this%z1 => z_tmp%x(1:n) + this%z2 => z_tmp%x(n+1:2*n) + this%z3 => z_tmp%x(2*n+1:3*n) + associate (p1 => this%p1, p2 => this%p2, p3 => this%p3, z1 => this%z1, & z2 => this%z2, z3 => this%z3, r1 => this%r1, r2 => this%r2, & r3 => this%r3, w1 => this%w1, w2 => this%w2, w3 => this%w3) @@ -262,6 +287,11 @@ function cg_cpld_solve(this, Ax, x, y, z, fx, fy, fz, & ksp_results%iter = 0 if (abscmp(rnorm, 0.0_rp)) then ksp_results%converged = .true. + nullify(this%w1, this%w2, this%w3) + nullify(this%r1, this%r2, this%r3) + nullify(this%p1, this%p2, this%p3) + nullify(this%z1, this%z2, this%z3) + call neko_scratch_registry%relinquish_host_array(temp_indices) return end if @@ -355,6 +385,11 @@ function cg_cpld_solve(this, Ax, x, y, z, fx, fy, fz, & end if end do end associate + nullify(this%w1, this%w2, this%w3) + nullify(this%r1, this%r2, this%r3) + nullify(this%p1, this%p2, this%p3) + nullify(this%z1, this%z2, this%z3) + call neko_scratch_registry%relinquish_host_array(temp_indices) call this%monitor_stop() ksp_results%res_final = rnorm ksp_results%iter = iter diff --git a/src/krylov/bcknd/cpu/gmres.f90 b/src/krylov/bcknd/cpu/gmres.f90 index fac6b7db769f..2eabf06ec491 100644 --- a/src/krylov/bcknd/cpu/gmres.f90 +++ b/src/krylov/bcknd/cpu/gmres.f90 @@ -43,6 +43,8 @@ module gmres use scalar_bc_projector, only : scalar_bc_projector_t use vector_bc_projector, only : vector_bc_projector_t, & vector_bc_projector_components + use host_array, only : host_array_t + use scratch_registry, only : neko_scratch_registry use math, only : glsc3, rzero, copy, sub2, cmult2, abscmp use neko_config, only : NEKO_BLK_SIZE use comm, only : NEKO_COMM, MPI_EXTRA_PRECISION @@ -50,30 +52,46 @@ module gmres implicit none private - !> Standard preconditioned generalized minimal residual method + !> CPU implementation of the preconditioned GMRES method. + !! + !! The large working vectors are associated with host arrays from the + !! scratch registry only for the duration of a solve. The small + !! extra-precision reduction arrays remain owned by the solver. type, public, extends(ksp_t) :: gmres_t + !> Maximum dimension of the Krylov basis before restarting. integer :: lgmres = 30 - real(kind=rp), allocatable :: w(:) - real(kind=rp), allocatable :: r(:) - real(kind=rp), allocatable :: z(:,:) - real(kind=rp), allocatable :: v(:,:) - !> reduction variables, extra prec + !> Operator action used during Arnoldi orthogonalisation. + real(kind=rp), pointer :: w(:) => null() + !> Residual at the start of a restart cycle. + real(kind=rp), pointer :: r(:) => null() + !> Preconditioned Krylov basis. + real(kind=rp), pointer :: z(:,:) => null() + !> Krylov basis. + real(kind=rp), pointer :: v(:,:) => null() + !> Upper-Hessenberg matrix in extra precision. real(kind=xp), allocatable :: h(:,:) + !> Sines of the Givens rotations in extra precision. real(kind=xp), allocatable :: s(:) + !> Rotated residual vector in extra precision. real(kind=xp), allocatable :: gam(:) + !> Cosines of the Givens rotations and solution coefficients. real(kind=xp), allocatable :: c(:) - !> Per-thread partial-sum buffer + !> Per-thread inner-product buffer in extra precision. real(kind=xp), allocatable :: hp(:,:) contains + !> Initialise a CPU GMRES solver. procedure, pass(this) :: init => gmres_init + !> Free a CPU GMRES solver. procedure, pass(this) :: free => gmres_free + !> Solve a linear system with the CPU GMRES method. procedure, pass(this) :: solve => gmres_solve + !> Solve three independent systems with the CPU GMRES method. procedure, pass(this) :: solve_coupled => gmres_solve_coupled end type gmres_t contains - !> Initialise a standard GMRES solver + !> Initialise a CPU GMRES solver. subroutine gmres_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) class(gmres_t), target, intent(inout) :: this integer, intent(in) :: n @@ -90,16 +108,10 @@ subroutine gmres_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) this%M => M end if - allocate(this%w(n)) - allocate(this%r(n)) - allocate(this%c(this%lgmres)) allocate(this%s(this%lgmres)) allocate(this%gam(this%lgmres + 1)) - allocate(this%z(n, this%lgmres)) - allocate(this%v(n, this%lgmres)) - allocate(this%h(this%lgmres, this%lgmres)) nthrds = 1 @@ -126,35 +138,26 @@ subroutine gmres_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) end subroutine gmres_init - !> Deallocate a standard GMRES solver + !> Free a CPU GMRES solver. subroutine gmres_free(this) class(gmres_t), intent(inout) :: this call this%ksp_free() - if (allocated(this%w)) then - deallocate(this%w) - end if + nullify(this%w) if (allocated(this%c)) then deallocate(this%c) end if - if (allocated(this%r)) then - deallocate(this%r) - end if - - if (allocated(this%z)) then - deallocate(this%z) - end if + nullify(this%r) + nullify(this%z) if (allocated(this%h)) then deallocate(this%h) end if - if (allocated(this%v)) then - deallocate(this%v) - end if + nullify(this%v) if (allocated(this%s)) then deallocate(this%s) @@ -173,7 +176,7 @@ subroutine gmres_free(this) end subroutine gmres_free - !> Standard GMRES solve + !> Solve a linear system with the CPU GMRES method. function gmres_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & result(ksp_results) class(gmres_t), intent(inout) :: this @@ -192,6 +195,8 @@ function gmres_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & real(kind=xp) :: alpha, lr, alpha2, norm_fac, tmp, acc real(kind=rp) :: temp, rnorm logical :: conv + type(host_array_t), pointer :: w_tmp, r_tmp, z_tmp, v_tmp + integer :: temp_indices(4) conv = .false. iter = 0 @@ -206,6 +211,20 @@ function gmres_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & nthrds = 1 !$ nthrds = omp_get_max_threads() + call neko_scratch_registry%request_host_array(w_tmp, temp_indices(1), & + n, .false.) + call neko_scratch_registry%request_host_array(r_tmp, temp_indices(2), & + n, .false.) + call neko_scratch_registry%request_host_array(z_tmp, temp_indices(3), & + n * this%lgmres, .false.) + call neko_scratch_registry%request_host_array(v_tmp, temp_indices(4), & + n * this%lgmres, .false.) + + this%w => w_tmp%x + this%r => r_tmp%x + this%z(1:n, 1:this%lgmres) => z_tmp%x + this%v(1:n, 1:this%lgmres) => v_tmp%x + associate(w => this%w, c => this%c, r => this%r, z => this%z, h => this%h, & v => this%v, s => this%s, gam => this%gam, hp => this%hp) @@ -405,6 +424,8 @@ function gmres_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & end do end associate + nullify(this%w, this%r, this%z, this%v) + call neko_scratch_registry%relinquish_host_array(temp_indices) call this%monitor_stop() ksp_results%res_final = rnorm ksp_results%iter = iter @@ -412,7 +433,7 @@ function gmres_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & end function gmres_solve - !> Standard GMRES coupled solve + !> Solve three independent systems with the CPU GMRES method. function gmres_solve_coupled(this, Ax, x, y, z, fx, fy, fz, & n, coef, bc_projector, gs_h, niter) result(ksp_results) class(gmres_t), intent(inout) :: this From 1e00a630615e57dbfb0a034b97e801e2161a69f4 Mon Sep 17 00:00:00 2001 From: Timofey Mukha Date: Wed, 2 Sep 2026 16:11:59 +0200 Subject: [PATCH 2/3] Update src/krylov/bcknd/cpu/cg.f90 Co-authored-by: Niclas Jansson --- src/krylov/bcknd/cpu/cg.f90 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/krylov/bcknd/cpu/cg.f90 b/src/krylov/bcknd/cpu/cg.f90 index 0e6f1776fe85..627586930f3c 100644 --- a/src/krylov/bcknd/cpu/cg.f90 +++ b/src/krylov/bcknd/cpu/cg.f90 @@ -132,7 +132,7 @@ subroutine cg_free(this) end subroutine cg_free - !> Solve a linear system with the CPU PCG method. + !> Solve a linear system with the Standard PCG method. function cg_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & result(ksp_results) class(cg_t), intent(inout) :: this From e2e24b2d4084ad861095622df9d2b770535bab5b Mon Sep 17 00:00:00 2001 From: Timofey Mukha Date: Wed, 2 Sep 2026 16:14:59 +0200 Subject: [PATCH 3/3] standard gmres --- src/krylov/bcknd/cpu/gmres.f90 | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/krylov/bcknd/cpu/gmres.f90 b/src/krylov/bcknd/cpu/gmres.f90 index 2eabf06ec491..84888c1c3f50 100644 --- a/src/krylov/bcknd/cpu/gmres.f90 +++ b/src/krylov/bcknd/cpu/gmres.f90 @@ -30,7 +30,7 @@ ! ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE ! POSSIBILITY OF SUCH DAMAGE. ! -!> Defines various GMRES methods +!> Implements `gmres_t`. module gmres !$ use omp_lib use krylov, only : ksp_t, ksp_monitor_t @@ -52,7 +52,7 @@ module gmres implicit none private - !> CPU implementation of the preconditioned GMRES method. + !> CPU implementation of the preconditioned standard GMRES method. !! !! The large working vectors are associated with host arrays from the !! scratch registry only for the duration of a solve. The small @@ -79,19 +79,19 @@ module gmres !> Per-thread inner-product buffer in extra precision. real(kind=xp), allocatable :: hp(:,:) contains - !> Initialise a CPU GMRES solver. + !> Initialise a standard GMRES solver. procedure, pass(this) :: init => gmres_init - !> Free a CPU GMRES solver. + !> Free a standard GMRES solver. procedure, pass(this) :: free => gmres_free - !> Solve a linear system with the CPU GMRES method. + !> Solve a linear system with the standard GMRES method. procedure, pass(this) :: solve => gmres_solve - !> Solve three independent systems with the CPU GMRES method. + !> Solve three independent systems with the standard GMRES method. procedure, pass(this) :: solve_coupled => gmres_solve_coupled end type gmres_t contains - !> Initialise a CPU GMRES solver. + !> Initialise a standard GMRES solver. subroutine gmres_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) class(gmres_t), target, intent(inout) :: this integer, intent(in) :: n @@ -138,7 +138,7 @@ subroutine gmres_init(this, n, max_iter, M, rel_tol, abs_tol, monitor) end subroutine gmres_init - !> Free a CPU GMRES solver. + !> Free a standard GMRES solver. subroutine gmres_free(this) class(gmres_t), intent(inout) :: this @@ -176,7 +176,7 @@ subroutine gmres_free(this) end subroutine gmres_free - !> Solve a linear system with the CPU GMRES method. + !> Solve a linear system with the standard GMRES method. function gmres_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & result(ksp_results) class(gmres_t), intent(inout) :: this @@ -433,7 +433,7 @@ function gmres_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) & end function gmres_solve - !> Solve three independent systems with the CPU GMRES method. + !> Solve three independent systems with the standard GMRES method. function gmres_solve_coupled(this, Ax, x, y, z, fx, fy, fz, & n, coef, bc_projector, gs_h, niter) result(ksp_results) class(gmres_t), intent(inout) :: this