amg.F90

Go to the documentation of this file.

! This module implements an aggregation-based algebraic multigrid method.
! See Sec. 6.1 in the following paper for a theoretical overview of similar methods:
!   K. Stuben. "A review of algebraic multigrid."
!   Journal of Computational and Applied Mathematics (2001)
!   https://doi.org/10.1016/S0377-0427(00)00516-1
module amg
 
    use constants
    use utils, only: echk
#include <petsc/finclude/petsc.h>
    use petsc
    implicit none
 
    ! Structure used for storing the interpolation indices
    type arr1int4
        integer, dimension(:), allocatable :: arr
    end type arr1int4
 
    ! Structure used for storing the interpolation indices
    type arr3int4
        integer, dimension(:, :, :), allocatable :: arr
    end type arr3int4
 
    ! Structure used for storing the interpolation indices
    type arr4int4
        integer, dimension(:, :, :, :), allocatable :: arr
    end type arr4int4
 
    ! The number of levels
    integer(kind=intType) amglevels
 
    ! The number of outer iterations
    integer(kind=intType) amgouterits
 
    ! The number of smoothing iterations
    integer(kind=intType) amgnsmooth
 
    ! The number of local ILU iterations
    integer(kind=intType) amglocalpreconits
    integer(kind=intType) amglocalpreconitsfine
    integer(kind=intType) amglocalpreconitscoarse
 
    ! ASM overlap for the solver/smoother
    integer(kind=intType) amgasmoverlap
    integer(kind=intType) amgasmoverlapfine
    integer(kind=intType) amgasmoverlapcoarse
 
    ! ILU fill for the solver/smoother
    integer(kind=intType) amgfilllevel
    integer(kind=intType) amgfilllevelfine
    integer(kind=intType) amgfilllevelcoarse
 
    ! Ordering
    character(len=maxStringLen) :: amgmatrixordering
 
    ! Arrays for matrices/vectors/solvers on each level.
    mat, dimension(:), allocatable :: a
    ksp, dimension(:), allocatable :: ksplevels
    vec, dimension(:), allocatable :: res, rhs, sol, sol2
    pc shellpc
 
    mat, pointer :: finemat
 
    ! The interpolation arrays
    type(arr1int4), dimension(:), allocatable :: interps
 
    ! The coarse level indices
    type(arr3int4), dimension(:, :), allocatable, target :: coarseindices
    type(arr4int4), dimension(:, :), allocatable, target :: coarseoversetindices
 
    logical :: amgsetup = .false.
    integer :: bs
