gs_tune.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.
!
submodule(gather_scatter) gs_tune
  use gs_shmem, only : gs_shmem_avail
  use gs_caf, only : gs_caf_avail, gs_caf_usable, gs_caf_signal_auto, &
       gs_caf_signal_modes, gs_caf_mode_name, gs_caf_mode_get, &
       gs_caf_set_mode
  use gs_utofu, only : gs_utofu_avail
  use gs_mpi_rma, only : gs_mpi_rma_avail
  use gs_device_nccl, only : gs_device_nccl_avail
  use gs_device_shmem, only : gs_device_shmem_avail
  use neko_config, only : neko_device_mpi
  implicit none
 
  integer, parameter :: GS_TUNE_BCKND(11) = [gs_comm_mpi, gs_comm_neighbour, &
       gs_comm_mpirma, gs_comm_openshmem, gs_comm_caf, gs_comm_utofu, &
       gs_comm_crystal, gs_comm_mpigpu, gs_comm_nccl, gs_comm_crystalgpu, &
       gs_comm_nvshmem]
 
  logical, parameter :: GS_TUNE_DEFAULT(11) = [.true., .true., .true., &
       .true., .false., .true., .true., .true., .true., .true., .false.]
 
  logical, parameter :: GS_TUNE_DEV_AVAIL = (neko_bcknd_hip .eq. 1) .or. &
       (neko_bcknd_cuda .eq. 1)
 
  logical, save :: caf_signal_tuned = .false.
 
