domain_decomposition_2d.f90 Source File

!> @file domain_decomposition_2d.f90
!> @brief Pure 2D block decomposition of the Cartesian grid across an MPI process grid.
!! Neighbour ranks are supplied by the caller (from MPI_Cart_shift) so this stays pure
!! and unit-testable with synthetic coordinates.
module domain_decomposition_2d
  use precision, only: wp
  use mpi_runtime, only: parallel_fatal
  use domain_decomposition, only: rank_local_count, rank_first_global
  implicit none
  private
  public :: decomp_2d_t, decompose_2d, is_decomp_2d_feasible, validate_decomp_2d

  type :: decomp_2d_t
    integer :: coord_x = -1, coord_y = -1
    integer :: dim_x = -1, dim_y = -1
    integer :: nx_global = -1, ny_global = -1
    integer :: nx_local = -1, ny_local = -1
    integer :: ix_first_global = -1, ix_last_global = -1
    integer :: iy_first_global = -1, iy_last_global = -1
    integer :: halo_width = -1
    integer :: x_lo_neighbour = -1, x_hi_neighbour = -1
    integer :: y_lo_neighbour = -1, y_hi_neighbour = -1
    logical :: periodic_x = .false., periodic_y = .false.
  end type decomp_2d_t

contains

  pure function decompose_2d(coords, dims, nx_global, ny_global, halo_width, &
                             periodic_x, periodic_y, x_lo, x_hi, y_lo, y_hi) result(d)
    integer, intent(in) :: coords(2), dims(2), nx_global, ny_global, halo_width
    logical, intent(in) :: periodic_x, periodic_y
    integer, intent(in) :: x_lo, x_hi, y_lo, y_hi
    type(decomp_2d_t) :: d

    d % coord_x = coords(1); d % coord_y = coords(2)
    d % dim_x = dims(1); d % dim_y = dims(2)
    d % nx_global = nx_global; d % ny_global = ny_global
    d % halo_width = halo_width
    d % periodic_x = periodic_x; d % periodic_y = periodic_y
    d % x_lo_neighbour = x_lo; d % x_hi_neighbour = x_hi
    d % y_lo_neighbour = y_lo; d % y_hi_neighbour = y_hi

    call axis_split(coords(1), dims(1), nx_global, d % nx_local, d % ix_first_global, d % ix_last_global)
    call axis_split(coords(2), dims(2), ny_global, d % ny_local, d % iy_first_global, d % iy_last_global)
  end function decompose_2d

  !> Split one Cartesian axis (`n_global` cells across `dim` process slots) for
  !! the slot at `coord`. Delegates to the shared 1D slab-partition helpers so
  !! the per-axis block layout matches the 1D `decompose()` exactly (single
  !! source of truth: domain_decomposition%rank_local_count/rank_first_global).
  pure subroutine axis_split(coord, dim, n_global, n_local, i_first, i_last)
    integer, intent(in) :: coord, dim, n_global
    integer, intent(out) :: n_local, i_first, i_last
    n_local = rank_local_count(n_global, dim, coord)
    i_first = rank_first_global(n_global, dim, coord)
    i_last = i_first + n_local - 1
  end subroutine axis_split

  pure function is_decomp_2d_feasible(d) result(ok)
    type(decomp_2d_t), intent(in) :: d
    logical :: ok
    ok = d % nx_local >= d % halo_width .and. d % ny_local >= d % halo_width
  end function is_decomp_2d_feasible

  !> Abort with a clear message if either local tile dimension is too small to
  !! hold the halo cells the chosen reconstruction scheme requires. Mirrors the
  !! 1D `validate_decomp`; called from the 2D driver right after the
  !! decomposition is built (an over-fine MPI process grid, e.g. a thin axis
  !! from MPI_Dims_create, would otherwise silently over-read halos).
  subroutine validate_decomp_2d(d)
    type(decomp_2d_t), intent(in) :: d
    character(len=256) :: msg

    if (.not. is_decomp_2d_feasible(d)) then
      write (msg, '(A,I0,A,I0,A,I0,A,I0,A,I0,A,I0,A,I0,A)') &
        'decomp_2d: local tile (', d % nx_local, ' x ', d % ny_local, &
        ') < halo_width (', d % halo_width, &
        ') at cart coord (', d % coord_x, ',', d % coord_y, &
        ') of process grid ', d % dim_x, ' x ', d % dim_y, &
        '; use fewer ranks or a larger nx/ny'
      call parallel_fatal(trim(msg))
    end if
  end subroutine validate_decomp_2d

end module domain_decomposition_2d