Add coupled_bicgstab on the CPU

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Develop
 
+- Added a coupled CPU BiCGStab solver for three-component vector systems.
 - The gather-scatter comm. backend autotuning now covers the device-resident
   backends. With `NEKO_GS_COMM` unset, a CUDA or HIP build benchmarks
   `MPIGPU`, `NCCL` and `CRYSTALGPU` (`NVSHMEM` only when asked for) alongside

doc/pages/user-guide/case-file.md

@@ -465,8 +465,9 @@ the governing equations to feature the full viscous stress tensor, as required
 for a variable viscosity field.
 
 Note that the full viscous stress tensor requires the equations for the 3
-velocity components to be solved in a coupled manner. Therefore, the `coupled_cg`
-(or `fused_coupled_cg`) solver should be used for velocity.
+velocity components to be solved in a coupled manner. Therefore,
+`coupled_bicgstab`, `coupled_cg`, or `fused_coupled_cg` should be used for
+velocity. The `coupled_bicgstab` solver is currently available on CPUs only.
 
 ### Schwarz iterations
 This feature is enabled by setting the `schwarz_iterations` keyword inside
@@ -1889,6 +1890,8 @@ The following keywords are used, with the corresponding options.
   - `cg`, a conjugate gradient solver.
   - `pipecg`, a pipelined conjugate gradient solver.
   - `bicgstab`, a bi-conjugate gradient stabilized solver.
+  - `coupled_bicgstab`, a coupled bi-conjugate gradient stabilized solver for
+    CPU backends. It can be used for velocity when viscosity varies in space.
   - `cacg`, a communication-avoiding conjugate gradient solver.
   - `coupled_cg`, a coupled conjugate gradient solver. Must be used for velocity
     when viscosity varies in space.
@@ -2028,7 +2031,7 @@ concisely directly in the table.
 | `source_terms`                                     | Array of JSON objects, defining additional source terms.                                          | See list of source terms above                              | -             |
 | `gradient_jump_penalty`                            | Array of JSON objects, defining additional gradient jump penalty.                                 | See list of gradient jump penalty above                     | -             |
 | `boundary_types`                                   | Boundary types/conditions labels.                                                                 | Array of strings                                            | -             |
-| `velocity_solver.type`                             | Linear solver for the momentum equation.                                                          | `cg`, `pipecg`, `bicgstab`, `cacg`, `gmres`                 | -             |
+| `velocity_solver.type`                             | Linear solver for the momentum equation.                                                          | `cg`, `pipecg`, `bicgstab`, `coupled_bicgstab`, `coupled_cg`, `cacg`, `gmres` | -             |
 | `velocity_solver.preconditioner.type`              | Linear solver preconditioner for the momentum equation.                                           | `ident`, `hsmg`, `jacobi`                                   | -             |
 | `velocity_solver.absolute_tolerance`               | Linear solver convergence criterion for the momentum equation.                                    | Positive real                                               | -             |
 | `velocity_solver.maxiter`                          | Linear solver max iteration count for the momentum equation.                                      | Positive real                                               | 800           |

src/.depends

@@ -226,12 +226,13 @@ krylov/precon.lo : krylov/precon.f90 config/num_types.lo
 krylov/krylov.lo : krylov/krylov.f90 common/log.lo config/neko_config.lo krylov/bcknd/device/pc_identity_device.lo krylov/pc_identity.lo bc/vector_bc_projector.lo bc/scalar_bc_projector.lo common/utils.lo field/field.lo mesh/mesh.lo sem/coef.lo krylov/precon.lo config/num_types.lo math/ax.lo gs/gather_scatter.lo 
 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/device/bicgstab_device.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/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/device/bicgstab_device.lo krylov/bcknd/cpu/bicgstab_coupled.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/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/bicgstab.lo : krylov/bcknd/cpu/bicgstab.f90 comm/comm.lo common/utils.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/cpu/bicgstab_coupled.lo : krylov/bcknd/cpu/bicgstab_coupled.f90 common/utils.lo math/math.lo math/operators.lo comm/comm.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/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/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 

src/Makefile.am

@@ -235,6 +235,7 @@ neko_fortran_SOURCES = \
 	krylov/bcknd/cpu/cheby.f90\
 	krylov/bcknd/cpu/pipecg.f90\
 	krylov/bcknd/cpu/bicgstab.f90\
