hecmw_precond_SAAMG_matrix.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_matrix
  use hecmw_precond_saamg_util
  use hecmw_precond_saamg_comm, only: hecmwst_saamg_comm, hecmw_saamg_comm_update, &
       hecmw_saamg_abort, hecmw_saamg_check_alloc
  implicit none
 
  private
  public :: hecmwst_saamg_bcsr
  public :: hecmw_saamg_bcsr_free
  public :: hecmw_saamg_bcsr_copy
  public :: hecmw_saamg_bcsr_move
  public :: hecmw_saamg_bcsr_from_dense
  public :: hecmw_saamg_bcsr_from_triplets
  public :: hecmw_saamg_bcsr_from_block_triplets
  public :: hecmw_saamg_bcsr_transpose
  public :: hecmw_saamg_bcsr_to_dense
  public :: hecmw_saamg_matvec
  public :: hecmw_saamg_matvec_d
  public :: hecmw_saamg_spgemm
  public :: hecmw_saamg_galerkin_local
  public :: hecmw_saamg_is_symmetric
  ! block-CSR kernels (also reachable directly; the names above delegate to these)
  public :: hecmw_saamg_to_dense_blk
  public :: hecmw_saamg_matvec_blk
  public :: hecmw_saamg_transpose_blk
  public :: hecmw_saamg_spgemm_blk
  public :: hecmw_saamg_galerkin_blk
  public :: hecmw_saamg_triple_blk
  public :: hecmw_saamg_transpose_blk_rows
 
  type hecmwst_saamg_bcsr
    integer(kind=kint) :: n     = 0
    integer(kind=kint) :: ncol  = 0
    integer(kind=kint) :: nb    = 0
    integer(kind=kint) :: mb    = 0
    integer(kind=kint) :: nnz   = 0
    integer(kind=kint) :: nbrow = 0
    integer(kind=kint) :: nbcol = 0
    integer(kind=kint) :: nnzb  = 0
    integer(kind=kint), allocatable :: browptr(:)
    integer(kind=kint), allocatable :: bcol(:)
    real(kind=kreal),   allocatable :: bval(:)     
  end type hecmwst_saamg_bcsr
 
