! ----------------------------------------------------------------------
! Testing matrix multiplication by the cuBLAS library
! Library functions: cublasSgemm_v2 & cublasDgemm_v2
! Defined operator cuBLASgemm
! Compilers: gfortran (10.2), ifort (21.1), nvfortran (20.11)
! ----------------------------------------------------------------------
! Command-line compilation:
! ----------------------------------------------------------------------
! Linux
! (set LIBRARY_PATH and LD_LIBRARY_PATH)
! gfortran testcuBLAS.f90 -lcudart -lcublas
! ifort -heap-arrays testcuBLAS.f90 -lcudart -lcublas
! nvfortran testcuBLAS.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 testcuBLAS.f90 -lcudart -lcublas
! 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 testcuBLAS.f90 cudart.lib cublas.lib
! nvfortran not available (03/2021)
! ----------------------------------------------------------------------

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

! single and double precision, size of arrays
integer,parameter :: sp=4,dp=8,nmax=10000,mmax=nmax,kmax=nmax
! matrix initialization: 1 constant elements, 2 random numbers
integer,parameter :: iselect=1

logical,parameter,private :: T=.TRUE.,F=.FALSE.
logical,parameter :: runcuBLASsp=T
logical,parameter :: runcuBLASdp=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

! ----------------------------------------------------------------------
! Matrix multiplication via the cuBLAS library
! Overloaded defined operator .cuBLASgemm.
! Usage: c4=a4.cuBLASgemm.b4; c8=a8.cuBLASgemm.b8
!        for real(4) and real(8) matrices
! ----------------------------------------------------------------------
MODULE mcuBLAS

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

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

interface operator (.cuBLASgemm.)
module procedure mm_sp_cuBLAS,mm_dp_CUBLAS
end interface

enum,bind(c) ! cublasStatus_t
enumerator :: CUBLAS_STATUS_SUCCESS=0,CUBLAS_STATUS_NOT_INITIALIZED=1,CUBLAS_STATUS_ALLOC_FAILED=3,CUBLAS_STATUS_INVALID_VALUE=7, &
              CUBLAS_STATUS_ARCH_MISMATCH=8,CUBLAS_STATUS_MAPPING_ERROR=11,CUBLAS_STATUS_EXECUTION_FAILED=13, &
              CUBLAS_STATUS_INTERNAL_ERROR=14,CUBLAS_STATUS_NOT_SUPPORTED=15,CUBLAS_STATUS_LICENSE_ERROR=16
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    ! cuBLAS routines

integer(c_int) function cublasCreate(handle) bind(c,name='cublasCreate_v2')
import c_int,c_ptr
type(c_ptr) :: handle ! cublasHandle_t
end function

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

integer(c_int) function cublasAlloc(n,elemSize,devicePtr) bind(c,name='cublasAlloc')
import c_int,c_ptr
integer(c_int),value,intent(in) :: n,elemSize
type(c_ptr) :: devicePtr
end function

integer(c_int) function cublasFree(devicePtr) bind(c,name='cublasFree')
import c_int,c_ptr
type(c_ptr),value :: devicePtr
end function

integer(c_int) function cublasSgemm(handle,transa,transb,m,n,k,alpha,ad,lda,bd,ldb,beta,cd,ldc) bind(c,name='cublasSgemm_v2')
import c_int,c_float,c_double,c_ptr,sp
type(c_ptr),value,intent(in) :: handle           ! cublasHandle
integer(c_int),value,intent(in) :: transa,transb ! cublasOperation_t
integer(c_int),value,intent(in) :: m,n,k         ! int    
integer(c_int),value,intent(in) :: lda,ldb,ldc   ! int
real(sp),intent(in) :: alpha,beta                ! double
type(c_ptr),value,intent(in) :: ad,bd            ! double
type(c_ptr),value :: cd                          ! double
end function

integer(c_int) function cublasDgemm(handle,transa,transb,m,n,k,alpha,ad,lda,bd,ldb,beta,cd,ldc) bind(c,name='cublasDgemm_v2')
import c_int,c_float,c_double,c_ptr,dp
type(c_ptr),value,intent(in) :: handle           ! cublasHandle
integer(c_int),value,intent(in) :: transa,transb ! cublasOperation_t
integer(c_int),value,intent(in) :: m,n,k         ! int    
integer(c_int),value,intent(in) :: lda,ldb,ldc   ! int
real(dp),intent(in) :: alpha,beta                ! double
type(c_ptr),value,intent(in) :: ad,bd            ! double
type(c_ptr),value :: cd                          ! double
end function

end interface

CONTAINS

