implicit_gather_scatter.F90 Source File

!> @file implicit_gather_scatter.F90
!> @brief Gather / scatter helpers used by the Phase E implicit-stepper
!!        gather-solve-scatter wrapper.
!!
!! Hybrid cpp split mirrors src/parallel_reductions.F90 and
!! src/solution_gather.F90: MS-MPI ships only the legacy `use mpi`
!! module; OpenMPI ships both and we prefer `mpi_f08` for strong typing.

module implicit_gather_scatter
  use precision, only: wp
  use solver_state, only: neq
  use domain_decomposition, only: decomp_t, rank_local_count
#ifdef CFD_SOLVER_LEGACY_MPI
  use mpi, only: MPI_Gatherv, MPI_Scatterv, MPI_DOUBLE_PRECISION
#else
  use mpi_f08, only: MPI_Gatherv, MPI_Scatterv, MPI_DOUBLE_PRECISION
#endif
  use mpi_runtime, only: mpi_world, my_rank, n_ranks, parallel_fatal
  implicit none
  private

  public :: gather_field_to_root
  public :: scatter_field_from_root

contains

  !> Gather (neq, n_local) slices from every rank into (neq, n_global)
  !! on rank 0 via MPI_Gatherv. On non-root ranks `gather_buf` is not
  !! accessed; pass a `(neq, 0)` sentinel allocation.
  subroutine gather_field_to_root(local_field, n_local, decomp, gather_buf)
    real(wp),       intent(in)    :: local_field(:, :)
    integer,        intent(in)    :: n_local
    type(decomp_t), intent(in)    :: decomp
    real(wp),       intent(inout) :: gather_buf(:, :)

    integer, allocatable :: recv_counts(:), displs(:)
    integer :: send_count, ierr, info

    call build_counts_displs(decomp, recv_counts, displs)
    send_count = neq * n_local

    call MPI_Gatherv(local_field, send_count, MPI_DOUBLE_PRECISION, &
                     gather_buf, recv_counts, displs, MPI_DOUBLE_PRECISION, &
                     0, mpi_world(), ierr)
    if (ierr /= 0) call parallel_fatal('gather_field_to_root: MPI_Gatherv failed')

    deallocate (recv_counts, displs, stat=info)
    if (info /= 0) call parallel_fatal('gather_field_to_root: deallocation failed')
  end subroutine gather_field_to_root

  !> Scatter (neq, n_global) on rank 0 into per-rank (neq, n_local)
  !! slices via MPI_Scatterv. On non-root ranks `gather_buf` is not
  !! accessed; pass a `(neq, 0)` sentinel allocation.
  subroutine scatter_field_from_root(gather_buf, decomp, local_field, n_local)
    real(wp),       intent(in)    :: gather_buf(:, :)
    type(decomp_t), intent(in)    :: decomp
    real(wp),       intent(inout) :: local_field(:, :)
    integer,        intent(in)    :: n_local

    integer, allocatable :: send_counts(:), displs(:)
    integer :: recv_count, ierr, info

    call build_counts_displs(decomp, send_counts, displs)
    recv_count = neq * n_local

    call MPI_Scatterv(gather_buf, send_counts, displs, MPI_DOUBLE_PRECISION, &
                      local_field, recv_count, MPI_DOUBLE_PRECISION, &
                      0, mpi_world(), ierr)
    if (ierr /= 0) call parallel_fatal('scatter_field_from_root: MPI_Scatterv failed')

    deallocate (send_counts, displs, stat=info)
    if (info /= 0) call parallel_fatal('scatter_field_from_root: deallocation failed')
  end subroutine scatter_field_from_root

  ! Build per-rank counts and displacements from the same slab partition
  ! arithmetic decompose() uses. Both gather and scatter consume this same
  ! layout; non-root passes zero-size sentinel allocations so MPI gets
  ! valid pointer arguments.
  subroutine build_counts_displs(decomp, counts, displs)
    type(decomp_t),       intent(in)  :: decomp
    integer, allocatable, intent(out) :: counts(:), displs(:)

    integer :: rk, nr, info

    nr = n_ranks()

    if (my_rank() == 0) then
      allocate (counts(0:nr - 1), displs(0:nr - 1), stat=info)
      if (info /= 0) call parallel_fatal('build_counts_displs: allocation failed')
      do rk = 0, nr - 1
        counts(rk) = neq * rank_local_count(decomp % n_global, nr, rk)
      end do
      displs(0) = 0
      do rk = 1, nr - 1
        displs(rk) = displs(rk - 1) + counts(rk - 1)
      end do
    else
      ! Non-root must still pass valid pointers to MPI_Gatherv /
      ! MPI_Scatterv; allocate zero-length so the addresses are valid
      ! even though they are unused.
      allocate (counts(0:0), displs(0:0), stat=info)
      if (info /= 0) call parallel_fatal('build_counts_displs: allocation failed')
      counts = 0
      displs = 0
    end if
  end subroutine build_counts_displs

end module implicit_gather_scatter