contains
 
  subroutine hecmw_saamg_bcsr_free(A)
    implicit none
    type(hecmwst_saamg_bcsr), intent(inout) :: a
    if (allocated(a%browptr)) deallocate(a%browptr)
    if (allocated(a%bcol))    deallocate(a%bcol)
    if (allocated(a%bval))    deallocate(a%bval)
    a%n = 0; a%ncol = 0; a%nb = 0; a%mb = 0; a%nnz = 0
    a%nbrow = 0; a%nbcol = 0; a%nnzb = 0
  end subroutine hecmw_saamg_bcsr_free
 
  subroutine hecmw_saamg_bcsr_copy(A, B)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a
    type(hecmwst_saamg_bcsr), intent(out) :: b
    integer(kind=kint) :: astat
    call hecmw_saamg_bcsr_free(b)
    b%n = a%n; b%ncol = a%ncol; b%nb = a%nb; b%mb = a%mb; b%nnz = a%nnz
    b%nbrow = a%nbrow; b%nbcol = a%nbcol; b%nnzb = a%nnzb
    if (allocated(a%browptr)) then
      allocate(b%browptr(a%nbrow+1), b%bcol(a%nnzb), b%bval(a%nb*a%mb*a%nnzb), stat=astat)
      call hecmw_saamg_check_alloc(astat, 'bcsr_copy (browptr/bcol/bval)')
      b%browptr = a%browptr; b%bcol = a%bcol; b%bval = a%bval
    end if
  end subroutine hecmw_saamg_bcsr_copy
 
  subroutine hecmw_saamg_bcsr_move(A, B)
    implicit none
    type(hecmwst_saamg_bcsr), intent(inout) :: a
    type(hecmwst_saamg_bcsr), intent(out)   :: b
    call hecmw_saamg_bcsr_free(b)
    b%n = a%n; b%ncol = a%ncol; b%nb = a%nb; b%mb = a%mb; b%nnz = a%nnz
    b%nbrow = a%nbrow; b%nbcol = a%nbcol; b%nnzb = a%nnzb
    if (allocated(a%browptr)) call move_alloc(a%browptr, b%browptr)
    if (allocated(a%bcol))    call move_alloc(a%bcol,    b%bcol)
    if (allocated(a%bval))    call move_alloc(a%bval,    b%bval)
    call hecmw_saamg_bcsr_free(a)
  end subroutine hecmw_saamg_bcsr_move
 
  subroutine hecmw_saamg_bcsr_from_dense(dense, n, nb, A, mb)
    implicit none
    integer(kind=kint), intent(in)  :: n, nb
    real(kind=kreal),   intent(in)  :: dense(n,n)
    type(hecmwst_saamg_bcsr),     intent(out) :: a
    integer(kind=kint), optional, intent(in) :: mb
    real(kind=kreal), allocatable :: blk(:)
    integer(kind=kint) :: mb_, nbrow, nbcol, bs, i, j, ii, jj, kout, rbase, cbase
    logical :: any_nz
 
    mb_ = nb; if (present(mb)) mb_ = mb
    call hecmw_saamg_bcsr_free(a)
    a%n = n; a%ncol = n; a%nb = nb; a%mb = mb_
    nbrow = n / nb; nbcol = n / mb_; bs = nb*mb_
    a%nbrow = nbrow; a%nbcol = nbcol
    allocate(a%browptr(nbrow+1), blk(bs))
 
    ! pass 1: count nonzero blocks per block row
    a%browptr(1) = 1; a%nnzb = 0
    do i = 1, nbrow
      do j = 1, nbcol
        if (block_nonzero(i, j)) a%nnzb = a%nnzb + 1
      end do
      a%browptr(i+1) = a%nnzb + 1
    end do
    allocate(a%bcol(a%nnzb), a%bval(bs*a%nnzb)); a%nnz = bs*a%nnzb
 
    ! pass 2: copy each nonzero block
    kout = 0
    do i = 1, nbrow
      rbase = (i-1)*nb
      do j = 1, nbcol
        cbase = (j-1)*mb_; any_nz = .false.
        do jj = 1, mb_
          do ii = 1, nb
            blk((jj-1)*nb+ii) = dense(rbase+ii, cbase+jj)
            if (blk((jj-1)*nb+ii) /= 0.0d0) any_nz = .true.
          end do
        end do
        if (any_nz) then
          kout = kout + 1; a%bcol(kout) = j
          a%bval((kout-1)*bs+1:kout*bs) = blk(1:bs)
        end if
      end do
    end do
    deallocate(blk)
  contains
    logical function block_nonzero(ib, jb) result(nz)
      integer(kind=kint), intent(in) :: ib, jb
      integer(kind=kint) :: a, b
      nz = .false.
      do b = 1, mb_
        do a = 1, nb
          if (dense((ib-1)*nb+a, (jb-1)*mb_+b) /= 0.0d0) then; nz = .true.; return; end if
        end do
      end do
    end function block_nonzero
  end subroutine hecmw_saamg_bcsr_from_dense
 
  subroutine hecmw_saamg_matvec(A, x, y)
    implicit none
    type(hecmwst_saamg_bcsr),   intent(in)  :: a
    real(kind=kreal), intent(in)  :: x(:)
    real(kind=kreal), intent(out) :: y(:)
    call hecmw_saamg_matvec_blk(a, x, y)
  end subroutine hecmw_saamg_matvec
 
  subroutine hecmw_saamg_matvec_d(cmt, A, x, y)
    implicit none
    type(hecmwst_saamg_comm), intent(in)    :: cmt
    type(hecmwst_saamg_bcsr),  intent(in)    :: a
    real(kind=kreal),       intent(inout) :: x(:)
    real(kind=kreal),       intent(out)   :: y(:)
    call hecmw_saamg_comm_update(cmt, a%nb, x)
    call hecmw_saamg_matvec(a, x, y)
  end subroutine hecmw_saamg_matvec_d
 
  function hecmw_saamg_is_symmetric(A) result(rel)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in) :: a
    real(kind=kreal) :: rel
    real(kind=kreal) :: nrm, dif, aij, aji
    integer(kind=kint) :: i, j, t, tt, ii, jj, m, boff, foff
    logical :: found
 
    m = a%nb                                   ! square operator: nb == mb
    nrm = 0.0d0; dif = 0.0d0
    do i = 1, a%nbrow
      do t = a%browptr(i), a%browptr(i+1)-1
        j = a%bcol(t); boff = (t-1)*m*m
        ! locate the transposed block A(j,i)
        found = .false.; foff = 0
        do tt = a%browptr(j), a%browptr(j+1)-1
          if (a%bcol(tt) == i) then; found = .true.; foff = (tt-1)*m*m; exit; end if
        end do
        do jj = 1, m
          do ii = 1, m
            aij = a%bval(boff + (jj-1)*m + ii)
            nrm = nrm + aij*aij
            aji = 0.0d0
            if (found) aji = a%bval(foff + (ii-1)*m + jj)   ! A(j,i) entry (jj,ii)
            dif = dif + (aij - aji)**2
          end do
        end do
      end do
    end do
    if (nrm > 0.0d0) then
      rel = sqrt(dif) / sqrt(nrm)
    else
      rel = 0.0d0
    end if
  end function hecmw_saamg_is_symmetric
 
  subroutine hecmw_saamg_bcsr_from_triplets(nrow, ncol, nb, ti, tj, tv, nt, A, mb)
    implicit none
    integer(kind=kint),    intent(in)  :: nrow, ncol, nb, nt
    integer(kind=kint),    intent(in)  :: ti(:), tj(:)
    real(kind=kreal),      intent(in)  :: tv(:)
    type(hecmwst_saamg_bcsr), intent(out) :: a
    integer(kind=kint), optional, intent(in) :: mb
    integer(kind=kint), allocatable :: rcnt(:), rpos(:), jtmp(:), otmp(:), mark(:), touched(:)
    real(kind=kreal),   allocatable :: vtmp(:)
    integer(kind=kint) :: mb_, nbrow, nbcol, bs, t, i, br, rr, bc, cc, j, p, q, ntouch, kout, blk, astat
 
    mb_ = nb; if (present(mb)) mb_ = mb
    call hecmw_saamg_bcsr_free(a)
    a%n = nrow; a%ncol = ncol; a%nb = nb; a%mb = mb_
    nbrow = nrow / nb; nbcol = ncol / mb_; bs = nb*mb_
    a%nbrow = nbrow; a%nbcol = nbcol
 
    ! bucket triplets by block row, storing (block col, in-block offset, value)
    allocate(rcnt(nbrow+1)); rcnt = 0
    do t = 1, nt
      br = (ti(t)-1)/nb + 1; rcnt(br+1) = rcnt(br+1) + 1
    end do
    do i = 1, nbrow
      rcnt(i+1) = rcnt(i+1) + rcnt(i)
    end do
    allocate(jtmp(nt), otmp(nt), vtmp(nt), rpos(nbrow), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'bcsr_from_triplets (jtmp/otmp/vtmp)')
    do i = 1, nbrow
      rpos(i) = rcnt(i)
    end do
    do t = 1, nt
      br = (ti(t)-1)/nb + 1; rr = ti(t) - (br-1)*nb
      bc = (tj(t)-1)/mb_ + 1; cc = tj(t) - (bc-1)*mb_
      rpos(br) = rpos(br) + 1
      jtmp(rpos(br)) = bc; otmp(rpos(br)) = (cc-1)*nb + rr; vtmp(rpos(br)) = tv(t)
    end do
 
    ! count distinct block columns per block row
    allocate(mark(nbcol), touched(nbcol)); mark = 0
    a%nnzb = 0
    do i = 1, nbrow
      ntouch = 0
      do p = rcnt(i)+1, rcnt(i+1)
        if (mark(jtmp(p)) == 0) then
          ntouch = ntouch + 1; touched(ntouch) = jtmp(p); mark(jtmp(p)) = 1
        end if
      end do
      a%nnzb = a%nnzb + ntouch
      do q = 1, ntouch
        mark(touched(q)) = 0
      end do
    end do
 
    allocate(a%browptr(nbrow+1), a%bcol(a%nnzb), a%bval(bs*a%nnzb), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'from_triplets (bcol/bval)')
    a%bval = 0.0d0; a%browptr(1) = 1; kout = 0
    do i = 1, nbrow
      ntouch = 0
      do p = rcnt(i)+1, rcnt(i+1)
        j = jtmp(p)
        if (mark(j) == 0) then
          ntouch = ntouch + 1; touched(ntouch) = j; mark(j) = ntouch
          kout = kout + 1; a%bcol(kout) = j
        end if
        blk = kout - ntouch + mark(j)
        a%bval((blk-1)*bs + otmp(p)) = a%bval((blk-1)*bs + otmp(p)) + vtmp(p)
      end do
      do q = 1, ntouch
        mark(touched(q)) = 0
      end do
      a%browptr(i+1) = kout + 1
    end do
    a%nnz = bs*a%nnzb
 
    deallocate(rcnt, rpos, jtmp, otmp, vtmp, mark, touched)
  end subroutine hecmw_saamg_bcsr_from_triplets
 
  subroutine hecmw_saamg_bcsr_from_block_triplets(nbrow, nbcol, nb, mb, bi, bj, bv, nt, A)
    implicit none
    integer(kind=kint),    intent(in)  :: nbrow, nbcol, nb, mb, nt
    integer(kind=kint),    intent(in)  :: bi(:), bj(:)
    real(kind=kreal),      intent(in)  :: bv(:,:)   ! (nb*mb, nt)
    type(hecmwst_saamg_bcsr), intent(out) :: a
    integer(kind=kint), allocatable :: rcnt(:), rpos(:), btmp(:), mark(:), touched(:)
    real(kind=kreal),   allocatable :: vtmp(:,:)
    integer(kind=kint) :: t, i, j, p, q, r, ntouch, kout, bs, astat
 
    bs = nb*mb
    call hecmw_saamg_bcsr_free(a)
    a%n = nbrow*nb; a%ncol = nbcol*mb; a%nb = nb; a%mb = mb
    a%nbrow = nbrow; a%nbcol = nbcol
 
    ! bucket triplets by block row (unsorted, with duplicates)
    allocate(rcnt(nbrow+1)); rcnt = 0
    do t = 1, nt
      rcnt(bi(t)+1) = rcnt(bi(t)+1) + 1
    end do
    do i = 1, nbrow
      rcnt(i+1) = rcnt(i+1) + rcnt(i)
    end do
    allocate(btmp(nt), vtmp(bs, nt), rpos(nbrow), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'bcsr_from_block_triplets (btmp/vtmp)')
    do i = 1, nbrow
      rpos(i) = rcnt(i)
    end do
    do t = 1, nt
      i = bi(t)
      rpos(i) = rpos(i) + 1
      btmp(rpos(i)) = bj(t)
      vtmp(1:bs, rpos(i)) = bv(1:bs, t)
    end do
 
    ! count distinct block columns per row
    allocate(mark(nbcol), touched(nbcol)); mark = 0
    a%nnzb = 0
    do i = 1, nbrow
      ntouch = 0
      do p = rcnt(i)+1, rcnt(i+1)
        if (mark(btmp(p)) == 0) then
          ntouch = ntouch + 1; touched(ntouch) = btmp(p); mark(btmp(p)) = 1
        end if
      end do
      a%nnzb = a%nnzb + ntouch
      do q = 1, ntouch
        mark(touched(q)) = 0
      end do
    end do
 
    allocate(a%browptr(nbrow+1), a%bcol(a%nnzb), a%bval(bs*a%nnzb), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'from_triplets (bcol/bval)')
    a%bval = 0.0d0; a%browptr(1) = 1; kout = 0
    do i = 1, nbrow
      ntouch = 0
      do p = rcnt(i)+1, rcnt(i+1)
        j = btmp(p)
        if (mark(j) == 0) then
          ntouch = ntouch + 1; touched(ntouch) = j; mark(j) = ntouch
          kout = kout + 1
          a%bcol(kout) = j
          a%bval((kout-1)*bs+1:kout*bs) = vtmp(1:bs, p)
        else
          r = kout - ntouch + mark(j)
          a%bval((r-1)*bs+1:r*bs) = a%bval((r-1)*bs+1:r*bs) + vtmp(1:bs, p)
        end if
      end do
      do q = 1, ntouch
        mark(touched(q)) = 0
      end do
      a%browptr(i+1) = kout + 1
    end do
    a%nnz = bs*a%nnzb
 
    deallocate(rcnt, rpos, btmp, vtmp, mark, touched)
  end subroutine hecmw_saamg_bcsr_from_block_triplets
 
  subroutine hecmw_saamg_bcsr_transpose(A, At)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a
    type(hecmwst_saamg_bcsr), intent(out) :: at
    call hecmw_saamg_transpose_blk(a, at)
  end subroutine hecmw_saamg_bcsr_transpose
 
  subroutine hecmw_saamg_spgemm(A, B, C)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a, b
    type(hecmwst_saamg_bcsr), intent(out) :: c
    call hecmw_saamg_spgemm_blk(a, b, c)
  end subroutine hecmw_saamg_spgemm
 
  subroutine hecmw_saamg_galerkin_local(A, P, mblk, Ac)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a, p
    integer(kind=kint),    intent(in)  :: mblk
    type(hecmwst_saamg_bcsr), intent(out) :: ac
    call hecmw_saamg_galerkin_blk(a, p, mblk, ac)
  end subroutine hecmw_saamg_galerkin_local
 
  subroutine hecmw_saamg_bcsr_to_dense(A, dense)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a
    real(kind=kreal),      intent(out) :: dense(:,:)
    call hecmw_saamg_to_dense_blk(a, dense)
  end subroutine hecmw_saamg_bcsr_to_dense
 
  !-----------------------------------------------------------------------------
  ! block-CSR kernels
  !-----------------------------------------------------------------------------
 
  subroutine hecmw_saamg_to_dense_blk(A, dense)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a
    real(kind=kreal),      intent(out) :: dense(:,:)
    integer(kind=kint) :: i, t, j, ii, jj, rbase, cbase, boff
    dense = 0.0d0
    do i = 1, a%nbrow
      rbase = (i-1)*a%nb
      do t = a%browptr(i), a%browptr(i+1)-1
        j = a%bcol(t); cbase = (j-1)*a%mb; boff = (t-1)*a%nb*a%mb
        do jj = 1, a%mb
          do ii = 1, a%nb
            dense(rbase+ii, cbase+jj) = dense(rbase+ii, cbase+jj) + a%bval(boff+(jj-1)*a%nb+ii)
          end do
        end do
      end do
    end do
  end subroutine hecmw_saamg_to_dense_blk
 
  subroutine hecmw_saamg_matvec_blk(A, x, y)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a
    real(kind=kreal),      intent(in)  :: x(:)
    real(kind=kreal),      intent(out) :: y(:)
    integer(kind=kint) :: i, t, j, ii, jj, rbase, cbase, boff, nb, mb
    real(kind=kreal) :: xj
    nb = a%nb; mb = a%mb
    ! GPU (OpenACC) vs CPU (OpenMP) selected by #ifdef _OPENACC (mutually exclusive,
    ! so a build with both -acc and -mp is safe -- matches upstream las_33.F90).
    ! Arrays are unified (-gpu=mem:managed) so no data clauses; inner block-GEMV loops
    ! stay seq (accumulation, bitwise-stable).
