! ----------------------------------------------------------------------
! Testing LAPACK solvers by the cuSOLVER library
! Library functions: cuSolverDnSgetrf & Sgetrs & Dgetrf & Dgetrs
! Defined operator cuSOLVERgesv
! Compilers: gfortran (10.2), ifort (21.1), nvfortran (20.11)
! ----------------------------------------------------------------------
! Command-line compilation:
! ----------------------------------------------------------------------
! Linux
! (set LIBRARY_PATH and LD_LIBRARY_PATH)
! gfortran testcuSOLVER.f90 -lcudart -lcusolver
! ifort -heap-arrays testcuSOLVER.f90 -lcudart -lcusolver
! nvfortran testcuSOLVER.f90 -cudalib
! ----------------------------------------------------------------------
! Windows
! set PATH=C:\Program Files\PGI\win64\2019\cuda\10.1\bin;%PATH%
! set LIBRARY_PATH=C:\Program Files\PGI\win64\2019\cuda\10.1\lib\x64;%LIBRARY_PATH%
! gfortran testcuSOLVER.f90 -lcudart -lcusolver
! set PATH=C:\Program Files\PGI\win64\2019\cuda\10.1\bin;%PATH%
! set LIB=C:\Program Files\PGI\win64\2019\cuda\10.1\lib\x64;%LIB%
! ifort -heap-arrays testcuSOLVER.f90 cudart.lib cusolver.lib
! nvfortran not available (03/2021)
! ----------------------------------------------------------------------

! ----------------------------------------------------------------------
! Global constants
! ----------------------------------------------------------------------
MODULE mConst

! real kind, physical size of arrays
integer,parameter :: sp=4,dp=8,nmax=10000

logical,parameter,private :: T=.TRUE.,F=.FALSE.
logical,parameter :: runcuSOLVERsp=T
logical,parameter :: runcuSOLVERdp=T

END MODULE

! ----------------------------------------------------------------------
! Wallclock timer
! Usage: call tic(); call task(); print *,toc()
! ----------------------------------------------------------------------
MODULE mTicToc

implicit none
private
public tic,toc
integer,parameter :: wp=8
real(wp) :: time

CONTAINS

subroutine tic()
integer count,count_rate
call system_clock(count,count_rate)
time=real(count,wp)/real(count_rate,wp)
end subroutine

real(wp) function toc()
integer count,count_rate
call system_clock(count,count_rate)
toc=real(count,wp)/real(count_rate,wp)-time
end function

END MODULE

! ----------------------------------------------------------------------
! LU decomposition via the cuSOLVER library
! Overloaded defined operator .cuSOLVERgesv.
! Usage: c4=a4.cuSOLVERgesv.b4; c8=a8.cuSOLVERgesv.b8
!        for real(4) and real(8) matrices and vectors
! ----------------------------------------------------------------------
MODULE mcuSOLVER

use iso_c_binding
implicit none
private
public operator(.cuSOLVERgesv.)

integer,parameter :: sp=4,dp=8

interface operator (.cuSOLVERgesv.)
module procedure solve_sp_cuSOLVER,solve_dp_cuSOLVER
end interface

