crystal_router.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 crystal_router
  use comm, only : pe_rank, pe_size, neko_comm
  use num_types, only : i8
  use stack, only : stack_i8_t
  use utils, only : neko_error
  use mpi_f08, only : mpi_sendrecv, mpi_status, mpi_integer, mpi_integer8, &
       mpi_proc_null
  implicit none
  private
 
  integer, parameter :: cr_tag = 0
 
  public :: crystal_router_transfer, crystal_router_pack
 
contains
 
  subroutine crystal_router_pack(out, dest, body)
    type(stack_i8_t), intent(inout) :: out
    integer, intent(in) :: dest
    integer(i8), intent(in) :: body(:)
    integer(i8) :: w
    integer :: i
 
    w = int(dest, i8)
    call out%push(w)
    w = int(size(body), i8)
    call out%push(w)
    do i = 1, size(body)
       w = body(i)
       call out%push(w)
    end do
  end subroutine crystal_router_pack
 
  subroutine crystal_router_transfer(buf, n)
    integer(i8), allocatable, intent(inout) :: buf(:)
    integer, intent(inout) :: n
    integer(i8), allocatable :: keep(:), snd(:), rcv(:), extra(:)
    integer :: nkeep, nsnd, nrcv, nextra
    integer :: lo, hi, m, half, mid, r, partner
    logical :: lower
 
    ! Validate destinations up front so a packing bug surfaces here rather
    ! than as a silently misrouted (lost) record at scale.
    call cr_check_dest(buf, n)
 
    lo = 0
    hi = pe_size
 
    do while (hi - lo > 1)
       m = hi - lo
       half = m / 2
       mid = lo + half
       lower = (pe_rank < mid)
 
       ! Split the local buffer: records staying in our half vs. those bound
       ! for the opposite half.
       call cr_partition(buf, n, mid, lower, keep, nkeep, snd, nsnd)
 
       ! Pick the partner in the opposite half. Lower ranks always have one;
       ! in an odd-sized range the last upper rank (r == half) is unpaired.
       if (lower) then
          partner = mid + (pe_rank - lo)
       else
          r = pe_rank - mid
          if (r .lt. half) then
             partner = lo + r
          else
             partner = mpi_proc_null
          end if
       end if
 
       ! Bidirectional exchange with the partner, then rebuild the local
       ! buffer as (kept records) ++ (received records).
       call cr_exchange(snd, nsnd, partner, rcv, nrcv, partner)
       call cr_concat(buf, n, keep, nkeep, rcv, nrcv)
 
       ! Odd-sized range: the unpaired upper rank (hi-1) hands its lower-half
       ! records to rank lo, which absorbs them with one extra receive.
       if (iand(m, 1) .eq. 1) then
          if ((.not. lower) .and. ((pe_rank - mid) .eq. half)) then
             call cr_exchange(snd, nsnd, lo, extra, nextra, mpi_proc_null)
          else if (pe_rank .eq. lo) then
             call cr_exchange(snd, 0, mpi_proc_null, extra, nextra, hi - 1)
             call cr_append(buf, n, extra, nextra)
          end if
       end if
 
       ! Narrow the active range to the half containing pe_rank.
       if (lower) then
          hi = mid
       else
          lo = mid
       end if
    end do
 
    ! Release the per-stage scratch buffers (buf is returned to the caller).
    if (allocated(keep)) deallocate(keep)
    if (allocated(snd)) deallocate(snd)
    if (allocated(rcv)) deallocate(rcv)
    if (allocated(extra)) deallocate(extra)
 
  end subroutine crystal_router_transfer
 
  subroutine cr_check_dest(buf, n)
    integer(i8), allocatable, intent(in) :: buf(:)
    integer, intent(in) :: n
    integer :: p, dest, rlen
 
    p = 1
    do while (p .le. n)
       dest = int(buf(p))
       rlen = int(buf(p + 1))
       if (dest .lt. 0 .or. dest .ge. pe_size) then
          call neko_error('crystal_router: record destination out of range')
       end if
       p = p + 2 + rlen
    end do
  end subroutine cr_check_dest
 
  subroutine cr_partition(buf, n, mid, lower, keep, nkeep, snd, nsnd)
    integer(i8), allocatable, intent(in) :: buf(:)
    integer, intent(in) :: n, mid
    logical, intent(in) :: lower
    integer(i8), allocatable, intent(out) :: keep(:), snd(:)
    integer, intent(out) :: nkeep, nsnd
    integer :: p, dest, rlen, reclen
    logical :: in_lower
 
    allocate(keep(max(n, 1)))
    allocate(snd(max(n, 1)))
    nkeep = 0
    nsnd = 0
 
    p = 1
    do while (p .le. n)
       dest = int(buf(p))
       rlen = int(buf(p + 1))
       reclen = 2 + rlen
       in_lower = dest .lt. mid
       if (in_lower .eqv. lower) then
          keep(nkeep + 1:nkeep + reclen) = buf(p:p + reclen - 1)
          nkeep = nkeep + reclen
       else
          snd(nsnd + 1:nsnd + reclen) = buf(p:p + reclen - 1)
          nsnd = nsnd + reclen
       end if
       p = p + reclen
    end do
  end subroutine cr_partition
 
  subroutine cr_exchange(sbuf, sn, dst, rbuf, rn, src)
    integer(i8), intent(in) :: sbuf(:)
    integer, intent(in) :: sn, dst, src
    integer(i8), allocatable, intent(out) :: rbuf(:)
    integer, intent(out) :: rn
    type(mpi_status) :: status
    integer :: ierr
 
    rn = 0
    call mpi_sendrecv(sn, 1, mpi_integer, dst, cr_tag, &
         rn, 1, mpi_integer, src, cr_tag, neko_comm, status, ierr)
 
    allocate(rbuf(max(rn, 1)))
    call mpi_sendrecv(sbuf, sn, mpi_integer8, dst, cr_tag, &
         rbuf, rn, mpi_integer8, src, cr_tag, neko_comm, status, ierr)
  end subroutine cr_exchange
 
  subroutine cr_concat(buf, n, a, na, b, nb)
    integer(i8), allocatable, intent(inout) :: buf(:)
    integer, intent(out) :: n
    integer(i8), allocatable, intent(in) :: a(:), b(:)
    integer, intent(in) :: na, nb
 
    if (allocated(buf)) deallocate(buf)
    allocate(buf(max(na + nb, 1)))
    if (na .gt. 0) buf(1:na) = a(1:na)
    if (nb .gt. 0) buf(na + 1:na + nb) = b(1:nb)
    n = na + nb
  end subroutine cr_concat
 
  subroutine cr_append(buf, n, e, ne)
    integer(i8), allocatable, intent(inout) :: buf(:)
    integer, intent(inout) :: n
    integer(i8), allocatable, intent(in) :: e(:)
    integer, intent(in) :: ne
    integer(i8), allocatable :: tmp(:)
 
    if (ne .le. 0) return
    allocate(tmp(n + ne))
    if (n .gt. 0) tmp(1:n) = buf(1:n)
    tmp(n + 1:n + ne) = e(1:ne)
    call move_alloc(tmp, buf)
    n = n + ne
  end subroutine cr_append
 
end module crystal_router