gs_shmem.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_shmem
  use num_types, only : rp, c_rp, i8
  use gs_comm, only : gs_comm_t
  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
  use mpi_f08, only : mpi_allreduce, mpi_alltoall, mpi_integer, mpi_max
  use utils, only : neko_error
#ifdef HAVE_OPENSHMEM
  use shmem, only : shmem_malloc, shmem_calloc, shmem_free, &
       shmem_putmem_signal_nbi, shmem_signal_wait_until, &
       shmem_uint64_atomic_set, shmem_uint64_wait_until, &
       shmem_barrier_all, shmem_signal_set, shmem_cmp_ge
#endif
  use, intrinsic :: iso_c_binding, only : c_ptr, c_null_ptr, c_loc, &
       c_f_pointer, c_associated, c_sizeof, c_size_t, c_int64_t
  implicit none
  private
 
  type :: gs_shmem_buf_t
     integer, allocatable :: ndofs(:)
     integer, allocatable :: offset(:)
     integer, allocatable :: remote_offset(:)
     integer :: total = 0
     integer :: max_total = 0
     type(c_ptr) :: buf_ptr = c_null_ptr
   contains
     procedure, pass(this) :: init => gs_shmem_buf_init
     procedure, pass(this) :: free => gs_shmem_buf_free
  end type gs_shmem_buf_t
 
  type, public, extends(gs_comm_t) :: gs_shmem_t
     type(gs_shmem_buf_t) :: send_buf
     type(gs_shmem_buf_t) :: recv_buf
     ! Symmetric data-arrival signals; data_signals[r] is set to `iter`
     ! by PE r via shmem_putmem_signal_nbi when it has put data into
     ! our recv buffer. Sized to pe_size on every PE.
     type(c_ptr) :: data_signals_ptr = c_null_ptr
     ! Symmetric ack signals; ack_signals[r] is set to `iter` by PE r
     ! via shmem_uint64_atomic_set after it has consumed our most
     ! recent put. Sized to pe_size on every PE.
     type(c_ptr) :: ack_signals_ptr = c_null_ptr
     ! Monotonically increasing iteration counter; sender writes this
     ! value into the receiver's data signal slot via SHMEM_SIGNAL_SET,
     ! and the receiver writes the same value into the sender's ack
     ! slot after consumption.
     integer(kind=i8) :: iter = 0
   contains
     procedure, pass(this) :: init => gs_shmem_init
     procedure, pass(this) :: free => gs_shmem_free
     procedure, pass(this) :: nbsend => gs_shmem_nbsend
     procedure, pass(this) :: nbrecv => gs_shmem_nbrecv
     procedure, pass(this) :: nbwait => gs_shmem_nbwait
  end type gs_shmem_t
 
contains
 
  subroutine gs_shmem_buf_init(this, pe_order, dof_stack)
    class(gs_shmem_buf_t), intent(inout) :: this
    integer, intent(in) :: pe_order(:)
    type(stack_i4_t), intent(inout) :: dof_stack(0:)
    integer :: i, ierr, n
    integer(c_size_t) :: sz
    real(c_rp) :: rp_dummy
#ifndef HAVE_OPENSHMEM
    call neko_error('Neko was not built with OpenSHMEM support')
#else
 
    n = size(pe_order)
 
    allocate(this%ndofs(n))
    allocate(this%offset(n))
    allocate(this%remote_offset(n))
 
    do i = 1, n
       this%remote_offset(i) = -1
    end do
 
    this%total = 0
    do i = 1, n
       this%ndofs(i) = dof_stack(pe_order(i))%size()
       this%offset(i) = this%total
       this%total = this%total + this%ndofs(i)
    end do
 
    ! OpenSHMEM symmetric memory must be allocated with the same size on
    ! every PE.
    call mpi_allreduce(this%total, this%max_total, 1, mpi_integer, mpi_max, &
         neko_comm, ierr)
 
    sz = c_sizeof(rp_dummy) * int(max(this%max_total, 1), c_size_t)
    this%buf_ptr = shmem_malloc(sz)
    if (.not. c_associated(this%buf_ptr)) then
       call neko_error('shmem_malloc failed for gs_shmem buffer')
    end if
#endif
  end subroutine gs_shmem_buf_init
 
  subroutine gs_shmem_buf_free(this)
    class(gs_shmem_buf_t), intent(inout) :: this
 
    if (allocated(this%ndofs)) deallocate(this%ndofs)
    if (allocated(this%offset)) deallocate(this%offset)
    if (allocated(this%remote_offset)) deallocate(this%remote_offset)
 
