halo_exchange_2d Subroutine

   subroutine halo_exchange_2d(fld, decomp, nghost, nx_local, ny_local)
      !! Exchange ghost-cell halos for a single 2D field
      !!
      !! The field has dimensions (nx_local + 2*nghost, ny_local + 2*nghost).
      !! Physical cells occupy indices (nghost+1 : nghost+nx_local, nghost+1 : nghost+ny_local).
      !! This routine fills the nghost-wide ghost strips on each side by
      !! sending/receiving from neighbouring ranks.
      real(wp), intent(inout) :: fld(:, :)
         !! 2D field with ghost cells
      type(decomp_t), intent(in) :: decomp
         !! Domain decomposition descriptor
      integer, intent(in) :: nghost
         !! Ghost cell width
      integer, intent(in) :: nx_local
         !! Local physical cells in x
      integer, intent(in) :: ny_local
         !! Local physical cells in y

      type(comm_t) :: comm
      integer :: nx_total, ny_total
      integer :: i, j, k, idx
      integer :: rank_west, rank_east, rank_south, rank_north
      real(wp), allocatable :: send_west(:), send_east(:)
      real(wp), allocatable :: recv_west(:), recv_east(:)
      real(wp), allocatable :: send_south(:), send_north(:)
      real(wp), allocatable :: recv_south(:), recv_north(:)
      type(request_t) :: reqs(MAX_REQS)
      type(MPI_Status) :: stats(MAX_REQS)
      integer :: nreq
      integer :: strip_ew, strip_sn

      comm = comm_env_compute_comm()

      nx_total = nx_local + 2*nghost
      ny_total = ny_local + 2*nghost

      ! East-West strips: nghost columns x ny_total rows
      strip_ew = nghost*ny_total
      ! South-North strips: nx_total columns x nghost rows
      strip_sn = nx_total*nghost

      nreq = 0

      ! --- East/West exchange ---
      if (.not. decomp%has_west) then
         rank_west = decomp_rank_from_coords(decomp%px, decomp%rx - 1, decomp%ry)
         allocate (send_west(strip_ew), recv_west(strip_ew))

         ! Pack west send buffer: columns nghost+1 .. 2*nghost (first nghost physical columns)
         idx = 0
         do j = 1, ny_total
            do k = 1, nghost
               idx = idx + 1
               send_west(idx) = fld(nghost + k, j)
            end do
         end do

         nreq = nreq + 1
         call isend(comm, send_west, rank_west, 1, reqs(nreq))
         nreq = nreq + 1
         call irecv(comm, recv_west, rank_west, 2, reqs(nreq))
      end if

      if (.not. decomp%has_east) then
         rank_east = decomp_rank_from_coords(decomp%px, decomp%rx + 1, decomp%ry)
         allocate (send_east(strip_ew), recv_east(strip_ew))

         ! Pack east send buffer: columns nx_local+1 .. nx_local+nghost (last nghost physical cols)
         idx = 0
         do j = 1, ny_total
            do k = 1, nghost
               idx = idx + 1
               send_east(idx) = fld(nghost + nx_local - nghost + k, j)
            end do
         end do

         nreq = nreq + 1
         call isend(comm, send_east, rank_east, 2, reqs(nreq))
         nreq = nreq + 1
         call irecv(comm, recv_east, rank_east, 1, reqs(nreq))
      end if

      ! --- South/North exchange ---
      if (.not. decomp%has_south) then
         rank_south = decomp_rank_from_coords(decomp%px, decomp%rx, decomp%ry - 1)
         allocate (send_south(strip_sn), recv_south(strip_sn))

         ! Pack south send buffer: rows nghost+1 .. 2*nghost (first nghost physical rows)
         idx = 0
         do k = 1, nghost
            do i = 1, nx_total
               idx = idx + 1
               send_south(idx) = fld(i, nghost + k)
            end do
         end do

         nreq = nreq + 1
         call isend(comm, send_south, rank_south, 3, reqs(nreq))
         nreq = nreq + 1
         call irecv(comm, recv_south, rank_south, 4, reqs(nreq))
      end if

      if (.not. decomp%has_north) then
         rank_north = decomp_rank_from_coords(decomp%px, decomp%rx, decomp%ry + 1)
         allocate (send_north(strip_sn), recv_north(strip_sn))

         ! Pack north send buffer: rows ny_local+1 .. ny_local+nghost (last nghost physical rows)
         idx = 0
         do k = 1, nghost
            do i = 1, nx_total
               idx = idx + 1
               send_north(idx) = fld(i, nghost + ny_local - nghost + k)
            end do
         end do

         nreq = nreq + 1
         call isend(comm, send_north, rank_north, 4, reqs(nreq))
         nreq = nreq + 1
         call irecv(comm, recv_north, rank_north, 3, reqs(nreq))
      end if

      ! Wait for all
      if (nreq > 0) then
         call waitall(reqs(1:nreq), stats(1:nreq))
      end if

      ! --- Unpack received data into ghost cells ---
      if (.not. decomp%has_west) then
         idx = 0
         do j = 1, ny_total
            do k = 1, nghost
               idx = idx + 1
               fld(k, j) = recv_west(idx)
            end do
         end do
         deallocate (send_west, recv_west)
      end if

      if (.not. decomp%has_east) then
         idx = 0
         do j = 1, ny_total
            do k = 1, nghost
               idx = idx + 1
               fld(nghost + nx_local + k, j) = recv_east(idx)
            end do
         end do
         deallocate (send_east, recv_east)
      end if

      if (.not. decomp%has_south) then
         idx = 0
         do k = 1, nghost
            do i = 1, nx_total
               idx = idx + 1
               fld(i, k) = recv_south(idx)
            end do
         end do
         deallocate (send_south, recv_south)
      end if

      if (.not. decomp%has_north) then
         idx = 0
         do k = 1, nghost
            do i = 1, nx_total
               idx = idx + 1
               fld(i, nghost + ny_local + k) = recv_north(idx)
            end do
         end do
         deallocate (send_north, recv_north)
      end if

   end subroutine halo_exchange_2d