gs_mpi_rma.f90

Go to the documentation of this file.

! Copyright (c) 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 gs_mpi_rma
  use num_types, only : rp, i8
  use gs_comm, only : gs_comm_t, gs_vec_nc
  use gs_ops, only : gs_op_add, gs_op_max, gs_op_min, gs_op_mul
  use stack, only : stack_i4_t
  use comm, only : neko_comm, pe_rank, pe_size, mpi_real_precision
  use utils, only : neko_error
  use mpi_f08, only : mpi_win, mpi_info, mpi_win_allocate, mpi_win_free, &
       mpi_win_lock_all, mpi_win_unlock_all, mpi_win_flush, &
       mpi_win_flush_all, &
       mpi_win_sync, mpi_win_get_attr, mpi_put, mpi_accumulate, &
       mpi_alltoall, mpi_barrier, mpi_iprobe, mpi_status, &
       mpi_info_create, mpi_info_set, mpi_info_free, &
       mpi_integer, mpi_integer8, mpi_replace, mpi_mode_nocheck, &
       mpi_address_kind, mpi_win_model, mpi_win_unified, &
       mpi_any_source, mpi_any_tag
  use, intrinsic :: iso_c_binding, only : c_ptr, c_null_ptr, c_f_pointer, &
       c_associated, c_sizeof
  implicit none
  private
 
  logical, parameter, public :: gs_mpi_rma_avail = .true.
 
  integer, parameter :: gs_mpi_rma_poke = 64
 
  logical, save :: gs_mpi_rma_flush_all = .true.
  logical, save :: gs_mpi_rma_flush_bound = .false.
 
  type, public, extends(gs_comm_t) :: gs_mpi_rma_t
     real(kind=rp), allocatable :: send_buf(:)
     integer, allocatable :: send_ndofs(:)
     integer, allocatable :: send_offset(:)
     integer, allocatable :: send_rdisp(:)
 
     type(mpi_win) :: win_data
     type(c_ptr) :: recv_ptr = c_null_ptr
     integer, allocatable :: recv_ndofs(:)
     integer, allocatable :: recv_offset(:)
     integer :: recv_total = 0
 
     type(mpi_win) :: win_sig
     type(c_ptr) :: sig_ptr = c_null_ptr
 
     integer(kind=i8) :: iter = 0
 
     logical :: unified = .true.
 
     logical :: win_alloc = .false.
   contains
     procedure, pass(this) :: init => gs_mpi_rma_init
     procedure, pass(this) :: free => gs_mpi_rma_free
     procedure, pass(this) :: nbsend => gs_mpi_rma_nbsend
     procedure, pass(this) :: nbrecv => gs_mpi_rma_nbrecv
     procedure, pass(this) :: nbwait => gs_mpi_rma_nbwait
     procedure, pass(this) :: nbsend_vec => gs_mpi_rma_nbsend_vec
     procedure, pass(this) :: nbrecv_vec => gs_mpi_rma_nbrecv_vec
     procedure, pass(this) :: nbwait_vec => gs_mpi_rma_nbwait_vec
  end type gs_mpi_rma_t
 