#ifdef HAVE_OPENSHMEM
    if (c_associated(this%buf_ptr)) then
       ! shmem_free is collective; gs_shmem_free issues a barrier before
       ! tearing down buffers.
       call shmem_free(this%buf_ptr)
    end if
#endif
    this%buf_ptr = c_null_ptr
    this%total = 0
    this%max_total = 0
 
  end subroutine gs_shmem_buf_free
 
  subroutine gs_shmem_init(this, send_pe, recv_pe)
    class(gs_shmem_t), intent(inout) :: this
    type(stack_i4_t), intent(inout) :: send_pe
    type(stack_i4_t), intent(inout) :: recv_pe
    integer :: i, ierr
    integer(c_size_t) :: i64_size
    integer(c_int64_t) :: i64_dummy
    integer, allocatable :: local_offsets(:), remote_offsets(:)
#ifndef HAVE_OPENSHMEM
    call neko_error('Neko was not built with OpenSHMEM support')
#else
 
    call this%free()
    call this%init_order(send_pe, recv_pe)
 
    call this%send_buf%init(this%send_pe, this%send_dof)
    call this%recv_buf%init(this%recv_pe, this%recv_dof)
 
    ! Allocate the per-rank symmetric signal arrays. Size pe_size on
    ! every PE -> no Allreduce needed and no slot handshake required;
    ! every PE writes/reads at offset = remote PE's rank.
    i64_size = c_sizeof(i64_dummy)
    this%data_signals_ptr = shmem_calloc(int(pe_size, c_size_t), i64_size)
    if (.not. c_associated(this%data_signals_ptr)) then
       call neko_error('shmem_calloc failed for gs_shmem data signals')
    end if
    this%ack_signals_ptr = shmem_calloc(int(pe_size, c_size_t), i64_size)
    if (.not. c_associated(this%ack_signals_ptr)) then
       call neko_error('shmem_calloc failed for gs_shmem ack signals')
    end if
 
    ! Offset exchange via Alltoall. send_pe and recv_pe are built in
    ! gs_schedule with a shifted-modulo neighbor pattern; pairwise
    ! Isend/Irecv keyed by neighbor rank turned out to deadlock at
    ! certain PE counts. Alltoall is a single collective call that
    ! cannot mismatch, and the payload is just pe_size ints per PE.
    allocate(local_offsets(0:pe_size - 1))
    allocate(remote_offsets(0:pe_size - 1))
    local_offsets = -1
    do i = 1, size(this%recv_pe)
       local_offsets(this%recv_pe(i)) = this%recv_buf%offset(i)
    end do
    call mpi_alltoall(local_offsets, 1, mpi_integer, &
         remote_offsets, 1, mpi_integer, neko_comm, ierr)
    do i = 1, size(this%send_pe)
       this%send_buf%remote_offset(i) = remote_offsets(this%send_pe(i))
    end do
    deallocate(local_offsets)
    deallocate(remote_offsets)
 
    this%iter = 0
 
    ! Ensure all PEs have completed symmetric allocation before any
    ! one-sided communication is issued.
    call shmem_barrier_all()
#endif
  end subroutine gs_shmem_init
 
  subroutine gs_shmem_free(this)
    class(gs_shmem_t), intent(inout) :: this
 
#ifdef HAVE_OPENSHMEM
    ! shmem_free is collective; synchronize first to ensure no in-flight
    ! puts target the buffers we are about to release.
    call shmem_barrier_all()
 
    if (c_associated(this%data_signals_ptr)) then
       call shmem_free(this%data_signals_ptr)
    end if
    if (c_associated(this%ack_signals_ptr)) then
       call shmem_free(this%ack_signals_ptr)
    end if
    this%data_signals_ptr = c_null_ptr
    this%ack_signals_ptr = c_null_ptr
    this%iter = 0