contains
 
  function gs_comm_cand() result(cand)
    integer, allocatable :: cand(:)
    character(len=LOG_SIZE) :: log_buf
    character(len=13) :: label
    logical :: sel(size(GS_TUNE_BCKND)), named(size(GS_TUNE_BCKND))
    integer :: c(size(GS_TUNE_BCKND)), i, n
 
    call gs_tune_select(sel, named)
 
    n = 0
    do i = 1, size(gs_tune_bcknd)
       if (.not. sel(i)) cycle
       if (gs_comm_tunable(gs_tune_bcknd(i))) then
          n = n + 1
          c(n) = gs_tune_bcknd(i)
       else if (named(i)) then
          ! Asked for by name, but this build or run cannot drive it: say so
          ! rather than leaving the candidate quietly missing from the
          ! comparison below
          label = adjustl(gs_comm_name(gs_tune_bcknd(i)))
          write(log_buf, '(A,A,A12)') label, ': ', 'unavailable'
          call neko_log%message(log_buf)
       end if
    end do
 
    allocate(cand(n))
    cand = c(1:n)
 
  end function gs_comm_cand
 
  function gs_comm_tunable(comm_bcknd) result(tunable)
    integer, intent(in) :: comm_bcknd
    logical :: tunable
 
    select case (comm_bcknd)
    case (gs_comm_mpi, gs_comm_neighbour)
       tunable = .true.
    case (gs_comm_crystal)
       ! Nothing beyond MPI-1, and it addresses its partners by their rank
       ! in NEKO_COMM, so a communicator split does not rule it out either.
       ! It aggregates the halo into one message per routing stage instead
       ! of one per peer, which is a win only where per-message overhead
       ! dominates; the benchmark is what decides that
       tunable = .true.
    case (gs_comm_mpirma)
       ! Nothing beyond MPI-3, and it addresses its peers by their rank in
       ! NEKO_COMM rather than by a global PE or image number, so unlike the
       ! other one-sided backends a communicator split (NEKO_COMM_ID) does
       ! not rule it out. What it does assume is that the implementation
       ! makes RMA progress without the target entering MPI; on one that
       ! does not (Open MPI osc/pt2pt) it is slow rather than wrong, so it
       ! loses the benchmark and is dropped, which is the outcome the
       ! autotuning exists to produce. Rule it out with NEKO_GS_TUNE=-MPIRMA
       ! if paying for that measurement is not worth it.
       tunable = gs_mpi_rma_avail
    case (gs_comm_openshmem)
       ! The one-sided backends that address their peers by global PE or
       ! image number (OpenSHMEM PEs, coarray images) are only correct when
       ! NEKO_COMM spans every process, so skip them when the run has been
       ! split into several communicators (NEKO_COMM_ID). uTofu exchanges
       ! its addresses over NEKO_COMM itself and is unaffected.
       tunable = gs_shmem_avail .and. (pe_size .eq. global_pe_size)
    case (gs_comm_caf)
       ! Coarray support at configure time says nothing about the job
       ! actually running more than one image, see gs_caf_usable
       tunable = gs_caf_avail .and. gs_caf_usable() .and. &
            (pe_size .eq. global_pe_size)
    case (gs_comm_utofu)
       tunable = gs_utofu_avail
    case (gs_comm_mpigpu, gs_comm_crystalgpu)
       ! Both hand device pointers straight to MPI, which only holds on a
       ! build configured for device-aware MPI (--enable-device-mpi). They
       ! address their peers by their rank in NEKO_COMM, so a communicator
       ! split does not rule them out.
       tunable = gs_tune_dev_avail .and. neko_device_mpi
    case (gs_comm_nccl)
       ! The NCCL communicator is built once at startup
       ! (neko_comm_nccl_init) over the ranks of NEKO_COMM, but from a
       ! unique id broadcast over MPI_COMM_WORLD, so it is only well defined
       ! when NEKO_COMM spans every process, as for the PE addressed
       ! backends above
       tunable = gs_tune_dev_avail .and. gs_device_nccl_avail .and. &
            (pe_size .eq. global_pe_size)
    case (gs_comm_nvshmem)
       ! Addresses its peers by NVSHMEM PE, as OpenSHMEM does, so the same
       ! restriction on communicator splits applies
       tunable = gs_tune_dev_avail .and. gs_device_shmem_avail .and. &
            (pe_size .eq. global_pe_size)
    case default
       tunable = .false.
    end select
 
  end function gs_comm_tunable
 
  subroutine gs_tune_select(sel, named)
    logical, intent(out) :: sel(:), named(:)
    character(len=255) :: env_val
    character(len=32) :: tok, name
    character(len=1) :: op
    integer :: env_len, i, j, k
    logical :: delta, first
 
    sel = gs_tune_default
    named = .false.
 
    call get_environment_variable("NEKO_GS_TUNE", env_val, env_len)
    if (env_len .eq. 0) return
    ! env_len is the length of the value, which may well exceed the buffer
    env_len = min(env_len, len(env_val))
 
    delta = .false.
    first = .true.
    i = 1
    do while (i .le. env_len)
       if (scan(env_val(i:i), ', ') .ne. 0) then
          i = i + 1
          cycle
       end if
 
       j = i
       do while (j .le. env_len)
          if (scan(env_val(j:j), ', ') .ne. 0) exit
          j = j + 1
       end do
       tok = gs_tune_upcase(env_val(i:j-1))
       i = j
 
       op = ' '
       name = tok
       if (tok(1:1) .eq. '+' .or. tok(1:1) .eq. '-') then
          op = tok(1:1)
          name = tok(2:)
       end if
 
       if (first) then
          ! The first name decides how the list is read; a plain one starts
          ! the set from nothing rather than from the default
          delta = (op .ne. ' ')
          if (.not. delta) sel = .false.
          first = .false.
       else if ((op .ne. ' ') .neqv. delta) then
          call neko_error('NEKO_GS_TUNE: plain backend names and +/- ' // &
               'prefixed ones cannot be mixed')
       end if
 
       k = gs_tune_index(name)
       if (k .eq. 0) then
          call neko_error('NEKO_GS_TUNE: not a tunable Gather-scatter ' // &
               'comm. backend: ' // trim(name))
       end if
       sel(k) = (op .ne. '-')
       named(k) = .true.
    end do
 
  end subroutine gs_tune_select
 
  function gs_tune_index(name) result(idx)
    character(len=*), intent(in) :: name
    integer :: idx, bcknd, i
 
    select case (trim(name))
    case ('MPI')
       bcknd = gs_comm_mpi
    case ('NEIGHBOUR', 'NEIGHBOR')
       bcknd = gs_comm_neighbour
    case ('SHMEM', 'OPENSHMEM')
       bcknd = gs_comm_openshmem
    case ('CAF')
       bcknd = gs_comm_caf
    case ('UTOFU')
       bcknd = gs_comm_utofu
    case ('MPIRMA', 'RMA')
       bcknd = gs_comm_mpirma
    case ('CRYSTAL')
       bcknd = gs_comm_crystal
    case ('MPIGPU')
       bcknd = gs_comm_mpigpu
    case ('NCCL', 'RCCL')
       bcknd = gs_comm_nccl
    case ('NVSHMEM')
       bcknd = gs_comm_nvshmem
    case ('CRYSTALGPU')
       bcknd = gs_comm_crystalgpu
    case default
       bcknd = 0
    end select
 
    idx = 0
    do i = 1, size(gs_tune_bcknd)
       if (gs_tune_bcknd(i) .eq. bcknd) idx = i
    end do
 
  end function gs_tune_index
 
  function gs_tune_upcase(str) result(upper)
    character(len=*), intent(in) :: str
    character(len=len(str)) :: upper
    integer :: i, c
 
    do i = 1, len(str)
       c = iachar(str(i:i))
       if (c .ge. iachar('a') .and. c .le. iachar('z')) then
          upper(i:i) = achar(c - (iachar('a') - iachar('A')))
       else
          upper(i:i) = str(i:i)
       end if
    end do
 
  end function gs_tune_upcase
 
  subroutine gs_comm_switch(gs, comm_bcknd)
    type(gs_t), intent(inout) :: gs
    integer, intent(in) :: comm_bcknd
    class(gs_comm_t), allocatable :: comm_new
 
    call gs_comm_alloc(comm_new, comm_bcknd)
    call comm_new%take_schedule(gs%comm)
    call gs%comm%free()
    deallocate(gs%comm)
    call move_alloc(comm_new, gs%comm)
    call gs%comm%init_schedule()
 
    gs%bcknd%shared_on_host = .not. gs_comm_on_device(comm_bcknd)
 
  end subroutine gs_comm_switch
 
  module function gs_time_ops(gs, u, n, op, ntrials) result(t)
    type(gs_t), intent(inout) :: gs
    integer, intent(in) :: n
    real(kind=rp), dimension(n), intent(inout) :: u
    integer, intent(in) :: op, ntrials
    real(kind=dp) :: t
    integer :: i
 
    do i = 1, gs_tune_nwarmup
       call gs_op_vector(gs, u, n, op)
    end do
 
    if (neko_bcknd_device .eq. 1) call device_sync
    call mpi_barrier(neko_comm)
 
    t = mpi_wtime()
    do i = 1, ntrials
       call gs_op_vector(gs, u, n, op)
    end do
    if (neko_bcknd_device .eq. 1) call device_sync
    t = (mpi_wtime() - t) / real(ntrials, dp)
 
  end function gs_time_ops
 
  module subroutine gs_tune_comm(gs, n, comm_bcknd)
    type(gs_t), intent(inout) :: gs
    integer, intent(in) :: n
    integer, intent(in) :: comm_bcknd
    character(len=LOG_SIZE) :: log_buf
    character(len=13) :: label
    character(len=6) :: strtgy_str
    integer, allocatable :: cand(:)
    real(kind=dp), allocatable :: cand_time(:)
    real(kind=rp), allocatable :: tmp(:)
    type(c_ptr) :: tmp_d
    integer :: i, best, nmin, cur, dev_strtgy, dev_strtgy_avg
    logical :: dev_strtgy_env
 
    allocate(cand, source = gs_comm_cand())
 
    ! Track what gs is running as the sweep proceeds. The candidate list
    ! need neither start with the backend the schedule was built with nor
    ! contain it at all, and switching is what costs, so every candidate is
    ! switched to exactly once
    cur = comm_bcknd
 
    ! A rank without any dofs skips the halo exchange in gs_op_vector, which
    ! would leave every other rank hanging in a collective based backend.
    ! Keep the current backend in that case. NEKO_GS_TUNE may also have left
    ! us with nothing to compare, in which case a lone candidate is simply
    ! switched to, as if it had been requested with NEKO_GS_COMM.
    nmin = n
    call mpi_allreduce(mpi_in_place, nmin, 1, mpi_integer, mpi_min, neko_comm)
    if (nmin .eq. 0 .or. size(cand) .lt. 2) then
       call neko_log%message('Comm tuning  :      skipped')
       if (nmin .gt. 0 .and. size(cand) .eq. 1) then
          if (cand(1) .ne. cur) call gs_comm_switch(gs, cand(1))
          cur = cand(1)
          ! Nothing to compare the backend against, but device MPI still has
          ! a strategy to pick
          if (cur .eq. gs_comm_mpigpu) call gs_tune_dev_strtgy(gs, n)
       end if
       call neko_log%message('Tuned comm   : ' // gs_comm_name(cur))
       return
    end if
 
    allocate(cand_time(size(cand)))
 
    tmp_d = c_null_ptr
    allocate(tmp(n))
    tmp = 1.0_rp
    if (neko_bcknd_device .eq. 1) then
       call device_map(tmp, tmp_d, n)
       call device_memcpy(tmp, tmp_d, n, host_to_device, sync = .false.)
    end if
 
    ! GS_OP_MIN leaves the working vector untouched, so the benchmark can
    ! run for any number of trials without the values growing out of range
    dev_strtgy = -1
    dev_strtgy_avg = -1
    dev_strtgy_env = .false.
    do i = 1, size(cand)
       if (cand(i) .eq. gs_comm_caf .and. gs_caf_signal_auto() .and. &
            .not. caf_signal_tuned) then
          ! Benchmark the signaling modes as well; this leaves the CAF
          ! backend allocated in the winning mode
          cand_time(i) = gs_tune_caf_signal(gs, tmp, n)
          caf_signal_tuned = .true.
       else
          if (cand(i) .ne. cur) call gs_comm_switch(gs, cand(i))
          if (cand(i) .eq. gs_comm_mpigpu) then
             ! Benchmark the synchronisation strategies as well; this leaves
             ! the fastest one bound, and it has to be remembered in case the
             ! sweep switches away from device MPI and back again
             call gs_tune_strtgy(gs, tmp, n, dev_strtgy_avg, dev_strtgy_env, &
                  cand_time(i))
             dev_strtgy = gs_get_strtgy(gs)
          else
             cand_time(i) = gs_time_ops(gs, tmp, n, gs_op_min, &
                  gs_tune_ntrials)
          end if
       end if
       cur = cand(i)
    end do
 
    if (neko_bcknd_device .eq. 1) call device_unmap(tmp, tmp_d)
    deallocate(tmp)
 
    call mpi_allreduce(mpi_in_place, cand_time, size(cand), &
         mpi_double_precision, mpi_sum, neko_comm)
    cand_time = cand_time / pe_size
 
    best = minloc(cand_time, 1)
 
    do i = 1, size(cand)
       label = adjustl(gs_comm_name(cand(i)))
       ! What the coarray backend costs depends on the signaling mode in
       ! force, so report the timing under the mode it was measured in
       if (cand(i) .eq. gs_comm_caf .and. gs_caf_mode_get() .ne. 0) then
          label = 'CAF (' // &
               trim(adjustl(gs_caf_mode_name(gs_caf_mode_get()))) // ')'
       end if
       ! Device MPI is likewise only as fast as the synchronisation strategy
       ! it ran under, and that strategy is what the run goes on to use, so
       ! name it on the line it was measured on rather than above the
       ! comparison. 'Dev. MPI [xx]' is exactly the width of the label
       ! field, so the column is unaffected
       strtgy_str = ''
       if (cand(i) .eq. gs_comm_mpigpu .and. dev_strtgy_avg .ge. 0) then
          write(label, '(A,B0.2,A)') 'Dev. MPI [', dev_strtgy_avg, ']'
          ! A strategy that was named rather than measured is not a property
          ! of this comparison, so it is flagged instead of quietly shown
          if (dev_strtgy_env) strtgy_str = ' (env)'
       end if
       ! ES10.3 + ' s' fills the same 12 columns as the right-adjusted
       ! backend names, so the unit lines up with their last letter
       write(log_buf, '(A,A,ES10.3,A,A)') label, ': ', cand_time(i), ' s', &
            trim(strtgy_str)
       call neko_log%message(log_buf)
    end do
 
    if (cand(best) .ne. cur) call gs_comm_switch(gs, cand(best))
 
    ! A switch back to device MPI allocates a fresh backend, which comes up
    ! with the default strategy rather than the one benchmarked above
    if (dev_strtgy .ge. 0) call gs_set_strtgy(gs, dev_strtgy)
 
    call neko_log%message('Tuned comm   : ' // gs_comm_name(cand(best)))
 
    deallocate(cand_time)
 
  end subroutine gs_tune_comm
 
  module subroutine gs_tune_dev_strtgy(gs, n)
    type(gs_t), intent(inout) :: gs
    integer, intent(in) :: n
    character(len=LOG_SIZE) :: log_buf
    real(kind=rp), allocatable :: tmp(:)
    type(c_ptr) :: tmp_d
    integer :: strtgy_avg
    logical :: from_env
 
    tmp_d = c_null_ptr
    allocate(tmp(n))
    tmp = 1.0_rp
    call device_map(tmp, tmp_d, n)
    call device_memcpy(tmp, tmp_d, n, host_to_device, sync = .false.)
 
    call gs_tune_strtgy(gs, tmp, n, strtgy_avg, from_env)
 
    call device_unmap(tmp, tmp_d)
    deallocate(tmp)
 
    ! Nothing else was tuned, so the strategy gets the line to itself
    if (from_env) then
       write(log_buf, '(A,B0.2,A)') 'Env. strtgy  :         [', strtgy_avg, ']'
    else
       write(log_buf, '(A,B0.2,A)') 'Avg. strtgy  :         [', strtgy_avg, ']'
    end if
    call neko_log%message(log_buf)
 
  end subroutine gs_tune_dev_strtgy
 
  subroutine gs_tune_strtgy(gs, u, n, strtgy_avg, from_env, t)
    type(gs_t), intent(inout) :: gs
    integer, intent(in) :: n
    real(kind=rp), dimension(n), intent(inout) :: u
    integer, intent(out) :: strtgy_avg
    logical, intent(out) :: from_env
    real(kind=dp), intent(out), optional :: t
    integer, parameter :: strtgy(4) = [int(b'00'), int(b'01'), int(b'10'), &
         int(b'11')]
    real(kind=dp) :: strtgy_time(size(strtgy))
    character(len=255) :: env_strtgy
    integer :: i, env_len, best, avg
 
    call get_environment_variable("NEKO_GS_STRTGY", env_strtgy, env_len)
    from_env = (env_len .gt. 0)
 
    if (.not. from_env) then
       do i = 1, size(strtgy)
          call gs_set_strtgy(gs, strtgy(i))
          strtgy_time(i) = gs_time_ops(gs, u, n, gs_op_min, gs_tune_ntrials)
       end do
 
       best = minloc(strtgy_time, 1)
       call gs_set_strtgy(gs, strtgy(best))
       if (present(t)) t = strtgy_time(best)
 
       avg = best
       call mpi_allreduce(mpi_in_place, avg, 1, mpi_integer, mpi_sum, &
            neko_comm)
       avg = avg / pe_size
    else
       read(env_strtgy(1:env_len), *) best
 
       if (best .lt. 1 .or. best .gt. size(strtgy)) then
          call neko_error('Invalid gs sync strtgy')
       end if
 
       call gs_set_strtgy(gs, strtgy(best))
       if (present(t)) then
          t = gs_time_ops(gs, u, n, gs_op_min, gs_tune_ntrials)
       end if
 
       ! Every rank was given the same one, so there is nothing to average
       avg = best
    end if
 
    strtgy_avg = strtgy(avg)
 
  end subroutine gs_tune_strtgy
 
  subroutine gs_set_strtgy(gs, strtgy)
    type(gs_t), intent(inout) :: gs
    integer, intent(in) :: strtgy
 
    select type (c => gs%comm)
    type is (gs_device_mpi_t)
       c%nb_strtgy = strtgy
    end select
 
  end subroutine gs_set_strtgy
 
  function gs_get_strtgy(gs) result(strtgy)
    type(gs_t), intent(in) :: gs
    integer :: strtgy
 
    strtgy = -1
 
    select type (c => gs%comm)
    type is (gs_device_mpi_t)
       strtgy = c%nb_strtgy
    end select
 
  end function gs_get_strtgy
 
  function gs_tune_caf_signal(gs, u, n) result(t)
    type(gs_t), intent(inout) :: gs
    integer, intent(in) :: n
    real(kind=rp), dimension(n), intent(inout) :: u
    real(kind=dp) :: t
    character(len=LOG_SIZE) :: log_buf
    character(len=13) :: label
    integer, allocatable :: mode(:)
    real(kind=dp), allocatable :: mode_time(:)
    integer :: i, best
 
    allocate(mode, source = gs_caf_signal_modes())
    allocate(mode_time(size(mode)))
 
    do i = 1, size(mode)
       ! Bind the mode before building the backend: gs_caf_init sets up its
       ! per-instance signaling state according to the mode in force
       call gs_caf_set_mode(mode(i))
       call gs_comm_switch(gs, gs_comm_caf)
       mode_time(i) = gs_time_ops(gs, u, n, gs_op_min, gs_tune_ntrials)
    end do
 
    call mpi_allreduce(mpi_in_place, mode_time, size(mode), &
         mpi_double_precision, mpi_sum, neko_comm)
    mode_time = mode_time / pe_size
 
    best = minloc(mode_time, 1)
 
    do i = 1, size(mode)
       label = 'CAF ' // adjustl(gs_caf_mode_name(mode(i)))
       write(log_buf, '(A,A,ES10.3,A)') label, ': ', mode_time(i), ' s'
       call neko_log%message(log_buf)
    end do
 
    ! The last mode is the one currently allocated
    if (best .ne. size(mode)) then
       call gs_caf_set_mode(mode(best))
       call gs_comm_switch(gs, gs_comm_caf)
    end if
 
    call neko_log%message('Tuned CAF sig: ' // gs_caf_mode_name(mode(best)))
 
    t = mode_time(best)
 
    deallocate(mode, mode_time)
 
  end function gs_tune_caf_signal
 
end submodule gs_tune