contains
 
  subroutine gs_mpi_rma_init(this, send_pe, recv_pe)
    class(gs_mpi_rma_t), intent(inout) :: this
    type(stack_i4_t), intent(inout) :: send_pe
    type(stack_i4_t), intent(inout) :: recv_pe
    type(mpi_info) :: info
    integer(kind=MPI_ADDRESS_KIND) :: win_size, attr_val
    integer, allocatable :: local_disp(:), remote_disp(:)
    integer(kind=i8), pointer :: sig(:)
    integer :: i, n, send_total, ierr, disp_unit, env_len
    integer(kind=i8) :: i8_dummy
    real(kind=rp) :: rp_dummy
    character(len=255) :: env_val
    logical :: flag
 
    ! No self-free here: gs_schedule has just filled send_dof/recv_dof and
    ! gs_mpi_rma_free would deallocate them again. Like every other comm
    ! backend, init assumes a freshly allocated object.
    call this%init_order(send_pe, recv_pe)
 
    ! Bind the flush strategy once; every instance shares it, so that an
    ! A/B run measures one protocol rather than a mixture.
    if (.not. gs_mpi_rma_flush_bound) then
       call get_environment_variable("NEKO_GS_RMA_FLUSH_ALL", env_val, env_len)
       if (env_len .gt. 0) then
          env_len = min(env_len, len(env_val))
          gs_mpi_rma_flush_all = (env_val(1:env_len) .ne. '0')
       end if
       gs_mpi_rma_flush_bound = .true.
    end if
 
    ! Send side: exact local sizing, plain memory
    n = size(this%send_pe)
    allocate(this%send_ndofs(n), this%send_offset(n), this%send_rdisp(n))
    send_total = 0
    do i = 1, n
       this%send_ndofs(i) = this%send_dof(this%send_pe(i))%size()
       this%send_offset(i) = send_total
       send_total = send_total + this%send_ndofs(i)
    end do
    this%send_rdisp = -1
    allocate(this%send_buf(max(gs_vec_nc*send_total, 1)))
 
    ! Receive side: exact local sizing, exposed in a window
    n = size(this%recv_pe)
    allocate(this%recv_ndofs(n), this%recv_offset(n))
    this%recv_total = 0
    do i = 1, n
       this%recv_ndofs(i) = this%recv_dof(this%recv_pe(i))%size()
       this%recv_offset(i) = this%recv_total
       this%recv_total = this%recv_total + this%recv_ndofs(i)
    end do
 
    ! same_disp_unit lets the implementation skip the per-target lookup on
    ! every operation.
    !
    ! The two accumulate assertions matter more than they look. By default
    ! MPI must preserve rar/raw/war/waw ordering between accumulates, which
    ! can rule out a hardware-atomic path for the signals in favour of a
    ! general, ordered one -- and the signals are two of the four RMA calls
    ! this backend makes per peer per round. Neither assertion costs
    ! anything here: a round issues at most one accumulate to any given
    ! location and consecutive rounds are separated by a flush, so no two
    ! are ever outstanding against the same slot, and every accumulate in
    ! the protocol is MPI_REPLACE.
    call mpi_info_create(info, ierr)
    call mpi_info_set(info, 'same_disp_unit', 'true', ierr)
    call mpi_info_set(info, 'accumulate_ordering', 'none', ierr)
    call mpi_info_set(info, 'accumulate_ops', 'same_op', ierr)
 
    disp_unit = int(c_sizeof(rp_dummy))
    win_size = int(max(gs_vec_nc*this%recv_total, 1), mpi_address_kind) * &
         int(disp_unit, mpi_address_kind)
    call mpi_win_allocate(win_size, disp_unit, info, neko_comm, &
         this%recv_ptr, this%win_data, ierr)
    if (.not. c_associated(this%recv_ptr)) then
       call neko_error('MPI_Win_allocate failed for gs_mpi_rma data window')
    end if
 
    disp_unit = int(c_sizeof(i8_dummy))
    win_size = int(2*pe_size, mpi_address_kind) * &
         int(disp_unit, mpi_address_kind)
    call mpi_win_allocate(win_size, disp_unit, info, neko_comm, &
         this%sig_ptr, this%win_sig, ierr)
    if (.not. c_associated(this%sig_ptr)) then
       call neko_error('MPI_Win_allocate failed for gs_mpi_rma signal window')
    end if
 
    call mpi_info_free(info, ierr)
    this%win_alloc = .true.
 
    ! MPI_Win_allocate does not zero the memory, and the protocol reads the
    ! counters before anyone has written them (iter == 1 waits for an ack
    ! of >= 0).
    call c_f_pointer(this%sig_ptr, sig, [2*pe_size])
    sig = 0_i8
 
    ! Both windows come from the same call sequence, so one query answers
    ! for both.
    call mpi_win_get_attr(this%win_sig, mpi_win_model, attr_val, flag, ierr)
    this%unified = flag .and. (int(attr_val) .eq. mpi_win_unified)
 
    call mpi_win_lock_all(mpi_mode_nocheck, this%win_data, ierr)
    call mpi_win_lock_all(mpi_mode_nocheck, this%win_sig, ierr)
 
    ! Tell each peer where in our receive window its slab belongs. Alltoall
    ! rather than pairwise Isend/Irecv for the same reason as in gs_shmem:
    ! send_pe and recv_pe follow a shifted-modulo pattern that deadlocked
    ! pairwise exchanges at certain rank counts, and the payload is one int
    ! per rank.
    allocate(local_disp(0:pe_size - 1), remote_disp(0:pe_size - 1))
    local_disp = -1
    do i = 1, size(this%recv_pe)
       local_disp(this%recv_pe(i)) = this%recv_offset(i)
    end do
    call mpi_alltoall(local_disp, 1, mpi_integer, &
         remote_disp, 1, mpi_integer, neko_comm, ierr)
    do i = 1, size(this%send_pe)
       this%send_rdisp(i) = remote_disp(this%send_pe(i))
    end do
    deallocate(local_disp, remote_disp)
 
    this%iter = 0
    this%vec_supported = .true.
 
    ! The zeroing above is a local store into window memory; no rank may
    ! issue RMA until every rank has done it.
    call mpi_barrier(neko_comm, ierr)
 
  end subroutine gs_mpi_rma_init
 
  subroutine gs_mpi_rma_free(this)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer :: ierr
 
    if (this%win_alloc) then
       ! MPI_Win_free is collective and completes the epoch; barrier first
       ! so no rank is still spinning on, or putting into, a window that
       ! another rank is about to tear down.
       call mpi_barrier(neko_comm, ierr)
       call mpi_win_unlock_all(this%win_sig, ierr)
       call mpi_win_unlock_all(this%win_data, ierr)
       call mpi_win_free(this%win_sig, ierr)
       call mpi_win_free(this%win_data, ierr)
       this%win_alloc = .false.
    end if
    this%recv_ptr = c_null_ptr
    this%sig_ptr = c_null_ptr
    this%recv_total = 0
    this%iter = 0
 
    if (allocated(this%send_buf)) deallocate(this%send_buf)
    if (allocated(this%send_ndofs)) deallocate(this%send_ndofs)
    if (allocated(this%send_offset)) deallocate(this%send_offset)
    if (allocated(this%send_rdisp)) deallocate(this%send_rdisp)
    if (allocated(this%recv_ndofs)) deallocate(this%recv_ndofs)
    if (allocated(this%recv_offset)) deallocate(this%recv_offset)
 
    call this%free_order()
    call this%free_dofs()
 
  end subroutine gs_mpi_rma_free
 
  function gs_mpi_rma_load(s, slot) result(v)
    ! The whole array is the dummy rather than the one element that is read:
    ! a VOLATILE dummy fed from a pointer array has to be assumed-shape or a
    ! pointer array itself. No INTENT on it either, since gfortran rejects
    ! INTENT(IN) together with VOLATILE. Read only.
    integer(kind=i8), volatile :: s(:)
    integer, intent(in) :: slot
    integer(kind=i8) :: v
 
    v = s(slot)
 
  end function gs_mpi_rma_load
 
  subroutine gs_mpi_rma_wait_ge(this, slot, val)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer, intent(in) :: slot
    integer(kind=i8), intent(in) :: val
    integer(kind=i8), pointer :: sig(:)
    type(mpi_status) :: status
    integer :: spin, ierr
    logical :: flag
 
    call c_f_pointer(this%sig_ptr, sig, [2*pe_size])
 
    spin = 0
    do
       ! In the separate memory model the public window copy is only made
       ! visible to local loads by a sync; in the unified model this is a
       ! cheap memory barrier.
       if (.not. this%unified) call mpi_win_sync(this%win_sig, ierr)
 
       if (gs_mpi_rma_load(sig, slot) .ge. val) exit
 
       ! Poke MPI so implementations without asynchronous progress still
       ! advance incoming RMA. Not needed on hardware-driven one-sided
       ! components, hence only every GS_MPI_RMA_POKE iterations.
       spin = spin + 1
       if (mod(spin, gs_mpi_rma_poke) .eq. 0) then
          call mpi_iprobe(mpi_any_source, mpi_any_tag, neko_comm, flag, &
               status, ierr)
       end if
    end do
 
  end subroutine gs_mpi_rma_wait_ge
 
  subroutine gs_mpi_rma_nbsend(this, u, n, tag, deps, strm)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer, intent(in) :: n
    real(kind=rp), dimension(n), intent(inout) :: u
    integer, intent(in) :: tag
    type(c_ptr), intent(inout) :: deps
    type(c_ptr), intent(inout) :: strm
    integer(kind=MPI_ADDRESS_KIND) :: rdisp
    integer, pointer :: sp(:)
    integer :: i, j, dst, base, ndst, ierr
 
    ! Entered from inside the gs_op_vector OpenMP parallel region. Every MPI
    ! call is funnelled through the master: RMA under MPI_THREAD_MULTIPLE is
    ! the least tuned path in most implementations.
    !$omp master
    this%iter = this%iter + 1
    !$omp end master
 
    ! Pack and put one peer at a time.
    do i = 1, size(this%send_pe)
       dst = this%send_pe(i)
       base = this%send_offset(i)
       ndst = this%send_ndofs(i)
       sp => this%send_dof(dst)%array()
       !OCL NORECURRENCE, NOVREC, NOALIAS
       !DIR$ CONCURRENT
       !DIR$ IVDEP
       !GCC$ ivdep
       !NEC$ IVDEP
       !$omp do
       do j = 1, ndst
          this%send_buf(base + j) = u(sp(j))
       end do
       !$omp end do
 
       ! Wait for dst to have consumed our previous round before we
       ! overwrite its receive window. The ack slots start at 0 and dst
       ! writes iter after consuming, so round 1 passes immediately.
       !$omp master
       call gs_mpi_rma_wait_ge(this, pe_size + dst + 1, this%iter - 1_i8)
 
       rdisp = int(this%send_rdisp(i), mpi_address_kind)
       call mpi_put(this%send_buf(base + 1), ndst, mpi_real_precision, &
            dst, rdisp, ndst, mpi_real_precision, this%win_data, ierr)
       !$omp end master
    end do
 
    ! Announce each peer as soon as its own put is remotely complete. MPI
    ! orders accumulates only against the same location, so the flush is
    ! what stands in for a put-with-notify.
    !$omp master
    if (gs_mpi_rma_flush_all) call mpi_win_flush_all(this%win_data, ierr)
    rdisp = int(pe_rank, mpi_address_kind)
    do i = 1, size(this%send_pe)
       dst = this%send_pe(i)
       if (.not. gs_mpi_rma_flush_all) then
          call mpi_win_flush(dst, this%win_data, ierr)
       end if
       call mpi_accumulate(this%iter, 1, mpi_integer8, dst, rdisp, 1, &
            mpi_integer8, mpi_replace, this%win_sig, ierr)
    end do
    call mpi_win_flush_all(this%win_sig, ierr)
    !$omp end master
    !$omp barrier
 
  end subroutine gs_mpi_rma_nbsend
 
  subroutine gs_mpi_rma_nbrecv(this, tag)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer, intent(in) :: tag
 
  end subroutine gs_mpi_rma_nbrecv
 
  subroutine gs_mpi_rma_nbwait(this, u, n, op, strm)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer, intent(in) :: n
    real(kind=rp), dimension(n), intent(inout) :: u
    type(c_ptr), intent(inout) :: strm
    integer :: op
    integer(kind=MPI_ADDRESS_KIND) :: rdisp
    integer, pointer :: sp(:)
    real(kind=rp), pointer :: recv_data(:)
    integer :: i, j, src, base, nsrc, ierr
 
    call c_f_pointer(this%recv_ptr, recv_data, [max(this%recv_total, 1)])
 
    ! Serial over peers: a dof shared by three or more ranks appears in
    ! several recv_dof lists, so reducing two slabs concurrently would race
    ! on it. The parallelism is taken within each slab instead.
    do i = 1, size(this%recv_pe)
       src = this%recv_pe(i)
       base = this%recv_offset(i)
       nsrc = this%recv_ndofs(i)
 
       !$omp master
       call gs_mpi_rma_wait_ge(this, src + 1, this%iter)
       !$omp end master
       !$omp barrier
 
       sp => this%recv_dof(src)%array()
       select case (op)
       case (gs_op_add)
          !OCL NORECURRENCE, NOVREC, NOALIAS
          !DIR$ CONCURRENT
          !DIR$ IVDEP
          !GCC$ ivdep
          !NEC$ IVDEP
          !$omp do
          do j = 1, nsrc
             u(sp(j)) = u(sp(j)) + recv_data(base + j)
          end do
          !$omp end do
       case (gs_op_mul)
          !OCL NORECURRENCE, NOVREC, NOALIAS
          !DIR$ CONCURRENT
          !DIR$ IVDEP
          !GCC$ ivdep
          !NEC$ IVDEP
          !$omp do
          do j = 1, nsrc
             u(sp(j)) = u(sp(j)) * recv_data(base + j)
          end do
          !$omp end do
       case (gs_op_min)
          !OCL NORECURRENCE, NOVREC, NOALIAS
          !DIR$ CONCURRENT
          !DIR$ IVDEP
          !GCC$ ivdep
          !NEC$ IVDEP
          !$omp do
          do j = 1, nsrc
             u(sp(j)) = min(u(sp(j)), recv_data(base + j))
          end do
          !$omp end do
       case (gs_op_max)
          !OCL NORECURRENCE, NOVREC, NOALIAS
          !DIR$ CONCURRENT
          !DIR$ IVDEP
          !GCC$ ivdep
          !NEC$ IVDEP
          !$omp do
          do j = 1, nsrc
             u(sp(j)) = max(u(sp(j)), recv_data(base + j))
          end do
          !$omp end do
       case default
          call neko_error("Unknown operation in gs_mpi_rma_nbwait")
       end select
 
       ! The implicit barrier at end do guarantees the whole team is done
       ! with the slab before it is released back to src.
       !$omp master
       rdisp = int(pe_size + pe_rank, mpi_address_kind)
       call mpi_accumulate(this%iter, 1, mpi_integer8, src, rdisp, 1, &
            mpi_integer8, mpi_replace, this%win_sig, ierr)
       !$omp end master
    end do
 
    ! An accumulate is only guaranteed to be issued at a flush, so the acks
    ! have to be pushed out before we go back to computing.
    !$omp master
    call mpi_win_flush_all(this%win_sig, ierr)
    !$omp end master
    !$omp barrier
 
  end subroutine gs_mpi_rma_nbwait
 
  subroutine gs_mpi_rma_nbsend_vec(this, u, n, nc, tag, deps, strm)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer, intent(in) :: n, nc
    real(kind=rp), dimension(nc*n), intent(inout) :: u
    integer, intent(in) :: tag
    type(c_ptr), intent(inout) :: deps
    type(c_ptr), intent(inout) :: strm
    integer(kind=MPI_ADDRESS_KIND) :: rdisp
    integer, pointer :: sp(:)
    integer :: i, j, c, dst, base, ndst, ierr
 
    !$omp master
    this%iter = this%iter + 1
    !$omp end master
 
    ! Pack and put per peer, as in the scalar path; see gs_mpi_rma_nbsend.
    do i = 1, size(this%send_pe)
       dst = this%send_pe(i)
       base = this%send_offset(i)
       ndst = this%send_ndofs(i)
       sp => this%send_dof(dst)%array()
       !$omp do
       do j = 1, ndst
          do c = 1, nc
             this%send_buf(nc*base + (c-1)*ndst + j) = u((c-1)*n + sp(j))
          end do
       end do
       !$omp end do
 
       !$omp master
       call gs_mpi_rma_wait_ge(this, pe_size + dst + 1, this%iter - 1_i8)
 
       rdisp = int(nc*this%send_rdisp(i), mpi_address_kind)
       call mpi_put(this%send_buf(nc*base + 1), nc*ndst, &
            mpi_real_precision, dst, rdisp, nc*ndst, mpi_real_precision, &
            this%win_data, ierr)
       !$omp end master
    end do
 
    !$omp master
    if (gs_mpi_rma_flush_all) call mpi_win_flush_all(this%win_data, ierr)
    rdisp = int(pe_rank, mpi_address_kind)
    do i = 1, size(this%send_pe)
       dst = this%send_pe(i)
       if (.not. gs_mpi_rma_flush_all) then
          call mpi_win_flush(dst, this%win_data, ierr)
       end if
       call mpi_accumulate(this%iter, 1, mpi_integer8, dst, rdisp, 1, &
            mpi_integer8, mpi_replace, this%win_sig, ierr)
    end do
    call mpi_win_flush_all(this%win_sig, ierr)
    !$omp end master
    !$omp barrier
 
  end subroutine gs_mpi_rma_nbsend_vec
 
  subroutine gs_mpi_rma_nbrecv_vec(this, tag, nc)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer, intent(in) :: tag, nc
 
  end subroutine gs_mpi_rma_nbrecv_vec
 
  subroutine gs_mpi_rma_nbwait_vec(this, u, n, nc, op, strm)
    class(gs_mpi_rma_t), intent(inout) :: this
    integer, intent(in) :: n, nc
    real(kind=rp), dimension(nc*n), intent(inout) :: u
    type(c_ptr), intent(inout) :: strm
    integer :: op
    integer(kind=MPI_ADDRESS_KIND) :: rdisp
    integer, pointer :: sp(:)
    real(kind=rp), pointer :: recv_data(:)
    integer :: i, j, c, src, base, nsrc, ierr
 
    call c_f_pointer(this%recv_ptr, recv_data, [max(nc*this%recv_total, 1)])
 
    do i = 1, size(this%recv_pe)
       src = this%recv_pe(i)
       base = this%recv_offset(i)
       nsrc = this%recv_ndofs(i)
 
       !$omp master
       call gs_mpi_rma_wait_ge(this, src + 1, this%iter)
       !$omp end master
       !$omp barrier
 
       sp => this%recv_dof(src)%array()
       select case (op)
       case (gs_op_add)
          !$omp do
          do j = 1, nsrc
             do c = 1, nc
                u((c-1)*n + sp(j)) = u((c-1)*n + sp(j)) + &
                     recv_data(nc*base + (c-1)*nsrc + j)
             end do
          end do
          !$omp end do
       case (gs_op_mul)
          !$omp do
          do j = 1, nsrc
             do c = 1, nc
                u((c-1)*n + sp(j)) = u((c-1)*n + sp(j)) * &
                     recv_data(nc*base + (c-1)*nsrc + j)
             end do
          end do
          !$omp end do
       case (gs_op_min)
          !$omp do
          do j = 1, nsrc
             do c = 1, nc
                u((c-1)*n + sp(j)) = min(u((c-1)*n + sp(j)), &
                     recv_data(nc*base + (c-1)*nsrc + j))
             end do
          end do
          !$omp end do
       case (gs_op_max)
          !$omp do
          do j = 1, nsrc
             do c = 1, nc
                u((c-1)*n + sp(j)) = max(u((c-1)*n + sp(j)), &
                     recv_data(nc*base + (c-1)*nsrc + j))
             end do
          end do
          !$omp end do
       case default
          call neko_error("Unknown operation in gs_mpi_rma_nbwait_vec")
       end select
 
       !$omp master
       rdisp = int(pe_size + pe_rank, mpi_address_kind)
       call mpi_accumulate(this%iter, 1, mpi_integer8, src, rdisp, 1, &
            mpi_integer8, mpi_replace, this%win_sig, ierr)
       !$omp end master
    end do
 
    !$omp master
    call mpi_win_flush_all(this%win_sig, ierr)
    !$omp end master
    !$omp barrier
 
  end subroutine gs_mpi_rma_nbwait_vec
 
end module gs_mpi_rma