#endif
 
    call this%send_buf%free()
    call this%recv_buf%free()
 
    call this%free_order()
    call this%free_dofs()
 
  end subroutine gs_shmem_free
 
  subroutine gs_shmem_nbsend(this, u, n, tag, deps, strm)
    class(gs_shmem_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 :: i, j, dst, base
    integer(c_size_t) :: nbytes
    integer , pointer :: sp(:)
    real(kind=rp), pointer :: send_data(:), recv_data(:)
    integer(c_int64_t), pointer :: data_signals(:), ack_signals(:)
    real(c_rp) :: rp_dummy
#ifdef HAVE_OPENSHMEM
 
    ! Each gs op gets a fresh signal value so receivers can distinguish
    ! puts from one call vs. the next.
    this%iter = this%iter + 1
 
    call c_f_pointer(this%send_buf%buf_ptr, send_data, &
         [max(this%send_buf%max_total, 1)])
    call c_f_pointer(this%recv_buf%buf_ptr, recv_data, &
         [max(this%recv_buf%max_total, 1)])
    call c_f_pointer(this%data_signals_ptr, data_signals, [pe_size])
    call c_f_pointer(this%ack_signals_ptr, ack_signals, [pe_size])
 
    do i = 1, size(this%send_pe)
       dst = this%send_pe(i)
 
       ! Wait for dst to have consumed our previous round's data before
       ! we overwrite their recv buffer. ack_signals[dst] is calloc'd
       ! to 0 and dst writes the iter value (via atomic_set) after
       ! consumption, so for iter == 1 the wait succeeds immediately.
       call shmem_uint64_wait_until(c_loc(ack_signals(dst + 1)), &
            shmem_cmp_ge, this%iter - 1_8)
 
       sp => this%send_dof(dst)%array()
       base = this%send_buf%offset(i)
       do concurrent(j = 1:this%send_dof(dst)%size())
          send_data(base + j) = u(sp(j))
       end do
 
       nbytes = int(this%send_buf%ndofs(i), c_size_t) * c_sizeof(rp_dummy)
       ! Put data + signal dst's data_signals[my_rank].
       call shmem_putmem_signal_nbi( &
            c_loc(recv_data(this%send_buf%remote_offset(i) + 1)), &
            c_loc(send_data(base + 1)), &
            nbytes, &
            c_loc(data_signals(pe_rank + 1)), &
            this%iter, shmem_signal_set, dst)
    end do
#endif
  end subroutine gs_shmem_nbsend
 
  subroutine gs_shmem_nbrecv(this, tag)
    class(gs_shmem_t), intent(inout) :: this
    integer, intent(in) :: tag
 
  end subroutine gs_shmem_nbrecv
 
  subroutine gs_shmem_nbwait(this, u, n, op, strm)
    class(gs_shmem_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 :: i, j, src, base
    integer(c_int64_t) :: dummy
    integer , pointer :: sp(:)
    real(kind=rp), pointer :: recv_data(:)
    integer(c_int64_t), pointer :: data_signals(:), ack_signals(:)
#ifdef HAVE_OPENSHMEM
 
    call c_f_pointer(this%recv_buf%buf_ptr, recv_data, &
         [max(this%recv_buf%max_total, 1)])
    call c_f_pointer(this%data_signals_ptr, data_signals, [pe_size])
    call c_f_pointer(this%ack_signals_ptr, ack_signals, [pe_size])
 
    do i = 1, size(this%recv_pe)
       src = this%recv_pe(i)
 
       ! Wait for data from src; data_signals[src] is set by src's put.
       dummy = shmem_signal_wait_until(c_loc(data_signals(src + 1)), &
            shmem_cmp_ge, this%iter)
 
       sp => this%recv_dof(src)%array()
       base = this%recv_buf%offset(i)
       select case (op)
       case (gs_op_add)
          !NEC$ IVDEP
          do concurrent(j = 1:this%recv_dof(src)%size())
             u(sp(j)) = u(sp(j)) + recv_data(base + j)
          end do
       case (gs_op_mul)
          !NEC$ IVDEP
          do concurrent(j = 1:this%recv_dof(src)%size())
             u(sp(j)) = u(sp(j)) * recv_data(base + j)
          end do
       case (gs_op_min)
          !NEC$ IVDEP
          do concurrent(j = 1:this%recv_dof(src)%size())
             u(sp(j)) = min(u(sp(j)), recv_data(base + j))
          end do
       case (gs_op_max)
          !NEC$ IVDEP
          do concurrent(j = 1:this%recv_dof(src)%size())
             u(sp(j)) = max(u(sp(j)), recv_data(base + j))
          end do
       case default
          call neko_error("Unknown operation in gs_nbwait_shmem")
       end select
 
       ! Tell src we are done with this round's buffer so they may
       ! overwrite it on the next round. We set our own rank's slot
       ! in src's ack_signals; src waits there on the next nbsend.
       call shmem_uint64_atomic_set( &
            c_loc(ack_signals(pe_rank + 1)), &
            this%iter, src)
    end do
#endif
  end subroutine gs_shmem_nbwait
 
end module gs_shmem