#ifdef _OPENACC
    !$acc parallel loop gang vector private(t,j,ii,jj,rbase,cbase,boff,xj)
#else
    !$omp parallel do default(none) shared(A, x, y, nb, mb) private(i,t,j,ii,jj,rbase,cbase,boff,xj)
#endif
    do i = 1, a%nbrow
      rbase = (i-1)*nb
      y(rbase+1:rbase+nb) = 0.0d0
      !$acc loop seq
      do t = a%browptr(i), a%browptr(i+1)-1
        j = a%bcol(t); cbase = (j-1)*mb; boff = (t-1)*nb*mb
        !$acc loop seq
        do jj = 1, mb
          xj = x(cbase+jj)
          !$acc loop seq
          do ii = 1, nb
            y(rbase+ii) = y(rbase+ii) + a%bval(boff+(jj-1)*nb+ii) * xj
          end do
        end do
      end do
    end do
#ifndef _OPENACC
    !$omp end parallel do
#endif
  end subroutine hecmw_saamg_matvec_blk
 
  subroutine hecmw_saamg_transpose_blk(A, At)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a
    type(hecmwst_saamg_bcsr), intent(out) :: at
    integer(kind=kint), allocatable :: pos(:)
    integer(kind=kint) :: i, t, j, ii, jj, nb, mb, p, boff, aoff, astat
    nb = a%nb; mb = a%mb
    call hecmw_saamg_bcsr_free(at)
    at%n = a%ncol; at%ncol = a%n; at%nb = mb; at%mb = nb
    at%nbrow = a%nbcol; at%nbcol = a%nbrow; at%nnzb = a%nnzb; at%nnz = a%nnz
    allocate(at%browptr(at%nbrow+1), at%bcol(at%nnzb), at%bval(nb*mb*at%nnzb), pos(at%nbrow), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'transpose_blk (At)')
    at%browptr = 0
    do t = 1, a%nnzb
      at%browptr(a%bcol(t)+1) = at%browptr(a%bcol(t)+1) + 1
    end do
    at%browptr(1) = 1
    do j = 1, at%nbrow
      at%browptr(j+1) = at%browptr(j+1) + at%browptr(j)
    end do
    do j = 1, at%nbrow
      pos(j) = at%browptr(j)
    end do
    do i = 1, a%nbrow
      do t = a%browptr(i), a%browptr(i+1)-1
        j = a%bcol(t); p = pos(j)
        at%bcol(p) = i
        boff = (p-1)*mb*nb; aoff = (t-1)*nb*mb
        do jj = 1, mb            ! A block (ii,jj) -> At block (jj,ii)
          do ii = 1, nb
            at%bval(boff + (ii-1)*mb + jj) = a%bval(aoff + (jj-1)*nb + ii)
          end do
        end do
        pos(j) = pos(j) + 1
      end do
    end do
    deallocate(pos)
  end subroutine hecmw_saamg_transpose_blk
 
  subroutine hecmw_saamg_spgemm_blk(A, B, C)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a, b
    type(hecmwst_saamg_bcsr), intent(out) :: c
    integer(kind=kint), allocatable :: mark(:), clist(:)
    real(kind=kreal),   allocatable :: spa(:,:)
    integer(kind=kint) :: rb, kb, cb, ncb, i, ta, kk, tb, l, ii, jj, pp, cnt, nci, t, pos
    integer(kind=kint) :: aoff, boff, coff, astat
    real(kind=kreal) :: bkl
 
    if (a%mb /= b%nb .or. a%nbcol /= b%nbrow) then
      write(*,*) 'hecmw_saamg_spgemm_blk: inner dim mismatch (mb,nbcol)=', a%mb, a%nbcol, &
           ' vs (nb,nbrow)=', b%nb, b%nbrow
      call hecmw_saamg_abort('spgemm_blk: inner dimension mismatch')
    end if
    rb = a%nb; kb = a%mb; cb = b%mb; ncb = b%nbcol
    call hecmw_saamg_bcsr_free(c)
    c%n = a%n; c%ncol = b%ncol; c%nb = rb; c%mb = cb
    c%nbrow = a%nbrow; c%nbcol = ncb
    allocate(c%browptr(a%nbrow+1), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'spgemm_blk (browptr)')
 
    ! pass 1: count nonzero blocks per row (rows independent; per-thread mark).  The
    ! mark(l)/=i trick avoids re-zeroing: mark(l) holds the last row that touched l
    ! (distinct per row), so it works unchanged with a thread-private mark.
    !$omp parallel default(none) shared(A, B, C, ncb) private(mark, i, cnt, ta, kk, tb, l) if(A%nbrow > 256)
    allocate(mark(ncb)); mark = 0
    !$omp do
    do i = 1, a%nbrow
      cnt = 0
      do ta = a%browptr(i), a%browptr(i+1)-1
        kk = a%bcol(ta)
        do tb = b%browptr(kk), b%browptr(kk+1)-1
          l = b%bcol(tb)
          if (mark(l) /= i) then; mark(l) = i; cnt = cnt + 1; end if
        end do
      end do
      c%browptr(i+1) = cnt
    end do
    !$omp end do
    deallocate(mark)
    !$omp end parallel
 
    c%browptr(1) = 1                                  ! prefix-sum counts -> row pointers
    do i = 1, a%nbrow
      c%browptr(i+1) = c%browptr(i) + c%browptr(i+1)
    end do
    c%nnzb = c%browptr(a%nbrow+1) - 1
    allocate(c%bcol(c%nnzb), c%bval(rb*cb*c%nnzb), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'spgemm_blk (C bcol/bval)')
    c%nnz = rb*cb*c%nnzb
 
    ! pass 2: accumulate block products (rows independent; each row writes its own
    ! browptr(i)..browptr(i+1)-1 slice; per-thread mark/clist/spa accumulator)
    !$omp parallel default(none) shared(A, B, C, rb, kb, cb, ncb) if(A%nbrow > 256) &
    !$omp&   private(mark, clist, spa, i, nci, ta, kk, aoff, tb, l, boff, jj, pp, bkl, ii, t, pos, coff)
    allocate(mark(ncb), clist(ncb), spa(rb*cb, ncb)); mark = 0
    !$omp do
    do i = 1, a%nbrow
      nci = 0
      do ta = a%browptr(i), a%browptr(i+1)-1
        kk = a%bcol(ta); aoff = (ta-1)*rb*kb
        do tb = b%browptr(kk), b%browptr(kk+1)-1
          l = b%bcol(tb); boff = (tb-1)*kb*cb
          if (mark(l) /= i) then
            mark(l) = i; nci = nci + 1; clist(nci) = l; spa(1:rb*cb, l) = 0.0d0
          end if
          do jj = 1, cb
            do pp = 1, kb
              bkl = b%bval(boff + (jj-1)*kb + pp)
              if (bkl == 0.0d0) cycle
              do ii = 1, rb
                spa((jj-1)*rb+ii, l) = spa((jj-1)*rb+ii, l) + a%bval(aoff+(pp-1)*rb+ii) * bkl
              end do
            end do
          end do
        end do
      end do
      pos = c%browptr(i) - 1
      do t = 1, nci
        pos = pos + 1
        l = clist(t)
        c%bcol(pos) = l
        coff = (pos-1)*rb*cb
        c%bval(coff+1:coff+rb*cb) = spa(1:rb*cb, l)
      end do
    end do
    !$omp end do
    deallocate(mark, clist, spa)
    !$omp end parallel
  end subroutine hecmw_saamg_spgemm_blk
 
  subroutine hecmw_saamg_galerkin_blk(A, P, mblk, Ac)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a, p
    integer(kind=kint),    intent(in)  :: mblk
    type(hecmwst_saamg_bcsr), intent(out) :: ac
    type(hecmwst_saamg_bcsr) :: pt, c
    call hecmw_saamg_transpose_blk(p, pt)   ! Pt = P^T
    call hecmw_saamg_spgemm_blk(a, p, c)    ! C  = A P
    call hecmw_saamg_spgemm_blk(pt, c, ac)  ! Ac = P^T C
    ac%nb = mblk; ac%mb = mblk
    call hecmw_saamg_bcsr_free(pt)
    call hecmw_saamg_bcsr_free(c)
  end subroutine hecmw_saamg_galerkin_blk
 
  subroutine hecmw_saamg_transpose_blk_rows(A, nbrow_keep, At)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: a
    integer(kind=kint),    intent(in)  :: nbrow_keep
    type(hecmwst_saamg_bcsr), intent(out) :: at
    integer(kind=kint), allocatable :: pos(:)
    integer(kind=kint) :: i, t, j, ii, jj, nb, mb, p, boff, aoff, nnzb, astat
    nb = a%nb; mb = a%mb
    nnzb = a%browptr(nbrow_keep+1) - 1
    call hecmw_saamg_bcsr_free(at)
    at%n = a%ncol; at%ncol = nbrow_keep*nb; at%nb = mb; at%mb = nb
    at%nbrow = a%nbcol; at%nbcol = nbrow_keep; at%nnzb = nnzb; at%nnz = nnzb*nb*mb
    allocate(at%browptr(at%nbrow+1), at%bcol(nnzb), at%bval(nb*mb*nnzb), pos(at%nbrow), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'transpose_blk_rows (At)')
    at%browptr = 0
    do t = 1, nnzb
      at%browptr(a%bcol(t)+1) = at%browptr(a%bcol(t)+1) + 1
    end do
    at%browptr(1) = 1
    do j = 1, at%nbrow
      at%browptr(j+1) = at%browptr(j+1) + at%browptr(j)
    end do
    do j = 1, at%nbrow
      pos(j) = at%browptr(j)
    end do
    do i = 1, nbrow_keep
      do t = a%browptr(i), a%browptr(i+1)-1
        j = a%bcol(t); p = pos(j)
        at%bcol(p) = i
        boff = (p-1)*mb*nb; aoff = (t-1)*nb*mb
        do jj = 1, mb
          do ii = 1, nb
            at%bval(boff + (ii-1)*mb + jj) = a%bval(aoff + (jj-1)*nb + ii)
          end do
        end do
        pos(j) = pos(j) + 1
      end do
    end do
    deallocate(pos)
  end subroutine hecmw_saamg_transpose_blk_rows
 
  subroutine hecmw_saamg_triple_blk(L, A, R, Ac)
    implicit none
    type(hecmwst_saamg_bcsr), intent(in)  :: l, a, r
    type(hecmwst_saamg_bcsr), intent(out) :: ac
    integer(kind=kint), allocatable :: mark(:), clist(:)
    real(kind=kreal),   allocatable :: spa(:,:), la(:)
    integer(kind=kint) :: rb, k1, k2, cb, ncr, i, jt, j, kt, k, lt, ll, p, q, rr, c
    integer(kind=kint) :: cnt, nci, t, pos, loff, aoff, roff, coff, astat
    real(kind=kreal) :: av, rv
 
    if (l%mb /= a%nb .or. a%mb /= r%nb .or. l%nbcol /= a%nbrow .or. a%nbcol /= r%nbrow) then
      call hecmw_saamg_abort('triple_blk: dimension mismatch')
    end if
    rb = l%nb; k1 = l%mb; k2 = a%mb; cb = r%mb; ncr = r%nbcol
    call hecmw_saamg_bcsr_free(ac)
    ac%n = l%n; ac%ncol = r%ncol; ac%nb = rb; ac%mb = cb
    ac%nbrow = l%nbrow; ac%nbcol = ncr
    allocate(ac%browptr(l%nbrow+1), stat=astat)
    call hecmw_saamg_check_alloc(astat, 'triple_blk (browptr)')
 
    ! pass 1: count nonzero blocks per row of Ac (rows independent; per-thread mark)
    !$omp parallel default(none) shared(L, A, R, Ac, ncr) private(mark, i, cnt, jt, j, kt, k, lt, ll) if(L%nbrow > 256)
    allocate(mark(ncr)); mark = 0
    !$omp do
    do i = 1, l%nbrow
      cnt = 0
      do jt = l%browptr(i), l%browptr(i+1)-1
        j = l%bcol(jt)
        do kt = a%browptr(j), a%browptr(j+1)-1
          k = a%bcol(kt)
          do lt = r%browptr(k), r%browptr(k+1)-1
            ll = r%bcol(lt)
            if (mark(ll) /= i) then; mark(ll) = i; cnt = cnt + 1; end if
          end do
        end do
      end do
      ac%browptr(i+1) = cnt
    end do
    !$omp end do
    deallocate(mark)
    !$omp end parallel
 
    ac%browptr(1) = 1                                 ! prefix-sum counts -> row pointers
    do i = 1, l%nbrow
      ac%browptr(i+1) = ac%browptr(i) + ac%browptr(i+1)
    end do
    ac%nnzb = ac%browptr(l%nbrow+1) - 1
    allocate(ac%bcol(ac%nnzb), ac%bval(rb*cb*ac%nnzb), stat=astat); ac%nnz = rb*cb*ac%nnzb
    call hecmw_saamg_check_alloc(astat, 'triple_blk (Ac bcol/bval)')
 
    ! pass 2: accumulate L(i,j) A(j,k) R(k,ll) into Ac(i,ll) (rows independent;
    ! each row writes its own slice; per-thread mark/clist/spa/la scratch)
    !$omp parallel default(none) shared(L, A, R, Ac, rb, k1, k2, cb, ncr) if(L%nbrow > 256) &
    !$omp&   private(mark, clist, spa, la, i, nci, jt, j, loff, kt, k, aoff, q, p, av, rr, &
    !$omp&           lt, ll, roff, c, rv, t, pos, coff)
    allocate(mark(ncr), clist(ncr), spa(rb*cb, ncr), la(rb*k2)); mark = 0
    !$omp do
    do i = 1, l%nbrow
      nci = 0
      do jt = l%browptr(i), l%browptr(i+1)-1
        j = l%bcol(jt); loff = (jt-1)*rb*k1
        do kt = a%browptr(j), a%browptr(j+1)-1
          k = a%bcol(kt); aoff = (kt-1)*k1*k2
          ! la = L(i,j) * A(j,k)   (rb x k1) * (k1 x k2) = rb x k2
          la(1:rb*k2) = 0.0d0
          do q = 1, k2
            do p = 1, k1
              av = a%bval(aoff + (q-1)*k1 + p)
              if (av == 0.0d0) cycle
              do rr = 1, rb
                la((q-1)*rb+rr) = la((q-1)*rb+rr) + l%bval(loff+(p-1)*rb+rr) * av
              end do
            end do
          end do
          ! spa(:,ll) += la * R(k,ll)   (rb x k2) * (k2 x cb) = rb x cb
          do lt = r%browptr(k), r%browptr(k+1)-1
            ll = r%bcol(lt); roff = (lt-1)*k2*cb
            if (mark(ll) /= i) then
              mark(ll) = i; nci = nci + 1; clist(nci) = ll; spa(1:rb*cb, ll) = 0.0d0
            end if
            do c = 1, cb
              do q = 1, k2
                rv = r%bval(roff + (c-1)*k2 + q)
                if (rv == 0.0d0) cycle
                do rr = 1, rb
                  spa((c-1)*rb+rr, ll) = spa((c-1)*rb+rr, ll) + la((q-1)*rb+rr) * rv
                end do
              end do
            end do
          end do
        end do
      end do
      pos = ac%browptr(i) - 1
      do t = 1, nci
        pos = pos + 1
        ll = clist(t); ac%bcol(pos) = ll
        coff = (pos-1)*rb*cb
        ac%bval(coff+1:coff+rb*cb) = spa(1:rb*cb, ll)
      end do
    end do
    !$omp end do
    deallocate(mark, clist, spa, la)
    !$omp end parallel
  end subroutine hecmw_saamg_triple_blk
 
end module hecmw_precond_saamg_matrix