function mm_sp_CUBLAS(a,b) RESULT (c)
real(sp),intent(in) :: a(:,:),b(:,:)
real(sp) c(size(a,1),size(b,2))
real(sp) alpha,beta
integer m,n,k,lda,ldb,ldc,istatus
integer transa,transb
type(c_ptr) ad,bd,cd
type(c_ptr) handle
target a,b,c

c(1,1)=0.
transa=CUBLAS_OP_N  ! no-op CUBLAS_OP_N, transpose CUBLAS_OP_T, Hermite conjugate CUBLAS_OP_C
transb=CUBLAS_OP_N
m=size(a,1)
n=size(b,2)
k=size(b,1)
if (size(a,2)/=size(b,1)) print *,'cuBLASgemm: size(a,2)/=size(b,1)' 
alpha=1.
beta=0.
lda=size(a,1)
ldb=size(b,1)
ldc=size(c,1)
istatus=cublasCreate(handle)
istatus=cudaMalloc(ad,int(lda*k*sp,c_size_t))
istatus=cudaMemcpy(ad,c_loc(a),int(lda*k*sp,c_size_t),cudaMemcpyHostToDevice)
istatus=cudaMalloc(bd,int(ldb*n*sp,c_size_t))
istatus=cudaMemcpy(bd,c_loc(b),int(ldb*n*sp,c_size_t),cudaMemcpyHostToDevice)
istatus=cudaMalloc(cd,int(ldc*n*sp,c_size_t))
istatus=cublasSgemm(handle,transa,transb,m,n,k,alpha,ad,lda,bd,ldb,beta,cd,ldc)
istatus=cudaMemcpy(c_loc(c),cd,int(ldc*n*sp,c_size_t),cudaMemcpyDeviceToHost)
istatus=cudaFree(ad)
istatus=cudaFree(bd)
istatus=cudaFree(cd)
istatus=cublasDestroy(handle)
end function

function mm_dp_CUBLAS(a,b) RESULT (c)
real(dp),intent(in) :: a(:,:),b(:,:)
real(dp) c(size(a,1),size(b,2))
real(dp) alpha,beta
integer m,n,k,lda,ldb,ldc,istatus
integer transa,transb
type(c_ptr) ad,bd,cd
type(c_ptr) handle
target a,b,c

c(1,1)=0.
transa=CUBLAS_OP_N  ! no-op CUBLAS_OP_N, transpose CUBLAS_OP_T, Hermite conjugate CUBLAS_OP_C
transb=CUBLAS_OP_N
m=size(a,1)
n=size(b,2)
k=size(b,1)
if (size(a,2)/=size(b,1)) print *,'cuBLASgemm: size(a,2)/=size(b,1)' 
alpha=1.
beta=0.
lda=size(a,1)
ldb=size(b,1)
ldc=size(c,1)
istatus=cublasCreate(handle)
istatus=cudaMalloc(ad,int(lda*k*dp,c_size_t))
istatus=cudaMemcpy(ad,c_loc(a),int(lda*k*dp,c_size_t),cudaMemcpyHostToDevice)
istatus=cudaMalloc(bd,int(ldb*n*dp,c_size_t))
istatus=cudaMemcpy(bd,c_loc(b),int(ldb*n*dp,c_size_t),cudaMemcpyHostToDevice)
istatus=cudaMalloc(cd,int(ldc*n*dp,c_size_t))
istatus=cublasDgemm(handle,transa,transb,m,n,k,alpha,ad,lda,bd,ldb,beta,cd,ldc)
istatus=cudaMemcpy(c_loc(c),cd,int(ldc*n*dp,c_size_t),cudaMemcpyDeviceToHost)
istatus=cudaFree(ad)
istatus=cudaFree(bd)
istatus=cudaFree(cd)
istatus=cublasDestroy(handle)
end function

END MODULE

! ----------------------------------------------------------------------
  PROGRAM TestcuBLAS
! ----------------------------------------------------------------------
use mConst
use mTicToc
use mcuBLAS
implicit none
real(dp),allocatable :: a(:,:),b(:,:),c(:,:)

ALLOCATE (a(mmax,kmax),b(kmax,nmax),c(mmax,nmax))
select case (iselect)
case (1)
  a=1.
  b=1.
case (2)
  call random_number(a)
  call random_number(b)
end select

if (runcuBLASsp) then
  print *,'cuBLAS SP (GPU)'
  call tic()
  c=real(a,sp).cuBLASgemm.real(b,sp)
  print *,'dt (SP) = ',toc()
  print '(1x,a,10e20.12)','c =',c(1:1,1:1),sum(c)
endif

if (runcuBLASdp) then
  print *,'cuBLAS DP (GPU)'
  call tic()
  c=a.cuBLASgemm.b
  print *,'dt (DP) = ',toc()
  print '(1x,a,10e20.12)','c =',c(1:1,1:1),sum(c)
endif

END PROGRAM