enum,bind(c) ! cusolverStatus_t
enumerator :: &
  CUSOLVER_STATUS_SUCCESS=0,CUSOLVER_STATUS_NOT_INITIALIZED=1,CUSOLVER_STATUS_ALLOC_FAILED=2, &
  CUSOLVER_STATUS_INVALID_VALUE=3,CUSOLVER_STATUS_ARCH_MISMATCH=4,CUSOLVER_STATUS_MAPPING_ERROR=5, &
  CUSOLVER_STATUS_EXECUTION_FAILED=6,CUSOLVER_STATUS_INTERNAL_ERROR=7,CUSOLVER_STATUS_MATRIX_TYPE_NOT_SUPPORTED=8, &
  CUSOLVER_STATUS_NOT_SUPPORTED=9,CUSOLVER_STATUS_ZERO_PIVOT=10,CUSOLVER_STATUS_INVALID_LICENSE=11, &
  CUSOLVER_STATUS_IRS_PARAMS_NOT_INITIALIZED=12,CUSOLVER_STATUS_IRS_PARAMS_INVALID=13,CUSOLVER_STATUS_IRS_PARAMS_INVALID_PREC=14, &
  CUSOLVER_STATUS_IRS_PARAMS_INVALID_REFINE=15,CUSOLVER_STATUS_IRS_PARAMS_INVALID_MAXITER=16,CUSOLVER_STATUS_IRS_INTERNAL_ERROR=20,&
  CUSOLVER_STATUS_IRS_NOT_SUPPORTED=21,CUSOLVER_STATUS_IRS_OUT_OF_RANGE=22, &
  CUSOLVER_STATUS_IRS_NRHS_NOT_SUPPORTED_FOR_REFINE_GMRES=23,CUSOLVER_STATUS_IRS_INFOS_NOT_INITIALIZED=25, &
  CUSOLVER_STATUS_IRS_INFOS_NOT_DESTROYED=26,CUSOLVER_STATUS_IRS_MATRIX_SINGULAR=30,CUSOLVER_STATUS_INVALID_WORKSPACE=31
end enum

enum,bind(c) ! cudaDataType_t
enumerator :: &
  CUDA_R_32F= 0,  & ! real as a float
  CUDA_C_32F= 4,  & ! complex as a pair of float numbers
  CUDA_R_64F= 1,  & ! real as a double
  CUDA_C_64F= 5     ! complex as a pair of double numbers
end enum

enum,bind(c) ! cublasOperation_t
enumerator :: CUBLAS_OP_N=0,CUBLAS_OP_T=1,CUBLAS_OP_C=2,CUBLAS_OP_HERMITAN=2,CUBLAS_OP_CONJG=3
end enum

enum,bind(c) ! cudaMemcpyKind
enumerator :: &
  cudaMemcpyHostToHost=0,cudaMemcpyHostToDevice=1,cudaMemcpyDeviceToHost=2,cudaMemcpyDeviceToDevice=3,cudaMemcpyDefault=4
end enum

interface    ! CUDA Runtime API

integer(c_int) function cudaMalloc(devPtr,size) bind(c,name='cudaMalloc')
import c_int,c_size_t,c_ptr
type(c_ptr) :: devPtr                    ! void** devPtr
integer(c_size_t),value,intent(in) :: size
end function

integer(c_int) function cudaFree(devPtr) bind(c,name='cudaFree')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: devPtr   ! void* devPtr
end function

integer(c_int) function cudaMemcpy(dst,src,count,kind) bind(c,name='cudaMemcpy')
import c_int,c_size_t,c_ptr
type(c_ptr),value :: dst
type(c_ptr),value,intent(in) :: src
integer(c_size_t),value,intent(in) :: count
integer(c_int),value,intent(in) :: kind
end function

end interface

interface    ! cuSOLVER routines

integer(c_int) function cusolverDnCreate(handle) bind(c,name='cusolverDnCreate')
import c_int,c_ptr
type(c_ptr) :: handle ! cusolverDnHandle_t
end function

integer(c_int) function cusolverDnDestroy(handle) bind(c,name='cusolverDnDestroy')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: handle ! cusolverDnHandle_t
end function

integer(c_int) function cusolverDnSgetrf_bufferSize(handle,m,n,A,lda,Lwork) bind(c,name='cusolverDnSgetrf_bufferSize')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: handle
integer(c_int),value,intent(in) :: m,n,lda
type(c_ptr),value :: A
integer(c_int) :: Lwork
end function

integer(c_int) function cusolverDnDgetrf_bufferSize(handle,m,n,A,lda,Lwork) bind(c,name='cusolverDnDgetrf_bufferSize')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: handle
integer(c_int),value,intent(in) :: m,n,lda
type(c_ptr),value :: A
integer(c_int) :: Lwork
end function

