Implement fused kernels for vegetation carbon compute_auxiliary! - #192
Open
bgroenks96 wants to merge 7 commits into
Open
bgroenks96 wants to merge 7 commits into
bgroenks96 wants to merge 7 commits into
Conversation
bgroenks96
force-pushed
the
bg/kernel-fusing
branch
from
September 20, 2026 12:09
4ca5b30 to
9c39cd3
Compare
compute_auxiliary!compute_auxiliary!
bgroenks96
marked this pull request as ready for review
September 20, 2026 12:30
maximilian-gelbrecht
left a comment
Collaborator
There was a problem hiding this comment.
Looks pretty straight forward to me, just one question about the output field assignment (see comments)
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.
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 (467a1e6). The vegetation carbon cycle fusion is kept. Removes the fused surface-hydrology integration test and restores the pre-fusion surface sources and evapotranspiration docs.
bgroenks96
force-pushed
the
bg/kernel-fusing
branch
from
October 1, 2026 16:41
afdb0a5 to
2aea8cb
Compare
This branch has not been deployed
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
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.
This PR adds kernel fusion for vegetation carbon processes. Description below will be updated after those changes are added.
The following is generated by an LLM 🤖
Fuse the vegetation carbon cycle auxiliary kernel
This PR fuses the vegetation carbon cycle
compute_auxiliary!chain into asingle kernel launch, and documents (via a reverted experiment) why the same
fusion does not work for surface hydrology.
What changed
Vegetation carbon cycle fusion (kept). The five-stage chain
carbon dynamics → phenology → photosynthesis → stomatal conductance → autotrophic respiration
previously ran as five separate
compute_auxiliary!launches. Each stage is apure producer/consumer of the next within a cell (
balanced_leaf_area_index→leaf_area_index/phenology_factor→net_assimilation/gross_primary_production→canopy_water_conductance→net_primary_production),so the whole chain now resolves inside one
XYlaunch(
compute_auxiliary_kernel!). Root distribution and plant-available water stillrun beforehand as separate launches: they depend on the soil state and produce
fields (
soil_moisture_limiting_factor) that the fused kernel reads.Surface hydrology fusion (tried, reverted). The same fusion was implemented
for
SurfaceHydrology(interception → {evapotranspiration, surface runoff}) ind39a54c and reverted in afdb0a5. Unlike the vegetation chain,
evapotranspiration and runoff are independent siblings, not a dependency
chain: fusing them forces the register union of two unrelated stages, which
regressed GPU performance by ~20% with no end-to-end benefit. The per-process
fan-out is restored byte-for-byte. This finding is recorded in the plan doc
(rev 7) as the boundary of the fusion rule: fuse dependency chains and single
accumulators, do not bundle independent outputs.
Docs fix. The
@docssignature forcompute_auxiliary!(::PALADYNCanopyEvapotranspiration, ...)had gone stale(it listed a
snowparameter that the method does not take), which broke theDocumenter build on this branch. Fixed in 955111b.
Benchmarking
Microbenchmark of the vegetation carbon-cycle auxiliary pass (fused single launch vs the
pre-fusion seven-launch fan-out), interleaved A/B, minimum of 30 repeats, Float64, nz=30,
on the
:landconfiguration at three resolutions:Neutral on CPU and neutral-to-slightly-positive on GPU, as expected for a genuine
dependency chain: the fused kernel resolves the chain in one launch without paying the
register-union penalty that independent-sibling fusion incurs (see the reverted surface
hydrology experiment, ~20% GPU regression).
Testing
test/vegetation/integration_tests.jl): the fusedkernel's auxiliary output (
leaf_area_index,phenology_factor,net_assimilation,gross_primary_production,canopy_water_conductance,net_primary_production, ...) matches the per-process fan-out exactly,across the vegetation carbon-cycle configurations.
Files
src/processes/vegetation/vegetation_carbon_cycle.jlcompute_auxiliary_kernel!, rewrittencompute_auxiliary!test/vegetation/integration_tests.jltest/vegetation/vegetation_model_tests.jldocs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md