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

   implicit none

   contains

   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(:)
      !
      integer :: ierr
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC data present(SX,SY)
      !DEV_ACC host_data use_device(SX,SY)
#if defined __DXL_CUBLAS
      call cublasSaxpy(N,SA,SX,INCX,SY,INCY)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(SX,SY)
      ierr=rocblas_saxpy(handle,N,c_loc(SA),c_loc(SX),INCX,c_loc(SY),INCY)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(SX,SY)
      !DEV_OMPGPU dispatch
      call Saxpy(N,SA,SX,INCX,SY,INCY)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call SAXPY(N,SA,SX,INCX,SY,INCY)
#endif
      !DEV_ACC end host_data
      !DEV_ACC end data
      !
      call acc_loc_dev_sync()
      !
   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(:)
      !
      integer :: ierr
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC data present(DX,DY)
      !DEV_ACC host_data use_device(DX,DY)
#if defined __DXL_CUBLAS
      call cublasDaxpy(N,DA,DX,INCX,DY,INCY)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(DX,DY)
      ierr=rocblas_daxpy(handle,N,c_loc(DA),c_loc(DX),INCX,c_loc(DY),INCY)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(DX,DY)
      !DEV_OMPGPU dispatch
      call Daxpy(N,DA,DX,INCX,DY,INCY)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call DAXPY(N,DA,DX,INCX,DY,INCY)
#endif
      !DEV_ACC end host_data
      !DEV_ACC end data
      !
      call acc_loc_dev_sync()
      !
   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(:)
      !
      integer :: ierr
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC data present(CX,CY)
      !DEV_ACC host_data use_device(CX,CY)
#if defined __DXL_CUBLAS
      call cublasCaxpy(N,CA,CX,INCX,CY,INCY)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(CX,CY)
      ierr=rocblas_caxpy(handle,N,c_loc(CA),c_loc(CX),INCX,c_loc(CY),INCY)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(CX,CY)
      !DEV_OMPGPU dispatch
      call Caxpy(N,CA,CX,INCX,CY,INCY)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call CAXPY(N,CA,CX,INCX,CY,INCY)
#endif
      !DEV_ACC end host_data
      !DEV_ACC end data
      !
      call acc_loc_dev_sync()
      !
   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(:)
      !
      integer :: ierr
      !
      call acc_loc_dev_sync()
      !
      !DEV_ACC data present(CX,CY)
      !DEV_ACC host_data use_device(CX,CY)
#if defined __DXL_CUBLAS
      call cublasZaxpy(N,CA,CX,INCX,CY,INCY)
#elif defined __DXL_ROCBLAS
      !DEV_OMPGPU target data use_device_ptr(CX,CY)
      ierr=rocblas_zaxpy(handle,N,c_loc(CA),c_loc(CX),INCX,c_loc(CY),INCY)
      !DEV_OMPGPU end target data
#elif defined __DXL_MKL_GPU
      !DEV_OMPGPU target data use_device_addr(CX,CY)
      !DEV_OMPGPU dispatch
      call Zaxpy(N,CA,CX,INCX,CY,INCY)
      !DEV_OMPGPU end dispatch
      !DEV_OMPGPU end target data
#else
      call ZAXPY(N,CA,CX,INCX,CY,INCY)
#endif
      !DEV_ACC end host_data
      !DEV_ACC end data
      !
      call acc_loc_dev_sync()
      !
   end subroutine dev_ZAXPY_gpu
   !
end submodule devxlib_linalg_axpy