integer(c_int) function cusolverDnSgetrf(handle,m,n,A,lda,Workspace,devIpiv,devInfo) bind(c,name='cusolverDnSgetrf')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: handle
integer(c_int),value,intent(in) :: m,n,lda
type(c_ptr),value :: A,Workspace,devIpiv,devInfo
end function

integer(c_int) function cusolverDnDgetrf(handle,m,n,A,lda,Workspace,devIpiv,devInfo) bind(c,name='cusolverDnDgetrf')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: handle
integer(c_int),value,intent(in) :: m,n,lda
type(c_ptr),value :: A,Workspace,devIpiv,devInfo
end function

integer(c_int) function cusolverDnSgetrs(handle,trans,n,nrhs,A,lda,devIpiv,b,ldb,devInfo) bind(c,name='cusolverDnSgetrs')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: handle
integer(c_int),value,intent(in) :: trans
integer(c_int),value,intent(in) :: n,nrhs,lda,ldb
type(c_ptr),value :: A,b,devIpiv,devInfo
end function

integer(c_int) function cusolverDnDgetrs(handle,trans,n,nrhs,A,lda,devIpiv,b,ldb,devInfo) bind(c,name='cusolverDnDgetrs')
import c_int,c_ptr
type(c_ptr),value,intent(in) :: handle
integer(c_int),value,intent(in) :: trans
integer(c_int),value,intent(in) :: n,nrhs,lda,ldb
type(c_ptr),value :: A,b,devIpiv,devInfo
end function

end interface

CONTAINS

function solve_sp_cuSOLVER(A,b) RESULT (x)
real(sp),intent(in) :: A(:,:),b(:)
real(sp) :: x(size(A,1))
integer :: n,lda,ldb,Lwork,nrhs,info
integer :: trans,istatus
type(c_ptr) :: Ad,bd,ipiv,Workspace,infod
type(c_ptr) :: handle
target :: A,b,x,info

x(1)=0.
n=size(A,1)
nrhs=1
lda=size(A,1)
ldb=size(b,1)
trans=CUBLAS_OP_N  ! CUBLAS_OP_N, CUBLAS_OP_T, CUBLAS_OP_C
istatus=cusolverDnCreate(handle)
istatus=cudaMalloc(Ad,int(lda*n*sp,c_size_t))
istatus=cudaMalloc(bd,int(ldb*sp,c_size_t))
istatus=cudaMalloc(ipiv,int(n*c_int,c_size_t))
istatus=cudaMemcpy(Ad,c_loc(A),int(lda*n*sp,c_size_t),cudaMemcpyHostToDevice)
istatus=cusolverDnSgetrf_bufferSize(handle,n,n,Ad,lda,Lwork)
istatus=cudaMalloc(Workspace,int(Lwork*sp,c_size_t))
istatus=cudaMalloc(infod,int(c_int,c_size_t))
istatus=cusolverDnSgetrf(handle,n,n,Ad,lda,Workspace,ipiv,infod)
istatus=cudaMemcpy(c_loc(info),infod,int(c_int,c_size_t),cudaMemcpyDeviceToHost)
if (istatus/=0 .or. info/=0) print *,'getrf status/info =',istatus,info
istatus=cudaMemcpy(bd,c_loc(b),int(n*sp,c_size_t),cudaMemcpyHostToDevice)
istatus=cuSolverDnSgetrs(handle,trans,n,nrhs,Ad,lda,ipiv,bd,ldb,infod)
istatus=cudaMemcpy(c_loc(info),infod,int(c_int,c_size_t),cudaMemcpyDeviceToHost)
if (istatus/=0 .or. info/=0) print *,'getrs status/info =',istatus,info
istatus=cudaMemcpy(c_loc(x),bd,int(n*sp,c_size_t),cudaMemcpyDeviceToHost)
istatus=cudaFree(Ad)
istatus=cudaFree(bd)
istatus=cudaFree(ipiv)
istatus=cudaFree(Workspace)
istatus=cusolverDnDestroy(handle)
end function