+	krylov/bcknd/cpu/bicgstab_coupled.f90\
 	krylov/bcknd/cpu/gmres.f90\
 	krylov/bcknd/cpu/pc_jacobi.f90\
 	krylov/bcknd/cpu/cg_coupled.f90\

src/krylov/bcknd/cpu/bicgstab.f90

@@ -42,6 +42,8 @@ module bicgstab
   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, copy, NEKO_EPS, add2s2, p_update
   use utils, only : neko_error
   use comm, only : NEKO_COMM, MPI_EXTRA_PRECISION
@@ -57,19 +59,17 @@ module bicgstab
   !! three components independently and does not apply a coupled operator.
   type, public, extends(ksp_t) :: bicgstab_t
      !> Search direction \f$p\f$.
-     real(kind=rp), allocatable :: p(:)
+     real(kind=rp), pointer :: p(:) => null()
      !> Preconditioned search direction \f$\hat{p} = M^{-1}p\f$.
-     real(kind=rp), allocatable :: p_hat(:)
+     real(kind=rp), pointer :: p_hat(:) => null()
      !> Residual \f$r\f$.
-     real(kind=rp), allocatable :: r(:)
-     !> Intermediate residual \f$s = r - \alpha v\f$.
-     real(kind=rp), allocatable :: s(:)
+     real(kind=rp), pointer :: r(:) => null()
      !> Preconditioned intermediate residual \f$\hat{s} = M^{-1}s\f$.
-     real(kind=rp), allocatable :: s_hat(:)
+     real(kind=rp), pointer :: s_hat(:) => null()
      !> Operator action \f$t = A\hat{s}\f$.
-     real(kind=rp), allocatable :: t(:)
+     real(kind=rp), pointer :: t(:) => null()
      !> Operator action \f$v = A\hat{p}\f$.
-     real(kind=rp), allocatable :: v(:)
+     real(kind=rp), pointer :: v(:) => null()
    contains
      !> Initialise a CPU BiCGStab solver.
      procedure, pass(this) :: init => bicgstab_init
@@ -103,13 +103,6 @@ subroutine bicgstab_init(this, n, max_iter, M, rel_tol, abs_tol, monitor)
 
     call this%free()
 
-    allocate(this%p(n))
-    allocate(this%p_hat(n))
-    allocate(this%r(n))
-    allocate(this%s(n))
-    allocate(this%s_hat(n))
-    allocate(this%t(n))
-    allocate(this%v(n))
     if (present(M)) then
        this%M => M
     end if
@@ -140,36 +133,8 @@ subroutine bicgstab_free(this)
 
     call this%ksp_free()
 
-    if (allocated(this%v)) then
-       deallocate(this%v)
-    end if
-
-    if (allocated(this%r)) then
-       deallocate(this%r)
-    end if
-
-    if (allocated(this%t)) then
-       deallocate(this%t)
-    end if
-
-    if (allocated(this%p)) then
-       deallocate(this%p)
-    end if
-
-    if (allocated(this%p_hat)) then
-       deallocate(this%p_hat)
-    end if
-
-    if (allocated(this%s)) then
-       deallocate(this%s)
-    end if
-
-    if (allocated(this%s_hat)) then
-       deallocate(this%s_hat)
-    end if
-
     nullify(this%M)
-
+    nullify(this%p, this%p_hat, this%r, this%s_hat, this%t, this%v)
 
   end subroutine bicgstab_free
 
@@ -206,6 +171,9 @@ function bicgstab_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) &
     real(kind=rp) :: beta, alpha, omega, rho_1, rho_2
     ! Extra-precision accumulator for the fused residual reductions
     real(kind=xp) :: res_sum
+    type(host_array_t), pointer :: p_tmp, p_hat_tmp, r_tmp
+    type(host_array_t), pointer :: s_hat_tmp, t_tmp, v_tmp
+    integer :: temp_indices(6)
 
     if (present(niter)) then
        max_iter = niter
@@ -214,8 +182,28 @@ function bicgstab_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) &
     end if
     norm_fac = 1.0_rp / sqrt(coef%volume)
 
