vibe kappa shear - #81
Draft
JorgeG94 wants to merge 15 commits into
Draft
Conversation
Record each clock's name as cpu_clock_id() registers it, then open an NVTX range in cpu_clock_begin() and close it in cpu_clock_end(). Every existing cpu_clock_id() name becomes a named range in an nsys timeline with no call-site changes. The range is opened before, and closed after, the mpp clock so that it encloses it. Begin and end test the same condition, so starts and ends stay balanced for handles that were never named or that fall outside the table. The code is guarded behind -DMOM_USE_NVTX, which is undefined by default: without it the preprocessed source is byte-identical to before and no extra library is linked. Profiling builds need -DMOM_USE_NVTX and -cudalib=nvtx. Applied to both the FMS1 and FMS2 infra shims, which are byte-identical to each other. This also covers the halo and message-passing clocks, which MOM_domain_infra begins and ends through this module directly rather than through the MOM_cpu_clock framework wrapper. Based on Edward Yang's ae67665 ("add nvtx markers to clocks", branch benchmark_ALE_nvtx_clocks), which introduced the name-table approach and the range nesting; this adds the build guard so non-NVHPC builds are unaffected. Co-authored-by: Edward Yang <edward.yang@anu.edu.au> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
First increment of the MOM_kappa_shear GPU port. Converts the two
thickness-to-vertex interpolation loops in Calc_kappa_shear_vertex
(both branches of vertex_shear_OBC_bug) from !$OMP parallel do to
do concurrent, offloading them to the device.
Residency:
- h is host-authoritative in the (host-only) diabatic stack and is
mutated on host during diabatic (applyBoundaryFluxesInOut), so the
globally-resident device copy (CS%h) is refreshed with update to(h)
before the kernels; a plain map(to:) does not refresh a present
object.
- G%mask2dCu/Cv/T are already device-resident (initialize_MOM).
- h_at_u/h_at_v are device workspace (map alloc), copied back with
update from for the still-host per-column solver below. enter/exit
data are mirrored and refcount-neutral (release, not delete).
The per-column solver (kappa_shear_column/find_kappa_tke) stays on
host for now; later increments move it to the device, at which point
the update from(h_at_u,h_at_v) round-trip is removed.
Verification:
- benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1).
- nsys --stats=true confirms both do concurrent loops run on the GPU
(calc_kappa_shear_vertex_{550,553}_gpu, 12 launches each).
- Reviewed; no code defect found. The >=2-rank bitwise run remains the
mandatory pre-merge gate and is still pending (single-GPU node).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Groundwork for whole-column GPU kernels (kappa_shear next) that must
compute density derivatives from INSIDE a device region, where the
polymorphic EOS%type%... dispatch is not usable.
- Expose the existing this-free _loc elementals
calculate_density_derivs_elem_{Roquet_rho,buggy_Wright}_loc as public
and !$omp declare target.
- Add calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure,
drho_dT, drho_dS): a declare-target subroutine that select-case's on
the integer form and calls the matching _loc kernel, returning mks
derivatives. Unit rescaling and any `scale` factor remain the
caller's responsibility (as in calculate_density_derivs_1d). Forms
without a device _loc kernel hit the default branch; a device-using
module must FATAL at init on a GPU build before reaching it.
This is the reusable, form-general enabler chosen over a Roquet-scoped
inline; it is aligned with doc 06 6.4's select-case direction but adds
only the device-callable point path (not the full polymorphism removal).
No functional change: nothing calls the dispatcher yet.
Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu
reference (np=1). Dispatcher is exercised (and Fable-reviewed) in the
follow-up that wires kappa_shear to it.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…r on GPU builds Wires kappa_shear_column's Boussinesq density-derivative computation to the device-callable path added in 47ac931, so the column solver can later run inside a target region (the polymorphic calculate_density_derivs interface is not device-callable). - MOM_EOS: add get_EOS_form_and_scaling(EOS, form, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) so callers can reproduce, host-side, the unit rescaling that calculate_density_derivs_1d applies (EOS_type components are private). - kappa_shear_column: under #ifdef __NVCOMPILER_OPENMP_GPU, replace the calculate_density_derivs array call with an inline K=2..nzc loop calling calculate_density_derivs_elem_loc + the exact scale math of calculate_density_derivs_1d(dom=(/2,nzc/), scale=-g_R0). Non-GPU builds keep the original polymorphic call verbatim for all EOS forms. On GPU builds only ROQUET_RHO and WRIGHT have device _loc kernels; others FATAL. Verification: benchmark_ALE (ROQUET_RHO, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1) with the GPU path active on host. Reviewed; bitwise equivalence confirmed operation-by-operation, verdict ship. Follow-ups required by the offload step (increment 3), per review: hoist get_EOS_form_and_scaling + the form FATAL out of kappa_shear_column to the driver (both are host-only and cannot run in a device region), passing the form/scaling in; and FATAL on non-Boussinesq on GPU builds (that branch still uses polymorphic EOS calls). Untested-by-benchmark but code-identical: the non-unity unit-rescale branch and the WRIGHT dispatcher case. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…d prep) Prerequisite for offloading the column solver (increment 3): get_EOS_form_and_scaling and MOM_error are host-only and cannot run inside a device region, so move them out of kappa_shear_column up into both drivers (Calculate_kappa_shear and Calc_kappa_shear_vertex). - Resolve eos_form + the four unit-scaling factors once per driver call (guarded by use_temperature), and thread them into kappa_shear_column as new intent(in) args. - Move the "unsupported EOS form on GPU" FATAL to the drivers, and add a FATAL on non-Boussinesq on GPU builds (that branch still uses polymorphic EOS calls). - Add the five hoisted scalars to the shared() clause of both drivers' !$OMP parallel do (they are loop-invariant and read-only in the region; default(private) would leave them uninitialized inside the loop -- caught by the bitwise check). No functional change: same values, computed once per call instead of once per column. Verification: benchmark_ALE (vertex driver, ROQUET_RHO, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1). Reviewed; verdict ship, bitwise safety confirmed by inspection for both drivers. Not runtime-exercised by the benchmark (inspection-only): the h-point driver, the non-unity unit-scaling branch. Notes for the offload step: the driver FATAL guards execution not compilation, so declare target on kappa_shear_column will require #ifdef-excluding the non-Boussinesq branch; CS (pointer dummy) and tv (pointer-laden) must be demoted/pre-extracted before declare target. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Prepares the per-column solver for offload: the three workers (kappa_shear_column, find_kappa_tke, calculate_projected_state) are now !$omp declare target, so the driver column loop can run inside a target region in the next step. Device code is generated for all three (confirmed by -Minfo). Execution is still on the host (the column loop is still !$OMP parallel do), so this is bitwise-neutral. To make them device-compilable: - Demote CS from `pointer` to `type(Kappa_shear_CS), intent(in)` in kappa_shear_column and find_kappa_tke (the type is all scalars bar a diag pointer that the column math never touches). - Hoist use_temperature to a kappa_shear_column argument (resolved in the drivers) so the device path no longer reads the tv%T pointer via associated(); tv is then untouched on GPU builds. - #ifndef __NVCOMPILER_OPENMP_GPU around the non-Boussinesq branch's polymorphic calculate_specific_vol_derivs/calculate_density calls so they are excluded from device compilation (unreachable on GPU anyway -- the driver FATALs on non-Boussinesq). Notably the 26 early exit/return sites in the two solvers compile for the device without error on NVHPC 26.5 (the 25.11 exit-under-parallel miscompile appears fixed); their runtime correctness on device is gated by the bitwise check in the offload step, not pre-emptive rewrites. Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1). Next step flips the driver !$OMP parallel do to a target teams loop and adds residency; the solver then runs on the GPU. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Promote the per-J 2-D vertex slabs (u_2d..rho_2d) in Calc_kappa_shear_vertex to 3-D arrays (u_slab..rho_slab) computed in device do-concurrent passes ahead of the column loop. The loop bodies are verbatim except that I_hwt is inlined as a reciprocal multiply and the use_temperature branch is hoisted so each pass is purely elementwise; answers are bitwise-identical (verified np=1 against repro_ocean.stats on benchmark_ALE, V100/NVHPC 26.5). The key residency fix relative to the earlier reverted attempt: u_in/v_in and T_in/S_in (tv%T, tv%S) are already device-resident (mapped in MOM.F90), so a map(to:) on them does not copy - and the diabatic stack modifies T/S on the host before set_diffusivity, leaving the device copies stale. They now get an explicit "target update to" refresh, which resolves the NaN-diffusivity crash the first attempt hit. dz_3d is a fresh local each call, so its map(to:) does copy. h_at_u/h_at_v are now consumed on the device by the slab passes, so increment 1's copy-back is dropped; the slabs are copied back for the (still host) per-column solver loop, which reads them as X_slab(I,J,k). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Convert Calc_kappa_shear_vertex's per-column loop from a host OMP parallel do over J into an "omp target teams loop collapse(2)" over (J,I), calling the declare-target kappa_shear_column on the device. Bitwise-identical (np=1, benchmark_ALE vs repro_ocean.stats, V100/NVHPC 26.5), which also demonstrates that the solver's 26 early exit/return sites are runtime-correct on 26.5. Restructuring, all bitwise-neutral: - The per-J staging slices kappa_2d/tke_2d become device-resident 3-D staging arrays kappa_3d/tke_3d (map(alloc:)), written per column and consumed by new elementwise do-concurrent write-back passes for kappa_vertex/tke_io/kv_io. - The surface pressure is interpolated to vertices on the host ahead of the region (surface_pres_2d, verbatim expressions), keeping the possibly-null p_surf pointer and the CS%psurf_bug branch out of device code. - The h_vert fill stays on the host (it is only consumed by the host tracer-point averaging when VS_ThicknessMean=true), fed by a guarded update from(h_slab). - All per-column scratch is private at the loop; loop-invariant scalars are firstprivate; CS gets a per-call shallow map(to:) in the driver (all scalars plus a diag pointer that device code never dereferences). kappa_vertex, tke_io and kv_io are mapped to: (not alloc:) so host-set halo/initial values survive the full-array update from. diag_N2/S2_init/mean are mapped after their conditional host zeroing and copied back under one guard matching their consumers (post_data/Bchksum). Known limitation, to be fixed next: the solver's automatic arrays (sized nzc or SZK_(GV)) are device-heap allocated, so runs currently need NVCOMPILER_ACC_CUDA_HEAPSIZE=1073741824 (the default heap is too small: "DEVICE FORTRAN AUTO ALLOCATION FAILED"). Fixed-size locals will remove this. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… allocs)
The vertex driver's per-column private scratch and the column routines'
automatic locals are sized by GV%ke/nzc, which is not a compile-time constant,
so nvfortran allocated them from the device heap on every column call
(NVFORTRAN-W-0155). That exhausted the default device heap ("DEVICE FORTRAN
AUTO ALLOCATION FAILED" unless NVCOMPILER_ACC_CUDA_HEAPSIZE was raised) and
serialized columns on the device allocator (~2x total model slowdown).
In GPU builds (#ifdef __NVCOMPILER_OPENMP_GPU) these arrays now use a
compile-time ceiling, the new module parameter GPU_nk_max=128, so each device
thread gets stack ("local memory") arrays; kappa_shear_init fatals if
GV%ke > GPU_nk_max when the scheme is active. CPU builds keep the exact-size
declarations. No code queries these arrays' sizes and all accesses are
indexed to nzc/nz, so the enlargement is bitwise-neutral: verified np=1
against repro_ocean.stats at the default device heap, wall time back to the
pre-offload baseline (98s -> 45s), and the W-0155 warnings are gone.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…residency Review of the column-offload series found one latent halo bug: kv_io's actual argument (visc%Kv_shear_Bu) is already persistently device-resident, mapped in set_visc_init BEFORE the restart-reproducibility pass_var halo update, so this routine's map(to: kv_io) copies nothing and the full-array "update from" wrote stale pre-pass_var device halos over the corrected host halos - silently resurrecting the restart non-reproducibility in non-symmetric mode that the pass_var was added to fix. Invisible to a cold-start np=1 run. Fixed with an explicit "update to(tke_io, kv_io)" after the maps (tke_io defensively, so the routine does not depend on it staying unmapped elsewhere). Also corrects the increment-2 residency comment and drops the redundant "update to(T_in, S_in)": the actuals are the caller's convection-filtered T_f/S_f, fresh host locals each call, so their map(to:) already copies (they are not tv%T/tv%S as the comment claimed). This also establishes that the earlier reverted slab attempt's NaN was not T/S staleness; the operative fix was restructuring the interpolation (hoisting the use_temperature branch out of the do concurrent). And the GPU_nk_max FATAL message now notes it covers the shared tracer-point column routines too. Bit-identical at np=1 (benchmark_ALE vs repro_ocean.stats, V100/NVHPC 26.5). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
JorgeG94
force-pushed
the
feat/test_kappa_shear
branch
from
July 27, 2026 03:11
4fe5195 to
eaf1762
Compare
JorgeG94
marked this pull request as draft
August 6, 2026 22:40
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Threw the skill at kappa shear, yolo (will read the code it is just the weekend)