gs_mpi.f90

Go to the documentation of this file.

! Copyright (c) 2020-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
  use num_types, only : rp
  use gs_comm, only : gs_comm_t, gs_comm_mpi, gs_comm_mpigpu, 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 mpi_f08, only : mpi_statuses_ignore, mpi_status, &
       mpi_request, mpi_isend, mpi_irecv, mpi_testsome, mpi_testall, &
       mpi_thread_multiple
  use comm, only : neko_comm, mpi_real_precision, neko_mpi_thread_provided
  use, intrinsic :: iso_c_binding
  use utils, only : neko_error
  implicit none
  private
 
  type, public, extends(gs_comm_t) :: gs_mpi_t
     real(kind=rp), allocatable :: send_buf(:)
     real(kind=rp), allocatable :: recv_buf(:)
     integer, allocatable :: send_len(:), recv_len(:)
     integer, allocatable :: send_offset(:), recv_offset(:)
     type(mpi_request), allocatable :: send_request(:), recv_request(:)
     integer, allocatable :: recv_indices(:)
     type(mpi_status), allocatable :: recv_statuses(:)
     integer :: ncompleted
     real(kind=rp), allocatable :: send_buf_v(:)
     real(kind=rp), allocatable :: recv_buf_v(:)
   contains
     procedure, pass(this) :: init => gs_mpi_init
     procedure, pass(this) :: free => gs_mpi_free
     procedure, pass(this) :: nbsend => gs_nbsend_mpi
     procedure, pass(this) :: nbrecv => gs_nbrecv_mpi
     procedure, pass(this) :: nbwait => gs_nbwait_mpi
     procedure, pass(this) :: nbsend_vec => gs_nbsend_vec_mpi
     procedure, pass(this) :: nbrecv_vec => gs_nbrecv_vec_mpi
     procedure, pass(this) :: nbwait_vec => gs_nbwait_vec_mpi
  end type gs_mpi_t
 
