!! Distributed tripolar north-fold exchange (`px > 1`) — MPI backend.
!!
!! ROUTING
!! =======
!! The routing is the pure plan of `rdb_ocean_fold_plan` (read its header
!! first): every north-ghost / fold-line value comes from the tile that OWNS
!! the mirror point, never from a halo copy, so the fold needs no preceding
!! x exchange and is correct even where no halo precedes it.  Only the north
!! rank row takes part (`decomp%has_north` and a folded north edge); every
!! rank of that row reaches the same call sites, so each call is collective
!! over exactly that row.  Other ranks never enter.
!!
!! MESSAGES
!! ========
!! One message per ordered (sender, receiver) pair per call; the self pair
!! is a local copy through the same buffers (no MPI).  A call may carry a
!! GROUP of fields: `ocean_fold_begin` → `ocean_fold_pack` per field →
!! `ocean_fold_exchange` → `ocean_fold_unpack` per field, in the SAME
!! order → `ocean_fold_end`.  Per-peer layout (field-major, then layer,
!! then row, then entry — the halo's layer-outer convention):
!!
!!   idx = (p-1)*cap + offT*n(p,T) + offU*n(p,U)
!!         + ((L-1)*nrow + (r-1))*n(p,fam) + e
!!
!! where `offT`/`offU` count the (row x layer) slabs already packed per
!! column family.  Both ends advance the same cursors over the same field
!! order, and the receiver's `n(p,·)` (its receive counts) equal the
!! sender's send counts for that pair, so no lengths are exchanged.
!!
!! SIGN AND THE FOLD-LINE ROW
!! ==========================
!! The sign is applied at UNPACK (`dst = -buf` for a true vector), so one
!! pack kernel serves every field; IEEE negation is exact, so where it is
!! applied cannot change bits.  For v and corners the fold-line row is
!! always sent (message row 1); the receiver writes it only on its WEST-half
!! columns, zeroes a self-conjugate column for a vector and leaves it for a
!! scalar, and never touches the EAST (authoritative) half — the serial
!! projection rule, as a COPY: one value moves, exactly, so no
!! floating-point expression has to agree across ranks.  If the projection
!! is ever changed to an AVERAGE, the operand order must be canonical (east
!! operand first) on every rank.
!!
!! No source of the fold is ever a destination (sources: rows below the
!! fold line and the east half of the fold row), so pack-all → MPI →
!! unpack-all is hazard-free in place.
!!
!! DEVICE vs HOST
!! ==============
!! Mirrors `rdb_ocean_halo` exactly: the `device_resident` switch selects
!! `!$acc parallel loop` pack/unpack over device-present arrays (default)
!! or plain host loops (setup-time host fills), and the MPI calls pass the
!! buffers through `!$acc host_data use_device` — like the halo, the device
!! path always hands MPI DEVICE addresses (`RDB_CUDA_AWARE_MPI` only selects
!! a log line today; a host-staged branch would be added to both modules
!! together).  The index lists are copied to the device once at init.
!!
!! TAG
!! ===
!!   TAG_OC_FOLD = 15  (halo uses 11-14, `rdb_halo` 1-4).  One message per
!!   ordered pair per call and every call completes with `waitall` before
!!   returning, so (source, tag) is unambiguous, px = 2 included.
#ifdef RDB_DOUBLE_PRECISION
#define FOLD_ISEND_N comm_isend_real_dp_array_n
#define FOLD_IRECV_N comm_irecv_real_dp_array_n
#else
#define FOLD_ISEND_N comm_isend_real_sp_array_n
#define FOLD_IRECV_N comm_irecv_real_sp_array_n
#endif
module rdb_ocean_fold_exchange
   !! Owner-routed north-fold exchange for an east-west split north rank
   !! row, plus the px-dispatching single-field fold entry points
   !! (`ocean_fold_north_*`): `px = 1` keeps the local kernels of
   !! `rdb_ocean_fold` verbatim, `px > 1` always uses the exchange (the
   !! self-mirror tile through a local copy).
   use rdb_constants, only: wp
   use rdb_decomp, only: decomp_t, decomp_rank_from_coords
   use pic_mpi_lib, only: comm_t, request_t, MPI_Status, waitall, &
                          FOLD_ISEND_N, FOLD_IRECV_N
   use rdb_comm_env, only: comm_env_compute_comm
   use pic_logger, only: logger => global_logger
   use pic_strings, only: to_string
   use rdb_ocean_status, only: OCEAN_STATUS_OK, OCEAN_STATUS_ERR_SETUP
   use rdb_error_ring, only: fail
   use rdb_ocean_fold, only: fold_north_centre, fold_north_u_face, &
                             fold_north_v_face, fold_north_corner
   use rdb_ocean_fold_plan, only: fold_plan_t, fold_plan_build, fold_stagger_family, &
                                  fold_stagger_nrows, FOLD_PLAN_OK, FOLD_FAM_T, &
                                  FOLD_FAM_U, FOLD_NFAM, FOLD_ROW_WEST, FOLD_ROW_SELF, &
                                  FOLD_STAG_T, FOLD_STAG_U, FOLD_STAG_V, FOLD_STAG_CORNER
   implicit none
   private

   public :: ocean_fold_exchange_init
   public :: ocean_fold_exchange_destroy
   public :: ocean_fold_exchange_reserve
   public :: ocean_fold_is_distributed
   public :: ocean_fold_begin
   public :: ocean_fold_pack
   public :: ocean_fold_exchange
   public :: ocean_fold_unpack
   public :: ocean_fold_end
   public :: ocean_fold_north_centre
   public :: ocean_fold_north_u_face
   public :: ocean_fold_north_v_face
   public :: ocean_fold_north_corner
   public :: FOLD_STAG_T, FOLD_STAG_U, FOLD_STAG_V, FOLD_STAG_CORNER

   ! Rank-generic group primitives: 2D or 3D by the rank of fld.
   interface ocean_fold_pack
      module procedure ocean_fold_pack_2d, ocean_fold_pack_3d
   end interface ocean_fold_pack

   interface ocean_fold_unpack
      module procedure ocean_fold_unpack_2d, ocean_fold_unpack_3d
   end interface ocean_fold_unpack

   ! px-dispatching single-field folds (local kernel at px = 1, exchange
   ! at px > 1).  Same stagger semantics as the `rdb_ocean_fold` kernels.
   interface ocean_fold_north_centre
      module procedure fold_centre_2d, fold_centre_3d
   end interface ocean_fold_north_centre

   interface ocean_fold_north_u_face
      module procedure fold_u_2d, fold_u_3d
   end interface ocean_fold_north_u_face

   interface ocean_fold_north_v_face
      module procedure fold_v_2d, fold_v_3d
   end interface ocean_fold_north_v_face

   interface ocean_fold_north_corner
      module procedure fold_corner_2d
   end interface ocean_fold_north_corner

   integer, parameter :: TAG_OC_FOLD = 15
      !! Fold message tag (halo: 11-14, `rdb_halo`: 1-4).

   ! -----------------------------------------------------------------
   ! Topology
   ! -----------------------------------------------------------------
   logical :: fx_initialised = .false.
      !! True after `ocean_fold_exchange_init`.
   logical :: fx_active = .false.
      !! This rank takes part in a distributed fold (px > 1, folded north
      !! edge, north rank row).
   integer :: fx_ng = 0
      !! Ghost width.
   integer :: fx_nxl = 0
      !! Tile physical width.
   integer :: fx_nyl = 0
      !! Tile physical height (identical on every north-row tile).
   integer :: fx_px = 1
      !! Tiles along x.
   integer :: fx_ry = 0
      !! This rank's y coordinate (the north row).

   ! -----------------------------------------------------------------
   ! Plan, flattened for the kernels (device-resident when active)
   ! -----------------------------------------------------------------
   type(fold_plan_t) :: fx_plan
      !! Host copy of the routing plan.
   integer :: fx_npeer = 0
      !! Peers (self included when it routes to itself).
   integer :: fx_self = 0
      !! Peer index of this tile (0 if not a peer).
   integer :: fx_nmax = 1
      !! Largest (peer, family) entry count.
   integer, allocatable :: fx_peer_rank(:)
      !! Compute-comm rank of each peer.
   integer, allocatable :: fx_sn(:, :)
      !! Send counts (npeer, FOLD_NFAM).
   integer, allocatable :: fx_rn(:, :)
      !! Receive counts (npeer, FOLD_NFAM).
   integer :: fx_s0(FOLD_NFAM) = 0
      !! Start of each family's block in the flat send arrays.
   integer :: fx_r0(FOLD_NFAM) = 0
      !! Start of each family's block in the flat receive arrays.
   integer :: fx_ns(FOLD_NFAM) = 0
      !! Send entries per family.
   integer :: fx_nr(FOLD_NFAM) = 0
      !! Receive entries per family.
   integer, allocatable :: fx_scol(:), fx_speer(:), fx_se(:)
      !! Send entries: source column, peer index, position in its list.
   integer, allocatable :: fx_rcol(:), fx_rpeer(:), fx_re(:), fx_rcls(:)
      !! Receive entries: destination column, peer, position, fold-row class.

   ! -----------------------------------------------------------------
   ! Buffers + group cursors
   ! -----------------------------------------------------------------
   real(wp), allocatable :: fx_sbuf(:), fx_rbuf(:)
      !! Send / receive buffers, `fx_cap` values per peer region.
   integer :: fx_cap = 0
      !! Per-peer capacity (values).
   integer :: fx_slab_cap = 0
      !! Per-peer capacity in (row x layer) slabs of `fx_nmax` entries.
   type(request_t), allocatable :: fx_reqs(:)
      !! MPI requests (2 per peer).
   type(MPI_Status), allocatable :: fx_stats(:)
      !! MPI statuses (2 per peer).
   integer :: fx_poff(FOLD_NFAM) = 0
      !! Pack cursor: slabs packed per family in the open group.
   integer :: fx_uoff(FOLD_NFAM) = 0
      !! Unpack cursor: slabs unpacked per family in the open group.
   logical :: fx_open = .false.
      !! A group is open (between begin and end).
   logical :: fx_sent = .false.
      !! The open group has been exchanged (packing closed).

contains

   ! ==================================================================
   ! Lifecycle
   ! ==================================================================

   subroutine ocean_fold_exchange_init(decomp, nghost, north_fold, ierr)
      !! Record the topology and, on a north-row rank of an east-west split
      !! folded grid, build the routing plan and map it to the device.
      !! Every rank may call it (it is not collective); a rank that does
      !! not fold, or a `px = 1` run, keeps the exchange inactive.
      type(decomp_t), intent(in) :: decomp
         !! Domain decomposition.
      integer, intent(in) :: nghost
         !! Ghost width.
      logical, intent(in) :: north_fold
         !! This rank folds its north edge (`bc%north_fold`, rank-local).
      integer, intent(out), optional :: ierr
         !! `OCEAN_STATUS_ERR_SETUP` when the plan cannot be built (more
         !! tiles than fold-row columns); absent ⇒ `error stop`.

      integer :: status, p, f

      if (present(ierr)) ierr = OCEAN_STATUS_OK
      if (fx_initialised) call ocean_fold_exchange_destroy()

      fx_ng = nghost
      fx_nxl = decomp%nx_local
      fx_nyl = decomp%ny_local
      fx_px = decomp%px
      fx_ry = decomp%ry
      fx_active = north_fold .and. decomp%px > 1
      fx_initialised = .true.
      if (.not. fx_active) return

      call fold_plan_build(fx_plan, decomp%nx_global, decomp%px, nghost, decomp%rx, status)
      if (status /= FOLD_PLAN_OK) then
         call fail("ocean_fold_exchange_init: cannot build the fold plan for nx = "// &
                   to_string(decomp%nx_global)//", px = "//to_string(decomp%px)// &
                   ", nghost = "//to_string(nghost)//" (need nx >= px, nghost >= 1).", &
                   ierr, OCEAN_STATUS_ERR_SETUP)
         fx_active = .false.
         return
      end if

      fx_npeer = fx_plan%npeer
      fx_self = fx_plan%self_peer
      fx_nmax = max(fx_plan%nmax, 1)
      allocate (fx_peer_rank(fx_npeer))
      do p = 1, fx_npeer
         fx_peer_rank(p) = decomp_rank_from_coords(decomp%px, fx_plan%peer_rx(p), decomp%ry)
      end do
      allocate (fx_sn(fx_npeer, FOLD_NFAM), fx_rn(fx_npeer, FOLD_NFAM))
      fx_sn = fx_plan%send_n
      fx_rn = fx_plan%recv_n
      do f = 1, FOLD_NFAM
         fx_ns(f) = fx_plan%nsend(f)
         fx_nr(f) = fx_plan%nrecv(f)
         fx_s0(f) = 0
         fx_r0(f) = 0
         if (fx_npeer > 0) then
            fx_s0(f) = fx_plan%send_start(1, f)
            fx_r0(f) = fx_plan%recv_start(1, f)
         end if
      end do
      fx_scol = fx_plan%send_col
      fx_speer = fx_plan%send_peer
      fx_se = fx_plan%send_e
      fx_rcol = fx_plan%recv_col
      fx_rpeer = fx_plan%recv_peer
      fx_re = fx_plan%recv_e
      fx_rcls = fx_plan%recv_cls
      allocate (fx_reqs(2*max(fx_npeer, 1)), fx_stats(2*max(fx_npeer, 1)))
      !$acc enter data copyin(fx_sn, fx_rn, fx_scol, fx_speer, fx_se, &
      !$acc&                  fx_rcol, fx_rpeer, fx_re, fx_rcls)

      ! Base capacity: one field of the widest stagger, one layer.
      call grow_buffers(fx_ng + 1)
   end subroutine ocean_fold_exchange_init

   subroutine ocean_fold_exchange_destroy()
      !! Release the plan, buffers and topology.  Idempotent.
      if (.not. fx_initialised) return
      if (fx_active) then
         !$acc exit data delete(fx_sn, fx_rn, fx_scol, fx_speer, fx_se, &
         !$acc&                 fx_rcol, fx_rpeer, fx_re, fx_rcls)
         if (allocated(fx_sbuf)) then
            !$acc exit data delete(fx_sbuf, fx_rbuf)
            deallocate (fx_sbuf, fx_rbuf)
         end if
         deallocate (fx_peer_rank, fx_sn, fx_rn, fx_scol, fx_speer, fx_se, &
                     fx_rcol, fx_rpeer, fx_re, fx_rcls, fx_reqs, fx_stats)
      end if
      call fx_plan%destroy()
      fx_cap = 0
      fx_slab_cap = 0
      fx_npeer = 0
      fx_self = 0
      fx_nmax = 1
      fx_active = .false.
      fx_open = .false.
      fx_sent = .false.
      fx_initialised = .false.
   end subroutine ocean_fold_exchange_destroy

   subroutine ocean_fold_exchange_reserve(nslab)
      !! Pre-size the buffers for the largest group of the run, in
      !! (row x layer) slabs summed over its fields — e.g. the
      !! `ml_state` group is `(3 + ntracer)*nz` fields of at most `ng+1`
      !! rows — so no device reallocation fires mid-run.  No-op when
      !! inactive or already large enough.
      integer, intent(in) :: nslab
         !! Total (rows x layers) over the fields of the largest group.
      if (.not. fx_active) return
      if (nslab > fx_slab_cap) call grow_buffers(nslab)
   end subroutine ocean_fold_exchange_reserve

   pure function ocean_fold_is_distributed() result(flag)
      !! True iff this rank folds through the exchange (`px > 1` and a
      !! folded north edge on this rank).  False on every `px = 1` run and
      !! on ranks off the north row.
      logical :: flag
      flag = fx_initialised .and. fx_active
   end function ocean_fold_is_distributed

   subroutine grow_buffers(nslab)
      !! (Re)allocate both buffers for `nslab` slabs per peer
      !! (exit-delete / dealloc / alloc / enter-create, as
      !! `ocean_halo_buffers_ensure_nz`).  Never while a group is open.
      integer, intent(in) :: nslab

      if (fx_open) error stop "rdb_ocean_fold_exchange: buffer growth inside an open group"
      if (allocated(fx_sbuf)) then
         !$acc exit data delete(fx_sbuf, fx_rbuf)
         deallocate (fx_sbuf, fx_rbuf)
      end if
      fx_slab_cap = nslab
      fx_cap = nslab*fx_nmax
      allocate (fx_sbuf(fx_cap*max(fx_npeer, 1)), fx_rbuf(fx_cap*max(fx_npeer, 1)))
      !$acc enter data create(fx_sbuf, fx_rbuf)
   end subroutine grow_buffers

   ! ==================================================================
   ! Group protocol
   ! ==================================================================

   subroutine ocean_fold_begin(nslab)
      !! Open a group of `nslab` (rows x layers) slabs in total (an upper
      !! bound is fine).  Grows the buffers, with a warning, if the run did
      !! not reserve enough.  No-op when inactive.
      integer, intent(in) :: nslab
         !! Sum over the group's fields of rows x layers.

      if (.not. fx_active) return
      if (fx_open) error stop "rdb_ocean_fold_exchange: ocean_fold_begin with a group already open"
      if (nslab > fx_slab_cap) then
         call logger%warning("ocean_fold_begin: growing the fold buffers mid-run ("// &
                             to_string(fx_slab_cap)//" -> "//to_string(nslab)// &
                             " slabs). Call ocean_fold_exchange_reserve at init.")
         call grow_buffers(nslab)
      end if
      fx_poff = 0
      fx_uoff = 0
      fx_open = .true.
      fx_sent = .false.
   end subroutine ocean_fold_begin

   subroutine ocean_fold_end()
      !! Close the group: every packed slab must have been unpacked.
      if (.not. fx_active) return
      if (.not. fx_open) error stop "rdb_ocean_fold_exchange: ocean_fold_end without begin"
      if (any(fx_uoff /= fx_poff)) then
         error stop "rdb_ocean_fold_exchange: group closed with fields left unpacked"
      end if
      fx_open = .false.
      fx_sent = .false.
   end subroutine ocean_fold_end

   subroutine ocean_fold_pack_2d(fld, nxa, nya, stagger, device_resident)
      !! Pack a 2D field of storage shape (nxa, nya) into the open group.
      integer, intent(in) :: nxa, nya
         !! Storage extents of `fld`.
      real(wp), intent(in) :: fld(nxa, nya)
         !! Field (T: nxt x nyt; u: nxt+1 x nyt; v: nxt x nyt+1;
         !! corner: nxt+1 x nyt+1).
      integer, intent(in) :: stagger
         !! `FOLD_STAG_*`.
      logical, intent(in), optional :: device_resident
         !! `.false.` ⇒ host arrays; default device-resident.

      call ocean_fold_pack_3d(fld, nxa, nya, 1, stagger, device_resident)
   end subroutine ocean_fold_pack_2d

   subroutine ocean_fold_pack_3d(fld, nxa, nya, nz, stagger, device_resident)
      !! Pack every layer of a 3D field into the open group: the sender's
      !! owned mirror-source points, rows below (and, for v / corner, on)
      !! the fold line.
      integer, intent(in) :: nxa, nya, nz
         !! Storage extents of `fld`.
      real(wp), intent(in) :: fld(nxa, nya, nz)
         !! Field; see `ocean_fold_pack_2d` for the per-stagger shape.
      integer, intent(in) :: stagger
         !! `FOLD_STAG_*`.
      logical, intent(in), optional :: device_resident
         !! `.false.` ⇒ host arrays; default device-resident.

      integer :: fam, nrow, sbase, g, g0, ng_e, L, r, p, cap, oT, oU
      logical :: on_device

      if (.not. fx_active) return
      if (.not. fx_open .or. fx_sent) then
         error stop "rdb_ocean_fold_exchange: ocean_fold_pack outside an open, unsent group"
      end if
      on_device = .true.
      if (present(device_resident)) on_device = device_resident

      fam = fold_stagger_family(stagger)
      nrow = fold_stagger_nrows(stagger, fx_ng)
      if (sum(fx_poff) + nrow*nz > fx_slab_cap) then
         error stop "rdb_ocean_fold_exchange: group larger than its ocean_fold_begin size"
      end if
      ! Source row of message row r: ng+nyl+1-r (T, u) / ng+nyl+2-r (v,
      ! corner; r = 1 is the fold-line row) — `fold_row_map`.
      sbase = fx_ng + fx_nyl + 1
      if (nrow == fx_ng + 1) sbase = sbase + 1
      g0 = fx_s0(fam)
      ng_e = fx_ns(fam)
      cap = fx_cap
      oT = fx_poff(FOLD_FAM_T)
      oU = fx_poff(FOLD_FAM_U)

      if (ng_e > 0) then
         if (on_device) then
            !$acc parallel loop collapse(3) private(p) &
            !$acc& present(fx_sbuf, fx_scol, fx_speer, fx_se, fx_sn, fld)
            do L = 1, nz
               do r = 1, nrow
                  do g = g0 + 1, g0 + ng_e
                     p = fx_speer(g)
                     fx_sbuf((p - 1)*cap + oT*fx_sn(p, FOLD_FAM_T) + oU*fx_sn(p, FOLD_FAM_U) &
                             + ((L - 1)*nrow + (r - 1))*fx_sn(p, fam) + fx_se(g)) = &
                        fld(fx_scol(g), sbase - r, L)
                  end do
               end do
            end do
         else
            do L = 1, nz
               do r = 1, nrow
                  do g = g0 + 1, g0 + ng_e
                     p = fx_speer(g)
                     fx_sbuf((p - 1)*cap + oT*fx_sn(p, FOLD_FAM_T) + oU*fx_sn(p, FOLD_FAM_U) &
                             + ((L - 1)*nrow + (r - 1))*fx_sn(p, fam) + fx_se(g)) = &
                        fld(fx_scol(g), sbase - r, L)
                  end do
               end do
            end do
         end if
      end if
      fx_poff(fam) = fx_poff(fam) + nrow*nz
   end subroutine ocean_fold_pack_3d

   subroutine ocean_fold_exchange(device_resident)
      !! Move every packed message: one isend + irecv per non-self peer
      !! with a non-empty message, the self pair as a local copy, then
      !! `waitall`.  Collective over the north rank row.
      logical, intent(in), optional :: device_resident
         !! `.false.` ⇒ host buffers; default device-resident.

      integer :: p, n_s, n_r, o, nreq, i
      logical :: on_device
      type(comm_t) :: comm

      if (.not. fx_active) return
      if (.not. fx_open .or. fx_sent) then
         error stop "rdb_ocean_fold_exchange: ocean_fold_exchange outside an open, unsent group"
      end if
      on_device = .true.
      if (present(device_resident)) on_device = device_resident
      comm = comm_env_compute_comm()

      nreq = 0
      do p = 1, fx_npeer
         n_s = fx_poff(FOLD_FAM_T)*fx_sn(p, FOLD_FAM_T) + fx_poff(FOLD_FAM_U)*fx_sn(p, FOLD_FAM_U)
         n_r = fx_poff(FOLD_FAM_T)*fx_rn(p, FOLD_FAM_T) + fx_poff(FOLD_FAM_U)*fx_rn(p, FOLD_FAM_U)
         o = (p - 1)*fx_cap
         if (p == fx_self) then
            ! Self pair: the receive region is the send region (same list).
            if (on_device) then
               !$acc parallel loop present(fx_sbuf, fx_rbuf)
               do i = o + 1, o + n_s
                  fx_rbuf(i) = fx_sbuf(i)
               end do
            else
               fx_rbuf(o + 1:o + n_s) = fx_sbuf(o + 1:o + n_s)
            end if
            cycle
         end if
         if (on_device) then
            !$acc host_data use_device(fx_sbuf, fx_rbuf)
            if (n_r > 0) then
               nreq = nreq + 1
               call FOLD_IRECV_N(comm, fx_rbuf(o + 1:o + n_r), n_r, fx_peer_rank(p), &
                                 TAG_OC_FOLD, fx_reqs(nreq))
            end if
            if (n_s > 0) then
               nreq = nreq + 1
               call FOLD_ISEND_N(comm, fx_sbuf(o + 1:o + n_s), n_s, fx_peer_rank(p), &
                                 TAG_OC_FOLD, fx_reqs(nreq))
            end if
            !$acc end host_data
         else
            if (n_r > 0) then
               nreq = nreq + 1
               call FOLD_IRECV_N(comm, fx_rbuf(o + 1:o + n_r), n_r, fx_peer_rank(p), &
                                 TAG_OC_FOLD, fx_reqs(nreq))
            end if
            if (n_s > 0) then
               nreq = nreq + 1
               call FOLD_ISEND_N(comm, fx_sbuf(o + 1:o + n_s), n_s, fx_peer_rank(p), &
                                 TAG_OC_FOLD, fx_reqs(nreq))
            end if
         end if
      end do
      if (nreq > 0) call waitall(fx_reqs(1:nreq), fx_stats(1:nreq))
      fx_sent = .true.
   end subroutine ocean_fold_exchange

   subroutine ocean_fold_unpack_2d(fld, nxa, nya, stagger, negate, device_resident)
      !! Unpack the next field of the exchanged group into a 2D field.
      integer, intent(in) :: nxa, nya
         !! Storage extents of `fld`.
      real(wp), intent(inout) :: fld(nxa, nya)
         !! Field; same stagger and order as at pack.
      integer, intent(in) :: stagger
         !! `FOLD_STAG_*`.
      logical, intent(in) :: negate
         !! `.true.` for a true-vector component (sign flips across the fold).
      logical, intent(in), optional :: device_resident
         !! `.false.` ⇒ host arrays; default device-resident.

      call ocean_fold_unpack_3d(fld, nxa, nya, 1, stagger, negate, device_resident)
   end subroutine ocean_fold_unpack_2d

   subroutine ocean_fold_unpack_3d(fld, nxa, nya, nz, stagger, negate, device_resident)
      !! Unpack the next field of the exchanged group: write every north
      !! ghost row (all storage columns) and, for v / corner, the fold-line
      !! row's west half (self-conjugate column → 0 for a vector).
      integer, intent(in) :: nxa, nya, nz
         !! Storage extents of `fld`.
      real(wp), intent(inout) :: fld(nxa, nya, nz)
         !! Field; same stagger and order as at pack.
      integer, intent(in) :: stagger
         !! `FOLD_STAG_*`.
      logical, intent(in) :: negate
         !! `.true.` for a true-vector component.
      logical, intent(in), optional :: device_resident
         !! `.false.` ⇒ host arrays; default device-resident.

      integer :: fam, nrow, dbase, g, g0, ng_e, L, r, p, cap, oT, oU, ib, cls
      logical :: on_device, has_row
      real(wp) :: sgn, val

      if (.not. fx_active) return
      if (.not. fx_open .or. .not. fx_sent) then
         error stop "rdb_ocean_fold_exchange: ocean_fold_unpack before ocean_fold_exchange"
      end if
      on_device = .true.
      if (present(device_resident)) on_device = device_resident

      fam = fold_stagger_family(stagger)
      nrow = fold_stagger_nrows(stagger, fx_ng)
      has_row = (nrow == fx_ng + 1)
      ! Destination row of message row r: ng+nyl+r for every stagger
      ! (`fold_row_map`: T/u ng+nyl+d, v/corner ng+nyl+1+d with d = r-1).
      dbase = fx_ng + fx_nyl
      g0 = fx_r0(fam)
      ng_e = fx_nr(fam)
      cap = fx_cap
      oT = fx_uoff(FOLD_FAM_T)
      oU = fx_uoff(FOLD_FAM_U)
      sgn = 1.0_wp
      if (negate) sgn = -1.0_wp

      if (ng_e > 0) then
         if (on_device) then
            !$acc parallel loop collapse(3) private(p, ib, cls, val) &
            !$acc& present(fx_rbuf, fx_rcol, fx_rpeer, fx_re, fx_rcls, fx_rn, fld)
            do L = 1, nz
               do r = 1, nrow
                  do g = g0 + 1, g0 + ng_e
                     p = fx_rpeer(g)
                     ib = (p - 1)*cap + oT*fx_rn(p, FOLD_FAM_T) + oU*fx_rn(p, FOLD_FAM_U) &
                          + ((L - 1)*nrow + (r - 1))*fx_rn(p, fam) + fx_re(g)
                     val = sgn*fx_rbuf(ib)
                     if (has_row .and. r == 1) then
                        cls = fx_rcls(g)
                        if (cls == FOLD_ROW_WEST) then
                           fld(fx_rcol(g), dbase + 1, L) = val
                        else if (cls == FOLD_ROW_SELF .and. negate) then
                           fld(fx_rcol(g), dbase + 1, L) = 0.0_wp
                        end if
                     else
                        fld(fx_rcol(g), dbase + r, L) = val
                     end if
                  end do
               end do
            end do
         else
            do L = 1, nz
               do r = 1, nrow
                  do g = g0 + 1, g0 + ng_e
                     p = fx_rpeer(g)
                     ib = (p - 1)*cap + oT*fx_rn(p, FOLD_FAM_T) + oU*fx_rn(p, FOLD_FAM_U) &
                          + ((L - 1)*nrow + (r - 1))*fx_rn(p, fam) + fx_re(g)
                     val = sgn*fx_rbuf(ib)
                     if (has_row .and. r == 1) then
                        cls = fx_rcls(g)
                        if (cls == FOLD_ROW_WEST) then
                           fld(fx_rcol(g), dbase + 1, L) = val
                        else if (cls == FOLD_ROW_SELF .and. negate) then
                           fld(fx_rcol(g), dbase + 1, L) = 0.0_wp
                        end if
                     else
                        fld(fx_rcol(g), dbase + r, L) = val
                     end if
                  end do
               end do
            end do
         end if
      end if
      fx_uoff(fam) = fx_uoff(fam) + nrow*nz
   end subroutine ocean_fold_unpack_3d

   ! ==================================================================
   ! px-dispatching single-field folds
   ! ==================================================================
   ! Each takes the field's storage extents and the tile geometry the
   ! local kernel needs.  px = 1 (or an inactive rank): the `rdb_ocean_fold`
   ! kernel, verbatim.  px > 1: a one-field group through the exchange.

   subroutine fold_centre_2d(fld, nxt, nyt, nx_phys, ny_phys, ng, device_resident)
      !! Cell-centred 2D north fold (copy).
      integer, intent(in) :: nxt, nyt, nx_phys, ny_phys, ng
      real(wp), intent(inout) :: fld(nxt, nyt)
         !! Cell-centred field (nx_total, ny_total).
      logical, intent(in), optional :: device_resident
         !! Exchange path only: `.false.` ⇒ host arrays.

      if (ocean_fold_is_distributed()) then
         call ocean_fold_begin(fold_stagger_nrows(FOLD_STAG_T, fx_ng))
         call ocean_fold_pack(fld, nxt, nyt, FOLD_STAG_T, device_resident)
         call ocean_fold_exchange(device_resident)
         call ocean_fold_unpack(fld, nxt, nyt, FOLD_STAG_T, .false., device_resident)
         call ocean_fold_end()
      else
         call fold_north_centre(fld, nxt, nyt, nx_phys, ny_phys, ng)
      end if
   end subroutine fold_centre_2d

   subroutine fold_centre_3d(fld, nxt, nyt, nz, nx_phys, ny_phys, ng, device_resident)
      !! Cell-centred 3D north fold (copy).
      integer, intent(in) :: nxt, nyt, nz, nx_phys, ny_phys, ng
      real(wp), intent(inout) :: fld(nxt, nyt, nz)
         !! Cell-centred field (nx_total, ny_total, nz).
      logical, intent(in), optional :: device_resident
         !! Exchange path only: `.false.` ⇒ host arrays.

      if (ocean_fold_is_distributed()) then
         call ocean_fold_begin(fold_stagger_nrows(FOLD_STAG_T, fx_ng)*nz)
         call ocean_fold_pack(fld, nxt, nyt, nz, FOLD_STAG_T, device_resident)
         call ocean_fold_exchange(device_resident)
         call ocean_fold_unpack(fld, nxt, nyt, nz, FOLD_STAG_T, .false., device_resident)
         call ocean_fold_end()
      else
         call fold_north_centre(fld, nxt, nyt, nz, nx_phys, ny_phys, ng)
      end if
   end subroutine fold_centre_3d

   subroutine fold_u_2d(fld, nxf, nyt, nx_phys, ny_phys, ng, device_resident)
      !! x-face (u) 2D north fold (negate).
      integer, intent(in) :: nxf, nyt, nx_phys, ny_phys, ng
      real(wp), intent(inout) :: fld(nxf, nyt)
         !! x-face field (nx_total+1, ny_total).
      logical, intent(in), optional :: device_resident
         !! Exchange path only: `.false.` ⇒ host arrays.

      if (ocean_fold_is_distributed()) then
         call ocean_fold_begin(fold_stagger_nrows(FOLD_STAG_U, fx_ng))
         call ocean_fold_pack(fld, nxf, nyt, FOLD_STAG_U, device_resident)
         call ocean_fold_exchange(device_resident)
         call ocean_fold_unpack(fld, nxf, nyt, FOLD_STAG_U, .true., device_resident)
         call ocean_fold_end()
      else
         call fold_north_u_face(fld, nxf, nyt, nx_phys, ny_phys, ng)
      end if
   end subroutine fold_u_2d

   subroutine fold_u_3d(fld, nxf, nyt, nz, nx_phys, ny_phys, ng, device_resident)
      !! x-face (u) 3D north fold (negate).
      integer, intent(in) :: nxf, nyt, nz, nx_phys, ny_phys, ng
      real(wp), intent(inout) :: fld(nxf, nyt, nz)
         !! x-face field (nx_total+1, ny_total, nz).
      logical, intent(in), optional :: device_resident
         !! Exchange path only: `.false.` ⇒ host arrays.

      if (ocean_fold_is_distributed()) then
         call ocean_fold_begin(fold_stagger_nrows(FOLD_STAG_U, fx_ng)*nz)
         call ocean_fold_pack(fld, nxf, nyt, nz, FOLD_STAG_U, device_resident)
         call ocean_fold_exchange(device_resident)
         call ocean_fold_unpack(fld, nxf, nyt, nz, FOLD_STAG_U, .true., device_resident)
         call ocean_fold_end()
      else
         call fold_north_u_face(fld, nxf, nyt, nz, nx_phys, ny_phys, ng)
      end if
   end subroutine fold_u_3d

   subroutine fold_v_2d(fld, nxt, nyf, nx_phys, ny_phys, ng, device_resident)
      !! y-face (v) 2D north fold (negate + fold-line projection).
      integer, intent(in) :: nxt, nyf, nx_phys, ny_phys, ng
      real(wp), intent(inout) :: fld(nxt, nyf)
         !! y-face field (nx_total, ny_total+1).
      logical, intent(in), optional :: device_resident
         !! Exchange path only: `.false.` ⇒ host arrays.

      if (ocean_fold_is_distributed()) then
         call ocean_fold_begin(fold_stagger_nrows(FOLD_STAG_V, fx_ng))
         call ocean_fold_pack(fld, nxt, nyf, FOLD_STAG_V, device_resident)
         call ocean_fold_exchange(device_resident)
         call ocean_fold_unpack(fld, nxt, nyf, FOLD_STAG_V, .true., device_resident)
         call ocean_fold_end()
      else
         call fold_north_v_face(fld, nxt, nyf, nx_phys, ny_phys, ng)
      end if
   end subroutine fold_v_2d

   subroutine fold_v_3d(fld, nxt, nyf, nz, nx_phys, ny_phys, ng, device_resident)
      !! y-face (v) 3D north fold (negate + fold-line projection).
      integer, intent(in) :: nxt, nyf, nz, nx_phys, ny_phys, ng
      real(wp), intent(inout) :: fld(nxt, nyf, nz)
         !! y-face field (nx_total, ny_total+1, nz).
      logical, intent(in), optional :: device_resident
         !! Exchange path only: `.false.` ⇒ host arrays.

      if (ocean_fold_is_distributed()) then
         call ocean_fold_begin(fold_stagger_nrows(FOLD_STAG_V, fx_ng)*nz)
         call ocean_fold_pack(fld, nxt, nyf, nz, FOLD_STAG_V, device_resident)
         call ocean_fold_exchange(device_resident)
         call ocean_fold_unpack(fld, nxt, nyf, nz, FOLD_STAG_V, .true., device_resident)
         call ocean_fold_end()
      else
         call fold_north_v_face(fld, nxt, nyf, nz, nx_phys, ny_phys, ng)
      end if
   end subroutine fold_v_3d

   subroutine fold_corner_2d(fld, nxf, nyf, nx_phys, ny_phys, ng, negate, device_resident)
      !! SW-corner 2D north fold (+ fold-line projection); `negate` for a
      !! true-vector component, copy for a scalar / vorticity.
      integer, intent(in) :: nxf, nyf, nx_phys, ny_phys, ng
      real(wp), intent(inout) :: fld(nxf, nyf)
         !! Corner field (nx_total+1, ny_total+1).
      logical, intent(in) :: negate
         !! `.true.` ⇒ negate (true vector); `.false.` ⇒ copy.
      logical, intent(in), optional :: device_resident
         !! Exchange path only: `.false.` ⇒ host arrays.

      if (ocean_fold_is_distributed()) then
         call ocean_fold_begin(fold_stagger_nrows(FOLD_STAG_CORNER, fx_ng))
         call ocean_fold_pack(fld, nxf, nyf, FOLD_STAG_CORNER, device_resident)
         call ocean_fold_exchange(device_resident)
         call ocean_fold_unpack(fld, nxf, nyf, FOLD_STAG_CORNER, negate, device_resident)
         call ocean_fold_end()
      else
         call fold_north_corner(fld, nxf, nyf, nx_phys, ny_phys, ng, negate)
      end if
   end subroutine fold_corner_2d

end module rdb_ocean_fold_exchange