Skip to content

Use scratch arrays in CG, coupled CG, and GMRES on the CPU #2754

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
timofeymukha wants to merge 3 commits into
base: develop
Choose a base branch
from
+223 −147
Open

Use scratch arrays in CG, coupled CG, and GMRES on the CPU #2754

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1f5fe16
Implement
Sep 2, 2026
1e00a63
Update src/krylov/bcknd/cpu/cg.f90
timofeymukha Sep 2, 2026
e2e24b2
standard gmres
Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions src/.depends
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
87 changes: 52 additions & 35 deletions src/krylov/bcknd/cpu/cg.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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 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
Expand All @@ -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
Expand All @@ -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

@njansson njansson Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes me a bit nervous. Can't we add a rank-2 registry as well (@timfelle )?

@timfelle timfelle Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Array reshaping like this seem to be completely standard in the fortran community, But sure we can easily add another host_array_t to the system. But explicitly for a host implementation the only difference between host_array_t and vector_t is a %x_d=C_NULL_PTR, so I actually don't think the host array thing makes a difference at all since it should never be active in a device build.

@timfelle timfelle Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So my point is, you should probably just use a matrix_t here.

@njansson njansson Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Array reshaping like this seem to be completely standard in the fortran community, But sure we can easily add another host_array_t to the system. But explicitly for a host implementation the only difference between host_array_t and vector_t is a %x_d=C_NULL_PTR, so I actually don't think the host array thing makes a difference at all since it should never be active in a device build.

standard, and working on compilers are different things, as we know ;) But sure let's give it a try

@timfelle timfelle Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, but we do it all over the place already, we just do it at subroutine interfaces, with out of bound indexing and such.

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)
Comment on lines +161 to 179

@timfelle timfelle Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
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)
call neko_scratch_registry%request(this%w, temp_indices(1), n, .false.)
call neko_scratch_registry%request(this%r, temp_indices(2), n, .false.)
call neko_scratch_registry%request(this%z, temp_indices(4), n, .false.)
call neko_scratch_registry%request(this%alpha, temp_indices(5), CG_P_SPACE, .false.)
call neko_scratch_registry%request(p_tmp, temp_indices(3), &
n * CG_P_SPACE, .false.)
this%p(1:n, 1:CG_P_SPACE) => p_tmp%x
associate(w => this%w%x, r => this%r%x, p => this%p%x, &
z => this%z%x, alpha => this%alpha%x)


Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Loading
Loading