From 682d80a48eb5c6e7efb6ab87f9e8b62061e35da7 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Wed, 16 Sep 2026 11:10:27 +0200 Subject: [PATCH 1/7] Implement fused surface hydrology compute_auxiliary! --- .../2026-08/2026-08-31-PLAN_fused_kernels.md | 177 ++++++++++++------ .../canopy_interception.jl | 3 + .../bare_ground_evaporation.jl | 43 ++++- .../canopy_evapotranspiration.jl | 36 +++- src/processes/surface/surface_hydrology.jl | 43 ++++- test/surface/hydrology/integration_tests.jl | 91 +++++++++ .../hydrology/surface_hydrology_tests.jl | 4 + 7 files changed, 321 insertions(+), 76 deletions(-) create mode 100644 test/surface/hydrology/integration_tests.jl diff --git a/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md b/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md index b3982d2a0..f2d65860d 100644 --- a/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md +++ b/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md @@ -1,13 +1,16 @@ # Fused kernel launches for coupled soil, surface, and vegetation processes -> Status: **in progress** (approved; Phase 1 implemented, Phases 2–4 planned). Replace the per-process -> `compute_auxiliary!` / `compute_tendencies!` kernel launches inside the coupled process types -> (`SoilEnergyWaterCarbon`, `SurfaceHydrology`, `VegetationCarbonCycle`) with a single fused kernel -> each, following the pattern established for `SurfaceEnergyBalance` (commit `2fc7af72a`). Per-process -> launches are kept for standalone use and testing. Phase 1 first relocates the `surface_excess_water` -> prognostic from soil hydrology to surface runoff so the soil fusion in Phase 2 is clean — rev 4 -> found this also connects a currently dead tendency path (the pool has no drain in the coupled -> `LandModel` today), not just a scaling fix, which changes what Phase 1's test coverage must include. +> Status: **in progress** (rev 7: **soil fusion abandoned** after benchmarking; scope narrowed to +> **auxiliary-only** fusion). Phase 1 (relocate `surface_excess_water` to runoff) is merged (PR #188). +> The original goal — fuse the per-process `compute_auxiliary!` **and** `compute_tendencies!` launches +> inside the coupled process types (`SoilEnergyWaterCarbon`, `SurfaceHydrology`, `VegetationCarbonCycle`) +> into one kernel each, following `SurfaceEnergyBalance` (commit `2fc7af72a`) — is now **restricted to +> `compute_auxiliary!` only**. Benchmarking and register analysis (rev 6) established that fusing +> *tendencies* bundles independent output fields and raises register pressure (an occupancy regression +> on GPU), whereas fusing *auxiliaries* chains dependent stages and is a genuine win. The soil +> auxiliary pass is a single launch (energy/bgc auxiliaries are no-ops), so **soil is dropped entirely**. +> **Phase 3 (`SurfaceHydrology` auxiliary fusion) is implemented** (rev 7). The remaining target is the +> auxiliary chain in `VegetationCarbonCycle` (Phase 4). Date of initial draft: 2026-08-31 @@ -102,6 +105,71 @@ Base revision: b005a836aa2c26b305cc0faa304ba7748e1e4476 > runoff routes the excess into the pool (Case 1b), while the standalone call still discards. > Plus standalone runoff-tendency tests (`-min(D, S)`, cap binding, empty pool) and the coupled > `LandModel` `timestep!` regression test (pool decays as `S₀(1 − Δt/τ)ⁿ`). +> +> 2026-09-16 (rev 6) — **Soil fusion (Phase 2) abandoned; scope narrowed to auxiliary-only fusion.** +> Phase 2 was implemented and benchmarked (exclusive H100 / compute node, min-of-5 sustained repeats; +> a null-vs-null control established that single-repeat timings are dominated by spikes and cold GPU +> clocks). Results and root-cause analysis: +> +> - **CPU: fusion wins** (soil `RichardsEq` +27–61% ms/step; null control `NoFlow` +4–8%). Fewer +> launches dominate on CPU. +> - **GPU: fusion wins at launch-bound sizes (+3–7%) but regresses ~18% at the largest grid** (73728 +> columns). The null control is flat there, so the regression is real, not noise. +> - **Root cause = register pressure (measured, not assumed).** `ncu` is blocked on this cluster +> (`ERR_NVGPUCTRPERM`, driver-level), so registers were read from the exact PTX each launch compiles +> (`CUDA.@device_code_ptx` → `ptxas -c -v`, sm_86): fused tendency = **80 registers** vs standalone +> energy = 72, standalone water = 44. 80 regs → 25 warps/SM vs 72 → 28 (64K regs/SM). At large grids +> the GPU is occupancy-bound, so ~11% fewer resident warps to hide memory latency ≈ the ~18% loss. +> The fused *auxiliary* kernel is a GPU no-op (36 regs = standalone hydraulics 36): it wraps a single +> launch, since energy/bgc auxiliaries are no-ops. +> - **Structural rule (validated against Oceananigans' tendency kernels).** Oceananigans launches one +> kernel per *output field* and fuses only terms that **sum into a single accumulator** +> (`u_velocity_tendency`); it never bundles two independent output equations. The rule that predicts +> every measurement here: +> > **Fuse a dependency chain or a single accumulator; do not bundle independent outputs.** +> A chain/single-accumulator has live set ≈ `max(stage)` (each stage's temporaries retire before the +> next); bundling independent outputs has live set ≈ `union` (neither half's temporaries are released) +> — exactly the soil water+energy 80 = union(44, 72). `@noinline` on one half recovers occupancy +> (80→72, confirmed) but is a compiler workaround for a structural mismatch, not the pattern. +> - **Soil water and energy are two independent 3D output fields** — the worst possible fusion target, +> and the one Phase 2 fused. The pre-fusion per-field structure is already the Oceananigans-aligned +> one. The CPU win does not justify a GPU regression rooted in bundling independent outputs. +> - **Consequence for scope.** Tendency fusion is dropped from Phases 3 and 4 as well: the surface +> hydrology tendencies (canopy water, surface excess water) and vegetation tendencies (C, ν, GDD) are +> likewise *independent* outputs, so they carry the same union risk (the earlier "clean win" label on +> Phase 4 tendencies was wrong). Auxiliary fusion is retained for Phases 3–4 because those passes are +> genuine dependency chains on launch-bound (2D) grids — the case where fusion is correct on both +> architectures. The soil Phase 2 design is retained below for reference; it was implemented and +> documented in a closed PR and is **not** to be merged. +> +> 2026-09-16 (rev 7) — **Phase 3 (surface hydrology auxiliary fusion) implemented.** The three XY +> auxiliary launches (canopy → ET → runoff) collapse into one fused `compute_auxiliary_kernel!` +> dispatching on `SurfaceHydrology`, calling the per-cell mutating variants in dependency order. +> Deviations / findings during implementation: +> - **Uniform ET entry point.** A new `compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, +> evtr, interception, constants, atmos, soil, vegetation, snow)` kernel function is defined for both +> ET schemes (PALADYN canopy + bare-ground) with a shared signature; the standalone +> `compute_auxiliary_kernel!` for each scheme now just calls it, so the fused and standalone paths +> cannot drift. `NoCanopyInterception` gains a no-op `compute_canopy_auxiliary!` so the fused kernel +> dispatches uniformly (its `rainfall_ground` is a lazy passthrough, never written). +> - **`out` excludes lazy fields.** The fused host collects `auxiliary_fields(state, interception, +> evapotranspiration, surface_runoff)` but `filter`s out `FunctionField`s — e.g. +> `NoCanopyInterception`'s `rainfall_ground` passthrough — which no kernel writes and which the debug +> hook (`checkfinite!` → `parent`) cannot inspect. Full `fields` (no `except`) is passed so the chain +> reads prior stages' outputs from the aliased `Field`s. +> - **Latent bug fixed: canopy ET never saw `snow`.** The per-process canopy-ET `compute_auxiliary!` +> accepted `snow` but its `launch!` forwarded only up to `vegetation`, so `snow` fell into `args...` +> and the kernel defaulted to `snow = nothing` — canopy ground/canopy evaporation and transpiration +> were **never** scaled by the snow-free fraction, unlike bare-ground ET (which did pass `snow`). The +> fused kernel passes `snow` to both schemes; the standalone PALADYN host method was fixed to match +> (merge the snow fields, forward `snow`). This is a deliberate numerics change for the +> vegetated-under-snow config (ET now suppressed over snow-covered ground, as it should be); the +> fused and per-process paths agree bit-for-bit after the fix. Covered by a new test asserting +> `evaporation_ground/transpiration/evaporation_canopy == (1 − f_snow) ·` (snow-free flux). +> - **Tests.** New `test/surface/hydrology/fused_auxiliary_tests.jl`: (i) fused vs per-process fan-out +> over all 12 surface-hydrology auxiliaries, vegetated + snow, to machine precision; (ii) the +> snow-free-fraction scaling above. Existing `test/surface/*`, `test/coupled_models/land_model_tests.jl` +> (vegetated, bare, snow, meltwater, latent-partition, thin-snow, excess-water) all green. ## Problem description @@ -240,7 +308,13 @@ being drained at all before this phase. - New `test/surface/runoff_tests.jl`: standalone `DirectSurfaceRunoff` tendency equals `min(∂S∂t, S)`; coupled soil oversaturation fills the runoff-owned pool via `closure!`. -## Phase 2 — Fuse soil (`SoilEnergyWaterCarbon`) +## Phase 2 — Fuse soil (`SoilEnergyWaterCarbon`) — **ABANDONED (rev 6)** + +> **Not to be merged.** Implemented and benchmarked, then abandoned for the reasons in rev 6: the soil +> tendency fusion bundles two independent 3D output fields (`saturation_water_ice`, `internal_energy`), +> raising register pressure 72→80 and costing ~18% at large grids on GPU; the soil auxiliary fusion is a +> GPU no-op (one launch, energy/bgc auxiliaries are no-ops). The design below is retained for reference +> only (it lives in a closed PR). Revisit fusion for the auxiliary *chains* in Phases 3–4 instead. After Phase 1, soil hydrology's only tendency is `saturation_water_ice` (Richards), and its only auxiliary is `hydraulic_conductivity` (Richards). Energy contributes one tendency (`internal_energy`) @@ -304,7 +378,7 @@ untouched (recomputed in `initialize!`/`closure!`, not `compute_auxiliary!`). Add `debughook!` methods for the two new fused kernels (`checkfinite!(out)`), mirroring the existing per-process hooks. -## Phase 3 — Fuse surface hydrology (`SurfaceHydrology`) +## Phase 3 — Fuse surface hydrology (`SurfaceHydrology`) — **IMPLEMENTED (rev 7)** `SurfaceHydrology` = `canopy_interception` + `evapotranspiration` + `surface_runoff`, all XY. @@ -318,54 +392,37 @@ previous one's outputs from the shared `fields` (which alias `out`): hydrology::SurfaceHydrology, constants, atmos, soil, vegetation, snow) i, j = @index(Global, NTuple) compute_canopy_auxiliary!(out, i, j, grid, fields, hydrology.canopy_interception, atmos) - compute_evapotranspiration_fluxes!(out, i, j, grid, fields, hydrology.evapotranspiration, constants, atmos, ...) + compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, hydrology.evapotranspiration, + hydrology.canopy_interception, constants, atmos, soil, vegetation, snow) compute_surface_runoff!(out, i, j, grid, fields, hydrology.surface_runoff, hydrology.canopy_interception, get_hydrology(soil), snow) - return nothing end ``` -The per-cell mutating variants already exist (`compute_canopy_auxiliary!`, -`compute_evapotranspiration_fluxes!`, `compute_surface_runoff!`); the ET variant may need a thin -tuple-mutating wrapper to match the `out`-first convention. Host `compute_auxiliary!` collects the union -of the three sub-processes' auxiliary outputs as `out` and passes full `fields`. - -### `compute_tendencies!` — one fused XY launch +The canopy (`compute_canopy_auxiliary!`) and runoff (`compute_surface_runoff!`) per-cell mutating +variants already existed. A new `compute_evapotranspiration_auxiliary!` (defined for both ET schemes with +a shared signature) wraps each scheme's conductance-then-flux sequence; the standalone +`compute_auxiliary_kernel!` for each scheme calls it too, so the fused and standalone paths cannot drift. +`NoCanopyInterception` gains a no-op `compute_canopy_auxiliary!`. Host `compute_auxiliary!` collects the +union of the three sub-processes' auxiliary outputs as `out` (dropping lazy `FunctionField`s) and passes +full `fields`. See rev 7 for the latent canopy-ET `snow` bug this surfaced and fixed. -After Phase 1, two independent XY tendencies: canopy water and surface excess water. +### `compute_tendencies!` — **out of scope (rev 6)** -```julia -@kernel inbounds = true function compute_tendencies_kernel!(tendencies, grid, fields, - hydrology::SurfaceHydrology, evapotranspiration) - i, j = @index(Global, NTuple) - compute_canopy_water_tendency!(tendencies, i, j, grid, fields, hydrology.canopy_interception, evapotranspiration) - compute_surface_excess_water_tendency!(tendencies, i, j, grid, fields, hydrology.surface_runoff) - return nothing -end -``` - -`compute_canopy_water_tendency!` already takes the `tendencies` tuple; the runoff variant (from Phase 1) -is adapted to the tuple form. ET has no tendency. +After Phase 1 there are two XY tendencies (canopy water, surface excess water). They are **independent +outputs** (neither feeds the other), so fusing them is the register-union anti-pattern from rev 6, not a +chain. Tendency fusion is therefore **dropped** for surface hydrology; the two per-process tendency +launches stay. (If ever revisited, capture registers first — see `scratch/capture_soil_regs.jl` — and +only fuse if the chain holds or `@noinline` recovers occupancy.) ## Phase 4 — Fuse vegetation carbon (`VegetationCarbonCycle`) -### `compute_tendencies!` — one fused XY launch (clean win) - -Three independent XY tendencies (`carbon_vegetation`, `vegetation_area_fraction`, `growing_degree_days`): - -```julia -@kernel inbounds = true function compute_tendencies_kernel!(tendencies, grid, fields, - veg::VegetationCarbonCycle, constants, atmos) - i, j = @index(Global, NTuple) - compute_veg_carbon_tendencies!(tendencies, i, j, grid, fields, veg.carbon_dynamics, veg.traits) - compute_ν_tendencies!(tendencies, i, j, grid, fields, veg.vegetation_dynamics, veg.carbon_dynamics, veg.traits) - compute_gdd_tendency!(tendencies, i, j, grid, fields, veg.phenology, atmos) - return nothing -end -``` +### `compute_tendencies!` — **out of scope (rev 6)** -`compute_veg_carbon_tendencies!`, `compute_ν_tendencies!`, and `compute_gdd_tendency!` already take the -`tend` tuple. `vegetation_dynamics` may be `nothing` (`PrescribedVegetation`) — a `nothing` no-op variant -handles it. +Three XY tendencies (`carbon_vegetation`, `vegetation_area_fraction`, `growing_degree_days`). An earlier +revision labelled this a "clean win"; rev 6 corrects that — these are **independent outputs** (three +separate prognostics, none feeding another), so fusing them is the register-union anti-pattern, not a +chain. Tendency fusion is **dropped** for vegetation; the three per-process tendency launches stay. Only +the auxiliary chain (below) is fused. ### `compute_auxiliary!` — partial fusion (in scope, rev 3) @@ -417,9 +474,10 @@ launches → 1 XY launch (PAW's XYZ launch and derived `compute!` unchanged). Enzyme-safety check for each fused kernel. - **Reactant.** `test/reactant/` (own env, not under `--check-bounds=yes`) still compiles every fused kernel — they inherit only throw-free, allocation-free kernel functions. -- **Launch-count spot check.** Default vegetated model: soil tendencies 2→1 XYZ, soil auxiliaries - 1→1 XYZ (NoFlow 1→0), surface hydrology auxiliaries 3→1 XY and tendencies 1→1 XY, vegetation - tendencies 3→1 XY and auxiliaries 6→1 XY (PAW's XYZ launch and derived `compute!` unchanged). +- **Launch-count spot check (auxiliary-only, rev 6).** Default vegetated model: surface hydrology + auxiliaries 3→1 XY; vegetation auxiliaries 6→1 XY (PAW's XYZ launch and derived `compute!` unchanged). + Soil is unchanged (Phase 2 abandoned); all `compute_tendencies!` launch counts are unchanged (tendency + fusion dropped). ## Documentation changes @@ -431,12 +489,10 @@ launches → 1 XY launch (PAW's XYZ launch and derived `compute!` unchanged). ## Known limitations -- Phase 2 auxiliary fusion currently only folds in hydrology (energy/bgc auxiliaries are no-ops today); - the auxiliary-side win is limited to dropping the `NoFlow` no-op launch until more processes diagnose - auxiliaries. -- Soil `compute_tendencies!` keeps the `(state, grid, soil, constants)` signature; ET enters the soil - water balance through boundary conditions, so no `evtr`/`runoff` is threaded into the fused soil - tendencies kernel. +- **Soil is not fused at all (rev 6).** Its auxiliary pass is a single launch (energy/bgc auxiliaries are + no-ops) and its tendency pass bundles independent outputs (register-union regression). See rev 6. +- **Tendency fusion is out of scope everywhere (rev 6).** Only dependency-chain auxiliary passes are + fused. Independent-output tendencies keep their per-process launches. - Phase 4 auxiliary fusion is partial by construction: `plant_available_water` remains a separate XYZ launch (it must precede the fused XY kernel, which reads its derived `soil_moisture_limiting_factor`), and `root_distribution` stays a lazy `FunctionField` with no launch. @@ -460,3 +516,10 @@ launches → 1 XY launch (PAW's XYZ launch and derived `compute!` unchanged). `plant_available_water` in separate launches from the fused XY kernel. The photosynthesis / stomatal-conductance ordering is not an obstacle (photosynthesis uses only `compute_λc(stomcond, vpd)`, a parameter call, while stomatal conductance reads photosynthesis's output). + +## Decisions (rev 6) + +4. **Auxiliary-only fusion.** Fuse dependency-chain `compute_auxiliary!` passes; do **not** fuse + `compute_tendencies!` (independent outputs → register-union occupancy regression on GPU). Soil is + dropped entirely (auxiliary is a single launch; tendency is the union anti-pattern). Rule: + *fuse a chain or a single accumulator, never independent outputs.* diff --git a/src/processes/surface/canopy_interception/canopy_interception.jl b/src/processes/surface/canopy_interception/canopy_interception.jl index 668b876c9..c4a0f563a 100644 --- a/src/processes/surface/canopy_interception/canopy_interception.jl +++ b/src/processes/surface/canopy_interception/canopy_interception.jl @@ -16,6 +16,9 @@ passthrough_rainfall(grid, clock, fields, ::NoCanopyInterception) = fields.rainf @inline compute_auxiliary!(state, grid, ::NoCanopyInterception, args...) = nothing +# No-op per-cell variant: `rainfall_ground` is a lazy passthrough of `rainfall`, never written. +@inline compute_canopy_auxiliary!(out, i, j, grid, fields, ::NoCanopyInterception, args...) = nothing + @inline compute_tendencies!(state, grid, ::NoCanopyInterception, args...) = nothing @propagate_inbounds canopy_water(i, j, grid, fields, ::NoCanopyInterception) = zero(eltype(grid)) diff --git a/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl b/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl index 909f334be..7da62bd78 100644 --- a/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl +++ b/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl @@ -54,7 +54,7 @@ variables(::BareGroundEvaporation) = ( function compute_auxiliary!( state, grid, evaporation::BareGroundEvaporation, - ::NoCanopyInterception, + interception::NoCanopyInterception, constants::PhysicalConstants, atmos::AbstractAtmosphere, soil::Optional{AbstractSoil} = nothing, @@ -63,7 +63,7 @@ function compute_auxiliary!( out = auxiliary_fields(state, evaporation) # merge the snow cover fraction so the ground evaporation can be scaled by the snow-free fraction fields = merge(get_fields(state, evaporation, atmos, soil; except = out), get_fields(state, snow)) - launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evaporation, constants, atmos, soil, snow) + launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evaporation, interception, constants, atmos, soil, nothing, snow) return nothing end @@ -139,24 +139,49 @@ end out.evaporation_ground[i, j, 1] = E_gnd return out end + # Kernels +""" + $TYPEDSIGNATURES -@kernel inbounds = true function compute_auxiliary_kernel!( - out, grid, fields, - evapotranspiration::BareGroundEvaporation, +Compute and store the skin-driven ground evaporation conductance for bare-ground evaporation, then +evaluate the ground evaporation flux from it. The conductance is merged into `fields` so the flux +step reads it back without a global round-trip. The `interception` and `vegetation` arguments are +unused by this scheme. +""" +@propagate_inbounds function compute_evapotranspiration_auxiliary!( + out, i, j, grid, fields, + evaporation::BareGroundEvaporation, + interception::NoCanopyInterception, constants::PhysicalConstants, atmos::AbstractAtmosphere, soil::Optional{AbstractSoil} = nothing, + vegetation::Optional{AbstractVegetation} = nothing, snow::Optional{AbstractSnow} = nothing, + args... ) - i, j = @index(Global, NTuple) - # First compute conductances - compute_evapotranspiration_conductances!(out, i, j, grid, fields, evapotranspiration, constants, atmos, soil) + compute_evapotranspiration_conductances!(out, i, j, grid, fields, evaporation, constants, atmos, soil) # TODO: Annoyingly, we need to explicitly add these to `fields`; need a better solution to this problem conductances = (ground_evaporation_conductance = out.ground_evaporation_conductance,) fields = merge(fields, conductances) # Compute ET fluxes from stored conductances; `snow` scales ground evaporation by the snow-free fraction - compute_evapotranspiration_fluxes!(out, i, j, grid, fields, evapotranspiration, constants, atmos, snow) + compute_evapotranspiration_fluxes!(out, i, j, grid, fields, evaporation, constants, atmos, snow) + return out +end + +@kernel inbounds = true function compute_auxiliary_kernel!( + out, grid, fields, + evapotranspiration::BareGroundEvaporation, + interception::NoCanopyInterception, + constants::PhysicalConstants, + atmos::AbstractAtmosphere, + soil::Optional{AbstractSoil} = nothing, + vegetation::Optional{AbstractVegetation} = nothing, + snow::Optional{AbstractSnow} = nothing, + args... + ) + i, j = @index(Global, NTuple) + compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, snow, args...) end diff --git a/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl b/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl index 8101f24e8..305b0ddb9 100644 --- a/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl +++ b/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl @@ -1,7 +1,7 @@ """ $TYPEDEF -Canopy evapotranspiration scheme from PALADYN ([willeitPALADYNV10Comprehensive2016; Eq. (5)](@cite)) +Canopy evapotranspiration scheme from PALADYN ([willeitPALADYNV10Comprehensive2016; Eq. (5)](@cite)) that includes a canopy evaporation term based on the saturation fraction of canopy water defined by the canopy hydrology scheme. @@ -120,11 +120,13 @@ function compute_auxiliary!( atmos::AbstractAtmosphere, soil::AbstractSoil, vegetation::AbstractVegetation, + snow::Optional{AbstractSnow} = nothing, args... ) out = auxiliary_fields(state, evapotranspiration) - fields = get_fields(state, evapotranspiration, interception, atmos, soil, vegetation; except = out) - launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evapotranspiration, interception, constants, atmos, soil, vegetation) + # merge the snow cover fraction so the ground/canopy evaporation fluxes can be scaled by the snow-free fraction + fields = merge(get_fields(state, evapotranspiration, interception, atmos, soil, vegetation; except = out), get_fields(state, snow)) + launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, snow) return nothing end @@ -275,8 +277,15 @@ end # Kernels -@kernel inbounds = true function compute_auxiliary_kernel!( - out, grid, fields, +""" + $TYPEDSIGNATURES + +Compute and store the skin-driven vapor conductances for canopy evapotranspiration, then evaluate +the partitioned humidity fluxes from them. The conductances are merged into `fields` so the flux +step reads them back without a global round-trip. +""" +@propagate_inbounds function compute_evapotranspiration_auxiliary!( + out, i, j, grid, fields, evapotranspiration::PALADYNCanopyEvapotranspiration, interception::AbstractCanopyInterception, constants::PhysicalConstants, @@ -286,7 +295,6 @@ end snow::Optional{AbstractSnow} = nothing, args... ) - i, j = @index(Global, NTuple) # First compute conductances compute_evapotranspiration_conductances!(out, i, j, grid, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, args...) # TODO: Annoyingly, we need to explicitly add these to `fields`; need a better solution to this problem @@ -298,4 +306,20 @@ end fields = merge(fields, conductances) # Compute ET fluxes from stored conductances compute_evapotranspiration_fluxes!(out, i, j, grid, fields, evapotranspiration, constants, atmos, snow) + return out +end + +@kernel inbounds = true function compute_auxiliary_kernel!( + out, grid, fields, + evapotranspiration::PALADYNCanopyEvapotranspiration, + interception::AbstractCanopyInterception, + constants::PhysicalConstants, + atmos::AbstractAtmosphere, + soil::AbstractSoil, + vegetation::AbstractVegetation, + snow::Optional{AbstractSnow} = nothing, + args... + ) + i, j = @index(Global, NTuple) + compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, snow, args...) end diff --git a/src/processes/surface/surface_hydrology.jl b/src/processes/surface/surface_hydrology.jl index 63b1b8c42..33d80d700 100644 --- a/src/processes/surface/surface_hydrology.jl +++ b/src/processes/surface/surface_hydrology.jl @@ -43,12 +43,47 @@ function compute_auxiliary!( snow::Optional{AbstractSnow} = nothing, args... ) - compute_auxiliary!(state, grid, hydrology.canopy_interception, atmos) + interception = hydrology.canopy_interception + evapotranspiration = hydrology.evapotranspiration + surface_runoff = hydrology.surface_runoff + # The fused kernel writes the union of the three sub-processes' auxiliaries. Drop lazy + # (`FunctionField`) auxiliaries — e.g. `NoCanopyInterception`'s `rainfall_ground` passthrough — + # which no kernel writes. + out = filter(v -> v isa Field, auxiliary_fields(state, interception, evapotranspiration, surface_runoff)) + # Full fields (no `except`): within a cell the kernel writes `out.foo` and a later sub-process + # reads `fields.foo` — the same `Field` object, so the write is visible. This is what lets the + # canopy → evapotranspiration → runoff dependency chain run in a single launch. + fields = get_fields(state, interception, evapotranspiration, surface_runoff, atmos, soil, vegetation, snow) + launch!(grid, XY, compute_auxiliary_kernel!, out, fields, hydrology, constants, atmos, soil, vegetation, snow) + return nothing +end + +""" + $TYPEDSIGNATURES + +Fused auxiliary kernel for the coupled surface hydrology processes. Each sub-process's per-cell +mutating variant runs in dependency order — canopy interception, then evapotranspiration (which reads +the canopy saturation fraction), then surface runoff (which reads the ground rainfall) — so the chain +resolves within a single launch. A "prescribed"/no-op scheme contributes a no-op variant. +""" +@kernel inbounds = true function compute_auxiliary_kernel!( + out, grid, fields, + hydrology::SurfaceHydrology, + constants::PhysicalConstants, + atmos::AbstractAtmosphere, + soil::Optional{AbstractSoil} = nothing, + vegetation::Optional{AbstractVegetation} = nothing, + snow::Optional{AbstractSnow} = nothing, + args... + ) + i, j = @index(Global, NTuple) + compute_canopy_auxiliary!(out, i, j, grid, fields, hydrology.canopy_interception, atmos) # `snow` lets the (bare-ground) evaporation scheme scale ground evaporation by the snow-free fraction - compute_auxiliary!(state, grid, hydrology.evapotranspiration, hydrology.canopy_interception, constants, atmos, soil, vegetation, snow) + compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, hydrology.evapotranspiration, + hydrology.canopy_interception, constants, atmos, soil, vegetation, snow) # `snow` makes the surface runoff scheme's water input snow-aware (meltwater + bare-ground throughfall) - compute_auxiliary!(state, grid, hydrology.surface_runoff, hydrology.canopy_interception, soil, snow) - return nothing + compute_surface_runoff!(out, i, j, grid, fields, hydrology.surface_runoff, + hydrology.canopy_interception, get_hydrology(soil), snow) end """ $TYPEDSIGNATURES """ diff --git a/test/surface/hydrology/integration_tests.jl b/test/surface/hydrology/integration_tests.jl new file mode 100644 index 000000000..7aa10152a --- /dev/null +++ b/test/surface/hydrology/integration_tests.jl @@ -0,0 +1,91 @@ +using Terrarium +using Test + +# The coupled `SurfaceHydrology.compute_auxiliary!` fuses the canopy-interception, +# evapotranspiration, and surface-runoff auxiliaries into a single launch. These tests check that the +# fused kernel reproduces the per-process fan-out (the pre-fusion path) to machine precision, and that +# the canopy evapotranspiration fluxes are scaled by the snow-free fraction (the snow coupling the +# per-process canopy-ET launch used to drop). + +# Surface hydrology auxiliary fields written by the fused kernel (canopy + ET + runoff). +const SURFACE_HYDROLOGY_AUXILIARIES = ( + :canopy_water_interception, :canopy_water_removal, :saturation_canopy_water, :rainfall_ground, + :ground_evaporation_conductance, :canopy_evaporation_conductance, :transpiration_conductance, + :evaporation_canopy, :evaporation_ground, :transpiration, + :surface_runoff, :infiltration, +) + +function vegetated_snow_land(NF) + grid = ColumnGrid(CPU(), ExponentialSpacing(Δz_max = 1.0, N = 50)) + soil = SoilEnergyWaterCarbon(NF; hydrology = SoilHydrology(NF, RichardsEq())) + vegetation = VegetationCarbonCycle(NF) + snow = SingleLayerSnow(NF) + land = LandModel(grid; soil, vegetation, snow) + initializers = ( + temperature = (x, z) -> 2.0 - 0.02 * z, + saturation_water_ice = (x, z) -> min(1, 0.8 - 0.05 * z), + carbon_vegetation = 0.1, + snow_water_equivalent = 0.2, + snow_temperature = -2.0, + ) + integrator = initialize(land; initializers) + set!(integrator.state.rainfall, 1.0e-7) + Terrarium.closure!(integrator.state, land) + # Populate atmosphere / soil / snow / vegetation auxiliaries (surface hydrology runs last). + compute_auxiliary!(integrator.state, land) + return land, integrator.state +end + +@testset "Fused surface hydrology auxiliary matches per-process fan-out" begin + land, state = vegetated_snow_land(Float64) + grid = get_grid(land) + hydrology = land.surface_hydrology + constants, atmos, soil, vegetation, snow = land.constants, land.atmosphere, land.soil, land.vegetation, land.snow + + # Fused path: a single launch over the coupled type. + fused = deepcopy(state) + compute_auxiliary!(fused, grid, hydrology, constants, atmos, soil, vegetation, snow) + + # Per-process path: the three standalone launches the fan-out used to issue, in dependency order. + per_process = deepcopy(state) + compute_auxiliary!(per_process, grid, hydrology.canopy_interception, atmos) + compute_auxiliary!(per_process, grid, hydrology.evapotranspiration, hydrology.canopy_interception, + constants, atmos, soil, vegetation, snow) + compute_auxiliary!(per_process, grid, hydrology.surface_runoff, hydrology.canopy_interception, soil, snow) + + @testset "$(name)" for name in SURFACE_HYDROLOGY_AUXILIARIES + fused_vals = Array(interior(getproperty(fused, name))) + per_process_vals = Array(interior(getproperty(per_process, name))) + @test all(isfinite.(fused_vals)) + @test fused_vals ≈ per_process_vals + end +end + +@testset "Canopy ET fluxes are scaled by the snow-free fraction" begin + # With a snow-covered surface, the canopy scheme's ground/canopy evaporation and transpiration must + # be scaled by (1 − f_snow) — the coupling the per-process canopy-ET launch used to drop. Recomputing + # the auxiliaries with the snow cover fraction forced to zero must recover the unscaled fluxes. + land, state = vegetated_snow_land(Float64) + grid = get_grid(land) + hydrology = land.surface_hydrology + constants, atmos, soil, vegetation, snow = land.constants, land.atmosphere, land.soil, land.vegetation, land.snow + + @test all(Array(interior(state.snow_cover_fraction)) .> 0) # snow is present + + covered = deepcopy(state) + compute_auxiliary!(covered, grid, hydrology, constants, atmos, soil, vegetation, snow) + + # Same state with the snow cover fraction zeroed: the fused kernel must then leave the ET fluxes + # unscaled, i.e. the covered fluxes equal the bare fluxes times (1 − f_snow). + bare = deepcopy(state) + set!(bare.snow_cover_fraction, 0.0) + compute_auxiliary!(bare, grid, hydrology, constants, atmos, soil, vegetation, snow) + + f = Array(interior(state.snow_cover_fraction)) + @test all(f .< 1) # partially covered, so the scaling is nontrivial + for name in (:evaporation_ground, :transpiration, :evaporation_canopy) + covered_vals = Array(interior(getproperty(covered, name))) + bare_vals = Array(interior(getproperty(bare, name))) + @test all(isapprox.(covered_vals, (1 .- f) .* bare_vals; rtol = 1.0e-9)) + end +end diff --git a/test/surface/hydrology/surface_hydrology_tests.jl b/test/surface/hydrology/surface_hydrology_tests.jl index bff0fda77..44885b70c 100644 --- a/test/surface/hydrology/surface_hydrology_tests.jl +++ b/test/surface/hydrology/surface_hydrology_tests.jl @@ -12,3 +12,7 @@ end @testset "Surface runoff" begin include("surface_runoff_tests.jl") end + +@testset "Fused surface hydrology" begin + include("integration_tests.jl") +end From d483a7a90f418523c7f2f60218009ec7c8d5beb2 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Wed, 16 Sep 2026 17:24:04 +0200 Subject: [PATCH 2/7] Fuse vegetation carbon cycle auxiliaries Collapse the five XY auxiliary stages of VegetationCarbonCycle (carbon dynamics -> phenology -> photosynthesis -> stomatal conductance -> autotrophic respiration) into a single fused compute_auxiliary_kernel\!, following the SurfaceHydrology pattern. Each stage's per-cell mutating variant already existed; the fused kernel calls them in dependency order with full fields so a stage's write is visible to the next. PAW stays a preceding XYZ launch (it materializes soil_moisture_limiting_factor that photosynthesis and stomatal conductance read). New test/vegetation/integration_tests.jl checks fused vs per-process fan-out to machine precision across all nine written auxiliaries. --- .../2026-08/2026-08-31-PLAN_fused_kernels.md | 49 +++++++++++-- .../vegetation/vegetation_carbon_cycle.jl | 67 ++++++++++++------ test/vegetation/integration_tests.jl | 68 +++++++++++++++++++ test/vegetation/vegetation_model_tests.jl | 4 ++ 4 files changed, 163 insertions(+), 25 deletions(-) create mode 100644 test/vegetation/integration_tests.jl diff --git a/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md b/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md index f2d65860d..b25179dcf 100644 --- a/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md +++ b/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md @@ -9,8 +9,11 @@ > *tendencies* bundles independent output fields and raises register pressure (an occupancy regression > on GPU), whereas fusing *auxiliaries* chains dependent stages and is a genuine win. The soil > auxiliary pass is a single launch (energy/bgc auxiliaries are no-ops), so **soil is dropped entirely**. -> **Phase 3 (`SurfaceHydrology` auxiliary fusion) is implemented** (rev 7). The remaining target is the -> auxiliary chain in `VegetationCarbonCycle` (Phase 4). +> **Phase 3 (`SurfaceHydrology` auxiliary fusion) is implemented** (rev 7). **Phase 4 +> (`VegetationCarbonCycle` auxiliary fusion) is implemented** (rev 9): the five XY carbon-cycle stages +> (carbon dynamics → phenology → photosynthesis → stomatal conductance → autotrophic respiration) collapse +> to one launch, with `plant_available_water` kept as a preceding XYZ launch. Snow (rev 8) was considered +> and left as-is — its auxiliary pass is already a single launch. Date of initial draft: 2026-08-31 @@ -170,6 +173,39 @@ Base revision: b005a836aa2c26b305cc0faa304ba7748e1e4476 > over all 12 surface-hydrology auxiliaries, vegetated + snow, to machine precision; (ii) the > snow-free-fraction scaling above. Existing `test/surface/*`, `test/coupled_models/land_model_tests.jl` > (vegetated, bare, snow, meltwater, latent-partition, thin-snow, excess-water) all green. +> +> 2026-09-16 (rev 8) — **Snow considered for fusion; left as-is.** An earlier revision folded the snow +> auxiliary diagnosis into the fused surface-hydrology launch (a leading `compute_snow_properties!` stage). +> That was reverted: snow's auxiliaries should not be computed inside surface hydrology. Re-examining the +> snow side, there is **nothing to fuse**: `SingleLayerSnow.compute_auxiliary!` is already a single XY launch +> (`compute_snow_properties!` → `snow_depth`, `snow_cover_fraction`). The only other snow XY launch is the +> energy closure `energy_to_temperature!` (→ `snow_temperature`, `snow_liquid_fraction`), which (a) runs in a +> **different pass** (`closure!`, at the start of each timestep, vs `compute_auxiliary!` at finalize) and +> (b) produces **independent outputs** — the two launches share no dependency chain, so fusing them would be +> the register-union anti-pattern the rev 6 rule forbids. Snow is therefore unchanged. +> +> 2026-09-16 (rev 9) — **Phase 4 (vegetation carbon auxiliary fusion) implemented.** The five real XY +> auxiliary stages of `VegetationCarbonCycle` — carbon dynamics → phenology → photosynthesis → stomatal +> conductance → autotrophic respiration — collapse from five launches to one fused XY launch, following the +> Phase 3 pattern exactly. Each stage's per-cell mutating variant already existed (they are the bodies of +> the standalone `compute_auxiliary_kernel!`s), so the fused kernel just calls them in dependency order; the +> host collects `out = filter(v -> v isa Field, auxiliary_fields(state, ))` and passes full +> `fields` (no `except`) so a stage's write to `out.foo` is visible to a later stage reading `fields.foo`. +> - **`plant_available_water` stays separate and first.** It is an XYZ launch that then materializes the +> derived XY `soil_moisture_limiting_factor` via `compute!`; photosynthesis and stomatal conductance read +> that factor, so PAW must run before the fused kernel — the same ordering the fan-out already used. +> `root_distribution` (lazy `FunctionField`, no-op) and `vegetation_dynamics` (no-op) calls are unchanged. +> - **No `soil` in the fused kernel.** None of the five stages reads a soil field; the only soil-derived +> input is the PAW-produced limiting factor, read as a plain XY input. The kernel signature is +> `(out, grid, fields, veg, constants, atmos, args...)`. +> - **No `PrescribedPhenology` no-op variant.** `PrescribedPhenology` is only ever paired with +> `PrescribedVegetation` (which has its own `compute_auxiliary!` and no autotrophic respiration), never +> with `VegetationCarbonCycle`, so the fused kernel never dispatches on it — adding a no-op would guard an +> unreachable state. +> - **Tests.** New `test/vegetation/integration_tests.jl` (included from `vegetation_model_tests.jl`) checks +> the fused path against the per-process fan-out to machine precision across all nine written auxiliaries. +> Vegetation suite 148/148, land_model 45/45 green. Net: vegetation auxiliary pass 6 XY → 1 XY (PAW's XYZ +> launch and derived `compute!` unchanged). ## Problem description @@ -414,7 +450,7 @@ chain. Tendency fusion is therefore **dropped** for surface hydrology; the two p launches stay. (If ever revisited, capture registers first — see `scratch/capture_soil_regs.jl` — and only fuse if the chain holds or `@noinline` recovers occupancy.) -## Phase 4 — Fuse vegetation carbon (`VegetationCarbonCycle`) +## Phase 4 — Fuse vegetation carbon (`VegetationCarbonCycle`) — **IMPLEMENTED (rev 9)** ### `compute_tendencies!` — **out of scope (rev 6)** @@ -440,7 +476,7 @@ soil-dependent ones kept separate (rev 3): ```julia @kernel inbounds = true function compute_auxiliary_kernel!(out, grid, fields, - veg::VegetationCarbonCycle, constants, atmos, soil) + veg::VegetationCarbonCycle, constants, atmos, args...) i, j = @index(Global, NTuple) compute_veg_carbon_auxiliary!(out, i, j, grid, fields, veg.carbon_dynamics, veg.traits) compute_phenology!(out, i, j, grid, fields, veg.phenology, atmos) @@ -449,10 +485,13 @@ soil-dependent ones kept separate (rev 3): compute_photosynthesis!(out, i, j, grid, fields, veg.photosynthesis, veg.stomatal_conductance, veg.traits, constants, atmos) compute_stomatal_conductance!(out, i, j, grid, fields, veg.stomatal_conductance, veg.traits, constants, atmos) compute_autotrophic_respiration!(out, i, j, grid, fields, veg.autotrophic_respiration, veg.carbon_dynamics, veg.phenology, veg.traits, atmos) - return nothing end ``` +(The kernel takes no `soil`: none of the five stages reads a soil field — the only soil-derived input, +`soil_moisture_limiting_factor`, is materialized by PAW beforehand and read as a plain XY input. A `return` +is not permitted inside a `@kernel`, so the body ends on the last stage call.) + The photosynthesis → stomatal-conductance ordering is a clean forward dependency (rev 3): photosynthesis does not read stomatal's output field, so chaining them in one kernel is correct. `vegetation_dynamics` has no auxiliary (`compute_auxiliary!` is a no-op) and is omitted. The host `compute_auxiliary!` collects diff --git a/src/processes/vegetation/vegetation_carbon_cycle.jl b/src/processes/vegetation/vegetation_carbon_cycle.jl index 960d5ee5c..9d5900d32 100644 --- a/src/processes/vegetation/vegetation_carbon_cycle.jl +++ b/src/processes/vegetation/vegetation_carbon_cycle.jl @@ -87,36 +87,63 @@ function compute_auxiliary!( soil::Optional{AbstractSoil} = nothing, args... ) - # Compute auxiliary variables for each component - # Roots: need soil state and computes root_fraction + # Roots: need soil state and computes root_fraction (a lazy `FunctionField`, no launch) compute_auxiliary!(state, grid, veg.root_distribution, soil) - # PAW: needs soil saturation profile and computes soil_moisture_limiting_factor + # PAW: needs soil saturation profile and materializes soil_moisture_limiting_factor (a derived field). + # It runs before the fused kernel below, whose photosynthesis and stomatal conductance stages read it. compute_auxiliary!(state, grid, veg.plant_available_water, soil) - # Veg. carbon dynamics: needs C_veg(t) and computes LAI_b(t) - compute_auxiliary!(state, grid, veg.carbon_dynamics, veg.traits) - - # Phenology: needs LAI_b(t) and air temperature(t) and computes LAI(t) and phen(t) - compute_auxiliary!(state, grid, veg.phenology, veg.carbon_dynamics, atmos) - - # Photosynthesis: needs atm. inputs(t), LAI(t), and computes Rd(t) and GPP(t) - # N.B. We break the usual dependency pattern here to resolve the tight coupling between photosynthesis and - # stomatal conductance. Photosynthesis does *not* depend on the auxiliary state of stomatal conductance but - # requires its parameters to compute λc. - compute_auxiliary!(state, grid, veg.photosynthesis, veg.stomatal_conductance, veg.traits, constants, atmos) - - # Stomatal conductance: needs atm. inputs(t) and computes g_can(t) - compute_auxiliary!(state, grid, veg.stomatal_conductance, veg.traits, constants, atmos) - - # Autotrophic respiration: needs atm. inputs(t), GPP(t), Rd(t), C_veg(t), phen(t) and computes Ra(t) and NPP(t) - compute_auxiliary!(state, grid, veg.autotrophic_respiration, veg.carbon_dynamics, veg.phenology, veg.traits, atmos) + # The carbon-cycle chain — carbon dynamics → phenology → photosynthesis → stomatal conductance → + # autotrophic respiration — is fused into a single launch. Each stage reads the previous stage's + # output within a cell, so the dependency chain resolves without returning to the host. + carbon_dynamics = veg.carbon_dynamics + phenology = veg.phenology + photosynthesis = veg.photosynthesis + stomatal_conductance = veg.stomatal_conductance + autotrophic_respiration = veg.autotrophic_respiration + out = filter(v -> v isa Field, auxiliary_fields(state, carbon_dynamics, phenology, photosynthesis, + stomatal_conductance, autotrophic_respiration)) + # Full fields (no `except`): within a cell the kernel writes `out.foo` and a later stage reads + # `fields.foo` — the same `Field` object, so the write is visible to the stages below. + fields = get_fields(state, carbon_dynamics, phenology, photosynthesis, stomatal_conductance, + autotrophic_respiration, atmos) + launch!(grid, XY, compute_auxiliary_kernel!, out, fields, veg, constants, atmos) # Note: vegetation_dynamics compute_auxiliary! does nothing for now compute_auxiliary!(state, grid, veg.vegetation_dynamics) return nothing end +""" + $TYPEDSIGNATURES + +Fused auxiliary kernel for the vegetation carbon cycle. Each component's per-cell mutating variant runs in +dependency order — carbon dynamics (a pure producer of `balanced_leaf_area_index` from the vegetation carbon +pool), then phenology (which reads it to set `leaf_area_index` and `phenology_factor`), then photosynthesis +(which reads `leaf_area_index` and the soil moisture limiting factor to set `net_assimilation` and +`gross_primary_production`), then stomatal conductance (which reads `net_assimilation` to set +`canopy_water_conductance`), then autotrophic respiration (which reads `gross_primary_production` and +`phenology_factor` to set `net_primary_production`) — so the chain resolves within a single launch. +""" +@kernel inbounds = true function compute_auxiliary_kernel!( + out, grid, fields, + veg::VegetationCarbonCycle, + constants::PhysicalConstants, + atmos::AbstractAtmosphere, + args... + ) + i, j = @index(Global, NTuple) + compute_veg_carbon_auxiliary!(out, i, j, grid, fields, veg.carbon_dynamics, veg.traits) + compute_phenology!(out, i, j, grid, fields, veg.phenology, atmos) + # Photosynthesis reads stomatal conductance's parameters (λc) but not its auxiliary state + compute_photosynthesis!(out, i, j, grid, fields, veg.photosynthesis, veg.stomatal_conductance, + veg.traits, constants, atmos) + compute_stomatal_conductance!(out, i, j, grid, fields, veg.stomatal_conductance, veg.traits, constants, atmos) + compute_autotrophic_respiration!(out, i, j, grid, fields, veg.autotrophic_respiration, + veg.carbon_dynamics, veg.phenology, veg.traits, atmos) +end + """ $TYPEDSIGNATURES diff --git a/test/vegetation/integration_tests.jl b/test/vegetation/integration_tests.jl new file mode 100644 index 000000000..a1a8cd8bb --- /dev/null +++ b/test/vegetation/integration_tests.jl @@ -0,0 +1,68 @@ +using Terrarium +using Test + +# The coupled `VegetationCarbonCycle.compute_auxiliary!` fuses the carbon-dynamics, phenology, +# photosynthesis, stomatal-conductance, and autotrophic-respiration auxiliaries into a single launch. +# These tests check that the fused kernel reproduces the per-process fan-out (the pre-fusion path) to +# machine precision. + +# Auxiliary fields written by the fused kernel (carbon dynamics → phenology → photosynthesis → +# stomatal conductance → autotrophic respiration). +const VEGETATION_CARBON_CYCLE_AUXILIARIES = ( + :balanced_leaf_area_index, + :phenology_factor, :leaf_area_index, + :net_assimilation, :leaf_respiration, :gross_primary_production, + :canopy_water_conductance, + :autotrophic_respiration, :net_primary_production, +) + +function vegetated_land(NF) + grid = ColumnGrid(CPU(), ExponentialSpacing(Δz_max = 1.0, N = 50)) + swrc = VanGenuchten(α = 2.0, n = 2.0) + hydraulic_properties = ConstantSoilHydraulics(NF; swrc, unsat_hydraulic_cond = UnsatKVanGenuchten(NF)) + hydrology = SoilHydrology(NF, RichardsEq(); hydraulic_properties) + soil = SoilEnergyWaterCarbon(NF; hydrology) + vegetation = VegetationCarbonCycle(NF) + land = LandModel(grid; soil, vegetation) + initializers = ( + temperature = (x, z) -> 15.0 - 0.02 * z, + saturation_water_ice = (x, z) -> min(1, 0.8 - 0.05 * z), + carbon_vegetation = 0.5, + ) + integrator = initialize(land; initializers) + set!(integrator.state.rainfall, 1.0e-7) + Terrarium.closure!(integrator.state, land) + # Populate atmosphere / soil auxiliaries (vegetation reads air temperature, CO₂, radiation). + compute_auxiliary!(integrator.state, land) + return land, integrator.state +end + +@testset "Fused vegetation carbon cycle auxiliary matches per-process fan-out" begin + land, state = vegetated_land(Float64) + grid = get_grid(land) + vegetation = land.vegetation + constants, atmos, soil = land.constants, land.atmosphere, land.soil + + # Fused path: a single launch over the coupled type (plus the root/PAW pre-stages it keeps). + fused = deepcopy(state) + compute_auxiliary!(fused, grid, vegetation, constants, atmos, soil) + + # Per-process path: the standalone launches the fan-out used to issue, in dependency order. + per_process = deepcopy(state) + compute_auxiliary!(per_process, grid, vegetation.root_distribution, soil) + compute_auxiliary!(per_process, grid, vegetation.plant_available_water, soil) + compute_auxiliary!(per_process, grid, vegetation.carbon_dynamics, vegetation.traits) + compute_auxiliary!(per_process, grid, vegetation.phenology, vegetation.carbon_dynamics, atmos) + compute_auxiliary!(per_process, grid, vegetation.photosynthesis, vegetation.stomatal_conductance, + vegetation.traits, constants, atmos) + compute_auxiliary!(per_process, grid, vegetation.stomatal_conductance, vegetation.traits, constants, atmos) + compute_auxiliary!(per_process, grid, vegetation.autotrophic_respiration, vegetation.carbon_dynamics, + vegetation.phenology, vegetation.traits, atmos) + + @testset "$(name)" for name in VEGETATION_CARBON_CYCLE_AUXILIARIES + fused_vals = Array(interior(getproperty(fused, name))) + per_process_vals = Array(interior(getproperty(per_process, name))) + @test all(isfinite.(fused_vals)) + @test fused_vals ≈ per_process_vals + end +end diff --git a/test/vegetation/vegetation_model_tests.jl b/test/vegetation/vegetation_model_tests.jl index 5c36dc98d..ec5aa3a43 100644 --- a/test/vegetation/vegetation_model_tests.jl +++ b/test/vegetation/vegetation_model_tests.jl @@ -32,3 +32,7 @@ end @testset "Plant Available Water" begin include("plant_available_water_tests.jl") end + +@testset "Vegetation carbon cycle integration" begin + include("integration_tests.jl") +end From a868156a8b2e4c04c1963ee778fc4e90589d3fec Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Wed, 16 Sep 2026 17:52:57 +0200 Subject: [PATCH 3/7] Fix formatting errors --- src/processes/surface/surface_hydrology.jl | 12 ++++++--- .../vegetation/vegetation_carbon_cycle.jl | 26 +++++++++++++------ test/surface/hydrology/integration_tests.jl | 6 +++-- test/vegetation/integration_tests.jl | 12 ++++++--- 4 files changed, 38 insertions(+), 18 deletions(-) diff --git a/src/processes/surface/surface_hydrology.jl b/src/processes/surface/surface_hydrology.jl index 33d80d700..2b6886596 100644 --- a/src/processes/surface/surface_hydrology.jl +++ b/src/processes/surface/surface_hydrology.jl @@ -79,11 +79,15 @@ resolves within a single launch. A "prescribed"/no-op scheme contributes a no-op i, j = @index(Global, NTuple) compute_canopy_auxiliary!(out, i, j, grid, fields, hydrology.canopy_interception, atmos) # `snow` lets the (bare-ground) evaporation scheme scale ground evaporation by the snow-free fraction - compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, hydrology.evapotranspiration, - hydrology.canopy_interception, constants, atmos, soil, vegetation, snow) + compute_evapotranspiration_auxiliary!( + out, i, j, grid, fields, hydrology.evapotranspiration, + hydrology.canopy_interception, constants, atmos, soil, vegetation, snow + ) # `snow` makes the surface runoff scheme's water input snow-aware (meltwater + bare-ground throughfall) - compute_surface_runoff!(out, i, j, grid, fields, hydrology.surface_runoff, - hydrology.canopy_interception, get_hydrology(soil), snow) + compute_surface_runoff!( + out, i, j, grid, fields, hydrology.surface_runoff, + hydrology.canopy_interception, get_hydrology(soil), snow + ) end """ $TYPEDSIGNATURES """ diff --git a/src/processes/vegetation/vegetation_carbon_cycle.jl b/src/processes/vegetation/vegetation_carbon_cycle.jl index 9d5900d32..49a557c10 100644 --- a/src/processes/vegetation/vegetation_carbon_cycle.jl +++ b/src/processes/vegetation/vegetation_carbon_cycle.jl @@ -102,12 +102,18 @@ function compute_auxiliary!( photosynthesis = veg.photosynthesis stomatal_conductance = veg.stomatal_conductance autotrophic_respiration = veg.autotrophic_respiration - out = filter(v -> v isa Field, auxiliary_fields(state, carbon_dynamics, phenology, photosynthesis, - stomatal_conductance, autotrophic_respiration)) + out = filter( + v -> v isa Field, auxiliary_fields( + state, carbon_dynamics, phenology, photosynthesis, + stomatal_conductance, autotrophic_respiration + ) + ) # Full fields (no `except`): within a cell the kernel writes `out.foo` and a later stage reads # `fields.foo` — the same `Field` object, so the write is visible to the stages below. - fields = get_fields(state, carbon_dynamics, phenology, photosynthesis, stomatal_conductance, - autotrophic_respiration, atmos) + fields = get_fields( + state, carbon_dynamics, phenology, photosynthesis, stomatal_conductance, + autotrophic_respiration, atmos + ) launch!(grid, XY, compute_auxiliary_kernel!, out, fields, veg, constants, atmos) # Note: vegetation_dynamics compute_auxiliary! does nothing for now @@ -137,11 +143,15 @@ pool), then phenology (which reads it to set `leaf_area_index` and `phenology_fa compute_veg_carbon_auxiliary!(out, i, j, grid, fields, veg.carbon_dynamics, veg.traits) compute_phenology!(out, i, j, grid, fields, veg.phenology, atmos) # Photosynthesis reads stomatal conductance's parameters (λc) but not its auxiliary state - compute_photosynthesis!(out, i, j, grid, fields, veg.photosynthesis, veg.stomatal_conductance, - veg.traits, constants, atmos) + compute_photosynthesis!( + out, i, j, grid, fields, veg.photosynthesis, veg.stomatal_conductance, + veg.traits, constants, atmos + ) compute_stomatal_conductance!(out, i, j, grid, fields, veg.stomatal_conductance, veg.traits, constants, atmos) - compute_autotrophic_respiration!(out, i, j, grid, fields, veg.autotrophic_respiration, - veg.carbon_dynamics, veg.phenology, veg.traits, atmos) + compute_autotrophic_respiration!( + out, i, j, grid, fields, veg.autotrophic_respiration, + veg.carbon_dynamics, veg.phenology, veg.traits, atmos + ) end """ diff --git a/test/surface/hydrology/integration_tests.jl b/test/surface/hydrology/integration_tests.jl index 7aa10152a..041f39667 100644 --- a/test/surface/hydrology/integration_tests.jl +++ b/test/surface/hydrology/integration_tests.jl @@ -49,8 +49,10 @@ end # Per-process path: the three standalone launches the fan-out used to issue, in dependency order. per_process = deepcopy(state) compute_auxiliary!(per_process, grid, hydrology.canopy_interception, atmos) - compute_auxiliary!(per_process, grid, hydrology.evapotranspiration, hydrology.canopy_interception, - constants, atmos, soil, vegetation, snow) + compute_auxiliary!( + per_process, grid, hydrology.evapotranspiration, hydrology.canopy_interception, + constants, atmos, soil, vegetation, snow + ) compute_auxiliary!(per_process, grid, hydrology.surface_runoff, hydrology.canopy_interception, soil, snow) @testset "$(name)" for name in SURFACE_HYDROLOGY_AUXILIARIES diff --git a/test/vegetation/integration_tests.jl b/test/vegetation/integration_tests.jl index a1a8cd8bb..a3db89de5 100644 --- a/test/vegetation/integration_tests.jl +++ b/test/vegetation/integration_tests.jl @@ -53,11 +53,15 @@ end compute_auxiliary!(per_process, grid, vegetation.plant_available_water, soil) compute_auxiliary!(per_process, grid, vegetation.carbon_dynamics, vegetation.traits) compute_auxiliary!(per_process, grid, vegetation.phenology, vegetation.carbon_dynamics, atmos) - compute_auxiliary!(per_process, grid, vegetation.photosynthesis, vegetation.stomatal_conductance, - vegetation.traits, constants, atmos) + compute_auxiliary!( + per_process, grid, vegetation.photosynthesis, vegetation.stomatal_conductance, + vegetation.traits, constants, atmos + ) compute_auxiliary!(per_process, grid, vegetation.stomatal_conductance, vegetation.traits, constants, atmos) - compute_auxiliary!(per_process, grid, vegetation.autotrophic_respiration, vegetation.carbon_dynamics, - vegetation.phenology, vegetation.traits, atmos) + compute_auxiliary!( + per_process, grid, vegetation.autotrophic_respiration, vegetation.carbon_dynamics, + vegetation.phenology, vegetation.traits, atmos + ) @testset "$(name)" for name in VEGETATION_CARBON_CYCLE_AUXILIARIES fused_vals = Array(interior(getproperty(fused, name))) From 93563fbed472c4bac641301e02ee94ee70617257 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Wed, 16 Sep 2026 20:49:18 +0200 Subject: [PATCH 4/7] Fix outdated doc ref --- docs/src/processes/surface_hydrology/evapotranspiration.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/processes/surface_hydrology/evapotranspiration.md b/docs/src/processes/surface_hydrology/evapotranspiration.md index 1b058b372..0560a747d 100644 --- a/docs/src/processes/surface_hydrology/evapotranspiration.md +++ b/docs/src/processes/surface_hydrology/evapotranspiration.md @@ -271,7 +271,7 @@ compute_auxiliary!(state, grid, ::BareGroundEvaporation, ::NoCanopyInterception, ``` ```@docs; canonical = false -compute_auxiliary!(state, grid, ::PALADYNCanopyEvapotranspiration, ::AbstractCanopyInterception, ::PhysicalConstants, ::AbstractAtmosphere, ::AbstractSoil, ::AbstractVegetation, args...) +compute_auxiliary!(state, grid, ::PALADYNCanopyEvapotranspiration, ::AbstractCanopyInterception, ::PhysicalConstants, ::AbstractAtmosphere, ::AbstractSoil, ::AbstractVegetation, ::Optional{AbstractSnow}, args...) ``` ## Coupling to soil hydrology From 8923663e1b2d1ebd6966c4bc8e83576ca5003cb8 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Wed, 16 Sep 2026 20:49:37 +0200 Subject: [PATCH 5/7] Update plan doc --- docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md b/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md index b25179dcf..c292791bf 100644 --- a/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md +++ b/docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md @@ -447,7 +447,7 @@ full `fields`. See rev 7 for the latent canopy-ET `snow` bug this surfaced and f After Phase 1 there are two XY tendencies (canopy water, surface excess water). They are **independent outputs** (neither feeds the other), so fusing them is the register-union anti-pattern from rev 6, not a chain. Tendency fusion is therefore **dropped** for surface hydrology; the two per-process tendency -launches stay. (If ever revisited, capture registers first — see `scratch/capture_soil_regs.jl` — and +launches stay. (If ever revisited, capture registers first — see `scratch/benchmarks/soil/capture_soil_regs.jl` — and only fuse if the chain holds or `@noinline` recovers occupancy.) ## Phase 4 — Fuse vegetation carbon (`VegetationCarbonCycle`) — **IMPLEMENTED (rev 9)** From 9335e03fdfd865c0355c852a454317132241f468 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Sun, 20 Sep 2026 14:27:14 +0200 Subject: [PATCH 6/7] Revert fused surface hydrology compute_auxiliary! The fused single-launch kernel for SurfaceHydrology regressed GPU performance by ~20% (register-union of the independent evapotranspiration and surface-runoff stages) with no end-to-end benefit. Restore the per-process fan-out (467a1e689). The vegetation carbon cycle fusion is kept. Removes the fused surface-hydrology integration test and restores the pre-fusion surface sources and evapotranspiration docs. --- .../surface_hydrology/evapotranspiration.md | 2 +- .../canopy_interception.jl | 3 - .../bare_ground_evaporation.jl | 43 ++------- .../canopy_evapotranspiration.jl | 36 ++----- src/processes/surface/surface_hydrology.jl | 47 +--------- test/surface/hydrology/integration_tests.jl | 93 ------------------- .../hydrology/surface_hydrology_tests.jl | 4 - 7 files changed, 20 insertions(+), 208 deletions(-) delete mode 100644 test/surface/hydrology/integration_tests.jl diff --git a/docs/src/processes/surface_hydrology/evapotranspiration.md b/docs/src/processes/surface_hydrology/evapotranspiration.md index 0560a747d..1b058b372 100644 --- a/docs/src/processes/surface_hydrology/evapotranspiration.md +++ b/docs/src/processes/surface_hydrology/evapotranspiration.md @@ -271,7 +271,7 @@ compute_auxiliary!(state, grid, ::BareGroundEvaporation, ::NoCanopyInterception, ``` ```@docs; canonical = false -compute_auxiliary!(state, grid, ::PALADYNCanopyEvapotranspiration, ::AbstractCanopyInterception, ::PhysicalConstants, ::AbstractAtmosphere, ::AbstractSoil, ::AbstractVegetation, ::Optional{AbstractSnow}, args...) +compute_auxiliary!(state, grid, ::PALADYNCanopyEvapotranspiration, ::AbstractCanopyInterception, ::PhysicalConstants, ::AbstractAtmosphere, ::AbstractSoil, ::AbstractVegetation, args...) ``` ## Coupling to soil hydrology diff --git a/src/processes/surface/canopy_interception/canopy_interception.jl b/src/processes/surface/canopy_interception/canopy_interception.jl index c4a0f563a..668b876c9 100644 --- a/src/processes/surface/canopy_interception/canopy_interception.jl +++ b/src/processes/surface/canopy_interception/canopy_interception.jl @@ -16,9 +16,6 @@ passthrough_rainfall(grid, clock, fields, ::NoCanopyInterception) = fields.rainf @inline compute_auxiliary!(state, grid, ::NoCanopyInterception, args...) = nothing -# No-op per-cell variant: `rainfall_ground` is a lazy passthrough of `rainfall`, never written. -@inline compute_canopy_auxiliary!(out, i, j, grid, fields, ::NoCanopyInterception, args...) = nothing - @inline compute_tendencies!(state, grid, ::NoCanopyInterception, args...) = nothing @propagate_inbounds canopy_water(i, j, grid, fields, ::NoCanopyInterception) = zero(eltype(grid)) diff --git a/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl b/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl index 7da62bd78..909f334be 100644 --- a/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl +++ b/src/processes/surface/evapotranspiration/bare_ground_evaporation.jl @@ -54,7 +54,7 @@ variables(::BareGroundEvaporation) = ( function compute_auxiliary!( state, grid, evaporation::BareGroundEvaporation, - interception::NoCanopyInterception, + ::NoCanopyInterception, constants::PhysicalConstants, atmos::AbstractAtmosphere, soil::Optional{AbstractSoil} = nothing, @@ -63,7 +63,7 @@ function compute_auxiliary!( out = auxiliary_fields(state, evaporation) # merge the snow cover fraction so the ground evaporation can be scaled by the snow-free fraction fields = merge(get_fields(state, evaporation, atmos, soil; except = out), get_fields(state, snow)) - launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evaporation, interception, constants, atmos, soil, nothing, snow) + launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evaporation, constants, atmos, soil, snow) return nothing end @@ -139,49 +139,24 @@ end out.evaporation_ground[i, j, 1] = E_gnd return out end - # Kernels -""" - $TYPEDSIGNATURES -Compute and store the skin-driven ground evaporation conductance for bare-ground evaporation, then -evaluate the ground evaporation flux from it. The conductance is merged into `fields` so the flux -step reads it back without a global round-trip. The `interception` and `vegetation` arguments are -unused by this scheme. -""" -@propagate_inbounds function compute_evapotranspiration_auxiliary!( - out, i, j, grid, fields, - evaporation::BareGroundEvaporation, - interception::NoCanopyInterception, +@kernel inbounds = true function compute_auxiliary_kernel!( + out, grid, fields, + evapotranspiration::BareGroundEvaporation, constants::PhysicalConstants, atmos::AbstractAtmosphere, soil::Optional{AbstractSoil} = nothing, - vegetation::Optional{AbstractVegetation} = nothing, snow::Optional{AbstractSnow} = nothing, - args... ) + i, j = @index(Global, NTuple) + # First compute conductances - compute_evapotranspiration_conductances!(out, i, j, grid, fields, evaporation, constants, atmos, soil) + compute_evapotranspiration_conductances!(out, i, j, grid, fields, evapotranspiration, constants, atmos, soil) # TODO: Annoyingly, we need to explicitly add these to `fields`; need a better solution to this problem conductances = (ground_evaporation_conductance = out.ground_evaporation_conductance,) fields = merge(fields, conductances) # Compute ET fluxes from stored conductances; `snow` scales ground evaporation by the snow-free fraction - compute_evapotranspiration_fluxes!(out, i, j, grid, fields, evaporation, constants, atmos, snow) - return out -end - -@kernel inbounds = true function compute_auxiliary_kernel!( - out, grid, fields, - evapotranspiration::BareGroundEvaporation, - interception::NoCanopyInterception, - constants::PhysicalConstants, - atmos::AbstractAtmosphere, - soil::Optional{AbstractSoil} = nothing, - vegetation::Optional{AbstractVegetation} = nothing, - snow::Optional{AbstractSnow} = nothing, - args... - ) - i, j = @index(Global, NTuple) - compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, snow, args...) + compute_evapotranspiration_fluxes!(out, i, j, grid, fields, evapotranspiration, constants, atmos, snow) end diff --git a/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl b/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl index 305b0ddb9..8101f24e8 100644 --- a/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl +++ b/src/processes/surface/evapotranspiration/canopy_evapotranspiration.jl @@ -1,7 +1,7 @@ """ $TYPEDEF -Canopy evapotranspiration scheme from PALADYN ([willeitPALADYNV10Comprehensive2016; Eq. (5)](@cite)) +Canopy evapotranspiration scheme from PALADYN ([willeitPALADYNV10Comprehensive2016; Eq. (5)](@cite)) that includes a canopy evaporation term based on the saturation fraction of canopy water defined by the canopy hydrology scheme. @@ -120,13 +120,11 @@ function compute_auxiliary!( atmos::AbstractAtmosphere, soil::AbstractSoil, vegetation::AbstractVegetation, - snow::Optional{AbstractSnow} = nothing, args... ) out = auxiliary_fields(state, evapotranspiration) - # merge the snow cover fraction so the ground/canopy evaporation fluxes can be scaled by the snow-free fraction - fields = merge(get_fields(state, evapotranspiration, interception, atmos, soil, vegetation; except = out), get_fields(state, snow)) - launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, snow) + fields = get_fields(state, evapotranspiration, interception, atmos, soil, vegetation; except = out) + launch!(grid, XY, compute_auxiliary_kernel!, out, fields, evapotranspiration, interception, constants, atmos, soil, vegetation) return nothing end @@ -277,15 +275,8 @@ end # Kernels -""" - $TYPEDSIGNATURES - -Compute and store the skin-driven vapor conductances for canopy evapotranspiration, then evaluate -the partitioned humidity fluxes from them. The conductances are merged into `fields` so the flux -step reads them back without a global round-trip. -""" -@propagate_inbounds function compute_evapotranspiration_auxiliary!( - out, i, j, grid, fields, +@kernel inbounds = true function compute_auxiliary_kernel!( + out, grid, fields, evapotranspiration::PALADYNCanopyEvapotranspiration, interception::AbstractCanopyInterception, constants::PhysicalConstants, @@ -295,6 +286,7 @@ step reads them back without a global round-trip. snow::Optional{AbstractSnow} = nothing, args... ) + i, j = @index(Global, NTuple) # First compute conductances compute_evapotranspiration_conductances!(out, i, j, grid, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, args...) # TODO: Annoyingly, we need to explicitly add these to `fields`; need a better solution to this problem @@ -306,20 +298,4 @@ step reads them back without a global round-trip. fields = merge(fields, conductances) # Compute ET fluxes from stored conductances compute_evapotranspiration_fluxes!(out, i, j, grid, fields, evapotranspiration, constants, atmos, snow) - return out -end - -@kernel inbounds = true function compute_auxiliary_kernel!( - out, grid, fields, - evapotranspiration::PALADYNCanopyEvapotranspiration, - interception::AbstractCanopyInterception, - constants::PhysicalConstants, - atmos::AbstractAtmosphere, - soil::AbstractSoil, - vegetation::AbstractVegetation, - snow::Optional{AbstractSnow} = nothing, - args... - ) - i, j = @index(Global, NTuple) - compute_evapotranspiration_auxiliary!(out, i, j, grid, fields, evapotranspiration, interception, constants, atmos, soil, vegetation, snow, args...) end diff --git a/src/processes/surface/surface_hydrology.jl b/src/processes/surface/surface_hydrology.jl index 2b6886596..63b1b8c42 100644 --- a/src/processes/surface/surface_hydrology.jl +++ b/src/processes/surface/surface_hydrology.jl @@ -43,51 +43,12 @@ function compute_auxiliary!( snow::Optional{AbstractSnow} = nothing, args... ) - interception = hydrology.canopy_interception - evapotranspiration = hydrology.evapotranspiration - surface_runoff = hydrology.surface_runoff - # The fused kernel writes the union of the three sub-processes' auxiliaries. Drop lazy - # (`FunctionField`) auxiliaries — e.g. `NoCanopyInterception`'s `rainfall_ground` passthrough — - # which no kernel writes. - out = filter(v -> v isa Field, auxiliary_fields(state, interception, evapotranspiration, surface_runoff)) - # Full fields (no `except`): within a cell the kernel writes `out.foo` and a later sub-process - # reads `fields.foo` — the same `Field` object, so the write is visible. This is what lets the - # canopy → evapotranspiration → runoff dependency chain run in a single launch. - fields = get_fields(state, interception, evapotranspiration, surface_runoff, atmos, soil, vegetation, snow) - launch!(grid, XY, compute_auxiliary_kernel!, out, fields, hydrology, constants, atmos, soil, vegetation, snow) - return nothing -end - -""" - $TYPEDSIGNATURES - -Fused auxiliary kernel for the coupled surface hydrology processes. Each sub-process's per-cell -mutating variant runs in dependency order — canopy interception, then evapotranspiration (which reads -the canopy saturation fraction), then surface runoff (which reads the ground rainfall) — so the chain -resolves within a single launch. A "prescribed"/no-op scheme contributes a no-op variant. -""" -@kernel inbounds = true function compute_auxiliary_kernel!( - out, grid, fields, - hydrology::SurfaceHydrology, - constants::PhysicalConstants, - atmos::AbstractAtmosphere, - soil::Optional{AbstractSoil} = nothing, - vegetation::Optional{AbstractVegetation} = nothing, - snow::Optional{AbstractSnow} = nothing, - args... - ) - i, j = @index(Global, NTuple) - compute_canopy_auxiliary!(out, i, j, grid, fields, hydrology.canopy_interception, atmos) + compute_auxiliary!(state, grid, hydrology.canopy_interception, atmos) # `snow` lets the (bare-ground) evaporation scheme scale ground evaporation by the snow-free fraction - compute_evapotranspiration_auxiliary!( - out, i, j, grid, fields, hydrology.evapotranspiration, - hydrology.canopy_interception, constants, atmos, soil, vegetation, snow - ) + compute_auxiliary!(state, grid, hydrology.evapotranspiration, hydrology.canopy_interception, constants, atmos, soil, vegetation, snow) # `snow` makes the surface runoff scheme's water input snow-aware (meltwater + bare-ground throughfall) - compute_surface_runoff!( - out, i, j, grid, fields, hydrology.surface_runoff, - hydrology.canopy_interception, get_hydrology(soil), snow - ) + compute_auxiliary!(state, grid, hydrology.surface_runoff, hydrology.canopy_interception, soil, snow) + return nothing end """ $TYPEDSIGNATURES """ diff --git a/test/surface/hydrology/integration_tests.jl b/test/surface/hydrology/integration_tests.jl deleted file mode 100644 index 041f39667..000000000 --- a/test/surface/hydrology/integration_tests.jl +++ /dev/null @@ -1,93 +0,0 @@ -using Terrarium -using Test - -# The coupled `SurfaceHydrology.compute_auxiliary!` fuses the canopy-interception, -# evapotranspiration, and surface-runoff auxiliaries into a single launch. These tests check that the -# fused kernel reproduces the per-process fan-out (the pre-fusion path) to machine precision, and that -# the canopy evapotranspiration fluxes are scaled by the snow-free fraction (the snow coupling the -# per-process canopy-ET launch used to drop). - -# Surface hydrology auxiliary fields written by the fused kernel (canopy + ET + runoff). -const SURFACE_HYDROLOGY_AUXILIARIES = ( - :canopy_water_interception, :canopy_water_removal, :saturation_canopy_water, :rainfall_ground, - :ground_evaporation_conductance, :canopy_evaporation_conductance, :transpiration_conductance, - :evaporation_canopy, :evaporation_ground, :transpiration, - :surface_runoff, :infiltration, -) - -function vegetated_snow_land(NF) - grid = ColumnGrid(CPU(), ExponentialSpacing(Δz_max = 1.0, N = 50)) - soil = SoilEnergyWaterCarbon(NF; hydrology = SoilHydrology(NF, RichardsEq())) - vegetation = VegetationCarbonCycle(NF) - snow = SingleLayerSnow(NF) - land = LandModel(grid; soil, vegetation, snow) - initializers = ( - temperature = (x, z) -> 2.0 - 0.02 * z, - saturation_water_ice = (x, z) -> min(1, 0.8 - 0.05 * z), - carbon_vegetation = 0.1, - snow_water_equivalent = 0.2, - snow_temperature = -2.0, - ) - integrator = initialize(land; initializers) - set!(integrator.state.rainfall, 1.0e-7) - Terrarium.closure!(integrator.state, land) - # Populate atmosphere / soil / snow / vegetation auxiliaries (surface hydrology runs last). - compute_auxiliary!(integrator.state, land) - return land, integrator.state -end - -@testset "Fused surface hydrology auxiliary matches per-process fan-out" begin - land, state = vegetated_snow_land(Float64) - grid = get_grid(land) - hydrology = land.surface_hydrology - constants, atmos, soil, vegetation, snow = land.constants, land.atmosphere, land.soil, land.vegetation, land.snow - - # Fused path: a single launch over the coupled type. - fused = deepcopy(state) - compute_auxiliary!(fused, grid, hydrology, constants, atmos, soil, vegetation, snow) - - # Per-process path: the three standalone launches the fan-out used to issue, in dependency order. - per_process = deepcopy(state) - compute_auxiliary!(per_process, grid, hydrology.canopy_interception, atmos) - compute_auxiliary!( - per_process, grid, hydrology.evapotranspiration, hydrology.canopy_interception, - constants, atmos, soil, vegetation, snow - ) - compute_auxiliary!(per_process, grid, hydrology.surface_runoff, hydrology.canopy_interception, soil, snow) - - @testset "$(name)" for name in SURFACE_HYDROLOGY_AUXILIARIES - fused_vals = Array(interior(getproperty(fused, name))) - per_process_vals = Array(interior(getproperty(per_process, name))) - @test all(isfinite.(fused_vals)) - @test fused_vals ≈ per_process_vals - end -end - -@testset "Canopy ET fluxes are scaled by the snow-free fraction" begin - # With a snow-covered surface, the canopy scheme's ground/canopy evaporation and transpiration must - # be scaled by (1 − f_snow) — the coupling the per-process canopy-ET launch used to drop. Recomputing - # the auxiliaries with the snow cover fraction forced to zero must recover the unscaled fluxes. - land, state = vegetated_snow_land(Float64) - grid = get_grid(land) - hydrology = land.surface_hydrology - constants, atmos, soil, vegetation, snow = land.constants, land.atmosphere, land.soil, land.vegetation, land.snow - - @test all(Array(interior(state.snow_cover_fraction)) .> 0) # snow is present - - covered = deepcopy(state) - compute_auxiliary!(covered, grid, hydrology, constants, atmos, soil, vegetation, snow) - - # Same state with the snow cover fraction zeroed: the fused kernel must then leave the ET fluxes - # unscaled, i.e. the covered fluxes equal the bare fluxes times (1 − f_snow). - bare = deepcopy(state) - set!(bare.snow_cover_fraction, 0.0) - compute_auxiliary!(bare, grid, hydrology, constants, atmos, soil, vegetation, snow) - - f = Array(interior(state.snow_cover_fraction)) - @test all(f .< 1) # partially covered, so the scaling is nontrivial - for name in (:evaporation_ground, :transpiration, :evaporation_canopy) - covered_vals = Array(interior(getproperty(covered, name))) - bare_vals = Array(interior(getproperty(bare, name))) - @test all(isapprox.(covered_vals, (1 .- f) .* bare_vals; rtol = 1.0e-9)) - end -end diff --git a/test/surface/hydrology/surface_hydrology_tests.jl b/test/surface/hydrology/surface_hydrology_tests.jl index 44885b70c..bff0fda77 100644 --- a/test/surface/hydrology/surface_hydrology_tests.jl +++ b/test/surface/hydrology/surface_hydrology_tests.jl @@ -12,7 +12,3 @@ end @testset "Surface runoff" begin include("surface_runoff_tests.jl") end - -@testset "Fused surface hydrology" begin - include("integration_tests.jl") -end From 2aea8cb6d685941f1761e7d0c949195b4d5a2d98 Mon Sep 17 00:00:00 2001 From: Brian Groenke Date: Thu, 1 Oct 2026 18:40:56 +0200 Subject: [PATCH 7/7] Remove unnecessary filtering in fused veg aux --- .../vegetation/vegetation_carbon_cycle.jl | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/src/processes/vegetation/vegetation_carbon_cycle.jl b/src/processes/vegetation/vegetation_carbon_cycle.jl index 49a557c10..73ee62398 100644 --- a/src/processes/vegetation/vegetation_carbon_cycle.jl +++ b/src/processes/vegetation/vegetation_carbon_cycle.jl @@ -95,25 +95,14 @@ function compute_auxiliary!( compute_auxiliary!(state, grid, veg.plant_available_water, soil) # The carbon-cycle chain — carbon dynamics → phenology → photosynthesis → stomatal conductance → - # autotrophic respiration — is fused into a single launch. Each stage reads the previous stage's - # output within a cell, so the dependency chain resolves without returning to the host. + # autotrophic respiration — is fused into a single launch. carbon_dynamics = veg.carbon_dynamics phenology = veg.phenology photosynthesis = veg.photosynthesis stomatal_conductance = veg.stomatal_conductance autotrophic_respiration = veg.autotrophic_respiration - out = filter( - v -> v isa Field, auxiliary_fields( - state, carbon_dynamics, phenology, photosynthesis, - stomatal_conductance, autotrophic_respiration - ) - ) - # Full fields (no `except`): within a cell the kernel writes `out.foo` and a later stage reads - # `fields.foo` — the same `Field` object, so the write is visible to the stages below. - fields = get_fields( - state, carbon_dynamics, phenology, photosynthesis, stomatal_conductance, - autotrophic_respiration, atmos - ) + out = auxiliary_fields(state, carbon_dynamics, phenology, photosynthesis, stomatal_conductance, autotrophic_respiration) + fields = get_fields(state, carbon_dynamics, phenology, photosynthesis, stomatal_conductance, autotrophic_respiration, atmos) launch!(grid, XY, compute_auxiliary_kernel!, out, fields, veg, constants, atmos) # Note: vegetation_dynamics compute_auxiliary! does nothing for now