Compiler Explorer

Source code

module stdlib_krylov
implicit none

type :: linop
    procedure(mxm), nopass, pointer :: times => null()
    procedure(dot), nopass, pointer :: inner => dot_co_sp ! default to be changed
end type

abstract interface
    ! evaluate y := op(x)
    subroutine mxm(x,y)
        real, intent(in)  :: x(:)
        real, intent(out) :: y(:)
    end subroutine
    function dot(x,y) result(r)
        real, intent(in) :: x(:)
        real, intent(in) :: y(:)
        real :: r
    end function
end interface

interface cgsolve
    module procedure :: cgsolve_generic
    module procedure :: cgsolve_dense_sp
end interface

contains

function dot_co_sp(x,y) result(r)
    real, intent(in) :: x(:)
    real, intent(in) :: y(:)
    real :: r
    r = dot_product(x,y)
#ifdef CAF
    call co_sum(x)
#endif
end function

subroutine cgsolve_generic(A,b,x,tol)
    class(linop), intent(in) :: A
    real, intent(in) :: b(:), tol
    real, intent(out) :: x(:)

real :: rtr, rtrold, alpha, beta
    real, allocatable :: r(:), p(:), Ap(:)

allocate(Ap,mold=b)
    allocate(r,source=b)
    x = 0.0
    rtr = A%inner(r, r)
    p = b
    beta = 0

do while( norm2(r) > tol*norm2(b))
        p = r + beta * p
        call A%times(p,Ap) 
        alpha = rtr / A%inner(p, Ap)
        x = x + alpha * p
        r = r - alpha * Ap
        rtrold = rtr
        rtr = A%inner(r, r)
        beta = rtr / rtrold
    end do

end subroutine

subroutine cgsolve_dense_sp(A,b,x,tol)
    real, intent(in) :: A(:,:)
    real, intent(in) :: b(:), tol
    real, intent(out) :: x(:)
    !-------------------------
    type(linop) :: op
    !-------------------------
    op%times => custom_linop
    call cgsolve_generic(op,b,x,tol)
    contains
    subroutine custom_linop(x,y)
        real, intent(in)  :: x(:)
        real, intent(out) :: y(:)
        y = matmul(A,x)
    end subroutine
end subroutine

end module

subroutine example_dense()
    use stdlib_krylov
    implicit none
    real :: A(2,2) = reshape([4.0,1.0,1.0,3.0],shape=[2,2])
    real :: b(2) = [1.0,2.0]
    real :: x(2)
    real :: tol = 0.0001

call cgsolve(A,b,x,tol)
    print *, x

end subroutine

program main
implicit none

call example_dense()

end program