!
! Copyright (C) 2021, MaX CoE
! Distributed under the MIT License
! (license terms are at http://opensource.org/licenses/MIT).
!
!--
!
! Utility functions to perform memcpy and memset on the device with CUDA Fortran
! cuf_memXXX contains a CUF KERNEL to perform the selected operation
! cu_memsync are wrappers for cuda_memcpy functions
!
module devxlib_environment
   !
   use iso_c_binding, only : c_int, c_intptr_t, c_ptr
#if defined __DXL_OPENACC
   use openacc
#elif defined __DXL_CUDAF
   use cudafor
#endif

   implicit none

#if defined __DXL_OPENACC

   interface
      !
      function cudaDeviceSynchronize() &
         result(ierr) bind(c, name="cudaDeviceSynchronize")
         import c_int
         integer(c_int) :: ierr
      end function cudaDeviceSynchronize
      !
      function cudaStreamCreate(pstream) &
         result(ierr) bind(c, name="cudaStreamCreate")
         import c_ptr, c_int
         type(c_ptr) :: pstream
         integer(c_int) :: ierr
      end function cudaStreamCreate
      !
      function cudaStreamDestroy(pstream) &
         result(ierr) bind(c, name="cudaStreamDestroy")
         import c_ptr, c_int
         type(c_ptr), value :: pstream
         integer(c_int) :: ierr
      end function cudaStreamDestroy
      !
      function cudaStreamSynchronize(pstream) &
         result(ierr) bind(c, name="cudaStreamSynchronize")
         import c_ptr, c_int
         type(c_ptr), value :: pstream
         integer(c_int) :: ierr
      end function cudaStreamSynchronize
      !
      function cudaGetDeviceCount(count) &
         result(ierr) bind(c, name="cudaGetDeviceCount")
         import :: c_int
         implicit none
         integer(c_int) :: count
         integer(c_int) :: ierr
      end function cudaGetDeviceCount
      !
      function cudaSetDevice(device) &
         result(ierr) bind(c, name = "cudaSetDevice")
         import :: c_int
         implicit none
         integer(c_int), value :: device
         integer(c_int) :: ierr
      end function cudaSetDevice

      function cudaGetLastError() &
         result(ierr) bind(c, name = "cudaGetLastError")
         import :: c_int
         implicit none
         integer(c_int) :: ierr
      end function cudaGetLastError
      !
      function cudaGetDevice(device) &
         result(ierr) bind(c, name = "cudaGetDevice")
         import :: c_int
         implicit none
         integer(c_int) :: device
         integer(c_int) :: ierr
      end function cudaGetDevice
      !
   end interface

   interface
      function acc_malloc_f(total_byte_dim) bind(c, name="acc_malloc")
         use iso_c_binding, only : c_ptr, c_size_t
         implicit none
         type(c_ptr)                          :: acc_malloc_f
         integer(c_size_t), value, intent(in) :: total_byte_dim
      endfunction acc_malloc_f

      subroutine acc_free_f(dev_ptr) bind(c, name="acc_free")
         use iso_c_binding, only : c_ptr
         implicit none
         type(c_ptr), value :: dev_ptr
      endsubroutine acc_free_f

      function acc_get_cuda_stream_f(async) result(pstream) bind(c,name="acc_get_cuda_stream")
         import c_int, c_intptr_t
         integer(c_int), value :: async
         integer(c_intptr_t)   :: pstream
         !type(c_ptr)          :: pstream
      end function acc_get_cuda_stream_f

      subroutine acc_memcpy_to_device_f(dev_ptr, host_ptr, total_byte_dim) bind(c, name="acc_memcpy_to_device")
         use iso_c_binding, only : c_ptr, c_size_t
         implicit none
         type(c_ptr),       value :: dev_ptr
         type(c_ptr),       value :: host_ptr
         integer(c_size_t), value :: total_byte_dim
      endsubroutine acc_memcpy_to_device_f

      subroutine acc_memcpy_from_device_f(host_ptr, dev_ptr, total_byte_dim) bind(c, name="acc_memcpy_from_device")
         use iso_c_binding, only : c_ptr, c_size_t
         implicit none
         type(c_ptr),       value :: host_ptr
         type(c_ptr),       value :: dev_ptr
         integer(c_size_t), value :: total_byte_dim
      endsubroutine acc_memcpy_from_device_f

      subroutine acc_memcpy_device_f(dev_ptr_dst, dev_ptr_src, total_byte_dim) bind(c, name="acc_memcpy_device")
         use iso_c_binding, only : c_ptr, c_size_t
         implicit none
         type(c_ptr),       value :: dev_ptr_dst
         type(c_ptr),       value :: dev_ptr_src
         integer(c_size_t), value :: total_byte_dim
      endsubroutine acc_memcpy_device_f

   endinterface
#endif

#if defined __DXL_CUDAF
   interface
      integer (c_int) function cudaMalloc_f(buffer, size)  bind(C,name="cudaMalloc")
         use iso_c_binding
         implicit none
         type (c_ptr)              :: buffer
         integer (c_size_t), value :: size
      end function cudaMalloc_f

      integer (c_int) function cudaFree_f(buffer)  bind(C,name="cudaFree")
         use iso_c_binding
         implicit none
         type (c_ptr), value :: buffer
      end function cudaFree_f

      integer(c_int) function cudaMemcpy_f(dst, src, count, kind) bind(C, name='cudaMemcpy')
         use iso_c_binding
         type(c_ptr), value :: dst
         type(c_ptr), value :: src
         integer(c_size_t), value :: count
         integer(c_int), value :: kind
      end function cudaMemcpy_f
   endinterface
#endif

   contains

      subroutine devxlib_error( calling_routine, message, ierr )
         implicit none
         character(len=*), intent(in) :: calling_routine, message
         ! the name of the calling calling_routine
         ! the output message
         integer,          intent(in) :: ierr
         !
         write(0,*) trim(calling_routine), trim(message), ierr
         !
      end subroutine devxlib_error

endmodule devxlib_environment

