run_stage Subroutine

   subroutine run_stage(grid, metrics, dyn, eos, cor, ct, pgf, hv, bd, ss, va, hd, vd, vmix, ms, dt, stage, sf, geo, lateral_mix, epbl, kshear, slopes, vmix_tidal, td, cav)
      !! One FE stage of the multilayer step.  Order of operations:
      !!
      !!   1. EOS: rho_layer <- linear(T, S)
      !!   2. Velocity-tendency computes — Coriolis, PGF, hvisc,
      !!      bottom drag, surface stress.  All read h_old, u, v
      !!      and write to their own tendency buffers; none touch
      !!      h yet.
      !!   3. continuity_tracer_step_split — interleaved
      !!      Lie-split horizontal step: zonal flux + tracer zonal +
      !!      apply zonal + meridional flux + tracer meridional +
      !!      apply meridional.  Updates h AND hTr together; the
      !!      CWC discrete theorem holds (uniform T stays uniform).
      !!   4. tracer_hdiff — Laplacian diffusion on hTr.
      !!   5. compute_w_from_continuity — diagnose w from the total
      !!      horizontal divergence stored in `flux_h_layer`.
      !!   6. tracer_advect_vertical — upwind-in-z + h update that
      !!      cancels the horizontal h change (Eulerian-z mode).
      !!   7. Velocity-tendency applies (Coriolis, PGF, hvisc, drag,
      !!      surface stress) — all additive.
      !!   8. Surface tracer fluxes (heat, salt).
      !!   9. Vertical mixing closure + vdiff.
      type(hgrid_t), intent(in) :: grid
      type(ocean_metrics_t), intent(in) :: metrics
         !! Curvilinear horizontal metrics — forwarded to the PGF.
      type(ocean_dyn_t), intent(in) :: dyn
      type(eos_t), intent(in) :: eos
      type(coriolis_adv_t), intent(inout) :: cor
      type(continuity_t), intent(inout) :: ct
      type(ocean_pressure_force_t), intent(inout) :: pgf
      type(ocean_horizontal_viscosity_t), intent(inout) :: hv
      type(ocean_bottom_drag_t), intent(inout) :: bd
      type(ocean_top_drag_t), intent(inout), optional :: td
         !! Ice-shelf TOP-drag slot (`&ocean_tdrag_nml`).  OPTIONAL so the
         !! many direct `ocean_dyn_step*` / `run_stage*` call sites in the
         !! test suite need no churn; the production driver always passes
         !! it.  Absent, or present and disabled, => no kernel launch and
         !! a bit-identical step.
      type(ocean_cavity_flux_t), intent(inout), optional :: cav
         !! Ice-shelf basal-melt slot (`&ocean_cavity_melt_nml`).
         !! OPTIONAL for the same reason `td` is: the direct
         !! `ocean_dyn_step*` / `run_stage*` call sites in the test suite
         !! need no churn, and the production driver always passes it.
         !! Absent, disabled, or `freshwater="virtual"` => no kernel
         !! launch and a bit-identical step.  It is threaded down here
         !! rather than acted on in `engine_step_finalize` because the
         !! real-freshwater volume must be spent in the SAME stage, at
         !! the SAME stage weight and from the SAME `melt` value as the
         !! salt and heat halves the surface-flux apply spends -- see
         !! `ocean_cavity_mass_step`'s docstring.
      type(ocean_surface_stress_t), intent(inout) :: ss
      type(ocean_vertical_advection_t), intent(inout) :: va
      type(ocean_hdiff_tracer_t), intent(inout) :: hd
      type(ocean_vdiff_t), intent(inout) :: vd
      type(ocean_vmix_t), intent(inout) :: vmix
      type(multilayer_state_t), intent(inout) :: ms
      real(wp), intent(in) :: dt
      integer, intent(in) :: stage
         !! RK2 stage (1 or 2); see vmix_apply_in_stage.
      type(ocean_surface_flux_t), intent(in), optional :: sf
      type(ocean_geothermal_t), intent(in), optional :: geo
         !! Geothermal bottom-heat-flux slot.  See `ocean_dyn_step`.
      type(ocean_lateral_mix_t), intent(inout), optional :: lateral_mix
      type(ocean_epbl_t), intent(inout), optional :: epbl
      type(ocean_kappa_shear_t), intent(inout), optional :: kshear
      type(ocean_slopes_t), intent(inout), optional :: slopes
      type(ocean_tidal_mixing_t), intent(inout), optional :: vmix_tidal

      logical :: therm_active
      logical :: fold_top
         !! `.true.` when the ice-shelf top drag is folded into the vdiff
         !! `k = nz` diagonal — see the gate note at the `vmix_apply_in_stage`
         !! call below for why the test is `implicit_fold`, not `present(td)`.
      logical :: publish_shelf
         !! `.true.` when the ice-shelf top-drag slot is live and its
         !! `stress_top` is therefore full-sized and freshly written.
      integer :: i_ss, j_ss, nx_ss, ny_ss
         !! Loop/extent locals for the inline `stress_shelf` publish.
      real(wp) :: therm_dt

      therm_active = dyn%enable_thermodynamics .and. dyn%is_thermo_step()
      therm_dt = dyn%therm_dt(dt)

      ! EOS is DYNAMICS, not thermodynamics: rho_layer feeds the baroclinic
      ! PGF every step. Gating it on `is_thermo_step()` zero-order-holds the
      ! restoring force of the internal-gravity-wave oscillator for
      ! tau = (dt_therm_ratio-1)*dt, which is a delayed-restoring-force
      ! instability: sigma ~ omega^2*tau/2, unstable for ANY tau > 0, maximal
      ! at the grid scale. Measured on eady: sigma ∝ alpha_T ∝ N^2, sigma ∝ k^2,
      ! the mode is the diagonal 2-delta checkerboard, and the rate recovers
      ! the mode-2 internal wave speed. Survival was only ever by viscosity
      ! margin (nu*k^2 > sigma) — eady at nu_h=100 is marginal and blows up;
      ! double_gyre at nu_h=10000 merely looks fine.
      !
      ! So recompute rho EVERY dynamics step whenever thermodynamics is on.
      ! The EXPENSIVE thermo (mixing, ALE remap, tracer advection/hdiff,
      ! surface fluxes) stays on the slow `therm_active` cadence below, which
      ! is where the cost actually is. Bit-identical at dt_therm_ratio = 1,
      ! where is_thermo_step() is always true.
      call ocean_eos_compute(eos, ms, active=dyn%enable_thermodynamics)
      call coriolis_adv_compute_tendencies(grid, metrics, cor, ms)
      call ocean_pressure_force_compute(grid, metrics, pgf, ms, eos=eos)
      call ocean_lateral_mix_compute(grid, metrics, lateral_mix, ms)
      call ocean_horizontal_viscosity_compute_tendencies(grid, metrics, hv, ms, &
                                                         lateral_mix=lateral_mix, dt=dt)
      ! KE dissipation rate for the MEKE frictional source — captured here
      ! (du_visc fresh, u_face still pre-viscous) before the apply below.
      ! No-op unless MEKE's frictional source is enabled.
      call ocean_horizontal_viscosity_compute_ke_diss(hv, ms)
      call ocean_bottom_drag_compute_tendencies(grid, bd, ms, dt)
      call ocean_channel_drag_compute_tendencies(grid, metrics, bd, ms)
      ! Ice-shelf top drag (`&ocean_tdrag_nml`).  Sits with the other slow
      ! velocity-tendency computes; the kernel returns immediately when the
      ! slot is disabled, so an ordinary run pays one host branch.
      if (present(td)) call ocean_top_drag_compute_tendencies(td, ms, dt)
      ! Phase 4b: publish the ice-shelf base stress that BOTH boundary-
      ! layer schemes take `u_*` from.  Under a shelf the wind has been
      ! masked out of `tau` (so `stress_mag` is exactly 0 there) and the
      ! turbulent boundary layer is driven by the ice-ocean stress
      ! instead — `u_*^2 = |tau_top|/rho_0`.
      !
      ! Written INLINE as a `do concurrent`, not as a call handing
      ! `ss%stress_shelf` to an external subroutine: a host-gated call
      ! with a state array as an actual makes nvfortran treat the array
      ! as escaping and pessimises every `do concurrent` in this routine
      ! (CLAUDE.md, measured at +4.8%% on an inert porous pass).
      !
      ! Gated on `td%enable`, not `present(td)`: a DISABLED slot carries
      ! a `(1,1)` placeholder `stress_top`.
      !
      ! Placed here — after the top-drag compute, before
      ! `vmix_apply_in_stage` below — so KPP/EPBL read THIS stage's
      ! stress.  No lag.
      publish_shelf = .false.
      if (present(td)) publish_shelf = td%enable
      if (publish_shelf) then
         nx_ss = size(ss%stress_shelf, 1)
         ny_ss = size(ss%stress_shelf, 2)
         do concurrent(j_ss=1:ny_ss, i_ss=1:nx_ss)
            ss%stress_shelf(i_ss, j_ss) = td%stress_top(i_ss, j_ss)
         end do
      end if
      call ocean_surface_stress_compute_tendencies(grid, ss, ms)

      ! Horizontal step + tracer chain.  Continuity-tracer is
      ! unconditional (the h advection lives here regardless of
      ! thermodynamics); the tracer-only kernels self-gate on
      ! `therm_active`.
      !
      ! Phase 2 (6b) horizontal-tracer-advect cadence dispatch
      ! (`dt_tracer_advect_ratio`):
      !   ratio == 1 (default) → the existing every-step fused
      !     continuity+tracer kernel runs verbatim ⇒ bit-identical.
      !   ratio  > 1 → TR_MODE_ACCUMULATE: advance h + accumulate
      !     0.5·mass_flux·dt into ct%uhtr/vhtr per RK2 stage (hTr frozen);
      !     the boundary drain in `ocean_dyn_step` spends them.
      if (dyn%dt_tracer_advect_ratio <= 1) then
         call continuity_tracer_step_split(grid, metrics, ct, ms, dt)
      else
         call continuity_tracer_step_split(grid, metrics, ct, ms, dt, &
                                           tracer_mode=TR_MODE_ACCUMULATE)
      end if
      call tracer_hdiff(grid, metrics, hd, ms, therm_dt, active=therm_active)
      ! Mass budget: flux_h_layer now holds the total horizontal divergence
      ! (the same field compute_w_from_continuity reads).  Accumulate the
      ! boundary mass outflux for this RK2 stage (weight 0.5, since rk2_average
      ! halves each stage's contribution); vertical advection + ALE remap only
      ! redistribute within a column, so they don't affect the column-mass
      ! budget.  Closes the ocean mass Error to round-off with open BCs.
      call ocean_accumulate_mass_out(ms, ms%flux_h_layer, metrics%areaT, &
                                     grid%nghost, dt, 0.5_wp)
      call compute_w_from_continuity(grid, va, ms)
      call tracer_advect_vertical(grid, va, ms, therm_dt, active=therm_active)

      ! Velocity-side applies (no thermodynamics gate).  Async-chained on
      ! OpenACC queue 1 (additive accumulations onto u_face/v_face, FIFO-
      ! ordered).  The interleaved tracer applies below run on the default
      ! queue but touch disjoint arrays (hTr), so there is no cross-queue
      ! race.  ONE `!$acc wait(1)` before vmix — the first device consumer
      ! that reads the applied u_face/v_face (shear).
      call coriolis_adv_apply_tendencies(cor, ms, dt, no_wait=.true.)
      call ocean_pressure_force_apply(pgf, ms, dt, no_wait=.true.)
      call ocean_horizontal_viscosity_apply_tendencies(hv, ms, dt, no_wait=.true.)
      ! Double-count guard: when the bottom drag / wind stress are folded
      ! into the implicit vdiff tridiagonal (`&ocean_vdiff_nml implicit_*`),
      ! their explicit pre-solve applies are SKIPPED here so the forcing is
      ! not applied twice.  Channel (side-wall) drag is a distinct lateral
      ! term and always applies.  Defaults (both off) ⇒ both applies run ⇒
      ! bit-identical to the prior path.  The MOM6 BBL glue (`bbl_glue`)
      ! is the same kind of fold: its piston IS the bed drag, so the
      ! explicit (or `&ocean_bdrag_nml implicit` split-apply) one is skipped.
      if (.not. (vd%implicit_drag .or. vd%bbl_glue)) then
         call ocean_bottom_drag_apply_tendencies(bd, ms, dt, no_wait=.true.)
      end if
      call ocean_channel_drag_apply_tendencies(bd, ms, dt, no_wait=.true.)
      ! Top drag: same double-count guard as the bed.  `implicit_fold`
      ! (`&ocean_vdiff_nml implicit_top_drag`) folds the rate into the
      ! vdiff `k = nz` diagonal instead, so the explicit apply is skipped
      ! there.
      if (present(td)) then
         if (.not. td%implicit_fold) then
            call ocean_top_drag_apply_tendencies(td, ms, dt, no_wait=.true.)
         end if
      end if
      if (.not. vd%implicit_stress) then
         call ocean_surface_stress_apply_tendencies(ss, ms, dt, no_wait=.true.)
      end if

      ! Surface tracer fluxes, geothermal bottom flux, ideal-age, vmix.
      call ocean_surface_flux_apply_tracers(grid, sf, ms, therm_dt, active=therm_active)
      ! Ice-shelf real freshwater MASS (&ocean_cavity_melt_nml
      ! freshwater='mass').  Immediately after the surface-flux apply,
      ! because it REPLACES that apply's virtual cavity salt increment
      ! with the real advective one and adds the meltwater volume + its
      ! enthalpy in the same stage, at the same weight, from the same
      ! `melt`.  Absent / disabled / 'virtual' => immediate return.
      ! Weight 0.5 per SSP-RK2 stage, matching ocean_accumulate_mass_out.
      if (present(cav)) then
         call ocean_cavity_mass_step(grid, metrics, cav, ms, therm_dt, 0.5_wp, &
                                     active=therm_active)
      end if
      call apply_sw_and_restore(grid, metrics, sf, ms, therm_dt, therm_active)
      call ocean_geothermal_apply_tracers(grid, geo, ms, therm_dt, active=therm_active)
      ! Ideal-age interior aging only (PR-7): thermo-cadence gated, mirrors
      ! its tracer-kernel neighbours above.  The surface Dirichlet reset is
      ! NOT here — it runs once per outer step, after rk2_average, in
      ! `ocean_dyn_step` (a per-stage reset is halved by the RK2 average).
      call ocean_ideal_age_apply(grid, ms, therm_dt, active=dyn%is_thermo_step())

      ! Isopycnal-slope diagnostics — purely diagnostic, refreshed at
      ! thermo cadence (same gate as the lateral closures it feeds).
      ! No-op when absent / disabled (bit-identical).
      if (present(slopes) .and. therm_active) then
         if (slopes%enable) then
            call ocean_slopes_compute(grid, metrics, eos, slopes, ms, therm_dt)
         end if
      end if

      !$acc wait(1)
      ! No `bt_work` here: the unsplit path has no BT correction at all,
      ! so there is no consumer for visc_rem — do_remnant stays .false.
      ! The top-drag fold arrays are handed over as ARRAYS, not as the
      ! slot: `td` is optional here, and dereferencing an absent
      ! derived-type dummy is not allowed, whereas forwarding an absent
      ! optional ARRAY dummy on to another optional dummy is.
      !
      ! GATED ON `implicit_fold`, NOT on `present(td)`.  A DISABLED
      ! top-drag slot carries PLACEHOLDER-sized arrays, and the
      ! explicit-shape `(nu, nv)` dummy down in
      ! `diffuse_velocity_columns_impl` is mapped by nvfortran
      ! UNCONDITIONALLY — the `if (do_top)` guard inside the kernel is a
      ! runtime branch the compiler cannot see.  Handing over a `(2,1)`
      ! placeholder therefore aborts the GPU build with "variable in data
      ! clause is partially present", which is exactly what it did before
      ! this gate.  The fold requires `&ocean_tdrag_nml enable`
      ! (validate_config), so when it is on the arrays are full size.
      fold_top = .false.
      if (present(td)) fold_top = td%implicit_fold
      if (fold_top) then
         call vmix_apply_in_stage(grid, dyn, vmix, vd, ss, bd, ms, dt, stage, sf, epbl=epbl, kshear=kshear, &
                                  vmix_tidal=vmix_tidal, metrics=metrics, &
                                  lambda_top_u=td%lambda_top_u, lambda_top_v=td%lambda_top_v, &
                                  cover_u=td%cover_u, cover_v=td%cover_v)
      else
         call vmix_apply_in_stage(grid, dyn, vmix, vd, ss, bd, ms, dt, stage, sf, epbl=epbl, kshear=kshear, &
                                  vmix_tidal=vmix_tidal, metrics=metrics)
      end if
   end subroutine run_stage