Conversation
szy21
reviewed
Jul 13, 2026
szy21
approved these changes
Jul 13, 2026
szy21
left a comment
szy21
left a comment
Member
There was a problem hiding this comment.
I don't know too much about the gpu performance, but overall it looks good to me and I like the speedup:)
…d (ncol, nlev) layout for GPU speedups while preserving public API By switching internal storage layouts across optical properties (TwoStream, OneScalar), thermal and solar source arrays (SourceLW2Str, SourceLWNoScat, SourceSW2Str), and radiative fluxes (FluxLW, FluxSW, FluxBand) from (nlev, ncol) to contiguous column-major (ncol, nlev) F-order, adjacent GPU warp threads perform consecutive, coalesced memory reads and writes (`stride-1` over `gcol`). To strictly preserve the Layer-2 host data-exchange contract without adding runtime allocations or latency, all public getters (`lw_flux_up`, `lw_flux_dn`, `lw_flux_net`, `sw_flux_up`, `sw_flux_dn`, `sw_direct_flux_dn`, `sw_flux_net`, `net_flux`, `clear_*`, and all `spectral_*_flux_*` arrays) wrap internal buffers via `PermutedDimsArray(buffer, (2, 1))`. Verified allocation-free (@allocated == 0) via standalone API tests. Measured GPU speedups across uniform DYAMOND grids (86,400 horizontal columns × 63 vertical layers, 480 spectral g-points) on NVIDIA A100-SXM4-40GB: Per-kernel execution time comparison against baseline: - clear_sky_lw (Float32): 489.92 ms (down from 771.22 ms -> 0.64x / -36.4% reduction) - clear_sky_lw (Float64): 682.07 ms (down from 1049.30 ms -> 0.65x / -35.0% reduction) - clear_sky_sw (Float32): 455.41 ms (down from 464.95 ms -> 0.98x / -2.1% reduction) - clear_sky_sw (Float64): 639.97 ms (down from 653.72 ms -> 0.98x / -2.1% reduction) - all_sky_lw (Float32): 726.15 ms (down from 976.15 ms -> 0.74x / -25.6% reduction) - all_sky_lw (Float64): 992.34 ms (down from 1331.26 ms -> 0.75x / -25.5% reduction) - all_sky_sw (Float32): 671.15 ms (down from 843.34 ms -> 0.80x / -20.4% reduction) - all_sky_sw (Float64): 940.08 ms (down from 1097.44 ms -> 0.86x / -14.3% reduction) - all_sky_with_aerosols_lw (Float32): 797.79 ms (down from 1054.67 ms -> 0.76x / -24.4% reduction) - all_sky_with_aerosols_lw (Float64): 1096.99 ms (down from 1412.72 ms -> 0.78x / -22.3% reduction) - all_sky_with_aerosols_sw (Float32): 728.34 ms (down from 894.53 ms -> 0.81x / -18.6% reduction) - all_sky_with_aerosols_sw (Float64): 966.28 ms (down from 1173.98 ms -> 0.82x / -17.7% reduction) Combined solve (LW + SW step) speedup vs 12-thread Intel Xeon CPU: - Clear-sky (Float32): 0.95 s total GPU vs 77.95 s CPU -> 82.5x speedup (previously 63.1x) - All-sky McICA (Float32): 1.39 s total GPU vs 98.42 s CPU -> 70.8x speedup (previously 54.1x) - All-sky + aerosols (Float32): 1.52 s total GPU vs 99.39 s CPU -> 65.5x speedup (previously 51.0x) - Clear-sky (Float64): 1.32 s total GPU vs 79.64 s CPU -> 60.1x speedup (previously 46.8x) - All-sky McICA (Float64): 1.93 s total GPU vs 101.60 s CPU -> 52.6x speedup (previously 41.8x) - All-sky + aerosols (Float64): 2.07 s total GPU vs 106.33 s CPU -> 51.4x speedup (previously 41.1x)
… add transposed state cache for coalesced gas-optics reads Follow-ups to the (ncol, nlev) layout commit (298d688), all verified by the full CPU test suite on Julia 1.10.10 (previously failing there): 1. Flux presentation buffers (GPU getter contract restored). The Layer-2 flux getters returned doubly-wrapped lazy views (SubArray of PermutedDimsArray of the device array), which fall outside the single-level wrapper unions CUDA.jl/GPUArrays dispatch on: on GPU, Array(getter), collect, reductions, getter-only broadcasts, and scalar-RHS writes threw scalar-indexing errors (verified with the JLArrays reference backend; test/scalar_indexing.jl would fail on GPU CI). Getters now return plain (nlev, ncol) views of Fluxes.FluxPresentation arrays that update_lw_fluxes!/update_sw_fluxes! fill from the column-first compute buffers with fused transposing copies (allocation-free host loops / one device kernel); _flux_domain_view is deleted. The clear-sky snapshot buffers are stored directly in presentation layout (no extra memory), net_flux_buffer returns to (nlev, ncol) via a fused transpose-sum, and the opt-in per-band FluxBand buffers revert to host-facing (nlev, ncol, n_bnd) rather than doubling their memory with presentation copies. Extra memory: 7 broadband arrays (~310 MB at DYAMOND Float64); extra time: ~ms of transposes per update_fluxes! call. 2. Type stability on Julia 1.9-1.11. PermutedDimsArray(x, perm) does not constant-fold perm before Julia 1.12; construct with explicit type parameters (Fluxes.lazy_transpose). Also: broadcasts with a permuted- wrapper operand allocate 64 B on Julia <= 1.11 even when type-stable, so apply_metric_scaling! and the transposing copies use explicit loops on host arrays (broadcast on device arrays). The @inferred getter tests and the clear-sky @allocated == 0 assertions now pass on 1.10. 3. TransposedStateCache: column-first copies of the hot AtmosphericState arrays (layerdata: col_dry/p_lay/t_lay/rel_hum, and t_lev), refreshed once per spectral solve, read by all four spectral compute_optical_props! methods, so the g-point loop's repeated state reads are coalesced across GPU threads like the flux writes. One cache is shared by the longwave and shortwave workspaces (RRTMGPSolver); gray radiation uses none; opt out with state_cache = nothing. Verified bit-identical fluxes with and without the cache (pinned by a new clear-sky test) and zero-allocation solves. Also: cloudy-sky allocation tests measured call-site dispatch boxing of the type-unstable test setup rather than solver allocations (pristine values on 1.10: 880/1216 B against bounds of 448/368); a function barrier now gives the measurement concrete types and the bounds tighten to == 0. Stale layout docstrings (FluxLW, SourceLWNoScat, net_flux_buffer, band_flux) and docs (functional_core.md's live example read the surface flux where it claimed TOA) corrected; NEWS.md documents the layout change as breaking for Layer-1 raw-field access. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The `state_cache = nothing` opt-out's fallback accessor built a
`PermutedDimsArray` inside the gas-optics kernel to present the vertical-first
`AtmosphericState` arrays in the cache's transposed orientation.
`PermutedDimsArray`'s inner constructor builds an error string, which is not
GPU-compilable, so kernel compilation failed with
InvalidIRError: unsupported call to ... jl_string_to_genericmemory
on the P100 clear-sky test (exercised via the cache-equivalence test, which
builds a `state_cache = nothing` solver).
Replace the transposing 2D `layerdata_view` `[glay, fld]` access with per-field
1D accessors (`col_dry_view`/`p_lay_view`/`t_lay_view`/`rel_hum_view`, plus the
existing `t_lev_view`), each a plain `view(array, field, :, gcol)` — the same
form as the pre-transpose `getview_p_lay(as, gcol)` that has always compiled on
the GPU, with no `PermutedDimsArray` in any kernel. Coalescing is unchanged
(the cache's `view(cache.layerdata, gcol, :, f)` is still stride-1 across a
warp's columns).
Also switch the cache-refresh device path from a `dest .= _lazy_permute(src)`
broadcast to the native `permutedims!` primitive and delete `_lazy_permute`;
host arrays keep their explicit zero-allocation loops.
Full CPU suite green (correctness + zero-alloc + cache-equivalence).
`Fluxes.lazy_transpose` is unaffected: it is only ever called host-side in
Layer-2 orchestration, never in a kernel.
Corrections
Docbuild fixes
Test improvement
Updated performance table
- De-duplicate _coalesced_3d_zeros (build on _coalesced_3d + fill! of the physical storage via a new _physical helper), removing the Float64 intermediate of the old zeros(...) path. - Type both _default_state_cache device methods and document why the CPU default is `nothing`. - Mark the plain-Array _scale_by_transpose! method as a defensive fallback for hand-constructed Layer-1 buffers; put the CPU dual-layout method first. - Tighten helper docstrings: state the DA === Array dispatch contract and the indexed-(ncol, nlev) vs device-dependent-storage distinction; update_presentation! notes GPU transpose vs CPU copyto!. - Update narrative docs that still claimed buffers are stored column-first unconditionally (gpu.md parallelization paragraph, functional_core.md, getters.md, both NEWS.md bullets), and note in gpu.md's checkpointing section that adapting a CPU-constructed solver to the GPU keeps the CPU layout — construct on the device you run on. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
`docs/make.jl` uses `modules = [RRTMGP]` with the default strict `checkdocs = :all`, so every docstring must appear in a `@docs` block. Four internal helpers added on this branch had docstrings but no manual entry: `Fluxes.lazy_transpose`, `Fluxes._coalesced_2d`, `Fluxes._coalesced_3d`, and `AtmosphericStates.refresh_transposed_state!`. These are implementation plumbing, not public API, so demote their docstrings to plain comments (content preserved) rather than expose them on the API pages. No `@ref` targeted them. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
2 participants
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.
Stacked on
ts/tests-docs(docs branch); this PR is the performance work only.Transition internal buffers to coalesced
(ncol, nlev)layoutBy switching internal storage layouts across optical properties (
TwoStream,OneScalar), thermal and solar source arrays (SourceLW2Str,SourceLWNoScat,SourceSW2Str), and radiative fluxes (FluxLW,FluxSW,FluxBand) from(nlev, ncol)to contiguous column-major(ncol, nlev)order, adjacent GPU warp threads perform consecutive, coalesced memory reads and writes (stride-1overgcol).This preserves the public API and host data-exchange contract without adding runtime allocations or latency. All public getters (
lw_flux_up,lw_flux_dn,lw_flux_net,sw_flux_up,sw_flux_dn,sw_direct_flux_dn,sw_flux_net,net_flux,clear_*, and allspectral_*_flux_*arrays) wrap internal buffers viaPermutedDimsArray(buffer, (2, 1)).Measured GPU speedups across uniform DYAMOND grids (86,400 horizontal columns × 63 vertical layers, 480 spectral g-points) on NVIDIA A100-SXM4-40GB:
Per-kernel execution time comparison against baseline: