hecmw_precond_SAAMG_smoother.F90

Go to the documentation of this file.

!-------------------------------------------------------------------------------
! Copyright (c) 2026 FrontISTR Commons
! This software is released under the MIT License, see License.txt
!-------------------------------------------------------------------------------
module hecmw_precond_saamg_smoother
  use hecmw_precond_saamg_util
  use hecmw_precond_saamg_matrix
  use hecmw_precond_saamg_comm, only: hecmw_saamg_abort, hecmw_saamg_check_alloc
  implicit none
 
  private
  public :: hecmwst_saamg_blockdiag
  public :: hecmw_saamg_blockdiag_setup
  public :: hecmw_saamg_blockdiag_apply
  public :: hecmw_saamg_blockdiag_apply_bcsr
  public :: hecmw_saamg_blockdiag_apply_bcsr_inplace
  public :: hecmw_saamg_blockdiag_matvec
  public :: hecmw_saamg_blockdiag_free
  public :: hecmw_saamg_lanczos_lambda_max
  public :: hecmw_saamg_tridiag_max_eig
  public :: hecmwst_saamg_chebyshev
  public :: hecmw_saamg_cheb_setup
  public :: hecmw_saamg_cheb_apply_local
 
  type hecmwst_saamg_blockdiag
    integer(kind=kint) :: nnode = 0
    integer(kind=kint) :: nb    = 0
    real(kind=kreal),   allocatable :: dblk(:)  
    real(kind=kreal),   allocatable :: inv(:)   
  end type hecmwst_saamg_blockdiag
 
  type hecmwst_saamg_chebyshev
    integer(kind=kint) :: deg   = 2
    real(kind=kreal)   :: theta = 0.0d0 
    real(kind=kreal)   :: delta = 0.0d0 
    real(kind=kreal)   :: sigma = 0.0d0 
  end type hecmwst_saamg_chebyshev
 
contains
 
  !-----------------------------------------------------------------------------
  ! Node-block diagonal D
  !-----------------------------------------------------------------------------
 
  subroutine hecmw_saamg_blockdiag_setup(A, D)
    implicit none
    type(hecmwst_saamg_bcsr),      intent(in)  :: a
    type(hecmwst_saamg_blockdiag), intent(out) :: d
    integer(kind=kint) :: nb, nnode, inode, off, t, info, astat
    logical :: square_blocked
 
    nb = a%nb
    square_blocked = nb > 0 .and. a%mb == nb
    if (square_blocked) square_blocked = mod(a%n, nb) == 0
    if (.not. square_blocked) then
      write(*,*) 'hecmw_saamg_blockdiag_setup: A not square-blocked (nb,mb)=', nb, a%mb
      call hecmw_saamg_abort('blockdiag_setup: operator not square-blocked')
    end if
    nnode = a%nbrow
 
    call hecmw_saamg_blockdiag_free(d)
    d%nb = nb; d%nnode = nnode
    allocate(d%dblk(nb*nb*nnode), d%inv(nb*nb*nnode), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'blockdiag_setup (dblk/inv)')
    d%dblk = 0.0d0
 
    ! each node's block is inverted independently with the BLAS-free block_inverse
    ! (no per-thread workspace, no LAPACK -> no nested BLAS threading, OpenACC-ready)
    !$omp parallel do default(none) shared(A, D, nb, nnode) private(inode, off, t, info)
    do inode = 1, nnode
      off = (inode-1)*nb*nb
      do t = a%browptr(inode), a%browptr(inode+1)-1
        if (a%bcol(t) == inode) then              ! the diagonal block
          d%dblk(off+1:off+nb*nb) = a%bval((t-1)*nb*nb+1:t*nb*nb)
          exit
        end if
      end do
      call block_inverse(nb, d%dblk(off+1:off+nb*nb), d%inv(off+1:off+nb*nb), info)
      if (info /= 0) then
        !$omp critical
        write(*,'(a,i0)') 'hecmw_saamg_blockdiag_setup: singular diagonal block at node ', inode
        !$omp end critical
        call hecmw_saamg_abort('blockdiag_setup: singular diagonal block')
      end if
    end do
    !$omp end parallel do
  end subroutine hecmw_saamg_blockdiag_setup
 
  subroutine block_inverse(nb, a, ainv, info)
    implicit none
    integer(kind=kint), intent(in)  :: nb
    real(kind=kreal),   intent(in)  :: a(nb*nb)        ! col-major a(i,j) = a((j-1)*nb+i)
    real(kind=kreal),   intent(out) :: ainv(nb*nb)     ! col-major
    integer(kind=kint), intent(out) :: info
    real(kind=kreal) :: det, idet, piv, f, aug(nb, 2*nb)
    integer(kind=kint) :: i, j, k, p
    info = 0
    select case (nb)
    case (1)
      if (a(1) == 0.0d0) then; info = 1; return; end if
      ainv(1) = 1.0d0 / a(1)
    case (2)
      det = a(1)*a(4) - a(2)*a(3)
      if (det == 0.0d0) then; info = 1; return; end if
      idet = 1.0d0 / det
      ainv(1) =  a(4)*idet; ainv(2) = -a(2)*idet
      ainv(3) = -a(3)*idet; ainv(4) =  a(1)*idet
    case (3)
      det = a(1)*(a(5)*a(9)-a(8)*a(6)) - a(4)*(a(2)*a(9)-a(8)*a(3)) + a(7)*(a(2)*a(6)-a(5)*a(3))
      if (det == 0.0d0) then; info = 1; return; end if
      idet = 1.0d0 / det
      ainv(1) = (a(5)*a(9)-a(8)*a(6))*idet
      ainv(2) = (a(8)*a(3)-a(2)*a(9))*idet
      ainv(3) = (a(2)*a(6)-a(5)*a(3))*idet
      ainv(4) = (a(7)*a(6)-a(4)*a(9))*idet
      ainv(5) = (a(1)*a(9)-a(7)*a(3))*idet
      ainv(6) = (a(4)*a(3)-a(1)*a(6))*idet
      ainv(7) = (a(4)*a(8)-a(7)*a(5))*idet
      ainv(8) = (a(7)*a(2)-a(1)*a(8))*idet
      ainv(9) = (a(1)*a(5)-a(4)*a(2))*idet
    case default
      ! partial-pivoted Gauss-Jordan on the augmented [a | I] -> [I | a^{-1}]
      do j = 1, nb
        do i = 1, nb
          aug(i, j)    = a((j-1)*nb + i)
          aug(i, nb+j) = 0.0d0
        end do
        aug(j, nb+j) = 1.0d0
      end do
      do k = 1, nb
        p = k                                          ! partial pivot
        do i = k+1, nb
          if (abs(aug(i,k)) > abs(aug(p,k))) p = i
        end do
        if (aug(p,k) == 0.0d0) then; info = 1; return; end if
        if (p /= k) then
          do j = 1, 2*nb
            f = aug(k,j); aug(k,j) = aug(p,j); aug(p,j) = f
          end do
        end if
        piv = aug(k,k)
        do j = 1, 2*nb
          aug(k,j) = aug(k,j) / piv
        end do
        do i = 1, nb
          if (i == k) cycle
          f = aug(i,k)
          if (f == 0.0d0) cycle
          do j = 1, 2*nb
            aug(i,j) = aug(i,j) - f*aug(k,j)
          end do
        end do
      end do
      do j = 1, nb
        do i = 1, nb
          ainv((j-1)*nb + i) = aug(i, nb+j)
        end do
      end do
    end select
  end subroutine block_inverse
 
  subroutine hecmw_saamg_blockdiag_apply(D, x, y)
    implicit none
    type(hecmwst_saamg_blockdiag), intent(in)  :: d
    real(kind=kreal),      intent(in)  :: x(:)
    real(kind=kreal),      intent(out) :: y(:)
    integer(kind=kint) :: inode, base, ii, jj, off, nb
    real(kind=kreal) :: s
 
    nb = d%nb
    ! GPU (OpenACC) vs CPU (OpenMP) via #ifdef _OPENACC (mutually exclusive; safe with
    ! both -acc and -mp).  Managed memory => no data clauses; inner nb x nb seq.
#ifdef _OPENACC
    !$acc parallel loop gang vector private(base, off, ii, jj, s)
#else
    !$omp parallel do default(none) shared(D, x, y, nb) private(inode, base, off, ii, jj, s)
#endif
    do inode = 1, d%nnode
      base = (inode-1)*nb
      off  = (inode-1)*nb*nb
      !$acc loop seq
      do ii = 1, nb                              ! y_i = sum_j Dinv(i,j) x_j  (col-major)
        s = 0.0d0
        !$acc loop seq
        do jj = 1, nb
          s = s + d%inv(off + (jj-1)*nb + ii) * x(base+jj)
        end do
        y(base+ii) = s
      end do
    end do
#ifndef _OPENACC
    !$omp end parallel do
#endif
  end subroutine hecmw_saamg_blockdiag_apply
 
  subroutine hecmw_saamg_blockdiag_apply_bcsr(D, C, DC)
    implicit none
    type(hecmwst_saamg_blockdiag), intent(in)  :: d
    type(hecmwst_saamg_bcsr),      intent(in)  :: c
    type(hecmwst_saamg_bcsr),      intent(out) :: dc
    integer(kind=kint) :: nb, mb, inode, off, t, b0, ii, jj, cc, astat
    real(kind=kreal) :: s
 
    nb = d%nb; mb = c%mb
    if (c%nb /= nb .or. c%nbrow /= d%nnode) then
      write(*,*) 'hecmw_saamg_blockdiag_apply_bcsr: block-row mismatch', c%nb, nb, &
           c%nbrow, d%nnode
      call hecmw_saamg_abort('blockdiag_apply_bcsr: block-row mismatch')
    end if
 
    ! DC = D^{-1} C : same block structure as C, each block DC_IK = D_I^{-1} C_IK
    ! (nb x nb inverse times the nb x mb block, small dense matmul)
    call hecmw_saamg_bcsr_free(dc)
    dc%n = c%n; dc%ncol = c%ncol; dc%nb = nb; dc%mb = mb
    dc%nbrow = c%nbrow; dc%nbcol = c%nbcol; dc%nnzb = c%nnzb
    allocate(dc%browptr(c%nbrow+1), dc%bcol(c%nnzb), dc%bval(nb*mb*c%nnzb), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'blockdiag_apply_bcsr (DC)')
    dc%browptr = c%browptr; dc%bcol = c%bcol
 
    !$omp parallel do default(none) shared(C, DC, D, nb, mb) private(inode, off, t, b0, cc, ii, jj, s)
    do inode = 1, c%nbrow
      off = (inode-1)*nb*nb
      do t = c%browptr(inode), c%browptr(inode+1)-1
        b0 = (t-1)*nb*mb
        do cc = 1, mb                            ! DC(i,cc) = sum_j Dinv(i,j) C(j,cc)
          do ii = 1, nb
            s = 0.0d0
            do jj = 1, nb
              s = s + d%inv(off + (jj-1)*nb + ii) * c%bval(b0 + (cc-1)*nb + jj)
            end do
            dc%bval(b0 + (cc-1)*nb + ii) = s
          end do
        end do
      end do
    end do
    !$omp end parallel do
    ! block-only result: the sole consumer (smoothed prolongator) reads bval
  end subroutine hecmw_saamg_blockdiag_apply_bcsr
 
  subroutine hecmw_saamg_blockdiag_apply_bcsr_inplace(D, C)
    implicit none
    type(hecmwst_saamg_blockdiag), intent(in)    :: d
    type(hecmwst_saamg_bcsr),      intent(inout) :: c
    integer(kind=kint) :: nb, mb, inode, off, t, b0, ii, jj, cc
    real(kind=kreal) :: tmp(d%nb), s
    nb = d%nb; mb = c%mb
    if (c%nb /= nb .or. c%nbrow /= d%nnode) then
      write(*,*) 'hecmw_saamg_blockdiag_apply_bcsr_inplace: block-row mismatch', &
           c%nb, nb, c%nbrow, d%nnode
      call hecmw_saamg_abort('blockdiag_apply_bcsr_inplace: block-row mismatch')
    end if
    !$omp parallel do default(none) shared(C, D, nb, mb) private(inode, off, t, b0, cc, ii, jj, s, tmp)
    do inode = 1, c%nbrow
      off = (inode-1)*nb*nb
      do t = c%browptr(inode), c%browptr(inode+1)-1
        b0 = (t-1)*nb*mb
        do cc = 1, mb                            ! col cc <- Dinv * col cc (temp: read before overwrite)
          do jj = 1, nb
            tmp(jj) = c%bval(b0 + (cc-1)*nb + jj)
          end do
          do ii = 1, nb
            s = 0.0d0
            do jj = 1, nb
              s = s + d%inv(off + (jj-1)*nb + ii) * tmp(jj)
            end do
            c%bval(b0 + (cc-1)*nb + ii) = s
          end do
        end do
      end do
    end do
    !$omp end parallel do
  end subroutine hecmw_saamg_blockdiag_apply_bcsr_inplace
 
  subroutine hecmw_saamg_blockdiag_matvec(D, x, y)
    implicit none
    type(hecmwst_saamg_blockdiag), intent(in)  :: d
    real(kind=kreal),      intent(in)  :: x(:)
    real(kind=kreal),      intent(out) :: y(:)
    integer(kind=kint) :: inode, base, ii, jj, off
    real(kind=kreal) :: s
 
    !$omp parallel do default(none) shared(D, x, y) private(inode, base, off, ii, jj, s)
    do inode = 1, d%nnode
      base = (inode-1)*d%nb
      off  = (inode-1)*d%nb*d%nb
      do ii = 1, d%nb
        s = 0.0d0
        do jj = 1, d%nb
          s = s + d%dblk(off + (jj-1)*d%nb + ii) * x(base+jj)
        end do
        y(base+ii) = s
      end do
    end do
    !$omp end parallel do
  end subroutine hecmw_saamg_blockdiag_matvec
 
  subroutine hecmw_saamg_blockdiag_free(D)
    implicit none
    type(hecmwst_saamg_blockdiag), intent(inout) :: d
    if (allocated(d%dblk)) deallocate(d%dblk)
    if (allocated(d%inv))  deallocate(d%inv)
    d%nnode = 0; d%nb = 0
  end subroutine hecmw_saamg_blockdiag_free
 
  !-----------------------------------------------------------------------------
  ! Lanczos estimate of lambda_max(D^{-1} A)
  !-----------------------------------------------------------------------------
 
  function hecmw_saamg_tridiag_max_eig(diag, offd, m) result(lmax)
    implicit none
    real(kind=kreal),   intent(in) :: diag(:), offd(:)
    integer(kind=kint), intent(in) :: m
    real(kind=kreal) :: lmax
    real(kind=kreal), allocatable :: d(:), e(:)
    integer(kind=kint) :: info
#ifdef HECMW_WITH_LAPACK
    external :: dsterf
    if (m <= 1) then; lmax = diag(1); return; end if
    allocate(d(m), e(m-1))
    d(1:m) = diag(1:m); e(1:m-1) = offd(1:m-1)
    call dsterf(m, d, e, info)                 ! eigenvalues ascending in d
    if (info == 0) then; lmax = d(m); else; lmax = maxval(diag(1:m)); end if
    deallocate(d, e)
#else
    ! LAPACK-free fallback (largest diagonal = Rayleigh lower bound).  Unreached in
    ! practice: the backend aborts before setup when LAPACK is absent.
    lmax = maxval(diag(1:m))
#endif
  end function hecmw_saamg_tridiag_max_eig
 
  function hecmw_saamg_lanczos_lambda_max(A, D, niter) result(lambda)
    implicit none
    type(hecmwst_saamg_bcsr),      intent(in) :: a
    type(hecmwst_saamg_blockdiag), intent(in) :: d
    integer(kind=kint),    intent(in) :: niter
    real(kind=kreal) :: lambda
    real(kind=kreal), allocatable :: q(:), qp(:), w(:), aq(:), dv(:), alpha(:), beta(:)
    real(kind=kreal) :: qdq, wdw, bprev
    integer(kind=kint) :: j, i, m
    integer(kind=8) :: seed
 
    allocate(q(a%n), qp(a%n), w(a%n), aq(a%n), dv(a%n), alpha(niter), beta(niter))
 
    ! fixed-seed deterministic start (tiny LCG) mapped to (-1,1)
    seed = 88172645463325252_8
    do i = 1, a%n
      seed = seed*6364136223846793005_8 + 1442695040888963407_8
      q(i) = real(seed, kreal)/9.223372036854776d18
    end do
    ! D-normalize the start vector: <q,q>_D = q^T D q
    call hecmw_saamg_blockdiag_matvec(d, q, dv); qdq = dot_product(q, dv)
    if (qdq <= 0.0d0) then; lambda = 0.0d0; deallocate(q,qp,w,aq,dv,alpha,beta); return; end if
    q = q / sqrt(qdq); qp = 0.0d0; bprev = 0.0d0; m = 0
 
    do j = 1, niter
      call hecmw_saamg_matvec(a, q, aq)             ! aq = A q
      call hecmw_saamg_blockdiag_apply(d, aq, w)    ! w  = D^{-1} A q = M q
      alpha(j) = dot_product(aq, q)                 ! <M q, q>_D = (A q)^T q
      w = w - alpha(j)*q - bprev*qp                 ! orthogonalize (3-term recurrence)
      m = j
      call hecmw_saamg_blockdiag_matvec(d, w, dv); wdw = dot_product(w, dv)
      if (wdw <= 0.0d0) exit                        ! invariant subspace -> exact
      beta(j) = sqrt(wdw)
      if (j == niter) exit
      qp = q; bprev = beta(j); q = w / beta(j)
    end do
 
    lambda = hecmw_saamg_tridiag_max_eig(alpha, beta, m)
    deallocate(q, qp, w, aq, dv, alpha, beta)
  end function hecmw_saamg_lanczos_lambda_max
 
  !-----------------------------------------------------------------------------
  ! Chebyshev smoother
  !-----------------------------------------------------------------------------
 
  subroutine hecmw_saamg_cheb_setup(lambda_hat, deg, alpha, cheb)
    implicit none
    real(kind=kreal),      intent(in)  :: lambda_hat
    integer(kind=kint),    intent(in)  :: deg
    real(kind=kreal),      intent(in)  :: alpha
    type(hecmwst_saamg_chebyshev), intent(out) :: cheb
    real(kind=kreal) :: lmax, lmin
    lmax = 1.1d0 * lambda_hat
    lmin = lmax / alpha
    cheb%deg   = deg
    cheb%theta = 0.5d0*(lmax + lmin)
    cheb%delta = 0.5d0*(lmax - lmin)
    cheb%sigma = cheb%theta / cheb%delta
  end subroutine hecmw_saamg_cheb_setup
 
  subroutine hecmw_saamg_cheb_apply_local(A, D, cheb, b, x, zero_init)
    implicit none
    type(hecmwst_saamg_bcsr),      intent(in)    :: a
    type(hecmwst_saamg_blockdiag), intent(in)    :: d
    type(hecmwst_saamg_chebyshev), intent(in)    :: cheb
    real(kind=kreal),      intent(in)    :: b(:)
    real(kind=kreal),      intent(inout) :: x(:)
    logical,               intent(in)    :: zero_init
    real(kind=kreal), allocatable :: r(:), p(:), dr(:), ap(:)
    real(kind=kreal) :: rho, rho2, c1, c2
    integer(kind=kint) :: it, n
 
    n = a%n
    allocate(r(n), p(n), dr(n), ap(n))
 
    if (zero_init) then
      x = 0.0d0
      r = b
    else
      call hecmw_saamg_matvec(a, x, ap)
      r = b - ap
    end if
 
    call hecmw_saamg_blockdiag_apply(d, r, dr)
    p   = dr / cheb%theta
    rho = 1.0d0 / cheb%sigma
 
    do it = 1, cheb%deg
      x = x + p
      call hecmw_saamg_matvec(a, p, ap)
      r = r - ap
      rho2 = 1.0d0 / (2.0d0*cheb%sigma - rho)
      call hecmw_saamg_blockdiag_apply(d, r, dr)
      c1 = rho2 * rho
      c2 = 2.0d0 * rho2 / cheb%delta
      p  = c1 * p + c2 * dr
      rho = rho2
    end do
 
    deallocate(r, p, dr, ap)
  end subroutine hecmw_saamg_cheb_apply_local
 
end module hecmw_precond_saamg_smoother