comm.F90

Go to the documentation of this file.

module comm
  use mpi_f08, only : mpi_comm, mpi_datatype, mpi_initialized, mpi_init_thread, &
       mpi_init, mpi_thread_single, mpi_thread_funneled, &
       mpi_thread_serialized, mpi_thread_multiple, mpi_comm_rank, &
       mpi_comm_split, mpi_comm_dup, mpi_barrier, mpi_comm_free, mpi_finalize, &
       mpi_comm_world, mpi_double_precision, mpi_real, mpi_comm_size, &
       mpi_query_thread
  use utils, only : neko_error
  use neko_config
  use shmem
  !$ use omp_lib
  implicit none
  private
 
  interface
     subroutine neko_comm_wrapper_init(fcomm) &
          bind(c, name='neko_comm_wrapper_init')
       use, intrinsic :: iso_c_binding, only : c_int
       integer(c_int), value :: fcomm
     end subroutine neko_comm_wrapper_init
 
#ifdef HAVE_NVSHMEM
     subroutine neko_comm_nvshmem_init() &
          bind(c, name='neko_comm_nvshmem_init')
     end subroutine neko_comm_nvshmem_init
 
     subroutine neko_comm_nvshmem_finalize() &
          bind(c, name='neko_comm_nvshmem_finalize')
     end subroutine neko_comm_nvshmem_finalize
#endif
 
#if defined(HAVE_NCCL) || defined(HAVE_RCCL)
     subroutine neko_comm_nccl_init() &
          bind(c, name='neko_comm_nccl_init')
     end subroutine neko_comm_nccl_init
 
     subroutine neko_comm_nccl_finalize() &
          bind(c, name='neko_comm_nccl_finalize')
     end subroutine neko_comm_nccl_finalize
#endif
 
  end interface
 
 
  type(mpi_comm), public :: neko_comm
  type(mpi_comm), public :: neko_global_comm
 
#ifdef HAVE_MPI_PARAM_DTYPE
  type(mpi_datatype), public, parameter :: mpi_real_precision = mpi_double_precision
  type(mpi_datatype), public, parameter :: mpi_extra_precision = mpi_double_precision
#else
  type(mpi_datatype), public :: mpi_real_precision
  type(mpi_datatype), public :: mpi_extra_precision
#endif
 
  integer, public :: pe_rank
 
  integer, public :: pe_size
 
  logical, public :: nio
 
  integer, public :: global_pe_rank
 
  integer, public :: global_pe_size
 
  integer, public :: neko_mpi_thread_provided
 
  public :: comm_init, comm_free
 
contains
  subroutine comm_init
    integer :: ierr
    logical :: initialized
    integer :: provided, nthrds
    integer :: color = 0
    integer :: envvar_len
    character(len=255) :: color_str
    character(len=32) :: thread_str
    integer :: thread_envvar_len
    integer :: requested_thread_level
    logical :: user_thread_level
#ifdef HAVE_OPENSHMEM
    integer :: shmem_ierr
    integer :: shmem_requested
#endif
 
    pe_rank = -1
    pe_size = 0
    nio = .false.
 
    call mpi_initialized(initialized, ierr)
 
    call get_environment_variable("NEKO_COMM_ID", color_str, envvar_len)
    if (envvar_len .gt. 0) then
       read(color_str(1:envvar_len), *) color
    else
       color = 0
    end if
 
    nthrds = 1
    !$omp parallel
    !$omp master
    !$ nthrds = omp_get_num_threads()
    !$omp end master
    !$omp end parallel
 
    call get_environment_variable("NEKO_MPI_THREAD_LEVEL", thread_str, &
         thread_envvar_len)
    user_thread_level = thread_envvar_len .gt. 0
    if (user_thread_level) then
       select case (trim(adjustl(thread_str(1:thread_envvar_len))))
       case ("single", "SINGLE")
          requested_thread_level = mpi_thread_single
       case ("funneled", "FUNNELED")
          requested_thread_level = mpi_thread_funneled
       case ("serialized", "SERIALIZED")
          requested_thread_level = mpi_thread_serialized
       case ("multiple", "MULTIPLE")
          requested_thread_level = mpi_thread_multiple
       case default
          call neko_error('Unknown NEKO_MPI_THREAD_LEVEL: '// &
               trim(thread_str(1:thread_envvar_len)))
       end select
    end if
 
    if (.not.initialized) then
       if (user_thread_level) then
          if (requested_thread_level .eq. mpi_thread_single) then
             call mpi_init(ierr)
             provided = mpi_thread_single
          else
             call mpi_init_thread(requested_thread_level, provided, ierr)
             if (provided .lt. requested_thread_level) then
                call neko_error('Requested MPI thread level not provided')
             end if
          end if
       else if (nthrds .gt. 1) then
          call mpi_init_thread(mpi_thread_multiple, provided, ierr)
          if (provided .lt. mpi_thread_multiple) then
             ! MPI_THREAD_MULTIPLE is required for mt. device backends. For
             ! host backends the gather-scatter MPI calls are issued from the
             ! OpenMP master thread, so MPI_THREAD_SERIALIZED (or FUNNELED) is
             ! sufficient as a last resort.
             if (neko_bcknd_device .eq. 1) then
                call neko_error('Invalid thread support provided by MPI')
             else if (provided .lt. mpi_thread_funneled) then
                call neko_error('Invalid thread support provided by MPI')
             end if
          end if
       else
          call mpi_init(ierr)
       end if
    else
       call mpi_query_thread(provided, ierr)
    end if
 
    neko_mpi_thread_provided = provided
 
#ifndef HAVE_MPI_PARAM_DTYPE
    mpi_real_precision = mpi_double_precision
    mpi_extra_precision = mpi_double_precision
#endif
 
 
#ifdef HAVE_ADIOS2
    ! We split the communicator it to work asynchronously (MPMD)
    call mpi_comm_rank(mpi_comm_world, pe_rank, ierr)
    call mpi_comm_split(mpi_comm_world, 0, pe_rank, neko_global_comm, ierr)
#else
    ! Original version duplicates the communicator:
    call mpi_comm_dup(mpi_comm_world, neko_global_comm, ierr)
#endif
 
    call mpi_comm_rank(neko_global_comm, global_pe_rank, ierr)
    call mpi_comm_size(neko_global_comm, global_pe_size, ierr)
    if (envvar_len .gt. 0) then
       call mpi_comm_split(neko_global_comm, color, global_pe_rank, &
            neko_comm, ierr)
    else
       call mpi_comm_dup(neko_global_comm, neko_comm, ierr)
    end if
    call mpi_comm_rank(neko_comm, pe_rank, ierr)
    call mpi_comm_size(neko_comm, pe_size, ierr)
 
    ! Setup C/C++ wrapper
    call neko_comm_wrapper_init(neko_comm%mpi_val)
 
 
#ifdef HAVE_NVSHMEM
    ! Setup NVSHMEM (if requested)
    call neko_comm_nvshmem_init()
#endif
 
#if defined(HAVE_NCCL) | defined(HAVE_RCCL)
    ! Setup NCCL (if requested)
    call neko_comm_nccl_init()
#endif
 
#ifdef HAVE_OPENSHMEM
    ! Setup OpenSHMEM (if requested)
    if (user_thread_level) then
       select case (requested_thread_level)
       case (mpi_thread_single)
          shmem_requested = shmem_thread_single
       case (mpi_thread_funneled)
          shmem_requested = shmem_thread_funneled
       case (mpi_thread_serialized)
          shmem_requested = shmem_thread_serialized
       case (mpi_thread_multiple)
          shmem_requested = shmem_thread_multiple
       end select
       if (shmem_requested .eq. shmem_thread_single) then
          call shmem_init()
       else
          shmem_ierr = shmem_init_thread(shmem_requested, provided)
          if (provided .lt. shmem_requested) then
             call neko_error('Requested SHMEM thread level not provided')
          end if
       end if
    else if (nthrds .gt. 1) then
       shmem_ierr = shmem_init_thread(shmem_thread_multiple, provided)
       if (provided .ne. shmem_thread_multiple) then
          if (neko_bcknd_device .eq. 1) then
             call neko_error('Invalid thread support provided by SHMEM')
          else
             shmem_ierr = shmem_init_thread(shmem_thread_serialized, provided)
             if (provided .ne. shmem_thread_serialized) then
                call neko_error('Invalid thread support provided by SHMEM')
             end if
          end if
       end if
    else
       call shmem_init()
    end if
#endif
 
 
  end subroutine comm_init
 
  subroutine comm_free
    integer :: ierr
 
    call mpi_barrier(neko_comm, ierr)
    call mpi_comm_free(neko_comm, ierr)
    call mpi_comm_free(neko_global_comm, ierr)
 
#ifdef HAVE_NCCL
    call neko_comm_nccl_finalize()
#endif
 
#ifdef HAVE_NVSHMEM
    call neko_comm_nvshmem_finalize()
#endif
 
#ifdef HAVE_OPENSHMEM
    call shmem_finalize()
#endif
 
    call mpi_finalize(ierr)
 
  end subroutine comm_free
 
end module comm