contains
 
    subroutine setupamg(inputMat, nCell, blockSize, levels, nSmooth)
 
        use blockpointers
        use communication
        use utils, only: setpointers
        use haloexchange, only: whalo1to1intgeneric
        use adjointvars, only: ncellslocal
        integer(kind=intType), intent(in) :: nCell, blockSize, levels, nSmooth
        mat, target :: inputmat
 
        vecscatter :: scat
        vec :: recvvec, indexvec
        is :: is1
 
        integer(kind=intTYpe) :: nnx, nny, nnz, l, nn, i, j, k, ii, jj, kk, lvl, n, m
        integer(kind=intType) :: idim, count, cursize, ierr, coarseIndex
        integer(kind=intType) :: level, nVar, sps, nextLvl, ncoarse
        integer(kind=intType), dimension(:, :, :), allocatable :: sizes
        integer(kind=intType), dimension(:), allocatable :: offsets, procStarts, nnzOn, nnzOff
        real(kind=realtype), dimension(:), pointer :: indptr
        integer(kind=intType), dimension(:), allocatable :: indicesToGet
 
        if (amgsetup) then
            return
        end if
 
        ! Set some module variables
        amglevels = levels
        amgnsmooth = nsmooth
 
        ! Block size here refers to the number of states in each cell
        bs = blocksize
 
        ! Set the pointer to the fine level the AMG will be working on
        finemat => inputmat
 
        ! Allocate the list of the mats/vects/ksp
        allocate ( &
            a(2:amglevels), &
            ksplevels(1:amglevels), &
            res(1:amglevels), &
            rhs(1:amglevels), &
            sol(1:amglevels), &
            sol2(1:amglevels))
 
        ! First allocate the coarse level indices.
        allocate (coarseindices(ndom, 1:amglevels - 1))
        allocate (coarseoversetindices(ndom, 1:amglevels - 1))
        allocate (sizes(3, ndom, 1:amglevels))
        allocate (offsets(0:nproc), procstarts(1:amglevels))
        allocate (interps(1:amglevels - 1))
        do lvl = 1, amglevels
 
            do nn = 1, ndom
                call setpointers(nn, 1, 1)
                if (lvl == 1) then
 
                    ! Set the sizes for the finest level.
                    ! This is the number of owned cells.
                    sizes(1, nn, 1) = nx
                    sizes(2, nn, 1) = ny
                    sizes(3, nn, 1) = nz
                end if
            end do
 
            if (lvl > 1) then
                do nn = 1, ndom
                    do idim = 1, 3
                        cursize = sizes(idim, nn, lvl - 1)
                        if (cursize == 1) then
                            sizes(idim, nn, lvl) = cursize
                        else if (mod(cursize, 2) == 0) then
                            ! Evenly divides
                            sizes(idim, nn, lvl) = cursize / 2
                        else if (mod(cursize, 2) == 1) then
                            ! Odd, so have an extra
                            sizes(idim, nn, lvl) = (cursize - 1) / 2 + 1
                        end if
                    end do
                end do
            end if
 
            ! Get the number of cells on this proc:
            i = 0
            do nn = 1, ndom
                i = i + sizes(1, nn, lvl) * sizes(2, nn, lvl) * sizes(3, nn, lvl)
            end do
 
            call mpi_allgather(i, 1, mpi_integer, offsets(1:nproc), 1, mpi_integer, &
                               adflow_comm_world, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Prefix sum
            offsets(0) = 0
            do i = 1, nproc
                offsets(i) = offsets(i - 1) + offsets(i)
            end do
 
            ! Get my starting index.
            procstarts(lvl) = offsets(myid)
 
        end do
 
        ! Now set up the interp and the coarse-level indices. Note that this
        ! loop is the number of levels MINUS 1 since we are generating the
        ! interpolations between levels and we already have the 1st level
        ! of the  coarseIndices
 
        call veccreatempi(adflow_comm_world, ncellslocal(1_inttype), &
                          petsc_determine, indexvec, ierr)
        call echk(ierr, __file__, __line__)
 
        do lvl = 1, amglevels - 1
 
            do nn = 1, ndom
                call setpointers(nn, 1, 1)
                allocate (coarseindices(nn, lvl)%arr(0:ib, 0:jb, 0:kb))
                allocate (coarseoversetindices(nn, lvl)%arr(8, 0:ib, 0:jb, 0:kb))
                coarseindices(nn, lvl)%arr = -1
                coarseoversetindices(nn, lvl)%arr = -1
            end do
 
            ! Allocate the linear algebra interpolation array for this level
            ! (first count the number of nodes to be restricted on level lvl)
            n = 0
            do nn = 1, ndom
                n = n + sizes(1, nn, lvl) * sizes(2, nn, lvl) * sizes(3, nn, lvl)
            end do
            allocate (interps(lvl)%arr(n))
 
            ! Interps uses the local ordering so we always start at 0.
            n = 0
            count = 0
            ! Loop over the blocks
            do nn = 1, ndom
 
                ! Sizes for next level
                nnx = sizes(1, nn, lvl + 1)
                nny = sizes(2, nn, lvl + 1)
                nnz = sizes(3, nn, lvl + 1)
 
                ! Loop over the sizes of this level
                do k = 1, sizes(3, nn, lvl)
                    do j = 1, sizes(2, nn, lvl)
                        do i = 1, sizes(1, nn, lvl)
 
                            ! These are the indices on the next level
                            ii = (i - 1) / 2 + 1
                            jj = (j - 1) / 2 + 1
                            kk = (k - 1) / 2 + 1
 
                            coarseindex = (kk - 1) * nnx * nny + (jj - 1) * nnx + ii + count
 
                            ! Linear algebra info
                            n = n + 1
                            interps(lvl)%arr(n) = coarseindex
 
                        end do
                    end do
                end do
                count = count + nnx * nny * nnz
            end do
 
            ! We are not done yet; We need to fill in the block-based
            ! coarse indices and then do a halo-exchange on it so procs
            ! know where to put their off-proc entries on the coarser levels.
 
            ! Loop over the blocks.
            n = 0
            ii = 0
 
            call vecgetarrayf90(indexvec, indptr, ierr)
            call echk(ierr, __file__, __line__)
 
            do nn = 1, ndom
                call setpointers(nn, 1, 1)
                ii = ii + (kb + 1) * (jb + 1) * (ib + 1) ! Count maximum double halos
                ! Loop over the sizes of this level
                do k = 2, kl
                    do j = 2, jl
                        do i = 2, il
 
                            n = n + 1
 
                            ! Coarse Index: This is the first coarse index for
                            ! the current finest level element.
                            coarseindex = interps(1)%arr(n)
 
                            ! For levels higher than 2, we need to trace
                            ! through the subsequent levels to find what the
                            ! coarse index is for level lvl.
                            do nextlvl = 2, lvl
                                coarseindex = interps(nextlvl)%arr(coarseindex)
                            end do
 
                            ! Now we set the coarse level index into the
                            ! coarseIndices array. Note that we make the index
                            ! global here by adding procStarts. We also
                            ! subtract 1 to make it zero-based for petsc.
                            coarseindex = coarseindex + procstarts(lvl + 1) - 1
                            coarseindices(nn, lvl)%arr(i, j, k) = coarseindex
 
                            ! Now put that index into the array
                            indptr(n) = transfer(int(coarseindex, 8), one)
 
                        end do
                    end do
                end do
            end do
            call vecrestorearrayf90(indexvec, indptr, ierr)
            call echk(ierr, __file__, __line__)
 
            allocate (indicestoget(8 * ii))
            ii = 0
            do nn = 1, ndom
                call setpointers(nn, 1, 1)
                do k = 0, kb
                    do j = 0, jb
                        do i = 0, ib
                            ! If this cell is is an interpolated cell, record the indices we need to get
                            if (iblank(i, j, k) == -1) then
                                do m = 1, 8
                                    if (flowdoms(nn, 1, 1)%gInd(m, i, j, k) >= 0) then
                                        ii = ii + 1
                                        indicestoget(ii) = flowdoms(nn, 1, 1)%gInd(m, i, j, k)
                                    end if
                                end do
                            end if
                        end do
                    end do
                end do
            end do
 
            ! Now create the scatter to retrieve the "indicesToGet" from
            ! the indexVec. Petsc is always annoying for this.
 
            call iscreategeneral(adflow_comm_world, ii, indicestoget(1:ii), petsc_copy_values, &
                                 is1, ierr)
            call echk(ierr, __file__, __line__)
            deallocate (indicestoget)
 
            ! Create array to dump the result
            call veccreatempi(adflow_comm_world, ii, petsc_determine, recvvec, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Create the scatter
            call vecscattercreate(indexvec, is1, recvvec, petsc_null_is, scat, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Do the actual scatter
            call vecscatterbegin(scat, indexvec, recvvec, insert_values, scatter_forward, ierr)
            call echk(ierr, __file__, __line__)
            call vecscatterend(scat, indexvec, recvvec, insert_values, scatter_forward, ierr)
            call echk(ierr, __file__, __line__)
 
            call vecgetarrayf90(recvvec, indptr, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Loop back over at set the coarse indices
            ii = 0
            do nn = 1, ndom
                call setpointers(nn, 1, 1)
                ! Loop over the sizes of this level
                do k = 0, kb
                    do j = 0, jb
                        do i = 0, ib
                            if (iblank(i, j, k) == -1) then
                                do m = 1, 8
                                    if (flowdoms(nn, 1, 1)%gInd(m, i, j, k) >= 0) then
                                        ii = ii + 1
                                        coarseoversetindices(nn, lvl)%arr(m, i, j, k) = &
                                            int(transfer(indptr(ii), 1_8), inttype)
                                    end if
                                end do
                            end if
                        end do
                    end do
                end do
            end do
 
            call vecrestorearrayf90(recvvec, indptr, ierr)
            call echk(ierr, __file__, __line__)
 
            call vecscatterdestroy(scat, ierr)
            call echk(ierr, __file__, __line__)
 
            call vecdestroy(recvvec, ierr)
            call echk(ierr, __file__, __line__)
 
            call isdestroy(is1, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Now we need to exchange the coarse level indices. Set
            ! pointers to the coarseIndices arrays and then call the
            ! generic integer halo exchange.
            level = 1
            sps = 1
            nvar = 1
 
            do nn = 1, ndom
                flowdoms(nn, level, sps)%intCommVars(1)%var => coarseindices(nn, lvl)%arr(:, :, :)
            end do
 
            call whalo1to1intgeneric(nvar, level, sps, commpatterncell_2nd, internalcell_2nd)
 
        end do
 
        call vecdestroy(indexvec, ierr)
        call echk(ierr, __file__, __line__)
 
        ! Next we need to set up the matrices and vectors
 
        do lvl = 1, amglevels
 
            call kspcreate(adflow_comm_world, ksplevels(lvl), ierr)
            call echk(ierr, __file__, __line__)
 
            ! Create the coarse levels
            if (lvl >= 2) then
                ncoarse = maxval(interps(lvl - 1)%arr)
 
                allocate (nnzon(1:ncoarse), nnzoff(1:ncoarse))
                nnzon = 14
                nnzoff = 7
                call matcreatebaij(adflow_comm_world, bs, ncoarse * bs, ncoarse * bs, &
                                   petsc_determine, petsc_determine, 0, nnzon, 0, nnzoff, &
                                   a(lvl), ierr)
                call echk(ierr, __file__, __line__)
                deallocate (nnzon, nnzoff)
 
                call matsetoption(a(lvl), mat_row_oriented, petsc_false, ierr)
                call echk(ierr, __file__, __line__)
 
                call matsetoption(a(lvl), mat_new_nonzero_allocation_err, petsc_false, ierr)
                call echk(ierr, __file__, __line__)
 
                call matcreatevecs(a(lvl), res(lvl), sol(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
            else
                call veccreatempi(adflow_comm_world, ncell * bs, petsc_determine, res(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
                call vecduplicate(res(lvl), sol(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
            end if
            call vecduplicate(res(lvl), rhs(lvl), ierr)
            call echk(ierr, __file__, __line__)
 
            call vecduplicate(res(lvl), sol2(lvl), ierr)
            call echk(ierr, __file__, __line__)
 
        end do
        amgsetup = .true.
 
    end subroutine setupamg
 
    subroutine destroyamg
 
        integer(kind=intType) :: lvl, ierr, i, j
        if (amgsetup) then
            do lvl = 1, amglevels
                ! Destroy all of our vectors/matrices
                if (lvl > 1) then
                    call matdestroy(a(lvl), ierr)
                    call echk(ierr, __file__, __line__)
                end if
 
                call vecdestroy(res(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
                call vecdestroy(sol(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
                call vecdestroy(sol2(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
                call vecdestroy(rhs(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
                call kspdestroy(ksplevels(lvl), ierr)
                call echk(ierr, __file__, __line__)
            end do
 
            deallocate (a, res, rhs, sol, sol2, ksplevels)
            deallocate (coarseindices, coarseoversetindices, interps)
            amgsetup = .false.
 
        end if
    end subroutine destroyamg
 
    subroutine applyshellpc(pc, x, y, ierr)
        use communication
        ! Input/Output
        pc pc
        vec x, y
        integer(kind=intType) ierr
 
        ! Working
        integer(kind=intType) :: i
 
        if (amglevels > 1) then
 
            ! Copy the input vector x into rhs because x is read-only
            call veccopy(x, rhs(1), ierr)
            call echk(ierr, __file__, __line__)
 
            ! Zero the output vector y
            call vecset(y, zero, ierr)
            call echk(ierr, __file__, __line__)
 
            if (amgouterits == 1) then
                call mgprecon(rhs(1), y, 1) ! y is the new approximate solution
            else
 
                do i = 1, amgouterits
 
                    call mgprecon(rhs(1), sol(1), 1) ! The solution update is stored in sol(1)
 
                    ! Update the solution
                    call vecaypx(y, one, sol(1), ierr) ! y = y + sol(1)
                    call echk(ierr, __file__, __line__)
 
                    if (i < amgouterits) then
                        ! Compute new residual
                        call matmult(finemat, y, rhs(1), ierr)
                        call echk(ierr, __file__, __line__)
 
                        call vecaypx(rhs(1), -one, x, ierr)
                        call echk(ierr, __file__, __line__)
                    end if
 
                end do
            end if
        else
            ! Solve the fine level
            ! This is equivalent to not using multigrid
            call kspsolve(ksplevels(1), x, y, ierr)
            call echk(ierr, __file__, __line__)
        end if
 
    end subroutine applyshellpc
 
    subroutine setupshellpc(pc, ierr)
 
        use communication, only: adflow_comm_world
        use inputadjoint
 
        ! Input/Output
        pc pc
        integer(kind=intTYpe) :: ierr
 
        ! Working
        pc globalpc, subpc
        ksp subksp
        integer(kind=intType) :: lvl
        integer(kind=intType) :: nlocal, first
 
        do lvl = 1, amglevels
            if (lvl == 1) then
                call kspsetoperators(ksplevels(lvl), finemat, finemat, ierr)
                call echk(ierr, __file__, __line__)
 
                amgfilllevel = amgfilllevelfine
                amglocalpreconits = amglocalpreconitsfine
                amgasmoverlap = amgasmoverlapfine
            else
                call kspsetoperators(ksplevels(lvl), a(lvl), a(lvl), ierr)
                call echk(ierr, __file__, __line__)
 
                amgfilllevel = amgfilllevelcoarse
                amglocalpreconits = amglocalpreconitscoarse
                amgasmoverlap = amgasmoverlapcoarse
            end if
 
            call kspsetnormtype(ksplevels(lvl), ksp_norm_none, ierr)
            call echk(ierr, __file__, __line__)
 
            call kspsettype(ksplevels(lvl), 'richardson', ierr)
            call echk(ierr, __file__, __line__)
 
            call kspsettolerances(ksplevels(lvl), petsc_default_real, &
                                  petsc_default_real, petsc_default_real, &
                                  amgnsmooth, ierr)
            call echk(ierr, __file__, __line__)
 
            call kspgetpc(ksplevels(lvl), globalpc, ierr)
            call echk(ierr, __file__, __line__)
 
            call pcsettype(globalpc, 'asm', ierr)
            call echk(ierr, __file__, __line__)
 
            call pcasmsetoverlap(globalpc, amgasmoverlap, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Set up the main ksp context before extracting the subdomains
            call kspsetup(ksplevels(lvl), ierr)
            call echk(ierr, __file__, __line__)
 
            ! Extract the ksp objects for each subdomain
            call pcasmgetsubksp(globalpc, nlocal, first, subksp, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Since there is an extra matMult required when using the Richardson preconditioner
            ! with only 1 iteration, only use it when we need to do more than 1 iteration.
            if (amglocalpreconits > 1) then
                ! Set the subksp object to Richardson so we can do multiple iterations on the sub-domains
                call kspsettype(subksp, 'richardson', ierr)
                call echk(ierr, __file__, __line__)
 
                ! Set the number of iterations to do on local blocks. Tolerances are ignored.
                call kspsettolerances(subksp, petsc_default_real, petsc_default_real, petsc_default_real, &
                                      amglocalpreconits, ierr)
                call echk(ierr, __file__, __line__)
 
                ! normtype is NONE because we don't want to check error
                call kspsetnormtype(subksp, ksp_norm_none, ierr)
                call echk(ierr, __file__, __line__)
            else
                ! Set the subksp object to preonly because we are only doing one iteration
                call kspsettype(subksp, 'preonly', ierr)
                call echk(ierr, __file__, __line__)
            end if
 
            ! Extract the preconditioner for subksp object
            call kspgetpc(subksp, subpc, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Set the subpc type to ILU
            call pcsettype(subpc, 'ilu', ierr)
            call echk(ierr, __file__, __line__)
 
            ! Set up the matrix ordering for the subpc object
            call pcfactorsetmatorderingtype(subpc, amgmatrixordering, ierr)
            call echk(ierr, __file__, __line__)
 
            ! Set the ILU parameters
            call pcfactorsetlevels(subpc, amgfilllevel, ierr)
            call echk(ierr, __file__, __line__)
        end do
 
    end subroutine setupshellpc
 
    subroutine destroyshellpc(pc, ierr)
 
        ! Input/Ouput
        pc pc
        integer(kind=intType) :: ierr
 
        ! Working
        integer(kind=intType) :: lvl
 
    end subroutine destroyshellpc
 
    subroutine restrictvec(x, y, interp)
 
        ! Input/Output
        vec x, y
        integer(kind=intType), dimension(:), intent(in) :: interp
 
        ! Working
        real(kind=realtype), pointer :: yptr(:), xptr(:)
        real(kind=realtype), pointer :: xptrblk(:, :), yptrblk(:, :)
        integer(kind=intType) :: ierr, n, i, j
 
        ! Restrict x -> y
        call vecgetarrayf90(x, xptr, ierr)
        call echk(ierr, __file__, __line__)
 
        call vecgetarrayf90(y, yptr, ierr)
        call echk(ierr, __file__, __line__)
 
        ! Number of block nodes
        n = size(interp)
 
        ! Convenience block-based pointers
        xptrblk(1:bs, 1:size(xptr) / bs) => xptr
        yptrblk(1:bs, 1:size(yptr) / bs) => yptr
 
        ! Zero the output array
        yptr = zero
 
        ! Loop over the interpolation array, summing into the coarse array
        do i = 1, n
            j = interp(i)
            yptrblk(:, j) = yptrblk(:, j) + xptrblk(:, i)
        end do
 
        call vecrestorearrayf90(x, xptr, ierr)
        call echk(ierr, __file__, __line__)
 
        call vecrestorearrayf90(y, yptr, ierr)
        call echk(ierr, __file__, __line__)
 
    end subroutine restrictvec
 
    subroutine prolongvec(x, y, interp)
 
        ! Input/Output
        vec x, y
        integer(kind=intType), dimension(:), intent(in) :: interp
 
        ! Working
        real(kind=realtype), pointer :: yptr(:), xptr(:)
        real(kind=realtype), pointer :: xptrblk(:, :), yptrblk(:, :)
        integer(kind=intType) :: ierr, n, i, j
 
        ! Prolong vector x -> y
 
        call vecgetarrayf90(x, xptr, ierr)
        call echk(ierr, __file__, __line__)
 
        call vecgetarrayf90(y, yptr, ierr)
        call echk(ierr, __file__, __line__)
 
        ! Number of block nodes
        n = size(interp)
 
        ! Convenience pointers
        yptrblk(1:bs, 1:size(yptr) / bs) => yptr
        xptrblk(1:bs, 1:size(xptr) / bs) => xptr
 
        ! Loop over the interpolation array, injecting into the fine array
        do i = 1, n
            j = interp(i)
            yptrblk(:, i) = xptrblk(:, j)
        end do
 
        call vecrestorearrayf90(x, xptr, ierr)
        call echk(ierr, __file__, __line__)
 
        call vecrestorearrayf90(y, yptr, ierr)
        call echk(ierr, __file__, __line__)
 
    end subroutine prolongvec
 
    recursive subroutine mgprecon(r, y, k)
 
        ! r is the residual and y is the approximate solution
        ! k is the level
 
        ! Input/Output
        vec y, r
        integer(kind=intType), intent(in) :: k
 
        ! Working
        integer(kind=intType) :: ierr, i
 
        ! Step 1: Restrict the residual
        call restrictvec(r, rhs(k + 1), interps(k)%arr)
 
        ! Step 2: Solve the coarse level directly or recursively
        if (k == amglevels - 1) then
            ! The next level down is the bottom
            ! Break the recursion by solving
            call kspsolve(ksplevels(k + 1), rhs(k + 1), sol(k + 1), ierr)
            call echk(ierr, __file__, __line__)
        else
            ! Call the next level down recursively
            call mgprecon(rhs(k + 1), sol(k + 1), k + 1)
        end if
 
        ! Step 3: Prolongate the solution
        call prolongvec(sol(k + 1), y, interps(k)%arr)
 
        ! Step 4: Compute the new residual
        if (k == 1) then
            call matmult(finemat, y, res(k), ierr)  ! res = A(i) * z1(k)
        else
            call matmult(a(k), y, res(k), ierr)  ! res = A(i) * z1(k)
        end if
        call echk(ierr, __file__, __line__)
 
        call vecaypx(res(k), -one, r, ierr)    ! res = -res + r
        call echk(ierr, __file__, __line__)
 
        ! Step 5: Relax using the smoother
        call kspsolve(ksplevels(k), res(k), sol2(k), ierr)
        call echk(ierr, __file__, __line__)
 
        call vecaxpy(y, one, sol2(k), ierr)
        call echk(ierr, __file__, __line__)
 
    end subroutine mgprecon
 
end module amg