!
! 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

submodule (devxlib_linalg) devxlib_linalg_gemm

   implicit none

   contains

   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(:,:)
      !
#if defined __DXL_ROCBLAS
      integer(c_int) :: itransa, itransb
#else
      character(c_char) :: transa_,transb_
#endif
      integer :: ierr
      !
#if defined __DXL_ROCBLAS
      itransa = rocblas_get_operation(TRANSA)
      itransb = rocblas_get_operation(TRANSB)
#else
      transa_=TRANSA
      transb_=TRANSB
#endif
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC_DEBUG data present(A,B,C)
      !DEV_ACC host_data use_device(A,B,C)
#if defined __DXL_CUBLAS
      call cublasSgemm(transa_,transb_,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(A,B,C)
      ierr=rocblas_sgemm(handle,itransa,itransb,M,N,K,c_loc(ALPHA),c_loc(A),LDA,c_loc(B),LDB,c_loc(BETA),c_loc(C),LDC)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(A,B,C)
      !DEV_OMPGPU dispatch
      call SGEMM(transa,transb,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call SGEMM(transa,transb,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#endif
      !DEV_ACC end host_data
      !DEV_ACC_DEBUG end data
      !
      call acc_loc_dev_sync()
      !
   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(:,:)
      !
#if defined __DXL_ROCBLAS
      integer :: itransa, itransb
#else
      character(c_char) :: transa_,transb_
#endif
      integer :: ierr
      !
#if defined __DXL_ROCBLAS
      itransa = rocblas_get_operation(TRANSA)
      itransb = rocblas_get_operation(TRANSB)
#else
      transa_=TRANSA
      transb_=TRANSB
#endif
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC_DEBUG data present(A,B,C)
      !DEV_ACC host_data use_device(A,B,C)
#if defined __DXL_CUBLAS
      call cublasDgemm(transa_,transb_,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(A,B,C)
      ierr=rocblas_dgemm(handle,itransa,itransb,M,N,K,c_loc(ALPHA),c_loc(A),LDA,c_loc(B),LDB,c_loc(BETA),c_loc(C),LDC)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(A,B,C)
      !DEV_OMPGPU dispatch
      call DGEMM(transa,transb,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call DGEMM(transa,transb,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#endif
      !DEV_ACC end host_data
      !DEV_ACC_DEBUG end data
      !
      call acc_loc_dev_sync()
      !
   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(:,:)
      !
#if defined __DXL_ROCBLAS
      integer :: itransa, itransb
#else
      character(c_char) :: transa_,transb_
#endif

      integer :: ierr
      !
#if defined __DXL_ROCBLAS
      itransa = rocblas_get_operation(TRANSA)
      itransb = rocblas_get_operation(TRANSB)
#else
      transa_=TRANSA
      transb_=TRANSB
#endif
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC_DEBUG data present(A,B,C)
      !DEV_ACC host_data use_device(A,B,C)
#if defined __DXL_CUBLAS
      call cublasCgemm(transa_,transb_,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(A,B,C)
      ierr=rocblas_cgemm(handle,itransa,itransb,M,N,K,c_loc(ALPHA),c_loc(A),LDA,c_loc(B),LDB,c_loc(BETA),c_loc(C),LDC)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(A,B,C)
      !DEV_OMPGPU dispatch
      call CGEMM(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call CGEMM(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#endif
      !DEV_ACC end host_data
      !DEV_ACC_DEBUG end data
      !
      call acc_loc_dev_sync()
      !
   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(:,:)
      !
#if defined __DXL_ROCBLAS
      integer :: itransa, itransb
#else
      character(c_char) :: transa_,transb_
#endif

      integer :: ierr
      !
#if defined __DXL_ROCBLAS
      itransa = rocblas_get_operation(TRANSA)
      itransb = rocblas_get_operation(TRANSB)
#else
      transa_=TRANSA
      transb_=TRANSB
#endif
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC_DEBUG data present(A,B,C)
      !DEV_ACC host_data use_device(A,B,C)
#if defined __DXL_CUBLAS
      call cublasZgemm(transa_,transb_,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(A,B,C)
      ierr=rocblas_zgemm(handle,itransa,itransb,M,N,K,c_loc(ALPHA),c_loc(A),LDA,c_loc(B),LDB,c_loc(BETA),c_loc(C),LDC)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(A,B,C)
      !DEV_OMPGPU dispatch
      call ZGEMM(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call ZGEMM(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC)
#endif
      !DEV_ACC end host_data
      !DEV_ACC_DEBUG end data
      !
      call acc_loc_dev_sync()
      !
   end subroutine dev_ZGEMM_gpu
   !
end submodule devxlib_linalg_gemm