contains
 
  subroutine gs_mpi_init(this, send_pe, recv_pe)
    class(gs_mpi_t), intent(inout) :: this
    type(stack_i4_t), intent(inout) :: send_pe
    type(stack_i4_t), intent(inout) :: recv_pe
    integer :: i, nsend, nrecv, send_total, recv_total
 
    call this%init_order(send_pe, recv_pe)
 
    nsend = size(this%send_pe)
    nrecv = size(this%recv_pe)
 
    allocate(this%send_len(nsend), this%send_offset(nsend))
    allocate(this%send_request(nsend))
 
    allocate(this%recv_len(nrecv), this%recv_offset(nrecv))
    allocate(this%recv_request(nrecv))
    allocate(this%recv_indices(nrecv), this%recv_statuses(nrecv))
 
    send_total = 0
    do i = 1, nsend
       this%send_len(i) = this%send_dof(this%send_pe(i))%size()
       this%send_offset(i) = send_total
       send_total = send_total + this%send_len(i)
    end do
    allocate(this%send_buf(max(1, send_total)))
 
    recv_total = 0
    do i = 1, nrecv
       this%recv_len(i) = this%recv_dof(this%recv_pe(i))%size()
       this%recv_offset(i) = recv_total
       recv_total = recv_total + this%recv_len(i)
    end do
    allocate(this%recv_buf(max(1, recv_total)))
 
    ! Fused vector exchange buffers, sized for GS_VEC_NC components.
    allocate(this%send_buf_v(max(1, gs_vec_nc*send_total)))
    allocate(this%recv_buf_v(max(1, gs_vec_nc*recv_total)))
    this%vec_supported = .true.
 
  end subroutine gs_mpi_init
 
  subroutine gs_mpi_free(this)
    class(gs_mpi_t), intent(inout) :: this
 
    if (allocated(this%send_buf)) then
       deallocate(this%send_buf)
    end if
 
    if (allocated(this%recv_buf)) then
       deallocate(this%recv_buf)
    end if
 
    if (allocated(this%send_len)) then
       deallocate(this%send_len)
    end if
 
    if (allocated(this%recv_len)) then
       deallocate(this%recv_len)
    end if
 
    if (allocated(this%send_offset)) then
       deallocate(this%send_offset)
    end if
 
    if (allocated(this%recv_offset)) then
       deallocate(this%recv_offset)
    end if
 
    if (allocated(this%send_request)) then
       deallocate(this%send_request)
    end if
 
    if (allocated(this%recv_request)) then
       deallocate(this%recv_request)
    end if
 
    if (allocated(this%recv_indices)) then
       deallocate(this%recv_indices)
    end if
 
    if (allocated(this%recv_statuses)) then
       deallocate(this%recv_statuses)
    end if
 
    if (allocated(this%send_buf_v)) then
       deallocate(this%send_buf_v)
    end if
 
    if (allocated(this%recv_buf_v)) then
       deallocate(this%recv_buf_v)
    end if
 
    call this%free_order()
    call this%free_dofs()
 
  end subroutine gs_mpi_free
 
  subroutine gs_nbsend_mpi(this, u, n, tag, deps, strm)
    class(gs_mpi_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, ierr, dst, off, ndst
 
    ! Gather data from u into the per-peer slab of send_buf according
    ! to indices in send_dof. Slabs are contiguous so each MPI_Isend
    ! sends a single contiguous block.
 
    ! If the MPI library doesn't support MULTIPLE, use all threads to
    ! pack the send buffer. Otherwise, let each thread pack and send
    ! to a different neighbour
    if (neko_mpi_thread_provided .lt. mpi_thread_multiple) then
       do i = 1, size(this%send_pe)
          dst = this%send_pe(i)
          off = this%send_offset(i)
          ndst = this%send_len(i)
          select type (sp => this%send_dof(dst)%data)
          type is (integer)
             !$omp do
             do j = 1, ndst
                this%send_buf(off + j) = u(sp(j))
             end do
             !$omp end do
          end select
          !$omp master
          call mpi_isend(this%send_buf(off + 1), ndst, &
               mpi_real_precision, dst, tag, &
               neko_comm, this%send_request(i), ierr)
          !$omp end master
       end do
       !$omp barrier
    else
       !$omp do
       do i = 1, size(this%send_pe)
          dst = this%send_pe(i)
          off = this%send_offset(i)
          ndst = this%send_len(i)
          select type (sp => this%send_dof(dst)%data)
          type is (integer)
             !$omp simd
             do j = 1, ndst
                this%send_buf(off + j) = u(sp(j))
             end do
          end select
          call mpi_isend(this%send_buf(off + 1), ndst, &
               mpi_real_precision, dst, tag, &
               neko_comm, this%send_request(i), ierr)
       end do
       !$omp end do
    end if
 
  end subroutine gs_nbsend_mpi
 
  subroutine gs_nbrecv_mpi(this, tag)
    class(gs_mpi_t), intent(inout) :: this
    integer, intent(in) :: tag
    integer :: i, ierr, off, nsrc
 
    ! Issue recv requests, we will later check that these have finished
    ! in nbwait
 
 
    ! If the MPI library doesn't support MULTIPLE, the master thread
    ! will issue all Irecv's. Otherwise, threads will issue Irecv's
    ! concurrently.
    if (neko_mpi_thread_provided .lt. mpi_thread_multiple) then
       !$omp master
       do i = 1, size(this%recv_pe)
          off = this%recv_offset(i)
          nsrc = this%recv_len(i)
          call mpi_irecv(this%recv_buf(off + 1), nsrc, &
               mpi_real_precision, this%recv_pe(i), tag, &
               neko_comm, this%recv_request(i), ierr)
       end do
       !$omp end master
       !$omp barrier
    else
       !$omp do
       do i = 1, size(this%recv_pe)
          off = this%recv_offset(i)
          nsrc = this%recv_len(i)
          call mpi_irecv(this%recv_buf(off + 1), nsrc, &
               mpi_real_precision, this%recv_pe(i), tag, &
               neko_comm, this%recv_request(i), ierr)
       end do
       !$omp end do
    end if
  end subroutine gs_nbrecv_mpi
 
  subroutine gs_nbwait_mpi(this, u, n, op, strm)
    class(gs_mpi_t), intent(inout) :: this
    integer, intent(in) :: n
    real(kind=rp), dimension(n), intent(inout) :: u
    type(c_ptr), intent(inout) :: strm
    integer :: i, j, k, src, off, nsrc, ierr
    integer :: op
    integer :: nreqs
    logical :: sends_done
 
    ! Poll for any subset of the outstanding recv requests to complete,
    ! reduce each completed slab into u, and repeat until all are done.
    nreqs = size(this%recv_pe)
    do while (nreqs .gt. 0)
       !$omp master
       call mpi_testsome(size(this%recv_request), this%recv_request, &
            this%ncompleted, this%recv_indices, this%recv_statuses, ierr)
       !$omp end master
       !$omp barrier
       do k = 1, this%ncompleted
          i = this%recv_indices(k)
          src = this%recv_pe(i)
          off = this%recv_offset(i)
          nsrc = this%recv_len(i)
          select type (sp => this%recv_dof(src)%data)
          type is (integer)
             ! Do operation with data in buffer on dof specified by recv_dof
             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)) + this%recv_buf(off + 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)) * this%recv_buf(off + 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)), this%recv_buf(off + 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)), this%recv_buf(off + j))
                end do
                !$omp end do
             case default
                call neko_error("Unknown operation in gs_nbwait_mpi")
             end select
          end select
       end do
       nreqs = nreqs - this%ncompleted
       ! Synchronise before master can re-enter MPI_Testsome and overwrite
       ! this%ncompleted / this%recv_indices, which non-master threads are
       ! still reading in the unpack above.
       !$omp barrier
    end do
    !$omp master
    ! Finally, poll until all outstanding non-blocking sends have drained.
    if (size(this%send_request) .gt. 0) then
       sends_done = .false.
       do while (.not. sends_done)
          call mpi_testall(size(this%send_request), this%send_request, &
               sends_done, mpi_statuses_ignore, ierr)
       end do
    end if
    !$omp end master
    !$omp barrier
 
  end subroutine gs_nbwait_mpi
 
  subroutine gs_nbsend_vec_mpi(this, u, n, nc, tag, deps, strm)
    class(gs_mpi_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 :: i, j, c, ierr, dst, off, ndst
 
    if (neko_mpi_thread_provided .lt. mpi_thread_multiple) then
       do i = 1, size(this%send_pe)
          dst = this%send_pe(i)
          off = this%send_offset(i)
          ndst = this%send_len(i)
          select type (sp => this%send_dof(dst)%data)
          type is (integer)
             !$omp do
             do j = 1, ndst
                do c = 1, nc
                   this%send_buf_v(nc*off + (c-1)*ndst + j) = &
                        u((c-1)*n + sp(j))
                end do
             end do
             !$omp end do
          end select
          !$omp master
          call mpi_isend(this%send_buf_v(nc*off + 1), nc*ndst, &
               mpi_real_precision, dst, tag, &
               neko_comm, this%send_request(i), ierr)
          !$omp end master
       end do
       !$omp barrier
    else
       !$omp do
       do i = 1, size(this%send_pe)
          dst = this%send_pe(i)
          off = this%send_offset(i)
          ndst = this%send_len(i)
          select type (sp => this%send_dof(dst)%data)
          type is (integer)
             do c = 1, nc
                !$omp simd
                do j = 1, ndst
                   this%send_buf_v(nc*off + (c-1)*ndst + j) = &
                        u((c-1)*n + sp(j))
                end do
             end do
          end select
          call mpi_isend(this%send_buf_v(nc*off + 1), nc*ndst, &
               mpi_real_precision, dst, tag, &
               neko_comm, this%send_request(i), ierr)
       end do
       !$omp end do
    end if
 
  end subroutine gs_nbsend_vec_mpi
 
  subroutine gs_nbrecv_vec_mpi(this, tag, nc)
    class(gs_mpi_t), intent(inout) :: this
    integer, intent(in) :: tag, nc
    integer :: i, ierr, off, nsrc
 
    if (neko_mpi_thread_provided .lt. mpi_thread_multiple) then
       !$omp master
       do i = 1, size(this%recv_pe)
          off = this%recv_offset(i)
          nsrc = this%recv_len(i)
          call mpi_irecv(this%recv_buf_v(nc*off + 1), nc*nsrc, &
               mpi_real_precision, this%recv_pe(i), tag, &
               neko_comm, this%recv_request(i), ierr)
       end do
       !$omp end master
       !$omp barrier
    else
       !$omp do
       do i = 1, size(this%recv_pe)
          off = this%recv_offset(i)
          nsrc = this%recv_len(i)
          call mpi_irecv(this%recv_buf_v(nc*off + 1), nc*nsrc, &
               mpi_real_precision, this%recv_pe(i), tag, &
               neko_comm, this%recv_request(i), ierr)
       end do
       !$omp end do
    end if
  end subroutine gs_nbrecv_vec_mpi
 
  subroutine gs_nbwait_vec_mpi(this, u, n, nc, op, strm)
    class(gs_mpi_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 :: i, j, c, k, src, off, nsrc, ierr
    integer :: op
    integer :: nreqs
    logical :: sends_done
 
    nreqs = size(this%recv_pe)
    do while (nreqs .gt. 0)
       !$omp master
       call mpi_testsome(size(this%recv_request), this%recv_request, &
            this%ncompleted, this%recv_indices, this%recv_statuses, ierr)
       !$omp end master
       !$omp barrier
       do k = 1, this%ncompleted
          i = this%recv_indices(k)
          src = this%recv_pe(i)
          off = this%recv_offset(i)
          nsrc = this%recv_len(i)
          select type (sp => this%recv_dof(src)%data)
          type is (integer)
             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)) + &
                           this%recv_buf_v(nc*off + (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)) * &
                           this%recv_buf_v(nc*off + (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)), &
                           this%recv_buf_v(nc*off + (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)), &
                           this%recv_buf_v(nc*off + (c-1)*nsrc + j))
                   end do
                end do
                !$omp end do
             case default
                call neko_error("Unknown operation in gs_nbwait_vec_mpi")
             end select
          end select
       end do
       nreqs = nreqs - this%ncompleted
       !$omp barrier
    end do
    !$omp master
    if (size(this%send_request) .gt. 0) then
       sends_done = .false.
       do while (.not. sends_done)
          call mpi_testall(size(this%send_request), this%send_request, &
               sends_done, mpi_statuses_ignore, ierr)
       end do
    end if
    !$omp end master
    !$omp barrier
 
  end subroutine gs_nbwait_vec_mpi
 
end module gs_mpi