-    associate(r => this%r, t => this%t, s => this%s, v => this%v, &
-         p => this%p, s_hat => this%s_hat, p_hat => this%p_hat)
+    call neko_scratch_registry%request_host_array(p_tmp, temp_indices(1), &
+         n, .false.)
+    call neko_scratch_registry%request_host_array(p_hat_tmp, &
+         temp_indices(2), n, .false.)
+    call neko_scratch_registry%request_host_array(r_tmp, temp_indices(3), &
+         n, .false.)
+    call neko_scratch_registry%request_host_array(s_hat_tmp, &
+         temp_indices(4), n, .false.)
+    call neko_scratch_registry%request_host_array(t_tmp, temp_indices(5), &
+         n, .false.)
+    call neko_scratch_registry%request_host_array(v_tmp, temp_indices(6), &
+         n, .false.)
+
+    this%p => p_tmp%x
+    this%p_hat => p_hat_tmp%x
+    this%r => r_tmp%x
+    this%s_hat => s_hat_tmp%x
+    this%t => t_tmp%x
+    this%v => v_tmp%x
+
+    associate(p => this%p, p_hat => this%p_hat, r => this%r, &
+         s_hat => this%s_hat, t => this%t, v => this%v)
 
       res_sum = 0.0_xp
       !$omp parallel do reduction(+:res_sum)
@@ -246,6 +234,8 @@ function bicgstab_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) &
       if (r_norm .le. 0.0_rp .or. rnorm .lt. this%abs_tol .or. &
            rnorm .lt. gamma) then
          ksp_results%converged = .true.
+         nullify(this%p, this%p_hat, this%r, this%s_hat, this%t, this%v)
+         call neko_scratch_registry%relinquish_host_array(temp_indices)
          return
       end if
 
@@ -288,11 +278,13 @@ function bicgstab_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) &
             call neko_error('BiCGStab failure: non-finite alpha')
          end if
 
+         ! The previous residual is no longer needed after p has been formed,
+         ! so store the intermediate residual in r.
          res_sum = 0.0_xp
          !$omp parallel do reduction(+:res_sum)
          do i = 1, n
-            s(i) = r(i) - alpha * v(i)
-            res_sum = res_sum + s(i) * coef%mult(i,1,1,1) * s(i)
+            r(i) = r(i) - alpha * v(i)
+            res_sum = res_sum + r(i) * coef%mult(i,1,1,1) * r(i)
          end do
          !$omp end parallel do
 
@@ -308,12 +300,12 @@ function bicgstab_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) &
             exit
          end if
 
-         call this%M%solve(s_hat, s, n)
+         call this%M%solve(s_hat, r, n)
          call Ax%compute(t, s_hat, coef, x%msh, x%Xh)
          call gs_h%op(t, n, GS_OP_ADD)
          call bc_projector%apply(t, n)
 
-         call bicgstab_product_and_norm(stt, ttt, s, t, coef%mult, n)
+         call bicgstab_product_and_norm(stt, ttt, r, t, coef%mult, n)
          t_norm = bicgstab_sqrt(ttt, 'operator result t')
          if (t_norm .le. 0.0_rp) then
             call neko_error('BiCGStab breakdown: zero omega denominator')
@@ -330,7 +322,7 @@ function bicgstab_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) &
          !$omp parallel do reduction(+:res_sum)
          do i = 1, n
             x%x(i,1,1,1) = x%x(i,1,1,1) + alpha * p_hat(i) + omega * s_hat(i)
-            r(i) = s(i) - omega * t(i)
+            r(i) = r(i) - omega * t(i)
             res_sum = res_sum + r(i) * coef%mult(i,1,1,1) * r(i)
          end do
          !$omp end parallel do
@@ -354,11 +346,13 @@ function bicgstab_solve(this, Ax, x, f, n, coef, bc_projector, gs_h, niter) &
          rho_2 = rho_1
 
       end do
+      call this%monitor_stop()
+      ksp_results%res_final = rnorm
+      ksp_results%iter = iter
+      ksp_results%converged = this%is_converged(iter, rnorm)
     end associate
-    call this%monitor_stop()
-    ksp_results%res_final = rnorm
-    ksp_results%iter = iter
-    ksp_results%converged = this%is_converged(iter, rnorm)
+    nullify(this%p, this%p_hat, this%r, this%s_hat, this%t, this%v)
+    call neko_scratch_registry%relinquish_host_array(temp_indices)
   end function bicgstab_solve
 
   !> Check an inner product for a BiCGStab breakdown.