Skip to content

Implement fused kernels for vegetation carbon compute_auxiliary! - #192

Open
bgroenks96 wants to merge 7 commits into
mainfrom
bg/kernel-fusing
Open

bgroenks96 wants to merge 7 commits into
mainfrom
bg/kernel-fusing

Conversation

@bgroenks96

@bgroenks96 bgroenks96 commented Sep 16, 2026 •

Copy link
Copy Markdown
Collaborator

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 a
single 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 a
pure 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 XY launch
(compute_auxiliary_kernel!). Root distribution and plant-available water still
run 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}) in
d39a54c 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 @docs signature for
compute_auxiliary!(::PALADYNCanopyEvapotranspiration, ...) had gone stale
(it listed a snow parameter that the method does not take), which broke the
Documenter 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 :land configuration at three resolutions:

nlat_half columns CPU speedup (min) A40 GPU speedup (min / median)
8 512 1.00x 1.01x / 1.02x
24 4,608 1.01x 1.06x / 1.03x
48 18,432 1.00x 1.01x / 1.01x

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

  • New integration test (test/vegetation/integration_tests.jl): the fused
    kernel'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.
  • Existing vegetation, surface hydrology, and land model suites pass.

Files

File Change
src/processes/vegetation/vegetation_carbon_cycle.jl Fused compute_auxiliary_kernel!, rewritten compute_auxiliary!
test/vegetation/integration_tests.jl New fused-vs-fan-out equivalence test
test/vegetation/vegetation_model_tests.jl Include the new test file
docs/dev/2026-08/2026-08-31-PLAN_fused_kernels.md Rev 7: scope narrowed to vegetation; surface/soil fusion findings

@bgroenks96 bgroenks96 changed the title Implement fused surface hydrology compute_auxiliary! Implement fused kernels for vegetation carbon compute_auxiliary! Sep 20, 2026
@bgroenks96
bgroenks96 marked this pull request as ready for review September 20, 2026 12:30

@maximilian-gelbrecht maximilian-gelbrecht left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks pretty straight forward to me, just one question about the output field assignment (see comments)

Comment thread src/processes/vegetation/vegetation_carbon_cycle.jl Outdated
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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants