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_multiple, mpi_thread_serialized, mpi_comm_rank, &
       mpi_comm_split, mpi_comm_dup, mpi_barrier, mpi_comm_free, mpi_finalize, &
       mpi_comm_world, mpi_double_precision, mpi_real
  use utils, only : neko_error
  use neko_config
  !$ use omp_lib
  implicit none
  private
 
  interface
     subroutine neko_comm_wrapper_init(fcomm) &
          bind(c, name='neko_comm_wrapper_init')
       integer, 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
 
  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
 
    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
 
    if (.not.initialized) then
       if (nthrds .gt. 1) then
          call mpi_init_thread(mpi_thread_multiple, provided, ierr)
          if (provided .ne. mpi_thread_multiple) then
             ! MPI_THREAD_MULTIPLE is required for mt. device backends
             if (neko_bcknd_device .eq. 1) then
                call neko_error('Invalid thread support provided by MPI')
             else
                call mpi_init_thread(mpi_thread_serialized, provided, ierr)
                if (provided .ne. mpi_thread_serialized) then
                   call neko_error('Invalid thread support provided by MPI')
                end if
             end if
          end if
       else
          call mpi_init(ierr)
       end if
    end if
 
#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
 
 
  end subroutine comm_init
 
  subroutine comm_free
    integer :: ierr
 
    call mpi_barrier(neko_comm, ierr)
    call mpi_comm_free(neko_comm, ierr)
 
#ifdef HAVE_NCCL
    call neko_comm_nccl_finalize()
#endif
 
#ifdef HAVE_NVSHMEM
    call neko_comm_nvshmem_finalize()
#endif
 
    call mpi_finalize(ierr)
 
  end subroutine comm_free
 
end module comm