function solve_dp_cuSOLVER(A,b) RESULT (x)
real(dp),intent(in) :: A(:,:),b(:)
real(dp) :: x(size(A,1))
integer :: n,lda,ldb,Lwork,nrhs,info
integer :: trans,istatus
type(c_ptr) :: Ad,bd,ipiv,Workspace,infod
type(c_ptr) :: handle
target :: A,b,x,info

x(1)=0.
n=size(A,1)
nrhs=1
lda=size(A,1)
ldb=size(b,1)
trans=CUBLAS_OP_N  ! CUBLAS_OP_N, CUBLAS_OP_T, CUBLAS_OP_C
istatus=cusolverDnCreate(handle)
istatus=cudaMalloc(Ad,int(lda*n*dp,c_size_t))
istatus=cudaMalloc(bd,int(ldb*dp,c_size_t))
istatus=cudaMalloc(ipiv,int(n*c_int,c_size_t))
istatus=cudaMemcpy(Ad,c_loc(A),int(lda*n*dp,c_size_t),cudaMemcpyHostToDevice)
istatus=cusolverDnDgetrf_bufferSize(handle,n,n,Ad,lda,Lwork)
istatus=cudaMalloc(Workspace,int(Lwork*dp,c_size_t))
istatus=cudaMalloc(infod,int(c_int,c_size_t))
istatus=cusolverDnDgetrf(handle,n,n,Ad,lda,Workspace,ipiv,infod)
istatus=cudaMemcpy(c_loc(info),infod,int(c_int,c_size_t),cudaMemcpyDeviceToHost)
if (istatus/=0 .or. info/=0) print *,'getrf status/info =',istatus,info
istatus=cudaMemcpy(bd,c_loc(b),int(n*dp,c_size_t),cudaMemcpyHostToDevice)
istatus=cuSolverDnDgetrs(handle,trans,n,nrhs,Ad,lda,ipiv,bd,ldb,infod)
istatus=cudaMemcpy(c_loc(info),infod,int(c_int,c_size_t),cudaMemcpyDeviceToHost)
if (istatus/=0 .or. info/=0) print *,'getrs status/info =',istatus,info
istatus=cudaMemcpy(c_loc(x),bd,int(n*dp,c_size_t),cudaMemcpyDeviceToHost)
istatus=cudaFree(Ad)
istatus=cudaFree(bd)
istatus=cudaFree(ipiv)
istatus=cudaFree(Workspace)
istatus=cusolverDnDestroy(handle)
end function

END MODULE

! ----------------------------------------------------------------------
  PROGRAM TestcuSOLVER
! ----------------------------------------------------------------------
use mConst
use mTicToc
use mcuSOLVER
implicit none
real(dp),allocatable :: A(:,:),b(:),x(:)

ALLOCATE (A(nmax,nmax),b(nmax),x(nmax))
call random_number(A)

if (runcuSOLVERsp) then
  print *,'cuSOLVER SP (GPU)'
  b=1.
  call tic()
  x=real(A,sp).cuSOLVERgesv.real(b,sp)
  print *,'dt (SP) = ',toc()
  print '(1x,a,10e20.12)','x =',x(1:3)
  b=matmul(A,x)
  print '(1x,a,10e20.12)','b =',b(1:3)
endif

if (runcuSOLVERdp) then
  print *,'cuSOLVER DP (GPU)'
  b=1.
  call tic()
  x=A.cuSOLVERgesv.b
  print *,'dt (DP) = ',toc()
  print '(1x,a,10e20.12)','x =',x(1:3)
  b=matmul(A,x)
  print '(1x,a,10e20.12)','b =',b(1:3)
endif

END PROGRAM