!
! Copyright (C) 2021, MaX CoE
! Distributed under the MIT License
! (license terms are at http://opensource.org/licenses/MIT).
!
!--
!
! Wrappers for selected Linear Algebra subroutines
!
#include<devxlib_defs.h>
#include<devxlib_macros.h>
!
#ifdef __DXL_HAVE_DEVICE
  !
  ! checks
  !
#if ! defined __DXL_CUBLAS && ! defined __DXL_ROCBLAS && ! defined __DXL_MKL_GPU
# error
#endif
#endif

module devxlib_linalg

   use iso_c_binding,       only : c_loc, c_char
   use iso_fortran_env,     only : real32, real64
   use devxlib_environment, only : devxlib_error
   !
#if !defined __DXL_MKL_GPU
   use blas95_m
#endif
#if defined __DXL_OPENACC
   use devxlib_environment, only : cudaStreamSynchronize
   use openacc
#endif
#if !defined __DXL_MKL_GPU && !defined __DXL_ROCBLAS
   use devxlib_cublas
#elif defined __DXL_MKL_GPU
   use devxlib_mklgpu
#elif defined __DXL_ROCBLAS
   use devxlib_rocblas
#endif

   implicit none

   logical :: dev_linalg_init=.false.
   logical :: dev_linalg_close=.false.

   !
   ! HOST interfaces
   !
   interface devxlib_xDOT
     procedure :: SDOT, DDOT
   end interface devxlib_xDOT

   interface devxlib_xDOTU
     procedure :: SDOT, DDOT, CDOTU, ZDOTU
   end interface devxlib_xDOTU

   interface devxlib_xDOTC
     procedure :: SDOT, DDOT, CDOTC, ZDOTC
   end interface devxlib_xDOTC

   interface devxlib_xAXPY
     procedure :: SAXPY, DAXPY, CAXPY, ZAXPY
   end interface devxlib_xAXPY

   interface devxlib_xGEMV
     procedure :: SGEMV, DGEMV, CGEMV, ZGEMV
   end interface devxlib_xGEMV

   interface devxlib_xGEMM
     procedure :: SGEMM, DGEMM, CGEMM, ZGEMM
   end interface devxlib_xGEMM
   !
   interface devxlib_xAXPY_gpu
      module procedure &
         dev_SAXPY_gpu, dev_DAXPY_gpu, dev_CAXPY_gpu, dev_ZAXPY_gpu
   endinterface devxlib_xAXPY_gpu

   interface devxlib_xDOT_gpu
      module procedure &
         dev_SDOT_gpu, dev_DDOT_gpu
   endinterface devxlib_xDOT_gpu

   interface devxlib_xDOTU_gpu
      module procedure &
         dev_SDOT_gpu, dev_DDOT_gpu, dev_CDOTU_gpu, dev_ZDOTU_gpu
   endinterface devxlib_xDOTU_gpu

   interface devxlib_xDOTC_gpu
      module procedure &
         dev_SDOT_gpu, dev_DDOT_gpu, dev_CDOTC_gpu, dev_ZDOTC_gpu
   endinterface devxlib_xDOTC_gpu

   interface devxlib_xGEMV_gpu
      module procedure &
         dev_SGEMV_gpu, dev_DGEMV_gpu, dev_CGEMV_gpu, dev_ZGEMV_gpu
   endinterface devxlib_xGEMV_gpu

   interface devxlib_xGEMM_gpu
      module procedure &
         dev_SGEMM_gpu, dev_DGEMM_gpu, dev_CGEMM_gpu, dev_ZGEMM_gpu
   endinterface devxlib_xGEMM_gpu

   interface
      module subroutine dev_SAXPY_gpu(N,SA,SX,INCX,SY,INCY)
         implicit none
         integer,      intent(in) :: N
         integer,      intent(in) :: INCX,INCY
         real(real32), intent(in) TAR_ATTR  :: SA
         real(real32), intent(in) TAR_ATTR  DEV_ATTR :: SX(:)
         real(real32)             TAR_ATTR  DEV_ATTR :: SY(:)
      end subroutine dev_SAXPY_gpu
      !
      module subroutine dev_DAXPY_gpu(N,DA,DX,INCX,DY,INCY)
         implicit none
         integer,      intent(in) :: N
         integer,      intent(in) :: INCX,INCY
         real(real64), intent(in) TAR_ATTR :: DA
         real(real64), intent(in) TAR_ATTR DEV_ATTR :: DX(:)
         real(real64)             TAR_ATTR DEV_ATTR :: DY(:)
      end subroutine dev_DAXPY_gpu
      !
      module subroutine dev_CAXPY_gpu(N,CA,CX,INCX,CY,INCY)
         implicit none
         integer,         intent(in) :: N
         integer,         intent(in) :: INCX,INCY
         complex(real32), intent(in) TAR_ATTR :: CA
         complex(real32), intent(in) TAR_ATTR DEV_ATTR :: CX(:)
         complex(real32)             TAR_ATTR DEV_ATTR :: CY(:)
      end subroutine dev_CAXPY_gpu
      !
      module subroutine dev_ZAXPY_gpu(N,CA,CX,INCX,CY,INCY)
         implicit none
         integer,         intent(in) :: N
         integer,         intent(in) :: INCX,INCY
         complex(real64), intent(in) TAR_ATTR :: CA
         complex(real64), intent(in) TAR_ATTR DEV_ATTR :: CX(:)
         complex(real64)             TAR_ATTR DEV_ATTR :: CY(:)
      end subroutine dev_ZAXPY_gpu
   endinterface

   interface
      module function dev_SDOT_gpu(N,X,INCX,Y,INCY) result(res)
         implicit none
         real(real32) TAR_ATTR :: res
         integer,      intent(in) :: N
         integer,      intent(in) :: INCX,INCY
         real(real32), intent(in) TAR_ATTR DEV_ATTR :: X(:)
         real(real32)             TAR_ATTR DEV_ATTR :: Y(:)
      end function dev_SDOT_gpu
      !
      module function dev_DDOT_gpu(N,X,INCX,Y,INCY) result(res)
         implicit none
         real(real64) TAR_ATTR :: res
         integer,      intent(in) :: N
         integer,      intent(in) :: INCX,INCY
         real(real64), intent(in) TAR_ATTR DEV_ATTR :: X(:)
         real(real64)             TAR_ATTR DEV_ATTR :: Y(:)
      end function dev_DDOT_gpu
      !
      module function dev_CDOTU_gpu(N,X,INCX,Y,INCY) result(res)
         implicit none
         complex(real32) TAR_ATTR :: res
         integer,         intent(in) :: N
         integer,         intent(in) :: INCX,INCY
         complex(real32), intent(in) TAR_ATTR DEV_ATTR :: X(:)
         complex(real32)             TAR_ATTR DEV_ATTR :: Y(:)
      end function dev_CDOTU_gpu
      !
      module function dev_ZDOTU_gpu(N,X,INCX,Y,INCY) result(res)
         implicit none
         complex(real64) TAR_ATTR :: res
         integer,         intent(in) :: N
         integer,         intent(in) :: INCX,INCY
         complex(real64), intent(in) TAR_ATTR DEV_ATTR :: X(:)
         complex(real64)             TAR_ATTR DEV_ATTR :: Y(:)
      end function dev_ZDOTU_gpu
      !
      module function dev_CDOTC_gpu(N,X,INCX,Y,INCY) result(res)
         implicit none
         complex(real32) TAR_ATTR :: res
         integer,         intent(in) :: N
         integer,         intent(in) :: INCX,INCY
         complex(real32), intent(in) TAR_ATTR DEV_ATTR :: X(:)
         complex(real32)             TAR_ATTR DEV_ATTR :: Y(:)
      end function dev_CDOTC_gpu
      !
      module function dev_ZDOTC_gpu(N,X,INCX,Y,INCY) result(res)
         implicit none
         complex(real64) TAR_ATTR :: res
         integer,         intent(in) :: N
         integer,         intent(in) :: INCX,INCY
         complex(real64), intent(in) TAR_ATTR DEV_ATTR :: X(:)
         complex(real64)             TAR_ATTR DEV_ATTR :: Y(:)
      end function dev_ZDOTC_gpu
   endinterface

   interface
      module subroutine dev_SGEMV_gpu(TRANS,M,N,ALPHA,A,LDA,X,INCX,BETA,Y,INCY)
         implicit none
         real(real32),  intent(in) TAR_ATTR :: ALPHA,BETA
         integer,       intent(in) :: M,N,LDA,INCX,INCY
         character,     intent(in) :: TRANS
         real(real32),  intent(in) TAR_ATTR DEV_ATTR :: A(:,:),X(:)
         real(real32)              TAR_ATTR DEV_ATTR :: Y(:)
      end subroutine dev_SGEMV_gpu
      !
      module subroutine dev_DGEMV_gpu(TRANS,M,N,ALPHA,A,LDA,X,INCX,BETA,Y,INCY)
         implicit none
         real(real64),  intent(in) TAR_ATTR :: ALPHA,BETA
         integer,       intent(in) :: M,N,LDA,INCX,INCY
         character,     intent(in) :: TRANS
         real(real64),  intent(in) TAR_ATTR DEV_ATTR :: A(:,:),X(:)
         real(real64)              TAR_ATTR DEV_ATTR :: Y(:)
      end subroutine dev_DGEMV_gpu
      !
      module subroutine dev_CGEMV_gpu(TRANS,M,N,ALPHA,A,LDA,X,INCX,BETA,Y,INCY)
         implicit none
         complex(real32), intent(in) TAR_ATTR :: ALPHA,BETA
         integer,         intent(in) :: M,N,LDA,INCX,INCY
         character,       intent(in) :: TRANS
         complex(real32), intent(in) TAR_ATTR DEV_ATTR :: A(:,:),X(:)
         complex(real32)             TAR_ATTR DEV_ATTR :: Y(:)
      end subroutine dev_CGEMV_gpu
      !
      module subroutine dev_ZGEMV_gpu(TRANS,M,N,ALPHA,A,LDA,X,INCX,BETA,Y,INCY)
         implicit none
         complex(real64), intent(in) TAR_ATTR :: ALPHA,BETA
         integer,         intent(in) :: M,N,LDA,INCX,INCY
         character,       intent(in) :: TRANS
         complex(real64), intent(in) TAR_ATTR DEV_ATTR :: A(:,:),X(:)
         complex(real64)             TAR_ATTR DEV_ATTR :: Y(:)
      end subroutine dev_ZGEMV_gpu
   endinterface

   interface
      module subroutine dev_SGEMM_gpu(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
         implicit none
         real(real32), target, intent(in) :: ALPHA,BETA
         integer,              intent(in) :: K,LDA,LDB,LDC,M,N
         character,            intent(in) :: TRANSA,TRANSB
         real(real32), target, intent(in) DEV_ATTR :: A(:,:),B(:,:)
         real(real32), target             DEV_ATTR :: C(:,:)
      end subroutine dev_SGEMM_gpu
      !
      module subroutine dev_DGEMM_gpu(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
         implicit none
         real(real64), target, intent(in) :: ALPHA,BETA
         integer,              intent(in) :: K,LDA,LDB,LDC,M,N
         character,            intent(in) :: TRANSA,TRANSB
         real(real64), target, intent(in) DEV_ATTR :: A(:,:),B(:,:)
         real(real64), target             DEV_ATTR :: C(:,:)
      end subroutine dev_DGEMM_gpu
      !
      module subroutine dev_CGEMM_gpu(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
         implicit none
         complex(real32), target, intent(in) :: ALPHA,BETA
         integer,                 intent(in) :: K,LDA,LDB,LDC,M,N
         character,               intent(in) :: TRANSA,TRANSB
         complex(real32), target, intent(in) DEV_ATTR :: A(:,:),B(:,:)
         complex(real32), target             DEV_ATTR :: C(:,:)
      end subroutine dev_CGEMM_gpu
      !
      module subroutine dev_ZGEMM_gpu(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
         implicit none
         complex(real64), target, intent(in) :: ALPHA,BETA
         integer,                 intent(in) :: K,LDA,LDB,LDC,M,N
         character,               intent(in) :: TRANSA,TRANSB
         complex(real64), target, intent(in) DEV_ATTR :: A(:,:),B(:,:)
         complex(real64), target             DEV_ATTR :: C(:,:)
      end subroutine dev_ZGEMM_gpu
   endinterface

   contains
      !
      ! aux funcs
      !
      subroutine dev_linalg_setup()
         implicit none
         integer :: istat
         !
#ifdef __DXL_CUBLAS
         ! Is cublasInit needed?
         istat = cublasInit()
         if(istat/=0) call devxlib_error('dev_linalg_setup','cublasInit failed',10)
         !
         istat = cublasCreate(cublas_handle)
         if(istat/=0) call devxlib_error('dev_linalg_setup','cublasCreate failed',10)
         !
#if defined __DXL_OPENACC
         !istat = cublasSetStream(cublas_handle, acc_get_cuda_stream(acc_async_noval))
         !if(istat/=0) call devxlib_error('dev_linalg_setup','cublasSetStream failed',10)
         istat = cublasSetKernelStream(acc_get_cuda_stream(acc_async_noval))
         if(istat/=0) call devxlib_error('dev_linalg_setup','cublasSetKernelStream failed',10)
#endif
         !
         !istat = cusolverDnCreate(cusolv_h)
         !if(istat/=CUSOLVER_STATUS_SUCCESS) call error('cusolverDnCreate failed')
#endif
#ifdef __DXL_ROCBLAS
         !
         istat = rocblas_create_handle(handle)
         if(istat/=0) call devxlib_error('dev_linalg_setup','rocblas_create_handle failed',10)
         !
#endif
         !
         dev_linalg_init=.true.
         !
      end subroutine dev_linalg_setup

      subroutine dev_linalg_shutdown()
         implicit none
         integer :: istat
         !
#ifdef __DXL_CUBLAS
         !
         istat = cublasDestroy(cublas_handle)
         if(istat/=0) call devxlib_error('dev_linalg_shutdown','cublasDestroy failed',10)
#endif
#ifdef __DXL_ROCBLAS
         !
         istat = rocblas_destroy_handle(handle)
         if(istat/=0) call devxlib_error('dev_linalg_shutdown','rocblas_destroy_handle failed',10)
         !
#endif
         !
         dev_linalg_close=.true.
         !
      end subroutine dev_linalg_shutdown


      subroutine acc_loc_dev_sync()
         !
         ! AF: in ths subroutine we may want to make use of
         !     cublasSetStream when using OpenACC
         !
         implicit none
         integer :: ierr
         !
#if defined __DXL_OPENACC
         !ierr=cudaDeviceSynchronize()
         ierr=cudaStreamSynchronize(c_null_ptr)
#endif
         return
      end subroutine acc_loc_dev_sync
      !
endmodule devxlib_linalg

