cuda_intf.F90

Go to the documentation of this file.

! Copyright (c) 2021-2026, The Neko Authors
! All rights reserved.
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions
! are met:
!
!   * Redistributions of source code must retain the above copyright
!     notice, this list of conditions and the following disclaimer.
!
!   * Redistributions in binary form must reproduce the above
!     copyright notice, this list of conditions and the following
!     disclaimer in the documentation and/or other materials provided
!     with the distribution.
!
!   * Neither the name of the authors nor the names of its
!     contributors may be used to endorse or promote products derived
!     from this software without specific prior written permission.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
! "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
! LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
! FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
! COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
! INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
! BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
! LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
! CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
! LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
! ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
! POSSIBILITY OF SUCH DAMAGE.
!
module cuda_intf
  use utils, only : neko_error
  use, intrinsic :: iso_c_binding
  !$ use omp_lib
  implicit none
 
#ifdef HAVE_CUDA
 
  enum, bind(c)
    enumerator :: cudasuccess = 0
    enumerator :: cudaerrorinvalidvalue = 1
    enumerator :: cudaerrormemoryallocation = 2
    enumerator :: cudaerrorinitializationerror = 3
  end enum
 
  enum, bind(c)
    enumerator :: cudamemcpyhosttohost = 0
    enumerator :: cudamemcpyhosttodevice = 1
    enumerator :: cudamemcpydevicetohost = 2
    enumerator :: cudamemcpydevicetodevice = 3
    enumerator :: cudamemcpydefault = 4
  end enum
 
  interface
     integer(c_int) function cudamalloc(ptr_d, s) &
          bind(c, name = 'cudaMalloc')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr) :: ptr_d
       integer(c_size_t), value :: s
     end function cudamalloc
 
     integer(c_int) function cudafree(ptr_d) &
          bind(c, name = 'cudaFree')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: ptr_d
     end function cudafree
 
     integer(c_int) function cudamemcpy(ptr_dst, ptr_src, s, dir) &
          bind(c, name = 'cudaMemcpy')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: ptr_dst, ptr_src
       integer(c_size_t), value :: s
       integer(c_int), value :: dir
     end function cudamemcpy
 
     integer(c_int) function cudamemcpyasync(ptr_dst, ptr_src, s, dir, stream) &
          bind(c, name = 'cudaMemcpyAsync')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: ptr_dst, ptr_src, stream
       integer(c_size_t), value :: s
       integer(c_int), value :: dir
     end function cudamemcpyasync
 
     integer(c_int) function cudamemsetasync(ptr, v, s, stream) &
          bind(c, name = 'cudaMemsetAsync')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: ptr, stream
       integer(c_int), value :: v
       integer(c_size_t), value :: s
     end function cudamemsetasync
 
     integer(c_int) function cudadevicesynchronize() &
          bind(c, name = 'cudaDeviceSynchronize')
       use, intrinsic :: iso_c_binding
       implicit none
     end function cudadevicesynchronize
 
     integer(c_int) function cudadevicereset() &
          bind(c, name = 'cudaDeviceReset')
       use, intrinsic :: iso_c_binding
       implicit none
     end function cudadevicereset
 
     subroutine cuda_buffer_free_all() &
          bind(c, name = 'cuda_buffer_free_all')
     end subroutine cuda_buffer_free_all
 
     integer(c_int) function cudagetdeviceproperties(prop, device) &
          bind(c, name = 'cudaGetDeviceProperties')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: prop
       integer(c_int), value :: device
     end function cudagetdeviceproperties
 
     integer(c_int) function cudastreamcreate(stream) &
          bind(c, name = 'cudaStreamCreate')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr) :: stream
     end function cudastreamcreate
 
     integer(c_int) function cudastreamcreatewithflags(stream, flags) &
          bind(c, name = 'cudaStreamCreateWithFlags')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr) :: stream
       integer(c_int), value :: flags
     end function cudastreamcreatewithflags
 
     integer(c_int) function cudastreamcreatewithpriority(stream, flags, prio) &
          bind(c, name = 'cudaStreamCreateWithPriority')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr) :: stream
       integer(c_int), value :: flags, prio
     end function cudastreamcreatewithpriority
 
     integer(c_int) function cudastreamdestroy(steam) &
          bind(c, name = 'cudaStreamDestroy')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: steam
     end function cudastreamdestroy
 
     integer(c_int) function cudastreamsynchronize(stream) &
          bind(c, name = 'cudaStreamSynchronize')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: stream
     end function cudastreamsynchronize
 
     integer(c_int) function cudastreamwaitevent(stream, event, flags) &
          bind(c, name = 'cudaStreamWaitEvent')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: stream, event
       integer(c_int), value :: flags
     end function cudastreamwaitevent
 
     integer(c_int) function cudadevicegetstreampriorityrange &
          (low_prio, high_prio) &
          bind(c, name = 'cudaDeviceGetStreamPriorityRange')
       use, intrinsic :: iso_c_binding
       implicit none
       integer(c_int) :: low_prio, high_prio
     end function cudadevicegetstreampriorityrange
 
     integer(c_int) function cudaprofilerstart() &
          bind(c, name = 'cudaProfilerStart')
       use, intrinsic :: iso_c_binding
       implicit none
     end function cudaprofilerstart
 
     integer(c_int) function cudaprofilerstop() &
          bind(c, name = 'cudaProfilerStop')
       use, intrinsic :: iso_c_binding
       implicit none
     end function cudaprofilerstop
 
     integer(c_int) function cudaeventcreate(event) &
          bind(c, name = 'cudaEventCreate')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr) :: event
     end function cudaeventcreate
 
     integer(c_int) function cudaeventdestroy(event) &
          bind(c, name = 'cudaEventDestroy')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: event
     end function cudaeventdestroy
 
     integer(c_int) function cudaeventcreatewithflags(event, flags) &
          bind(c, name = 'cudaEventCreateWithFlags')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr) :: event
       integer(c_int), value :: flags
     end function cudaeventcreatewithflags
 
     integer(c_int) function cudaeventrecord(event, stream) &
          bind(c, name = 'cudaEventRecord')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: event, stream
     end function cudaeventrecord
 
     integer(c_int) function cudaeventsynchronize(event) &
          bind(c, name = 'cudaEventSynchronize')
       use, intrinsic :: iso_c_binding
       implicit none
       type(c_ptr), value :: event
     end function cudaeventsynchronize
 
     integer(c_int) function cudagetdevicecount(device_count) &
          bind(c, name = 'cudaGetDeviceCount')
       use, intrinsic :: iso_c_binding
       implicit none
       integer(c_int) :: device_count
     end function cudagetdevicecount
 
     integer(c_int) function cuda_device_mp_count() &
          bind(c, name = 'cuda_device_mp_count')
       use, intrinsic :: iso_c_binding
       implicit none
     end function cuda_device_mp_count
 
     integer(c_size_t) function cuda_device_total_mem() &
          bind(c, name = 'cuda_device_total_mem')
       use, intrinsic :: iso_c_binding
       implicit none
     end function cuda_device_total_mem
 
     integer(c_int) function cudagetdevice(device) &
          bind(c, name = 'cudaGetDevice')
       use, intrinsic :: iso_c_binding
       implicit none
       integer(c_int) :: device
     end function cudagetdevice
 
     integer(c_int) function cudasetdevice(device) &
          bind(c, name = 'cudaSetDevice')
       use, intrinsic :: iso_c_binding
       implicit none
       integer(c_int), value :: device
     end function cudasetdevice
  end interface
 
contains
 
  subroutine cuda_init(glb_cmd_queue, aux_cmd_queue, &
       STRM_HIGH_PRIO, STRM_LOW_PRIO)
    type(c_ptr), intent(inout) :: glb_cmd_queue
    type(c_ptr), intent(inout) :: aux_cmd_queue
    integer, intent(inout) :: STRM_HIGH_PRIO
    integer, intent(inout) :: STRM_LOW_PRIO
    integer(c_int) :: device_id
    integer :: nthrds = 1
 
    !$omp parallel
    !$omp master
    !$ nthrds = omp_get_num_threads()
    !$omp end master
    !$omp end parallel
 
    ! Ensure that all threads are assigned to the same device
    if (nthrds .gt. 1) then
       if (cudagetdevice(device_id) .ne. cudasuccess) then
          call neko_error('Error retrieving device id')
       end if
 
       !$omp parallel
       if (cudasetdevice(device_id) .ne. cudasuccess) then
          call neko_error('Error setting device id')
       end if
       !$omp end parallel
    end if
 
    if (cudadevicegetstreampriorityrange(strm_low_prio, strm_high_prio) &
         .ne. cudasuccess) then
       call neko_error('Error retrieving stream priority range')
    end if
 
    if (cudastreamcreatewithpriority(glb_cmd_queue, 1, strm_high_prio) &
         .ne. cudasuccess) then
       call neko_error('Error creating main stream')
    end if
 
    if (cudastreamcreatewithpriority(aux_cmd_queue, 1, strm_low_prio) &
         .ne. cudasuccess) then
       call neko_error('Error creating aux stream')
    end if
  end subroutine cuda_init
 
  subroutine cuda_finalize(glb_cmd_queue, aux_cmd_queue)
    type(c_ptr), intent(inout) :: glb_cmd_queue
    type(c_ptr), intent(inout) :: aux_cmd_queue
    integer :: ierr
 
    ! Release all device buffers held by the device layer
    call cuda_buffer_free_all()
 
    if (cudastreamdestroy(glb_cmd_queue) .ne. cudasuccess) then
       call neko_error('Error destroying main stream')
    end if
 
    if (cudastreamdestroy(aux_cmd_queue) .ne. cudasuccess) then
       call neko_error('Error destroying aux stream')
    end if
 
    ierr = cudadevicesynchronize()
 
    ! Best-effort context teardown to release runtime-owned allocations.
    ! Skipped in pFUnit-enabled builds: unit tests cycle device
    ! init/finalize with MPI still up, and device-aware communication
    ! backends (CUDA-aware MPI, NCCL, NVSHMEM) cache the primary
    ! context from first use; destroying it here would leave them with
    ! a dangling context
#ifndef HAVE_PFUNIT
    ierr = cudadevicereset()
#endif
  end subroutine cuda_finalize
 
  subroutine cuda_device_name(name)
    character(len=*), intent(inout) :: name
    character(kind=c_char, len=8192), target :: prop
    integer :: end_pos
 
    !
    ! Yes this is an ugly hack!
    ! Since we're only interested in the device name (first 256 bytes)
    ! we pass down a large enough chunk of memory to the cuda runtime
    ! and extract what we need later on
    !
    ! This will of course break if sizeof(cudaDeviceProp) > 8192
    !
 
    if (cudagetdeviceproperties(c_loc(prop), 0) .ne. cudasuccess) then
       call neko_error('Failed to query device')
    end if
 
    end_pos = scan(prop(1:256), c_null_char)
    if (end_pos .ge. 2) then
       name(1:end_pos-1) = prop(1:end_pos-1)
    end if
  end subroutine cuda_device_name
 
  integer function cuda_device_count()
    integer(c_int) :: num_devices
 
    if (cudagetdevicecount(num_devices) .ne. cudasuccess) then
       call neko_error('Error retrieving device count')
    end if
 
    cuda_device_count = num_devices
  end function cuda_device_count
 
#endif
 
end module cuda_intf