From e4c2bdd080e7407a0e321acd4d399d2f0c3a6cf6 Mon Sep 17 00:00:00 2001 From: Tapio Schneider Date: Sun, 12 Jul 2026 01:37:39 +0000 Subject: [PATCH 1/6] Performance: Transition internal scratch & output buffers to coalesced (ncol, nlev) layout for GPU speedups while preserving public API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- docs/src/howto/gpu.md | 14 ++-- .../nvidia_a100_sxm4_40gb.txt | 26 +++--- src/api/getters.jl | 51 +++++++----- src/api/update_fluxes.jl | 14 ++-- src/optics/Fluxes.jl | 80 ++++++++++--------- src/optics/Sources.jl | 42 +++++----- src/optics/compute_optical_props.jl | 31 ++++--- src/optics/gray_optics_kernels.jl | 28 +++---- src/optics/optical_props.jl | 28 +++---- src/rte/driver_utils.jl | 20 ++--- src/rte/longwave_2stream.jl | 30 +++---- src/rte/longwave_noscat.jl | 21 +++-- src/rte/shortwave_2stream.jl | 40 +++++----- src/rte/shortwave_noscat.jl | 16 ++-- ...sky_with_aerosols_dyamond_gpu_benchmark.jl | 2 +- test/all_sky_with_aerosols_utils.jl | 33 ++++---- test/clear_sky_dyamond_gpu_benchmark.jl | 2 +- test/clear_sky_utils.jl | 12 +-- test/cloudy_sky_dyamond_gpu_benchmark.jl | 2 +- test/cloudy_sky_utils.jl | 16 ++-- test/cos_zenith_edge_cases.jl | 16 ++-- test/float32_consistency.jl | 16 ++-- test/gray_atm_utils.jl | 8 +- test/partial_cloud_fraction.jl | 4 +- test/standalone.jl | 4 +- 25 files changed, 283 insertions(+), 273 deletions(-) diff --git a/docs/src/howto/gpu.md b/docs/src/howto/gpu.md index 188fa6e32..1f7327e88 100644 --- a/docs/src/howto/gpu.md +++ b/docs/src/howto/gpu.md @@ -105,13 +105,13 @@ Every configuration evaluated by the benchmark (see [`perf/benchmark_ratchet.jl` ### Single precision (`Float32`) — Primary target -At `Float32` precision, memory traffic across lookup-table interpolations is halved compared to `Float64`, enabling large GPU speedups: +At `Float32` precision, memory traffic across lookup-table interpolations and solver evaluations is halved compared to `Float64`, enabling massive GPU speedups when combined with contiguous `(ncol, nlev)` internal memory coalescing: | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup | |---|---|---|---| -| Clear-sky (two-stream) | 77.95 s | 1.24 s | 63.1× | -| All-sky, McICA clouds | 98.42 s | 1.82 s | 54.1× | -| All-sky with aerosols | 99.39 s | 1.95 s | 51.0× | +| Clear-sky (two-stream) | 77.95 s | 0.95 s | 82.5× | +| All-sky, McICA clouds | 98.42 s | 1.39 s | 70.8× | +| All-sky with aerosols | 99.39 s | 1.52 s | 65.5× | ### Double precision (`Float64`) @@ -119,6 +119,6 @@ For reference, running the identical uniform benchmark sweep in double precision | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup | |---|---|---|---| -| Clear-sky (two-stream) | 79.64 s | 1.70 s | 46.8× | -| All-sky, McICA clouds | 101.60 s | 2.43 s | 41.8× | -| All-sky with aerosols | 106.33 s | 2.59 s | 41.1× | +| Clear-sky (two-stream) | 79.64 s | 1.32 s | 60.1× | +| All-sky, McICA clouds | 101.60 s | 1.93 s | 52.6× | +| All-sky with aerosols | 106.33 s | 2.07 s | 51.4× | diff --git a/perf/benchmark_baselines/nvidia_a100_sxm4_40gb.txt b/perf/benchmark_baselines/nvidia_a100_sxm4_40gb.txt index 96f3cdd44..079f32dc0 100644 --- a/perf/benchmark_baselines/nvidia_a100_sxm4_40gb.txt +++ b/perf/benchmark_baselines/nvidia_a100_sxm4_40gb.txt @@ -1,17 +1,17 @@ -# RRTMGP DYAMOND GPU benchmark measurement (minimum ns) +# RRTMGP DYAMOND GPU benchmark baseline (minimum ns) # device: NVIDIA A100-SXM4-40GB # julia: 1.12.6 # recorded: local run # regenerate with RRTMGP_BENCH_WRITE_BASELINE=1 -all_sky_lw_Float32 9.76149447e8 -all_sky_lw_Float64 1.331256434e9 -all_sky_sw_Float32 8.43344196e8 -all_sky_sw_Float64 1.097443624e9 -all_sky_with_aerosols_lw_Float32 1.05467407e9 -all_sky_with_aerosols_lw_Float64 1.412722056e9 -all_sky_with_aerosols_sw_Float32 8.94524846e8 -all_sky_with_aerosols_sw_Float64 1.173976877e9 -clear_sky_lw_Float32 7.7122294e8 -clear_sky_lw_Float64 1.049297832e9 -clear_sky_sw_Float32 4.64954059e8 -clear_sky_sw_Float64 6.53718113e8 +all_sky_lw_Float32 7.26150458e8 +all_sky_lw_Float64 9.92342112e8 +all_sky_sw_Float32 6.71148384e8 +all_sky_sw_Float64 9.40083761e8 +all_sky_with_aerosols_lw_Float32 7.97787842e8 +all_sky_with_aerosols_lw_Float64 1.09698575e9 +all_sky_with_aerosols_sw_Float32 7.28343007e8 +all_sky_with_aerosols_sw_Float64 9.66275699e8 +clear_sky_lw_Float32 4.89919056e8 +clear_sky_lw_Float64 6.82068605e8 +clear_sky_sw_Float32 4.55409919e8 +clear_sky_sw_Float64 6.39967283e8 diff --git a/src/api/getters.jl b/src/api/getters.jl index f059e9b2d..67f04aa5a 100644 --- a/src/api/getters.jl +++ b/src/api/getters.jl @@ -45,6 +45,13 @@ _domain_view(isothermal_boundary_layer::Bool, x::AbstractArray) = _domain_view(isothermal_boundary_layer::Bool, x::AbstractArray{<:Any, 3}) = view(x, 1:(size(x, 1) - Int(isothermal_boundary_layer)), :, :) +_flux_domain_view(s::RRTMGPSolver, x) = _flux_domain_view(s.grid_params.isothermal_boundary_layer, x) +_flux_domain_view(::Bool, ::Nothing) = nothing +_flux_domain_view(isothermal_boundary_layer::Bool, x::AbstractArray{<:Any, 2}) = + view(PermutedDimsArray(x, (2, 1)), 1:(size(x, 2) - Int(isothermal_boundary_layer)), :) +_flux_domain_view(isothermal_boundary_layer::Bool, x::AbstractArray{<:Any, 3}) = + view(PermutedDimsArray(x, (2, 1, 3)), 1:(size(x, 2) - Int(isothermal_boundary_layer)), :, :) + # The clear-sky flux workspaces exist only on `AllSkyRadiationWithClearSkyDiagnostics` # solvers; give the `clear_*` getters an informative error elsewhere (instead of a # cryptic `getproperty(::Nothing, ...)`). @@ -57,21 +64,21 @@ _require_clear_sky(x) = x # Potentially views top_of_atmosphere_lw_flux_dn(s::RRTMGPSolver) = _maybe_transpose(s.lws.bcs.inc_flux) top_of_atmosphere_diffuse_sw_flux_dn(s::RRTMGPSolver) = _maybe_transpose(s.sws.bcs.inc_flux_diffuse) -lw_flux_up(s::RRTMGPSolver) = _domain_view(s, s.lws.flux.flux_up) -lw_flux_dn(s::RRTMGPSolver) = _domain_view(s, s.lws.flux.flux_dn) -lw_flux_net(s::RRTMGPSolver) = _domain_view(s, s.lws.flux.flux_net) -clear_lw_flux_up(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_up) -clear_lw_flux_dn(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_dn) -clear_lw_flux(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_net) +lw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, s.lws.flux.flux_up) +lw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, s.lws.flux.flux_dn) +lw_flux_net(s::RRTMGPSolver) = _flux_domain_view(s, s.lws.flux.flux_net) +clear_lw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_up) +clear_lw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_dn) +clear_lw_flux(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_net) surface_emissivity(s::RRTMGPSolver) = s.lws.bcs.sfc_emis -sw_flux_up(s::RRTMGPSolver) = _domain_view(s, s.sws.flux.flux_up) -sw_flux_dn(s::RRTMGPSolver) = _domain_view(s, s.sws.flux.flux_dn) -sw_flux_net(s::RRTMGPSolver) = _domain_view(s, s.sws.flux.flux_net) -sw_direct_flux_dn(s::RRTMGPSolver) = _domain_view(s, s.sws.flux.flux_dn_dir) -clear_sw_flux_up(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_up) -clear_sw_flux_dn(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn) -clear_sw_direct_flux_dn(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn_dir) -clear_sw_flux(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_net) +sw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_up) +sw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_dn) +sw_flux_net(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_net) +sw_direct_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_dn_dir) +clear_sw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_up) +clear_sw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn) +clear_sw_direct_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn_dir) +clear_sw_flux(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_net) cloud_liquid_effective_radius(s::RRTMGPSolver) = _domain_view(s, s.as.cloud_state.cld_r_eff_liq) cloud_ice_effective_radius(s::RRTMGPSolver) = _domain_view(s, s.as.cloud_state.cld_r_eff_ice) cloud_liquid_water_path(s::RRTMGPSolver) = _domain_view(s, s.as.cloud_state.cld_path_liq) @@ -94,8 +101,8 @@ layer_temperature(s::RRTMGPSolver) = _domain_view(s, getview_ layer_relative_humidity(s::RRTMGPSolver) = _domain_view(s, getview_rel_hum(s.as)) level_pressure(s::RRTMGPSolver) = _domain_view(s, s.as.p_lev) level_temperature(s::RRTMGPSolver) = _domain_view(s, s.as.t_lev) -net_flux(s::RRTMGPSolver) = _domain_view(s, s.net_flux_buffer) -clear_net_flux(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_net_flux_buffer)) +net_flux(s::RRTMGPSolver) = _flux_domain_view(s, s.net_flux_buffer) +clear_net_flux(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_net_flux_buffer)) deep_atmosphere_inverse_scaling(s::RRTMGPSolver) = _domain_view(s, s.deep_atmosphere_inverse_scaling) #! format: on @@ -155,17 +162,17 @@ wavenumber range. The `up`/`dn`/`net` getters are views into the retained per-band buffers. """ spectral_lw_flux_up(s::RRTMGPSolver) = - _domain_view(s, _solver_band_flux(s.lws).flux_up) + _flux_domain_view(s, _solver_band_flux(s.lws).flux_up) spectral_lw_flux_dn(s::RRTMGPSolver) = - _domain_view(s, _solver_band_flux(s.lws).flux_dn) + _flux_domain_view(s, _solver_band_flux(s.lws).flux_dn) spectral_lw_flux_net(s::RRTMGPSolver) = - _domain_view(s, _solver_band_flux(s.lws).flux_net) + _flux_domain_view(s, _solver_band_flux(s.lws).flux_net) spectral_sw_flux_up(s::RRTMGPSolver) = - _domain_view(s, _solver_band_flux(s.sws).flux_up) + _flux_domain_view(s, _solver_band_flux(s.sws).flux_up) spectral_sw_flux_dn(s::RRTMGPSolver) = - _domain_view(s, _solver_band_flux(s.sws).flux_dn) + _flux_domain_view(s, _solver_band_flux(s.sws).flux_dn) spectral_sw_flux_net(s::RRTMGPSolver) = - _domain_view(s, _solver_band_flux(s.sws).flux_net) + _flux_domain_view(s, _solver_band_flux(s.sws).flux_net) """ lw_band_bounds(s::RRTMGPSolver) diff --git a/src/api/update_fluxes.jl b/src/api/update_fluxes.jl index d216fe819..02b2554bf 100644 --- a/src/api/update_fluxes.jl +++ b/src/api/update_fluxes.jl @@ -46,9 +46,9 @@ function update_lw_fluxes!( lookups.lookup_lw_aero, ms, ) - parent(clear_lw_flux_up(s)) .= parent(lw_flux_up(s)) - parent(clear_lw_flux_dn(s)) .= parent(lw_flux_dn(s)) - parent(clear_lw_flux(s)) .= parent(lw_flux_net(s)) + s.clear_flux_lw.flux_up .= s.lws.flux.flux_up + s.clear_flux_lw.flux_dn .= s.lws.flux.flux_dn + s.clear_flux_lw.flux_net .= s.lws.flux.flux_net RTESolver.solve_lw!( lw_solver, as, @@ -103,10 +103,10 @@ function update_sw_fluxes!( lookups.lookup_sw_aero, ms, ) - parent(clear_sw_flux_up(s)) .= parent(sw_flux_up(s)) - parent(clear_sw_flux_dn(s)) .= parent(sw_flux_dn(s)) - parent(clear_sw_direct_flux_dn(s)) .= parent(sw_direct_flux_dn(s)) - parent(clear_sw_flux(s)) .= parent(sw_flux_net(s)) + s.clear_flux_sw.flux_up .= s.sws.flux.flux_up + s.clear_flux_sw.flux_dn .= s.sws.flux.flux_dn + s.clear_flux_sw.flux_dn_dir .= s.sws.flux.flux_dn_dir + s.clear_flux_sw.flux_net .= s.sws.flux.flux_net RTESolver.solve_sw!( sw_solver, diff --git a/src/optics/Fluxes.jl b/src/optics/Fluxes.jl index ea6e738bd..9c3abf2a2 100644 --- a/src/optics/Fluxes.jl +++ b/src/optics/Fluxes.jl @@ -47,9 +47,9 @@ function FluxLW(grid_params::RRTMGPGridParams) (; ncol, nlay) = grid_params DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) - flux_up = DA{FT}(undef, nlay + 1, ncol) - flux_dn = DA{FT}(undef, nlay + 1, ncol) - flux_net = DA{FT}(undef, nlay + 1, ncol) + flux_up = DA{FT}(undef, ncol, nlay + 1) + flux_dn = DA{FT}(undef, ncol, nlay + 1) + flux_net = DA{FT}(undef, ncol, nlay + 1) return FluxLW{FT, typeof(flux_net)}(flux_up, flux_dn, flux_net) end @@ -59,10 +59,10 @@ end Upward, downward and net shortwave fluxes at each level. # Fields -- `flux_up`: Upward flux [W/m²] `(nlev, ncol)`. -- `flux_dn`: Downward flux [W/m²] `(nlev, ncol)`. -- `flux_net`: Net flux [W/m²] `(nlev, ncol)`. -- `flux_dn_dir`: Direct downward flux [W/m²] `(nlev, ncol)`. +- `flux_up`: Upward flux [W/m²] `(ncol, nlev)`. +- `flux_dn`: Downward flux [W/m²] `(ncol, nlev)`. +- `flux_net`: Net flux [W/m²] `(ncol, nlev)`. +- `flux_dn_dir`: Direct downward flux [W/m²] `(ncol, nlev)`. """ struct FluxSW{FT <: AbstractFloat, FTA2D <: AbstractArray{FT, 2}} <: AbstractFlux{FT, FTA2D} @@ -84,10 +84,10 @@ function FluxSW(grid_params::RRTMGPGridParams) (; nlay, ncol) = grid_params FT = eltype(grid_params) DA = ClimaComms.array_type(grid_params) - flux_up = DA{FT}(undef, nlay + 1, ncol) - flux_dn = DA{FT}(undef, nlay + 1, ncol) - flux_net = DA{FT}(undef, nlay + 1, ncol) - flux_dn_dir = DA{FT}(undef, nlay + 1, ncol) + flux_up = DA{FT}(undef, ncol, nlay + 1) + flux_dn = DA{FT}(undef, ncol, nlay + 1) + flux_net = DA{FT}(undef, ncol, nlay + 1) + flux_dn_dir = DA{FT}(undef, ncol, nlay + 1) return FluxSW{FT, typeof(flux_net)}(flux_up, flux_dn, flux_net, flux_dn_dir) end @@ -95,14 +95,14 @@ end FluxBand{FT, FTA3D} Optional per-band upward, downward, and net radiative fluxes at each level, -`(nlev, ncol, n_bnd)`. Only allocated when spectrally-resolved fluxes are requested. -Band `b`'s slice `[:, :, b]` has the same `(nlev, ncol)` layout as the broadband +`(ncol, nlev, n_bnd)`. Only allocated when spectrally-resolved fluxes are requested. +Band `b`'s slice `[:, :, b]` has the same `(ncol, nlev)` layout as the broadband fluxes, and summing over the band dimension recovers the broadband fluxes. # Fields -- `flux_up`: upward flux per band [W/m²], `(nlev, ncol, n_bnd)`. -- `flux_dn`: downward flux per band [W/m²], `(nlev, ncol, n_bnd)`. -- `flux_net`: net flux per band (`flux_up - flux_dn`) [W/m²], `(nlev, ncol, n_bnd)`. +- `flux_up`: upward flux per band [W/m²], `(ncol, nlev, n_bnd)`. +- `flux_dn`: downward flux per band [W/m²], `(ncol, nlev, n_bnd)`. +- `flux_net`: net flux per band (`flux_up - flux_dn`) [W/m²], `(ncol, nlev, n_bnd)`. """ struct FluxBand{FT <: AbstractFloat, FTA3D <: AbstractArray{FT, 3}} flux_up::FTA3D @@ -115,9 +115,9 @@ function FluxBand(grid_params::RRTMGPGridParams, n_bnd::Int) (; nlay, ncol) = grid_params FT = eltype(grid_params) DA = ClimaComms.array_type(grid_params) - flux_up = DA{FT}(undef, nlay + 1, ncol, n_bnd) - flux_dn = DA{FT}(undef, nlay + 1, ncol, n_bnd) - flux_net = DA{FT}(undef, nlay + 1, ncol, n_bnd) + flux_up = DA{FT}(undef, ncol, nlay + 1, n_bnd) + flux_dn = DA{FT}(undef, ncol, nlay + 1, n_bnd) + flux_net = DA{FT}(undef, ncol, nlay + 1, n_bnd) return FluxBand{FT, typeof(flux_up)}(flux_up, flux_dn, flux_net) end @@ -131,7 +131,7 @@ function set_band_flux_to_zero!(band::FluxBand{FT}) where {FT} end # Add one g-point's up/down flux (column `gcol`, per-g-point scratch `flux_up`/`flux_dn` -# of shape `(nlev, ncol)`) into its band `ibnd`. No-op when spectral fluxes are off, so +# of shape `(ncol, nlev)`) into its band `ibnd`. No-op when spectral fluxes are off, so # the broadband path pays nothing (the branch is specialized away on `::Nothing`). @inline accumulate_band_flux!(::Nothing, flux_up, flux_dn, gcol, ibnd, nlev) = nothing @@ -145,8 +145,8 @@ end ) bu, bd = band.flux_up, band.flux_dn @inbounds for ilev in 1:nlev - bu[ilev, gcol, ibnd] += flux_up[ilev, gcol] - bd[ilev, gcol, ibnd] += flux_dn[ilev, gcol] + bu[gcol, ilev, ibnd] += flux_up[gcol, ilev] + bd[gcol, ilev, ibnd] += flux_dn[gcol, ilev] end return nothing end @@ -163,14 +163,14 @@ Compute the net flux for column `gcol` across `nlev` levels: (; flux_up, flux_dn, flux_net) = flux @inbounds begin for ilev in 1:nlev - flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol] + flux_net[gcol, ilev] = flux_up[gcol, ilev] - flux_dn[gcol, ilev] end end return nothing end @inline compute_net_flux!(flux::AbstractFlux, gcol) = - compute_net_flux!(flux, gcol, size(flux.flux_up, 1)) + compute_net_flux!(flux, gcol, size(flux.flux_up, 2)) """ set_flux_to_zero!(flux::FluxLW{FT}, gcol::Int, nlev::Int) where {FT<:AbstractFloat} @@ -185,15 +185,15 @@ Set longwave flux for column `gcol` to zero across `nlev` levels. ) where {FT <: AbstractFloat} (; flux_up, flux_dn, flux_net) = flux @inbounds for ilev in 1:nlev - flux_up[ilev, gcol] = FT(0) - flux_dn[ilev, gcol] = FT(0) - flux_net[ilev, gcol] = FT(0) + flux_up[gcol, ilev] = FT(0) + flux_dn[gcol, ilev] = FT(0) + flux_net[gcol, ilev] = FT(0) end return nothing end @inline set_flux_to_zero!(flux::FluxLW, gcol::Int) = - set_flux_to_zero!(flux, gcol, size(flux.flux_up, 1)) + set_flux_to_zero!(flux, gcol, size(flux.flux_up, 2)) """ set_flux_to_zero!(flux::FluxSW{FT}, gcol::Int, nlev::Int) where {FT<:AbstractFloat} @@ -208,16 +208,16 @@ Set shortwave flux for column `gcol` to zero across `nlev` levels. ) where {FT <: AbstractFloat} (; flux_up, flux_dn, flux_net, flux_dn_dir) = flux @inbounds for ilev in 1:nlev - flux_up[ilev, gcol] = FT(0) - flux_dn[ilev, gcol] = FT(0) - flux_net[ilev, gcol] = FT(0) - flux_dn_dir[ilev, gcol] = FT(0) + flux_up[gcol, ilev] = FT(0) + flux_dn[gcol, ilev] = FT(0) + flux_net[gcol, ilev] = FT(0) + flux_dn_dir[gcol, ilev] = FT(0) end return nothing end @inline set_flux_to_zero!(flux::FluxSW, gcol::Int) = - set_flux_to_zero!(flux, gcol, size(flux.flux_up, 1)) + set_flux_to_zero!(flux, gcol, size(flux.flux_up, 2)) """ apply_metric_scaling!(flux::Union{FluxLW, FluxSW}, metric_scaling) @@ -233,10 +233,11 @@ function apply_metric_scaling!( flux::AbstractFlux, metric_scaling::AbstractArray, ) - flux.flux_up .= flux.flux_up .* metric_scaling - flux.flux_dn .= flux.flux_dn .* metric_scaling - flux.flux_net .= flux.flux_net .* metric_scaling - flux isa FluxSW && (flux.flux_dn_dir .= flux.flux_dn_dir .* metric_scaling) + sc = PermutedDimsArray(metric_scaling, (2, 1)) + flux.flux_up .= flux.flux_up .* sc + flux.flux_dn .= flux.flux_dn .* sc + flux.flux_net .= flux.flux_net .* sc + flux isa FluxSW && (flux.flux_dn_dir .= flux.flux_dn_dir .* sc) return nothing end function apply_metric_scaling!( @@ -252,8 +253,9 @@ end apply_metric_scaling!(::Nothing, metric_scaling) = nothing apply_metric_scaling!(band::FluxBand, metric_scaling::Nothing) = nothing function apply_metric_scaling!(band::FluxBand, metric_scaling::AbstractArray) - band.flux_up .= band.flux_up .* metric_scaling - band.flux_dn .= band.flux_dn .* metric_scaling + sc = PermutedDimsArray(metric_scaling, (2, 1)) + band.flux_up .= band.flux_up .* sc + band.flux_dn .= band.flux_dn .* sc return nothing end diff --git a/src/optics/Sources.jl b/src/optics/Sources.jl index 6090af81f..ba55aafd1 100644 --- a/src/optics/Sources.jl +++ b/src/optics/Sources.jl @@ -41,8 +41,8 @@ function SourceLWNoScat(grid_params::RRTMGPGridParams; params::RP.ARP) DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) sfc_source = DA{FT, 1}(undef, ncol) - lay_source = DA{FT, 2}(undef, nlay, ncol) - lev_source = DA{FT, 2}(undef, nlay + 1, ncol) + lay_source = DA{FT, 2}(undef, ncol, nlay) + lev_source = DA{FT, 2}(undef, ncol, nlay + 1) return SourceLWNoScat{ typeof(sfc_source), @@ -66,10 +66,10 @@ and at the surface for 2-stream calculations. # Fields - `param_set`: Parameter set. - `sfc_source`: Surface source [W/m²] `(ncol)`. -- `leveldata`: Storage for level source, albedo and src `(3, nlay+1, ncol)`. -- `lev_source`: Level source [W/m²] `(nlay+1, ncol)`; used in 2-stream calculations. -- `albedo`: Temporary storage array `(nlay+1, ncol)`; used in 2-stream calculations. -- `src`: Temporary storage array `(nlay+1, ncol)`; used in 2-stream calculations. +- `leveldata`: Storage for level source, albedo and src `(ncol, nlay+1, 3)`. +- `lev_source`: Level source [W/m²] `(ncol, nlay+1)`; used in 2-stream calculations. +- `albedo`: Temporary storage array `(ncol, nlay+1)`; used in 2-stream calculations. +- `src`: Temporary storage array `(ncol, nlay+1)`; used in 2-stream calculations. """ struct SourceLW2Str{S, D, V, PS} <: AbstractSourceLW param_set::PS @@ -83,9 +83,9 @@ end function Adapt.adapt_structure(to, src::SourceLW2Str) sfc_source = Adapt.adapt(to, src.sfc_source) leveldata = Adapt.adapt(to, src.leveldata) - lev_source = view(leveldata, 1, :, :) - albedo = view(leveldata, 2, :, :) - src_view = view(leveldata, 3, :, :) + lev_source = view(leveldata, :, :, 1) + albedo = view(leveldata, :, :, 2) + src_view = view(leveldata, :, :, 3) return SourceLW2Str{ typeof(sfc_source), typeof(leveldata), @@ -106,10 +106,10 @@ function SourceLW2Str(grid_params::RRTMGPGridParams; params::RP.ARP) DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) sfc_source = DA{FT, 1}(undef, ncol) - leveldata = DA{FT, 3}(undef, 3, nlay + 1, ncol) - lev_source = view(leveldata, 1, :, :) - albedo = view(leveldata, 2, :, :) - src = view(leveldata, 3, :, :) + leveldata = DA{FT, 3}(undef, ncol, nlay + 1, 3) + lev_source = view(leveldata, :, :, 1) + albedo = view(leveldata, :, :, 2) + src = view(leveldata, :, :, 3) return SourceLW2Str{ typeof(sfc_source), @@ -135,9 +135,9 @@ and at the surface for 2-stream calculations. # Fields - `sfc_source`: Surface source `(ncol)`. -- `leveldata`: Storage for albedo and src `(2, nlay+1, ncol)`. -- `albedo`: Albedo `(nlay+1, ncol)`. -- `src`: Temporary storage array `(nlay+1, ncol)`; used in 2-stream calculations. +- `leveldata`: Storage for albedo and src `(ncol, nlay+1, 2)`. +- `albedo`: Albedo `(ncol, nlay+1)`. +- `src`: Temporary storage array `(ncol, nlay+1)`; used in 2-stream calculations. """ struct SourceSW2Str{S, D, V} sfc_source::S @@ -149,8 +149,8 @@ end function Adapt.adapt_structure(to, src::SourceSW2Str) sfc_source = Adapt.adapt(to, src.sfc_source) leveldata = Adapt.adapt(to, src.leveldata) - albedo = view(leveldata, 1, :, :) - src_view = view(leveldata, 2, :, :) + albedo = view(leveldata, :, :, 1) + src_view = view(leveldata, :, :, 2) return SourceSW2Str{typeof(sfc_source), typeof(leveldata), typeof(albedo)}( sfc_source, leveldata, @@ -164,9 +164,9 @@ function SourceSW2Str(grid_params::RRTMGPGridParams) DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) sfc_source = DA{FT, 1}(undef, ncol) - leveldata = DA{FT, 3}(undef, 2, nlay + 1, ncol) - albedo = view(leveldata, 1, :, :) - src = view(leveldata, 2, :, :) + leveldata = DA{FT, 3}(undef, ncol, nlay + 1, 2) + albedo = view(leveldata, :, :, 1) + src = view(leveldata, :, :, 2) return SourceSW2Str{typeof(sfc_source), typeof(leveldata), typeof(albedo)}( sfc_source, diff --git a/src/optics/compute_optical_props.jl b/src/optics/compute_optical_props.jl index 67d233e0b..b56c60d20 100644 --- a/src/optics/compute_optical_props.jl +++ b/src/optics/compute_optical_props.jl @@ -32,7 +32,7 @@ Compute optical properties for the longwave problem. totplnk = view(lkp.planck.tot_planck, :, ibnd) as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) t_lev_col = view(as.t_lev, :, gcol) - τ = view(op.τ, :, gcol) + τ = view(op.τ, gcol, :) lev_src_inc_prev = zero(t_sfc) lev_src_dec_prev = zero(t_sfc) @@ -57,7 +57,7 @@ Compute optical properties for the longwave problem. # compute longwave source terms t_lev_inc = t_lev_col[glay + 1] - lay_source[glay, gcol] = + lay_source[gcol, glay] = interp1d_equispaced(t_lay, t_planck, totplnk) * planckfrac lev_src_inc = interp1d_equispaced(t_lev_inc, t_planck, totplnk) * planckfrac @@ -66,15 +66,15 @@ Compute optical properties for the longwave problem. if glay == 1 sfc_source[gcol] = interp1d_equispaced(t_sfc, t_planck, totplnk) * planckfrac - lev_source[glay, gcol] = lev_src_dec + lev_source[gcol, glay] = lev_src_dec else - lev_source[glay, gcol] = sqrt(lev_src_inc_prev * lev_src_dec) + lev_source[gcol, glay] = sqrt(lev_src_inc_prev * lev_src_dec) end lev_src_dec_prev = lev_src_dec lev_src_inc_prev = lev_src_inc t_lev_dec = t_lev_inc end - lev_source[nlay + 1, gcol] = lev_src_inc_prev + lev_source[gcol, nlay + 1] = lev_src_inc_prev if !isnothing(lkp_cld) cloud_state = as.cloud_state cld_r_eff_liq = view(cloud_state.cld_r_eff_liq, :, gcol) @@ -141,9 +141,9 @@ end totplnk = view(lkp.planck.tot_planck, :, ibnd) as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) t_lev_col = view(as.t_lev, :, gcol) - τ = view(op.τ, :, gcol) - ssa = view(op.ssa, :, gcol) - g = view(op.g, :, gcol) + τ = view(op.τ, gcol, :) + ssa = view(op.ssa, gcol, :) + g = view(op.g, gcol, :) end lev_src_inc_prev = zero(t_sfc) @@ -177,15 +177,15 @@ end if glay == 1 sfc_source[gcol] = interp1d_equispaced(t_sfc, t_planck, totplnk) * planckfrac - lev_source[glay, gcol] = lev_src_dec + lev_source[gcol, glay] = lev_src_dec else - lev_source[glay, gcol] = sqrt(lev_src_inc_prev * lev_src_dec) + lev_source[gcol, glay] = sqrt(lev_src_inc_prev * lev_src_dec) end lev_src_dec_prev = lev_src_dec lev_src_inc_prev = lev_src_inc t_lev_dec = t_lev_inc end - lev_source[nlay + 1, gcol] = lev_src_inc_prev + lev_source[gcol, nlay + 1] = lev_src_inc_prev end if !isnothing(lkp_cld) # clouds need TwoStream optics cloud_state = as.cloud_state @@ -257,7 +257,6 @@ Compute optical properties for the shortwave problem. igpt::Int, lkp::LookUpSW, ::Nothing, - ::Nothing, ) nlay = AtmosphericStates.get_nlay(as) (; vmr) = as @@ -265,7 +264,7 @@ Compute optical properties for the shortwave problem. ibnd = lkp.band_data.major_gpt2bnd[igpt] t_sfc = as.t_sfc[gcol] as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) - τ = view(op.τ, :, gcol) + τ = view(op.τ, gcol, :) end @inbounds for glay in 1:nlay col_dry, p_lay, t_lay = @@ -301,9 +300,9 @@ end ibnd = lkp.band_data.major_gpt2bnd[igpt] t_sfc = as.t_sfc[gcol] as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) - τ = view(op.τ, :, gcol) - ssa = view(op.ssa, :, gcol) - g = view(op.g, :, gcol) + τ = view(op.τ, gcol, :) + ssa = view(op.ssa, gcol, :) + g = view(op.g, gcol, :) end @inbounds for glay in 1:nlay col_dry, p_lay, t_lay = diff --git a/src/optics/gray_optics_kernels.jl b/src/optics/gray_optics_kernels.jl index 860191ce1..dae38a7e7 100644 --- a/src/optics/gray_optics_kernels.jl +++ b/src/optics/gray_optics_kernels.jl @@ -37,28 +37,28 @@ function compute_optical_props!( p_lev_glayplus1 = p_lev[glay + 1, gcol] Δp = p_lev_glayplus1 - p_lev_glay p = p_lay[glay, gcol] - τ[glay, gcol] = + τ[gcol, glay] = compute_gray_optical_thickness_lw(otp, p0, Δp, p, lat) p_lev_glay = p_lev_glayplus1 # compute longwave source terms t_lev_inc = t_lev[glay + 1, gcol] t_lay_loc = t_lay[glay, gcol] - lay_source[glay, gcol] = + lay_source[gcol, glay] = sbc * (t_lay_loc * t_lay_loc * t_lay_loc * t_lay_loc) / FT(π) # computing lay_source lev_src_inc = sbc * (t_lev_inc * t_lev_inc * t_lev_inc * t_lev_inc) / FT(π) lev_src_dec = sbc * (t_lev_dec * t_lev_dec * t_lev_dec * t_lev_dec) / FT(π) if glay == 1 - lev_source[glay, gcol] = lev_src_dec + lev_source[gcol, glay] = lev_src_dec else - lev_source[glay, gcol] = sqrt(lev_src_inc_prev * lev_src_dec) + lev_source[gcol, glay] = sqrt(lev_src_inc_prev * lev_src_dec) end lev_src_dec_prev = lev_src_dec lev_src_inc_prev = lev_src_inc t_lev_dec = t_lev_inc end - lev_source[nlay + 1, gcol] = lev_src_inc_prev + lev_source[gcol, nlay + 1] = lev_src_inc_prev end return nothing end @@ -88,7 +88,7 @@ function compute_optical_props!( p_lev_glayplus1 = p_lev[glay + 1, gcol] Δp = p_lev_glayplus1 - p_lev_glay p = p_lay[glay, gcol] - τ[glay, gcol] = + τ[gcol, glay] = compute_gray_optical_thickness_lw(otp, p0, Δp, p, lat) p_lev_glay = p_lev_glayplus1 # compute longwave source terms @@ -98,19 +98,19 @@ function compute_optical_props!( lev_src_dec = sbc * (t_lev_dec * t_lev_dec * t_lev_dec * t_lev_dec) / FT(π) if glay == 1 - lev_source[glay, gcol] = lev_src_dec + lev_source[gcol, glay] = lev_src_dec else - lev_source[glay, gcol] = sqrt(lev_src_inc_prev * lev_src_dec) + lev_source[gcol, glay] = sqrt(lev_src_inc_prev * lev_src_dec) end lev_src_dec_prev = lev_src_dec lev_src_inc_prev = lev_src_inc t_lev_dec = t_lev_inc end - lev_source[nlay + 1, gcol] = lev_src_inc_prev + lev_source[gcol, nlay + 1] = lev_src_inc_prev end zeroval = zero(FT) - map!(x -> zeroval, view(ssa, :, gcol), view(ssa, :, gcol)) - map!(x -> zeroval, view(g, :, gcol), view(g, :, gcol)) + map!(x -> zeroval, view(ssa, gcol, :), view(ssa, gcol, :)) + map!(x -> zeroval, view(g, gcol, :), view(g, gcol, :)) return nothing end @@ -139,7 +139,7 @@ function compute_optical_props!( @inbounds p_lev_glayplus1 = p_lev[glay + 1, gcol] @inbounds Δp = p_lev_glayplus1 - p_lev_glay @inbounds p = p_lay[glay, gcol] - @inbounds τ[glay, gcol] = + @inbounds τ[gcol, glay] = compute_gray_optical_thickness_sw(otp, p0, Δp, p, lat) p_lev_glay = p_lev_glayplus1 end @@ -147,8 +147,8 @@ function compute_optical_props!( (; ssa, g) = op FT = eltype(τ) zeroval = zero(FT) - map!(x -> zeroval, view(ssa, :, gcol), view(ssa, :, gcol)) - map!(x -> zeroval, view(g, :, gcol), view(g, :, gcol)) + map!(x -> zeroval, view(ssa, gcol, :), view(ssa, gcol, :)) + map!(x -> zeroval, view(g, gcol, :), view(g, gcol, :)) end return nothing end diff --git a/src/optics/optical_props.jl b/src/optics/optical_props.jl index c3a3d1dde..854f5b021 100644 --- a/src/optics/optical_props.jl +++ b/src/optics/optical_props.jl @@ -15,7 +15,7 @@ end function Adapt.adapt_structure(to, op::OneScalar) layerdata = Adapt.adapt(to, op.layerdata) - τ = view(layerdata, 1, :, :) + τ = view(layerdata, :, :, 1) return OneScalar{typeof(layerdata), typeof(τ)}(layerdata, τ) end @@ -23,8 +23,8 @@ function OneScalar(grid_params::RRTMGPGridParams) (; ncol, nlay) = grid_params DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) - layerdata = DA{FT, 3}(undef, 1, nlay, ncol) - τ = view(layerdata, 1, :, :) + layerdata = DA{FT, 3}(undef, ncol, nlay, 1) + τ = view(layerdata, :, :, 1) V = typeof(τ) return OneScalar{typeof(layerdata), V}(layerdata, τ) end @@ -39,9 +39,9 @@ calculations accounting for extinction and emission. # Fields - `layerdata`: Storage for optical depth, single scattering albedo and asymmetry parameter. -- `τ`: View into optical depth. -- `ssa`: View into single scattering albedo. -- `g`: View into asymmetry parameter. +- `τ`: View into optical depth `(ncol, nlay)`. +- `ssa`: View into single scattering albedo `(ncol, nlay)`. +- `g`: View into asymmetry parameter `(ncol, nlay)`. """ struct TwoStream{D, V} <: AbstractOpticalProps layerdata::D @@ -52,9 +52,9 @@ end function Adapt.adapt_structure(to, op::TwoStream) layerdata = Adapt.adapt(to, op.layerdata) - τ = view(layerdata, 1, :, :) - ssa = view(layerdata, 2, :, :) - g = view(layerdata, 3, :, :) + τ = view(layerdata, :, :, 1) + ssa = view(layerdata, :, :, 2) + g = view(layerdata, :, :, 3) return TwoStream{typeof(layerdata), typeof(τ)}(layerdata, τ, ssa, g) end @@ -62,13 +62,13 @@ function TwoStream(grid_params::RRTMGPGridParams) (; ncol, nlay) = grid_params DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) - layerdata = DA{FT, 3}(zeros(3, nlay, ncol)) - V = typeof(view(layerdata, 1, :, :)) + layerdata = DA{FT, 3}(zeros(ncol, nlay, 3)) + V = typeof(view(layerdata, :, :, 1)) return TwoStream{typeof(layerdata), V}( layerdata, - view(layerdata, 1, :, :), - view(layerdata, 2, :, :), - view(layerdata, 3, :, :), + view(layerdata, :, :, 1), + view(layerdata, :, :, 2), + view(layerdata, :, :, 3), ) end diff --git a/src/rte/driver_utils.jl b/src/rte/driver_utils.jl index aac9fc0ed..6fef367d9 100644 --- a/src/rte/driver_utils.jl +++ b/src/rte/driver_utils.jl @@ -45,13 +45,13 @@ end up, dn = flux.flux_up, flux.flux_dn if igpt == 1 @inbounds for ilev in 1:nlev - acc_up[ilev, gcol] = up[ilev, gcol] - acc_dn[ilev, gcol] = dn[ilev, gcol] + acc_up[gcol, ilev] = up[gcol, ilev] + acc_dn[gcol, ilev] = dn[gcol, ilev] end else @inbounds for ilev in 1:nlev - acc_up[ilev, gcol] += up[ilev, gcol] - acc_dn[ilev, gcol] += dn[ilev, gcol] + acc_up[gcol, ilev] += up[gcol, ilev] + acc_dn[gcol, ilev] += dn[gcol, ilev] end end return nothing @@ -69,15 +69,15 @@ end up, dn, dir = flux.flux_up, flux.flux_dn, flux.flux_dn_dir if igpt == 1 @inbounds for ilev in 1:nlev - acc_up[ilev, gcol] = up[ilev, gcol] - acc_dn[ilev, gcol] = dn[ilev, gcol] - acc_dir[ilev, gcol] = dir[ilev, gcol] + acc_up[gcol, ilev] = up[gcol, ilev] + acc_dn[gcol, ilev] = dn[gcol, ilev] + acc_dir[gcol, ilev] = dir[gcol, ilev] end else @inbounds for ilev in 1:nlev - acc_up[ilev, gcol] += up[ilev, gcol] - acc_dn[ilev, gcol] += dn[ilev, gcol] - acc_dir[ilev, gcol] += dir[ilev, gcol] + acc_up[gcol, ilev] += up[gcol, ilev] + acc_dn[gcol, ilev] += dn[gcol, ilev] + acc_dir[gcol, ilev] += dir[gcol, ilev] end end return nothing diff --git a/src/rte/longwave_2stream.jl b/src/rte/longwave_2stream.jl index c67c89a36..846ababdd 100644 --- a/src/rte/longwave_2stream.jl +++ b/src/rte/longwave_2stream.jl @@ -257,21 +257,21 @@ Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08) FT = eltype(τ) @inbounds flux_dn_ilevplus1 = isnothing(inc_flux) ? FT(0) : inc_flux[gcol, igpt] - @inbounds flux_dn[nlev, gcol] = flux_dn_ilevplus1 + @inbounds flux_dn[gcol, nlev] = flux_dn_ilevplus1 # Albedo of lowest level is the surface albedo... @inbounds albedo_ilev = FT(1) - sfc_emis[ibnd, gcol] - @inbounds albedo[1, gcol] = albedo_ilev + @inbounds albedo[gcol, 1] = albedo_ilev # ... and source of diffuse radiation is surface emission @inbounds src_ilev = FT(π) * sfc_emis[ibnd, gcol] * sfc_source[gcol] - @inbounds src[1, gcol] = src_ilev#FT(π) * sfc_emis[ibnd, gcol] * sfc_source[gcol] + @inbounds src[gcol, 1] = src_ilev#FT(π) * sfc_emis[ibnd, gcol] * sfc_source[gcol] # From bottom to top of atmosphere -- # compute albedo and source of upward radiation - @inbounds lev_src_bot = lev_source[1, gcol] + @inbounds lev_src_bot = lev_source[gcol, 1] @inbounds for ilev in 1:nlay - lev_src_top = lev_source[ilev + 1, gcol] + lev_src_top = lev_source[gcol, ilev + 1] τ_lay, ssa_lay, g_lay = - τ[ilev, gcol], ssa[ilev, gcol], g[ilev, gcol] + τ[gcol, ilev], ssa[gcol, ilev], g[gcol, ilev] Rdif, Tdif, src_up, src_dn = lw_2stream_coeffs( τ_lay, ssa_lay, @@ -287,25 +287,25 @@ Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08) # transmitted through the layer and reflected from layers below (Tdiff*src*albedo) src_ilevplus1 = src_up + Tdif * denom * (src_ilev + albedo_ilev * src_dn) - albedo[ilev + 1, gcol], src[ilev + 1, gcol] = + albedo[gcol, ilev + 1], src[gcol, ilev + 1] = albedo_ilevplus1, src_ilevplus1 lev_src_bot = lev_src_top albedo_ilev, src_ilev = albedo_ilevplus1, src_ilevplus1 end # Eq 12, at the top of the domain upwelling diffuse is due to ... - @inbounds flux_up[nlev, gcol] = - flux_dn_ilevplus1 * albedo[nlev, gcol] + # ... reflection of incident diffuse and - src[nlev, gcol] # scattering by the direct beam below + @inbounds flux_up[gcol, nlev] = + flux_dn_ilevplus1 * albedo[gcol, nlev] + # ... reflection of incident diffuse and + src[gcol, nlev] # scattering by the direct beam below # From the top of the atmosphere downward -- compute fluxes - @inbounds lev_src_top = lev_source[nlay + 1, gcol] + @inbounds lev_src_top = lev_source[gcol, nlay + 1] ilev = nlay @inbounds while ilev ≥ 1 lev_src_bot, albedo_ilev, src_ilev = - lev_source[ilev, gcol], albedo[ilev, gcol], src[ilev, gcol] + lev_source[gcol, ilev], albedo[gcol, ilev], src[gcol, ilev] τ_lay, ssa_lay, g_lay = - τ[ilev, gcol], ssa[ilev, gcol], g[ilev, gcol] + τ[gcol, ilev], ssa[gcol, ilev], g[gcol, ilev] Rdif, Tdif, _, src_dn = lw_2stream_coeffs( τ_lay, ssa_lay, @@ -319,10 +319,10 @@ Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08) (Tdif * flux_dn_ilevplus1 + # Equation 13 Rdif * src_ilev + src_dn) * denom - flux_up[ilev, gcol] = + flux_up[gcol, ilev] = flux_dn_ilev * albedo_ilev + # Equation 12 src_ilev - flux_dn[ilev, gcol] = flux_dn_ilev + flux_dn[gcol, ilev] = flux_dn_ilev flux_dn_ilevplus1 = flux_dn_ilev lev_src_top = lev_src_bot ilev -= 1 diff --git a/src/rte/longwave_noscat.jl b/src/rte/longwave_noscat.jl index ea7329d8c..24b8e6758 100644 --- a/src/rte/longwave_noscat.jl +++ b/src/rte/longwave_noscat.jl @@ -240,26 +240,26 @@ Transport for no-scattering longwave problem. intensity_dn_ilevplus1 = isnothing(inc_flux) ? FT(0) : inc_flux[gcol, igpt] * flux_to_intensity - @inbounds flux_dn[nlev, gcol] = + @inbounds flux_dn[gcol, nlev] = intensity_dn_ilevplus1 * intensity_to_flux # Top of domain is index nlev # Downward propagation ilev = nlay @inbounds while ilev ≥ 1 - τ_loc = τ[ilev, gcol] * Ds + τ_loc = τ[gcol, ilev] * Ds trans = exp(-τ_loc) - lay_src = lay_source[ilev, gcol] + lay_src = lay_source[gcol, ilev] intensity_dn_ilev = trans * intensity_dn_ilevplus1 + lw_noscat_source_dn( - lev_source[ilev, gcol], + lev_source[gcol, ilev], lay_src, τ_loc, trans, τ_thresh, ) intensity_dn_ilevplus1 = intensity_dn_ilev - flux_dn[ilev, gcol] = intensity_dn_ilev * intensity_to_flux + flux_dn[gcol, ilev] = intensity_dn_ilev * intensity_to_flux ilev -= 1 end @@ -267,24 +267,23 @@ Transport for no-scattering longwave problem. @inbounds intensity_up_ilevminus1 = intensity_dn_ilevplus1 * (FT(1) - sfc_emis[ibnd, gcol]) + sfc_emis[ibnd, gcol] * sfc_source[gcol] - #flux_dn[1, gcol] * (FT(1) - sfc_emis[ibnd, gcol]) + sfc_emis[ibnd, gcol] * sfc_source[gcol] - @inbounds flux_up[1, gcol] = intensity_up_ilevminus1 * intensity_to_flux + @inbounds flux_up[gcol, 1] = intensity_up_ilevminus1 * intensity_to_flux # Upward propagation @inbounds for ilev in 2:(nlay + 1) - τ_loc = τ[ilev - 1, gcol] * Ds + τ_loc = τ[gcol, ilev - 1] * Ds trans = exp(-τ_loc) - lay_src = lay_source[ilev - 1, gcol] + lay_src = lay_source[gcol, ilev - 1] intensity_up_ilev = trans * intensity_up_ilevminus1 + lw_noscat_source_up( - lev_source[ilev, gcol], + lev_source[gcol, ilev], lay_src, τ_loc, trans, τ_thresh, ) intensity_up_ilevminus1 = intensity_up_ilev - flux_up[ilev, gcol] = intensity_up_ilev * intensity_to_flux + flux_up[gcol, ilev] = intensity_up_ilev * intensity_to_flux end return nothing end diff --git a/src/rte/shortwave_2stream.jl b/src/rte/shortwave_2stream.jl index b55898f69..233ef1ce4 100644 --- a/src/rte/shortwave_2stream.jl +++ b/src/rte/shortwave_2stream.jl @@ -315,27 +315,27 @@ Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08) # cumulative optical depth and stored for the two passes below. flux_dn_dir_top = toa_flux * solar_frac * μ₀ inv_μ₀ = FT(1) / max(μ₀, Numerics.μ₀_min(FT)) - @inbounds flux_dn_dir[nlev, gcol] = flux_dn_dir_top + @inbounds flux_dn_dir[gcol, nlev] = flux_dn_dir_top τ_cum = FT(0) @inbounds for ilev in nlay:-1:1 - τ_cum += τ[ilev, gcol] - flux_dn_dir[ilev, gcol] = flux_dn_dir_top * exp(-τ_cum * inv_μ₀) + τ_cum += τ[gcol, ilev] + flux_dn_dir[gcol, ilev] = flux_dn_dir_top * exp(-τ_cum * inv_μ₀) end - flux_dn_dir_bot = @inbounds flux_dn_dir[1, gcol] # surface value + flux_dn_dir_bot = @inbounds flux_dn_dir[gcol, 1] # surface value sfc_source = flux_dn_dir_bot * sfc_alb_direct - @inbounds flux_dn[nlev, gcol] = FT(0) # set to incoming flux when provided? + @inbounds flux_dn[gcol, nlev] = FT(0) # set to incoming flux when provided? # Albedo of lowest level is the surface albedo... @inbounds surface_albedo = - albedo[1, gcol] = bcs_sw.sfc_alb_diffuse[ibnd, gcol] + albedo[gcol, 1] = bcs_sw.sfc_alb_diffuse[ibnd, gcol] # ... and source of diffuse radiation is surface emission - @inbounds src[1, gcol] = sfc_source + @inbounds src[gcol, 1] = sfc_source # From bottom to top of atmosphere -- # compute albedo and source of upward radiation albedo_ilev, src_ilev = surface_albedo, sfc_source @inbounds for ilev in 1:nlay τ_ilev, ssa_ilev, g_ilev = - τ[ilev, gcol], ssa[ilev, gcol], g[ilev, gcol] + τ[gcol, ilev], ssa[gcol, ilev], g[gcol, ilev] (Rdir, Tdir, _, Rdif, Tdif) = sw_2stream_coeffs(τ_ilev, ssa_ilev, g_ilev, μ₀) denom = FT(1) / (FT(1) - Rdif * albedo_ilev) # Eq 10 @@ -344,43 +344,43 @@ Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08) # Equation 11 -- source is emitted upward radiation at top of layer plus # radiation emitted at bottom of layer, # transmitted through the layer and reflected from layers below (Tdiff*src*albedo) - flux_dn_dir_ilevplus1 = flux_dn_dir[ilev + 1, gcol] + flux_dn_dir_ilevplus1 = flux_dn_dir[gcol, ilev + 1] src_up_ilev = Rdir * flux_dn_dir_ilevplus1 src_dn_ilev = Tdir * flux_dn_dir_ilevplus1 src_ilevplus1 = src_up_ilev + Tdif * denom * (src_ilev + albedo_ilev * src_dn_ilev) - albedo[ilev + 1, gcol], src[ilev + 1, gcol] = + albedo[gcol, ilev + 1], src[gcol, ilev + 1] = albedo_ilevplus1, src_ilevplus1 albedo_ilev = albedo_ilevplus1 src_ilev = src_ilevplus1 end # Eq 12, at the top of the domain upwelling diffuse is due to ... - @inbounds flux_up[nlev, gcol] = - flux_dn[nlev, gcol] * albedo[nlev, gcol] + # ... reflection of incident diffuse and - src[nlev, gcol] # scattering by the direct beam below + @inbounds flux_up[gcol, nlev] = + flux_dn[gcol, nlev] * albedo[gcol, nlev] + # ... reflection of incident diffuse and + src[gcol, nlev] # scattering by the direct beam below # From the top of the atmosphere downward -- compute fluxes - @inbounds flux_dn_ilevplus1 = flux_dn[nlev, gcol] - @inbounds flux_dn[nlev, gcol] += flux_dn_dir_top + @inbounds flux_dn_ilevplus1 = flux_dn[gcol, nlev] + @inbounds flux_dn[gcol, nlev] += flux_dn_dir_top ilev = nlay @inbounds while ilev ≥ 1 τ_ilev, ssa_ilev, g_ilev = - τ[ilev, gcol], ssa[ilev, gcol], g[ilev, gcol] - albedo_ilev, src_ilev = albedo[ilev, gcol], src[ilev, gcol] + τ[gcol, ilev], ssa[gcol, ilev], g[gcol, ilev] + albedo_ilev, src_ilev = albedo[gcol, ilev], src[gcol, ilev] (_, Tdir, _, Rdif, Tdif) = sw_2stream_coeffs(τ_ilev, ssa_ilev, g_ilev, μ₀) denom = FT(1) / (FT(1) - Rdif * albedo_ilev) # Eq 10 - src_dn_ilev = Tdir * flux_dn_dir[ilev + 1, gcol] + src_dn_ilev = Tdir * flux_dn_dir[gcol, ilev + 1] flux_dn_ilev = (Tdif * flux_dn_ilevplus1 + # Equation 13 Rdif * src_ilev + src_dn_ilev) * denom - flux_up[ilev, gcol] = + flux_up[gcol, ilev] = flux_dn_ilev * albedo_ilev + # Equation 12 src_ilev - flux_dn[ilev, gcol] = flux_dn_ilev + flux_dn_dir[ilev, gcol] + flux_dn[gcol, ilev] = flux_dn_ilev + flux_dn_dir[gcol, ilev] flux_dn_ilevplus1 = flux_dn_ilev ilev -= 1 end diff --git a/src/rte/shortwave_noscat.jl b/src/rte/shortwave_noscat.jl index b0d10530f..bb1963c81 100644 --- a/src/rte/shortwave_noscat.jl +++ b/src/rte/shortwave_noscat.jl @@ -129,17 +129,17 @@ No-scattering solver for the shortwave problem. (; flux_up, flux_dn, flux_dn_dir) = flux FT = eltype(toa_flux) # downward propagation - @inbounds flux_dn_dir[nlev, gcol] = + @inbounds flux_dn_dir[gcol, nlev] = toa_flux[gcol] * solar_frac * cos_zenith[gcol] - @inbounds flux_dn[nlev, gcol] = flux_dn_dir[nlev, gcol] - @inbounds flux_up[nlev, gcol] = FT(0) + @inbounds flux_dn[gcol, nlev] = flux_dn_dir[gcol, nlev] + @inbounds flux_up[gcol, nlev] = FT(0) ilev = nlev - 1 @inbounds while ilev ≥ 1 - flux_dn_dir[ilev, gcol] = - flux_dn_dir[ilev + 1, gcol] * - exp(-τ[ilev, gcol] / max(cos_zenith[gcol], Numerics.μ₀_min(FT))) - flux_dn[ilev, gcol] = flux_dn_dir[ilev, gcol] - flux_up[ilev, gcol] = FT(0) + flux_dn_dir[gcol, ilev] = + flux_dn_dir[gcol, ilev + 1] * + exp(-τ[gcol, ilev] / max(cos_zenith[gcol], Numerics.μ₀_min(FT))) + flux_dn[gcol, ilev] = flux_dn_dir[gcol, ilev] + flux_up[gcol, ilev] = FT(0) ilev -= 1 end end diff --git a/test/all_sky_with_aerosols_dyamond_gpu_benchmark.jl b/test/all_sky_with_aerosols_dyamond_gpu_benchmark.jl index d5b13fb61..0f1a6fe8e 100644 --- a/test/all_sky_with_aerosols_dyamond_gpu_benchmark.jl +++ b/test/all_sky_with_aerosols_dyamond_gpu_benchmark.jl @@ -212,7 +212,7 @@ function benchmark_all_sky_with_aerosols( ) slv_sw = SLVSW(grid_params; swbcs...) #------calling solvers - metric_scaling = DA(one.(slv_sw.flux.flux_up)) + metric_scaling = DA(one.(as.p_lev)) solve_lw!( slv_lw, as, diff --git a/test/all_sky_with_aerosols_utils.jl b/test/all_sky_with_aerosols_utils.jl index b777d4b50..8af7e2783 100644 --- a/test/all_sky_with_aerosols_utils.jl +++ b/test/all_sky_with_aerosols_utils.jl @@ -288,9 +288,9 @@ function all_sky_with_aerosols( comp_flux_net_lw = comp_flux_up_lw .- comp_flux_dn_lw comp_flux_net_sw = comp_flux_up_sw .- comp_flux_dn_sw - flux_up_lw = Array(slv_lw.flux.flux_up) - flux_dn_lw = Array(slv_lw.flux.flux_dn) - flux_net_lw = Array(slv_lw.flux.flux_net) + flux_up_lw = Array(PermutedDimsArray(slv_lw.flux.flux_up, (2, 1))) + flux_dn_lw = Array(PermutedDimsArray(slv_lw.flux.flux_dn, (2, 1))) + flux_net_lw = Array(PermutedDimsArray(slv_lw.flux.flux_net, (2, 1))) max_err_flux_up_lw = maximum(abs.(flux_up_lw .- comp_flux_up_lw)) max_err_flux_dn_lw = maximum(abs.(flux_dn_lw .- comp_flux_dn_lw)) @@ -319,10 +319,10 @@ function all_sky_with_aerosols( "L∞ relative error in flux_net = $(max_rel_err_flux_net_lw * 100) %\n", ) - flux_up_sw = Array(slv_sw.flux.flux_up) - flux_dn_sw = Array(slv_sw.flux.flux_dn) - flux_dn_dir_sw = Array(slv_sw.flux.flux_dn_dir) - flux_net_sw = Array(slv_sw.flux.flux_net) + flux_up_sw = Array(PermutedDimsArray(slv_sw.flux.flux_up, (2, 1))) + flux_dn_sw = Array(PermutedDimsArray(slv_sw.flux.flux_dn, (2, 1))) + flux_dn_dir_sw = Array(PermutedDimsArray(slv_sw.flux.flux_dn_dir, (2, 1))) + flux_net_sw = Array(PermutedDimsArray(slv_sw.flux.flux_net, (2, 1))) max_err_flux_up_sw = maximum(abs.(flux_up_sw .- comp_flux_up_sw)) max_err_flux_dn_sw = maximum(abs.(flux_dn_sw .- comp_flux_dn_sw)) @@ -400,7 +400,7 @@ function all_sky_with_aerosols( slv_sw = SLVSW(grid_params; swbcs...) # Use simple array to test pointwise mult op - metric_scaling = DA(one.(slv_sw.flux.flux_up) * FT(2)) + metric_scaling = DA(one.(as.p_lev) * FT(2)) solve_lw!( slv_lw, as, @@ -425,14 +425,15 @@ function all_sky_with_aerosols( flux_dn_lw = DA(slv_lw.flux.flux_dn) flux_net_lw = DA(slv_lw.flux.flux_net) flux_dn_dir_sw = DA(slv_sw.flux.flux_dn_dir) - - @test all(test_flux_up_sw == flux_up_sw ./ metric_scaling) - @test all(test_flux_dn_sw == flux_dn_sw ./ metric_scaling) - @test all(test_flux_net_sw == flux_net_sw ./ metric_scaling) - @test all(test_flux_up_lw == flux_up_lw ./ metric_scaling) - @test all(test_flux_dn_lw == flux_dn_lw ./ metric_scaling) - @test all(test_flux_net_lw == flux_net_lw ./ metric_scaling) - @test all(test_flux_dn_dir_sw == flux_dn_dir_sw ./ metric_scaling) + sc = PermutedDimsArray(metric_scaling, (2, 1)) + + @test all(test_flux_up_sw == flux_up_sw ./ sc) + @test all(test_flux_dn_sw == flux_dn_sw ./ sc) + @test all(test_flux_net_sw == flux_net_sw ./ sc) + @test all(test_flux_up_lw == flux_up_lw ./ sc) + @test all(test_flux_dn_lw == flux_dn_lw ./ sc) + @test all(test_flux_net_lw == flux_net_lw ./ sc) + @test all(test_flux_dn_dir_sw == flux_dn_dir_sw ./ sc) # Exercise the full update_fluxes! orchestrator and net-flux getters, # and check they reproduce the reference all-sky net flux. diff --git a/test/clear_sky_dyamond_gpu_benchmark.jl b/test/clear_sky_dyamond_gpu_benchmark.jl index b6722f15c..274a4387c 100644 --- a/test/clear_sky_dyamond_gpu_benchmark.jl +++ b/test/clear_sky_dyamond_gpu_benchmark.jl @@ -188,7 +188,7 @@ function benchmark_clear_sky( ) slv_sw = SLVSW(grid_params; swbcs...) #------calling solvers - metric_scaling = DA(one.(slv_sw.flux.flux_up)) + metric_scaling = DA(one.(as.p_lev)) solve_lw!(slv_lw, as, lookup_lw, nothing, nothing, metric_scaling) trial_lw = if device isa ClimaComms.CUDADevice @benchmark CUDA.@sync solve_lw!( diff --git a/test/clear_sky_utils.jl b/test/clear_sky_utils.jl index 7e866ed65..5d3d0d7ea 100644 --- a/test/clear_sky_utils.jl +++ b/test/clear_sky_utils.jl @@ -152,9 +152,9 @@ function clear_sky( comp_flux_net_lw = comp_flux_up_lw .- comp_flux_dn_lw - flux_up_lw = Array(slv_lw.flux.flux_up) - flux_dn_lw = Array(slv_lw.flux.flux_dn) - flux_net_lw = Array(slv_lw.flux.flux_net) + flux_up_lw = Array(PermutedDimsArray(slv_lw.flux.flux_up, (2, 1))) + flux_dn_lw = Array(PermutedDimsArray(slv_lw.flux.flux_dn, (2, 1))) + flux_net_lw = Array(PermutedDimsArray(slv_lw.flux.flux_net, (2, 1))) max_err_flux_up_lw = maximum(abs.(flux_up_lw .- comp_flux_up_lw)) max_err_flux_dn_lw = maximum(abs.(flux_dn_lw .- comp_flux_dn_lw)) @@ -186,9 +186,9 @@ function clear_sky( # comparing shortwave fluxes with data from RRTMGP FORTRAN code comp_flux_net_sw = comp_flux_up_sw .- comp_flux_dn_sw - flux_up_sw = Array(slv_sw.flux.flux_up) - flux_dn_sw = Array(slv_sw.flux.flux_dn) - flux_net_sw = Array(slv_sw.flux.flux_net) + flux_up_sw = Array(PermutedDimsArray(slv_sw.flux.flux_up, (2, 1))) + flux_dn_sw = Array(PermutedDimsArray(slv_sw.flux.flux_dn, (2, 1))) + flux_net_sw = Array(PermutedDimsArray(slv_sw.flux.flux_net, (2, 1))) # Test if shortwave fluxes are zero if zenith angle is ≥ π/2 cos_zenith = Array(cos_zenith) diff --git a/test/cloudy_sky_dyamond_gpu_benchmark.jl b/test/cloudy_sky_dyamond_gpu_benchmark.jl index 6fc4d7517..e2ac03b9b 100644 --- a/test/cloudy_sky_dyamond_gpu_benchmark.jl +++ b/test/cloudy_sky_dyamond_gpu_benchmark.jl @@ -202,7 +202,7 @@ function benchmark_all_sky( ) slv_sw = SLVSW(grid_params; swbcs...) #------calling solvers - metric_scaling = DA(one.(slv_sw.flux.flux_up)) + metric_scaling = DA(one.(as.p_lev)) solve_lw!(slv_lw, as, lookup_lw, lookup_lw_cld, nothing, metric_scaling) trial_lw = if device isa ClimaComms.CUDADevice @benchmark CUDA.@sync solve_lw!( diff --git a/test/cloudy_sky_utils.jl b/test/cloudy_sky_utils.jl index 3a947ce91..1e386b4e1 100644 --- a/test/cloudy_sky_utils.jl +++ b/test/cloudy_sky_utils.jl @@ -107,7 +107,7 @@ function cloudy_sky( ) slv_sw = SLVSW(grid_params; swbcs...) #------calling solvers - metric_scaling = DA(one.(slv_sw.flux.flux_up)) + metric_scaling = DA(one.(as.p_lev)) solve_lw!(slv_lw, as, lookup_lw, lookup_lw_cld, nothing, metric_scaling) if device isa ClimaComms.CPUSingleThreaded JET.@test_opt solve_lw!( @@ -158,9 +158,9 @@ function cloudy_sky( comp_flux_net_lw = comp_flux_up_lw .- comp_flux_dn_lw comp_flux_net_sw = comp_flux_up_sw .- comp_flux_dn_sw - flux_up_lw = Array(slv_lw.flux.flux_up) - flux_dn_lw = Array(slv_lw.flux.flux_dn) - flux_net_lw = Array(slv_lw.flux.flux_net) + flux_up_lw = Array(PermutedDimsArray(slv_lw.flux.flux_up, (2, 1))) + flux_dn_lw = Array(PermutedDimsArray(slv_lw.flux.flux_dn, (2, 1))) + flux_net_lw = Array(PermutedDimsArray(slv_lw.flux.flux_net, (2, 1))) max_err_flux_up_lw = maximum(abs.(flux_up_lw .- comp_flux_up_lw)) max_err_flux_dn_lw = maximum(abs.(flux_dn_lw .- comp_flux_dn_lw)) @@ -189,10 +189,10 @@ function cloudy_sky( "L∞ relative error in flux_net = $(max_rel_err_flux_net_lw * 100) %\n", ) - flux_up_sw = Array(slv_sw.flux.flux_up) - flux_dn_sw = Array(slv_sw.flux.flux_dn) - flux_dn_dir_sw = Array(slv_sw.flux.flux_dn_dir) - flux_net_sw = Array(slv_sw.flux.flux_net) + flux_up_sw = Array(PermutedDimsArray(slv_sw.flux.flux_up, (2, 1))) + flux_dn_sw = Array(PermutedDimsArray(slv_sw.flux.flux_dn, (2, 1))) + flux_dn_dir_sw = Array(PermutedDimsArray(slv_sw.flux.flux_dn_dir, (2, 1))) + flux_net_sw = Array(PermutedDimsArray(slv_sw.flux.flux_net, (2, 1))) max_err_flux_up_sw = maximum(abs.(flux_up_sw .- comp_flux_up_sw)) max_err_flux_dn_sw = maximum(abs.(flux_dn_sw .- comp_flux_dn_sw)) diff --git a/test/cos_zenith_edge_cases.jl b/test/cos_zenith_edge_cases.jl index af29e3ae6..0e0b88e1c 100644 --- a/test/cos_zenith_edge_cases.jl +++ b/test/cos_zenith_edge_cases.jl @@ -176,14 +176,14 @@ function test_cos_zenith_edge_cases( ) # Retrieve fluxes - flux_up_sw = Array(slv_sw.flux.flux_up) - flux_dn_sw = Array(slv_sw.flux.flux_dn) - flux_dn_dir_sw = Array(slv_sw.flux.flux_dn_dir) - flux_net_sw = Array(slv_sw.flux.flux_net) - - flux_up_lw = Array(slv_lw.flux.flux_up) - flux_dn_lw = Array(slv_lw.flux.flux_dn) - flux_net_lw = Array(slv_lw.flux.flux_net) + flux_up_sw = Array(PermutedDimsArray(slv_sw.flux.flux_up, (2, 1))) + flux_dn_sw = Array(PermutedDimsArray(slv_sw.flux.flux_dn, (2, 1))) + flux_dn_dir_sw = Array(PermutedDimsArray(slv_sw.flux.flux_dn_dir, (2, 1))) + flux_net_sw = Array(PermutedDimsArray(slv_sw.flux.flux_net, (2, 1))) + + flux_up_lw = Array(PermutedDimsArray(slv_lw.flux.flux_up, (2, 1))) + flux_dn_lw = Array(PermutedDimsArray(slv_lw.flux.flux_dn, (2, 1))) + flux_net_lw = Array(PermutedDimsArray(slv_lw.flux.flux_net, (2, 1))) color2 = :cyan printstyled( diff --git a/test/float32_consistency.jl b/test/float32_consistency.jl index 888b35100..1ea933ef4 100644 --- a/test/float32_consistency.jl +++ b/test/float32_consistency.jl @@ -125,10 +125,10 @@ function clear_sky_fluxes(context, ::Type{FT}, ::Type{SLVLW}; ncol) where {FT, S solve_lw!(slv_lw, as, lookup_lw) solve_sw!(slv_sw, as, lookup_sw) return (; - lw_up = Array(slv_lw.flux.flux_up), - lw_dn = Array(slv_lw.flux.flux_dn), - sw_up = Array(slv_sw.flux.flux_up), - sw_dn = Array(slv_sw.flux.flux_dn), + lw_up = Array(PermutedDimsArray(slv_lw.flux.flux_up, (2, 1))), + lw_dn = Array(PermutedDimsArray(slv_lw.flux.flux_dn, (2, 1))), + sw_up = Array(PermutedDimsArray(slv_sw.flux.flux_up, (2, 1))), + sw_dn = Array(PermutedDimsArray(slv_sw.flux.flux_dn, (2, 1))), ) end @@ -187,10 +187,10 @@ function cloudy_sky_fluxes(context, ::Type{FT}, ::Type{SLVLW}; ncol) where {FT, solve_lw!(slv_lw, as, lookup_lw, lookup_lw_cld) solve_sw!(slv_sw, as, lookup_sw, lookup_sw_cld) return (; - lw_up = Array(slv_lw.flux.flux_up), - lw_dn = Array(slv_lw.flux.flux_dn), - sw_up = Array(slv_sw.flux.flux_up), - sw_dn = Array(slv_sw.flux.flux_dn), + lw_up = Array(PermutedDimsArray(slv_lw.flux.flux_up, (2, 1))), + lw_dn = Array(PermutedDimsArray(slv_lw.flux.flux_dn, (2, 1))), + sw_up = Array(PermutedDimsArray(slv_sw.flux.flux_up, (2, 1))), + sw_dn = Array(PermutedDimsArray(slv_sw.flux.flux_dn, (2, 1))), ) end diff --git a/test/gray_atm_utils.jl b/test/gray_atm_utils.jl index 25df12d93..56d2382c3 100644 --- a/test/gray_atm_utils.jl +++ b/test/gray_atm_utils.jl @@ -67,7 +67,9 @@ function gray_atmos_lw_equil( setup_gray_as_pr_grid(context, nlay, lat, p0, pe, otp, param_set, DA) slv_lw = SLVLW(grid_params; params = param_set, sfc_emis, inc_flux) - (; flux_up, flux_dn, flux_net) = slv_lw.flux + flux_up = view(PermutedDimsArray(slv_lw.flux.flux_up, (2, 1)), :, :) + flux_dn = view(PermutedDimsArray(slv_lw.flux.flux_dn, (2, 1)), :, :) + flux_net = view(PermutedDimsArray(slv_lw.flux.flux_net, (2, 1)), :, :) (; t_lay, p_lay, t_lev, p_lev) = gray_as sbc = FT(RRTMGP.Parameters.Stefan(param_set)) cp_d_ = FT(RRTMGP.Parameters.cp_d(param_set)) @@ -199,9 +201,9 @@ function gray_atmos_sw_test( solve_sw!(slv_sw, as) - τ = Array(slv_sw.op.τ) + τ = Array(PermutedDimsArray(slv_sw.op.τ, (2, 1))) cos_zenith = Array(cos_zenith) - flux_dn_dir = Array(slv_sw.flux.flux_dn_dir) + flux_dn_dir = Array(PermutedDimsArray(slv_sw.flux.flux_dn_dir, (2, 1))) toa_flux = Array(toa_flux) # testing with exact solution diff --git a/test/partial_cloud_fraction.jl b/test/partial_cloud_fraction.jl index 929d89ce2..f5c675dbc 100644 --- a/test/partial_cloud_fraction.jl +++ b/test/partial_cloud_fraction.jl @@ -97,8 +97,8 @@ function partial_cloud_fraction_test( solve_lw!(slv_lw, as, lookup_lw, lookup_lw_cld, nothing, nothing) solve_sw!(slv_sw, as, lookup_sw, lookup_sw_cld, nothing, nothing) return ( - lw_net = copy(Array(slv_lw.flux.flux_net)), - sw_net = copy(Array(slv_sw.flux.flux_net)), + lw_net = copy(Array(PermutedDimsArray(slv_lw.flux.flux_net, (2, 1)))), + sw_net = copy(Array(PermutedDimsArray(slv_sw.flux.flux_net, (2, 1)))), cld_cover_lw = copy(Array(as.cloud_state.cld_cover_lw)), cld_cover_sw = copy(Array(as.cloud_state.cld_cover_sw)), ) diff --git a/test/standalone.jl b/test/standalone.jl index d0c5051b8..4dc13a467 100644 --- a/test/standalone.jl +++ b/test/standalone.jl @@ -242,8 +242,8 @@ end RRTMGP.update_fluxes!(solver) - # internal buffers are boundary-extended; getters are domain-masked - @test size(solver.net_flux_buffer, 1) == nlay + 1 # full levels + # internal buffers are boundary-extended and coalesced (ncol, nlev); getters are domain-masked + @test size(solver.net_flux_buffer, 2) == nlay + 1 # full levels @test size(RRTMGP.net_flux(solver), 1) == domain_nlay + 1 # domain levels @test size(RRTMGP.level_pressure(solver), 1) == domain_nlay + 1 @test size(RRTMGP.lw_flux_net(solver), 1) == domain_nlay + 1 From 842a881bbaf0fc3f18be80b1c79a2086426af342 Mon Sep 17 00:00:00 2001 From: Tapio Schneider Date: Sat, 11 Jul 2026 22:05:53 -0700 Subject: [PATCH 2/6] Fix Julia 1.10/1.11 + GPU getter regressions of the layout transpose; add transposed state cache for coalesced gas-optics reads Follow-ups to the (ncol, nlev) layout commit (298d688ea), 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 --- NEWS.md | 25 ++++ docs/src/functional_core.md | 12 +- docs/src/getters.md | 8 +- docs/src/howto/gpu.md | 9 ++ docs/src/optics/Fluxes.md | 4 + ext/RRTMGPCUDAExt.jl | 1 + ext/cuda/rte_longwave_2stream.jl | 4 + ext/cuda/rte_longwave_noscat.jl | 4 + ext/cuda/rte_shortwave_2stream.jl | 4 + ext/cuda/rte_shortwave_noscat.jl | 5 +- src/api/getters.jl | 51 ++++---- src/api/solver.jl | 45 +++++-- src/api/update_fluxes.jl | 48 +++++--- src/optics/AtmosphericStates.jl | 102 ++++++++++++++++ src/optics/Fluxes.jl | 176 ++++++++++++++++++++++++---- src/optics/Sources.jl | 4 +- src/optics/compute_optical_props.jl | 48 +++++--- src/rte/RTE.jl | 66 +++++++++-- src/rte/RTESolver.jl | 26 +++- src/rte/longwave_2stream.jl | 4 + src/rte/longwave_noscat.jl | 4 + src/rte/shortwave_2stream.jl | 4 + src/rte/shortwave_noscat.jl | 5 +- test/clear_sky_utils.jl | 13 ++ test/cloudy_sky_utils.jl | 100 +++++++++++----- test/scalar_indexing.jl | 12 +- test/standalone.jl | 4 +- 27 files changed, 633 insertions(+), 155 deletions(-) diff --git a/NEWS.md b/NEWS.md index daf98159c..d4cc31bc0 100644 --- a/NEWS.md +++ b/NEWS.md @@ -3,6 +3,31 @@ RRTMGP.jl Release Notes main ------ +- Internal array layouts transposed to column-first `(ncol, nlay/nlev)` for + coalesced GPU memory access: optical properties (`OneScalar`, `TwoStream`), + source functions (`SourceLWNoScat`, `SourceLW2Str`, `SourceSW2Str`), and the + broadband flux buffers (`FluxLW`, `FluxSW`). On an A100 at DYAMOND scale + this cuts longwave kernel times by 25-36% and shortwave all-sky by 14-20%. + The Layer-2 flux getters (`lw_flux_up`, `net_flux`, ...) still present + plain `(nlev, ncol)` views: they read `Fluxes.FluxPresentation` arrays that + `update_fluxes!` fills from the compute buffers with one fused transposing + copy per call (so getters remain `Array`-materializable, broadcastable, and + reducible on the GPU). The opt-in per-band buffers (`FluxBand`) keep the + host-facing `(nlev, ncol, n_bnd)` layout. **Breaking for Layer-1 + (functional-core) users**: code reading the workspace fields directly + (`slv.flux.flux_up`, `op.τ`, ...) must swap index order to `[gcol, ilev]` + or wrap the buffer in `Fluxes.lazy_transpose`. State inputs + (`AtmosphericState` arrays) and all boundary-condition arrays are + unchanged. +- New `TransposedStateCache`: column-first copies of the hot + `AtmosphericState` arrays (layer pressure/temperature/dry-air column + amount/relative humidity, level temperatures), refreshed by one + `permutedims!` per spectral solve, so the gas-optics g-point loop reads + coalesced memory on GPUs. On by default (the RTE workspace constructors + build one; the `RRTMGPSolver` shares a single cache between the longwave + and shortwave workspaces); pass `state_cache = nothing` to a workspace + constructor to opt out, or pass the same cache to several workspaces to + share the storage. - Documentation reorganized along tutorial / how-to / explanation / reference lines: two new executable (Literate.jl) tutorials — "A first radiation calculation" and "Radiative-convective equilibrium", which reproduces diff --git a/docs/src/functional_core.md b/docs/src/functional_core.md index 948ce3c2c..8c60750b6 100644 --- a/docs/src/functional_core.md +++ b/docs/src/functional_core.md @@ -32,9 +32,15 @@ workspace's flux buffers in place, and you can read the results: ```julia solve_lw!(slv_lw, state, lookup_lw) # or solve_lw!(slv_lw, state) for gray -F = slv_lw.flux.flux_net # (nlev, ncol) [W/m²] +F = slv_lw.flux.flux_net # (ncol, nlev) [W/m²] ``` +Note the axis order: the raw flux buffers are stored column-first, `(ncol, +nlev)`, so that neighboring GPU threads (one per column) access consecutive +memory. Only the Layer-2 getters (`net_flux`, `lw_flux_up`, ...) present the +vertical-first `(nlev, ncol)` orientation, from presentation copies that +`update_fluxes!` refreshes at the end of each call. + This is what RRTMGP's own tests drive, and what [`RRTMGPSolver`](@ref RRTMGP.RRTMGPSolver) wraps behind [`update_fluxes!`](@ref RRTMGP.update_fluxes!) and the named getters. @@ -81,10 +87,10 @@ slv_sw = TwoStreamSWRTE(grid_params; cos_zenith, toa_flux, solve_lw!(slv_lw, as) solve_sw!(slv_sw, as) -lw_net = slv_lw.flux.flux_net # (nlev, ncol) [W/m²] +lw_net = slv_lw.flux.flux_net # (ncol, nlev) [W/m²] sw_net = slv_sw.flux.flux_net net = lw_net .+ sw_net -net[end, 1] # net flux at the top of the atmosphere [W/m²] +net[1, end] # net flux at the top of the atmosphere [W/m²] ``` The one-liner [`solve_gray`](@ref RRTMGP.solve_gray) wraps this. diff --git a/docs/src/getters.md b/docs/src/getters.md index 4a634cdc3..90dc3d37f 100644 --- a/docs/src/getters.md +++ b/docs/src/getters.md @@ -53,9 +53,13 @@ F = RRTMGP.net_flux(solver) # (nlev, ncol) view of the resu With ClimaCore, this forms the bridge: `array2field(getter(solver), space)` wraps an input view as a `Field` (writes flow into the solver's buffer), and `field2array(dst) .= getter(solver)` copies an output into a destination field. -On the GPU, the returned views are strided `SubArray`s of device arrays; to +On the GPU, the returned views are plain `SubArray`s of device arrays; to bring one to the CPU, use `Array(getter(solver))` (or the element-wise `.=` -copy above). +copy above). One subtlety: the solver's compute buffers are stored +column-first, `(ncol, nlev)`, for coalesced GPU access, and `update_fluxes!` +ends by transposing them into the `(nlev, ncol)` presentation arrays the flux +getters expose — so read fluxes after `update_fluxes!`, not between manual +Layer-1 solves. ## Responsibilities: what the host provides vs what RRTMGP derives diff --git a/docs/src/howto/gpu.md b/docs/src/howto/gpu.md index 1f7327e88..4e1dcf411 100644 --- a/docs/src/howto/gpu.md +++ b/docs/src/howto/gpu.md @@ -78,6 +78,15 @@ both backends: g-point's lookup-table slices hot in cache while sweeping the columns. With a single-threaded context, the same code runs serially. +The memory layout matches the thread mapping: the internal scratch and output +buffers are stored column-first, `(ncol, nlev)`, so consecutive GPU threads +write consecutive addresses (coalesced access), and the frequently re-read +state arrays (layer pressures and temperatures, level temperatures) are +copied into a column-first `TransposedStateCache` once per solve so the +g-point loop also reads coalesced memory. The caller-owned +`AtmosphericState` and the getter views keep their vertical-first +`(nlev, ncol)` presentation. + ## Performance expectations CI benchmarks the kernels on GPUs (see `perf/benchmark_ratchet.jl` diff --git a/docs/src/optics/Fluxes.md b/docs/src/optics/Fluxes.md index 514eb5e09..84abdb7aa 100644 --- a/docs/src/optics/Fluxes.md +++ b/docs/src/optics/Fluxes.md @@ -8,6 +8,10 @@ CurrentModule = RRTMGP.Fluxes AbstractFlux FluxLW FluxSW +FluxPresentation +update_presentation! +transpose_into! +transpose_sum_into! set_flux_to_zero! compute_net_flux! ``` diff --git a/ext/RRTMGPCUDAExt.jl b/ext/RRTMGPCUDAExt.jl index 53a240c6f..fb510f878 100644 --- a/ext/RRTMGPCUDAExt.jl +++ b/ext/RRTMGPCUDAExt.jl @@ -23,6 +23,7 @@ import RRTMGP.AtmosphericStates: GrayAtmosphericState import RRTMGP.AtmosphericStates: AtmosphericState import RRTMGP.AtmosphericStates: CloudState import RRTMGP.AtmosphericStates: AerosolState +import RRTMGP.AtmosphericStates: TransposedStateCache import RRTMGP.AtmosphericStates import RRTMGP.GrayAtmosphere: update_profile_lw! import RRTMGP.GrayAtmosphere: compute_gray_heating_rate! diff --git a/ext/cuda/rte_longwave_2stream.jl b/ext/cuda/rte_longwave_2stream.jl index a55eb2e77..739c98c97 100644 --- a/ext/cuda/rte_longwave_2stream.jl +++ b/ext/cuda/rte_longwave_2stream.jl @@ -54,6 +54,7 @@ function rte_lw_2stream_solve!( bcs_lw::LwBCs, op::TwoStream, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing} = nothing, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, @@ -71,6 +72,7 @@ function rte_lw_2stream_solve!( nlay, ncol, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, @@ -91,6 +93,7 @@ function rte_lw_2stream_solve_CUDA!( nlay, ncol, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing}, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing}, @@ -115,6 +118,7 @@ function rte_lw_2stream_solve_CUDA!( bcs_lw, op, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, diff --git a/ext/cuda/rte_longwave_noscat.jl b/ext/cuda/rte_longwave_noscat.jl index 0d7dd145c..800175987 100644 --- a/ext/cuda/rte_longwave_noscat.jl +++ b/ext/cuda/rte_longwave_noscat.jl @@ -60,6 +60,7 @@ function rte_lw_noscat_solve!( op::OneScalar, angle_disc::AngularDiscretization, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing} = nothing, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, @@ -76,6 +77,7 @@ function rte_lw_noscat_solve!( nlay, ncol, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, @@ -96,6 +98,7 @@ function rte_lw_noscat_solve_CUDA!( nlay, ncol, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing}, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing}, @@ -123,6 +126,7 @@ function rte_lw_noscat_solve_CUDA!( Ds, w_μ, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, diff --git a/ext/cuda/rte_shortwave_2stream.jl b/ext/cuda/rte_shortwave_2stream.jl index afc0bd9e4..5c82ebd61 100644 --- a/ext/cuda/rte_shortwave_2stream.jl +++ b/ext/cuda/rte_shortwave_2stream.jl @@ -64,6 +64,7 @@ function rte_sw_2stream_solve!( bcs_sw::SwBCs, src_sw::SourceSW2Str, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_sw::LookUpSW, lookup_sw_cld::Union{LookUpCld, Nothing} = nothing, lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, @@ -81,6 +82,7 @@ function rte_sw_2stream_solve!( nlay, ncol, as, + state_cache, lookup_sw, lookup_sw_cld, lookup_sw_aero, @@ -101,6 +103,7 @@ function rte_sw_2stream_solve_CUDA!( nlay, ncol, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_sw::LookUpSW, lookup_sw_cld::Union{LookUpCld, Nothing} = nothing, lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, @@ -126,6 +129,7 @@ function rte_sw_2stream_solve_CUDA!( bcs_sw, src_sw, as, + state_cache, lookup_sw, lookup_sw_cld, lookup_sw_aero, diff --git a/ext/cuda/rte_shortwave_noscat.jl b/ext/cuda/rte_shortwave_noscat.jl index b4098fb3a..5681143d3 100644 --- a/ext/cuda/rte_shortwave_noscat.jl +++ b/ext/cuda/rte_shortwave_noscat.jl @@ -58,11 +58,12 @@ function rte_sw_noscat_solve!( op::OneScalar, bcs_sw::SwBCs, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_sw::LookUpSW, ) nlay, ncol = AtmosphericStates.get_dims(as) tx, bx = _configure_threadblock(ncol) - args = (flux, flux_sw, op, bcs_sw, nlay, ncol, as, lookup_sw) + args = (flux, flux_sw, op, bcs_sw, nlay, ncol, as, state_cache, lookup_sw) @cuda always_inline = true threads = (tx) blocks = (bx) rte_sw_noscat_solve_CUDA!( args..., ) @@ -77,6 +78,7 @@ function rte_sw_noscat_solve_CUDA!( nlay, ncol, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_sw::LookUpSW, ) gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id @@ -95,6 +97,7 @@ function rte_sw_noscat_solve_CUDA!( op, bcs_sw, as, + state_cache, lookup_sw, n_gpt, nlev, diff --git a/src/api/getters.jl b/src/api/getters.jl index 67f04aa5a..8304ddef3 100644 --- a/src/api/getters.jl +++ b/src/api/getters.jl @@ -45,13 +45,6 @@ _domain_view(isothermal_boundary_layer::Bool, x::AbstractArray) = _domain_view(isothermal_boundary_layer::Bool, x::AbstractArray{<:Any, 3}) = view(x, 1:(size(x, 1) - Int(isothermal_boundary_layer)), :, :) -_flux_domain_view(s::RRTMGPSolver, x) = _flux_domain_view(s.grid_params.isothermal_boundary_layer, x) -_flux_domain_view(::Bool, ::Nothing) = nothing -_flux_domain_view(isothermal_boundary_layer::Bool, x::AbstractArray{<:Any, 2}) = - view(PermutedDimsArray(x, (2, 1)), 1:(size(x, 2) - Int(isothermal_boundary_layer)), :) -_flux_domain_view(isothermal_boundary_layer::Bool, x::AbstractArray{<:Any, 3}) = - view(PermutedDimsArray(x, (2, 1, 3)), 1:(size(x, 2) - Int(isothermal_boundary_layer)), :, :) - # The clear-sky flux workspaces exist only on `AllSkyRadiationWithClearSkyDiagnostics` # solvers; give the `clear_*` getters an informative error elsewhere (instead of a # cryptic `getproperty(::Nothing, ...)`). @@ -64,21 +57,21 @@ _require_clear_sky(x) = x # Potentially views top_of_atmosphere_lw_flux_dn(s::RRTMGPSolver) = _maybe_transpose(s.lws.bcs.inc_flux) top_of_atmosphere_diffuse_sw_flux_dn(s::RRTMGPSolver) = _maybe_transpose(s.sws.bcs.inc_flux_diffuse) -lw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, s.lws.flux.flux_up) -lw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, s.lws.flux.flux_dn) -lw_flux_net(s::RRTMGPSolver) = _flux_domain_view(s, s.lws.flux.flux_net) -clear_lw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_up) -clear_lw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_dn) -clear_lw_flux(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_net) +lw_flux_up(s::RRTMGPSolver) = _domain_view(s, s.presented_flux_lw.flux_up) +lw_flux_dn(s::RRTMGPSolver) = _domain_view(s, s.presented_flux_lw.flux_dn) +lw_flux_net(s::RRTMGPSolver) = _domain_view(s, s.presented_flux_lw.flux_net) +clear_lw_flux_up(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_up) +clear_lw_flux_dn(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_dn) +clear_lw_flux(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_net) surface_emissivity(s::RRTMGPSolver) = s.lws.bcs.sfc_emis -sw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_up) -sw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_dn) -sw_flux_net(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_net) -sw_direct_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, s.sws.flux.flux_dn_dir) -clear_sw_flux_up(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_up) -clear_sw_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn) -clear_sw_direct_flux_dn(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn_dir) -clear_sw_flux(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_net) +sw_flux_up(s::RRTMGPSolver) = _domain_view(s, s.presented_flux_sw.flux_up) +sw_flux_dn(s::RRTMGPSolver) = _domain_view(s, s.presented_flux_sw.flux_dn) +sw_flux_net(s::RRTMGPSolver) = _domain_view(s, s.presented_flux_sw.flux_net) +sw_direct_flux_dn(s::RRTMGPSolver) = _domain_view(s, s.presented_flux_sw.flux_dn_dir) +clear_sw_flux_up(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_up) +clear_sw_flux_dn(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn) +clear_sw_direct_flux_dn(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn_dir) +clear_sw_flux(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_net) cloud_liquid_effective_radius(s::RRTMGPSolver) = _domain_view(s, s.as.cloud_state.cld_r_eff_liq) cloud_ice_effective_radius(s::RRTMGPSolver) = _domain_view(s, s.as.cloud_state.cld_r_eff_ice) cloud_liquid_water_path(s::RRTMGPSolver) = _domain_view(s, s.as.cloud_state.cld_path_liq) @@ -101,8 +94,8 @@ layer_temperature(s::RRTMGPSolver) = _domain_view(s, getview_ layer_relative_humidity(s::RRTMGPSolver) = _domain_view(s, getview_rel_hum(s.as)) level_pressure(s::RRTMGPSolver) = _domain_view(s, s.as.p_lev) level_temperature(s::RRTMGPSolver) = _domain_view(s, s.as.t_lev) -net_flux(s::RRTMGPSolver) = _flux_domain_view(s, s.net_flux_buffer) -clear_net_flux(s::RRTMGPSolver) = _flux_domain_view(s, _require_clear_sky(s.clear_net_flux_buffer)) +net_flux(s::RRTMGPSolver) = _domain_view(s, s.net_flux_buffer) +clear_net_flux(s::RRTMGPSolver) = _domain_view(s, _require_clear_sky(s.clear_net_flux_buffer)) deep_atmosphere_inverse_scaling(s::RRTMGPSolver) = _domain_view(s, s.deep_atmosphere_inverse_scaling) #! format: on @@ -162,17 +155,17 @@ wavenumber range. The `up`/`dn`/`net` getters are views into the retained per-band buffers. """ spectral_lw_flux_up(s::RRTMGPSolver) = - _flux_domain_view(s, _solver_band_flux(s.lws).flux_up) + _domain_view(s, _solver_band_flux(s.lws).flux_up) spectral_lw_flux_dn(s::RRTMGPSolver) = - _flux_domain_view(s, _solver_band_flux(s.lws).flux_dn) + _domain_view(s, _solver_band_flux(s.lws).flux_dn) spectral_lw_flux_net(s::RRTMGPSolver) = - _flux_domain_view(s, _solver_band_flux(s.lws).flux_net) + _domain_view(s, _solver_band_flux(s.lws).flux_net) spectral_sw_flux_up(s::RRTMGPSolver) = - _flux_domain_view(s, _solver_band_flux(s.sws).flux_up) + _domain_view(s, _solver_band_flux(s.sws).flux_up) spectral_sw_flux_dn(s::RRTMGPSolver) = - _flux_domain_view(s, _solver_band_flux(s.sws).flux_dn) + _domain_view(s, _solver_band_flux(s.sws).flux_dn) spectral_sw_flux_net(s::RRTMGPSolver) = - _flux_domain_view(s, _solver_band_flux(s.sws).flux_net) + _domain_view(s, _solver_band_flux(s.sws).flux_net) """ lw_band_bounds(s::RRTMGPSolver) diff --git a/src/api/solver.jl b/src/api/solver.jl index 3abe81d76..676ec734c 100644 --- a/src/api/solver.jl +++ b/src/api/solver.jl @@ -1,6 +1,7 @@ using ..Fluxes using ..Sources -using ..AtmosphericStates: AtmosphericState, AbstractAtmosphericState +using ..AtmosphericStates: + AtmosphericState, AbstractAtmosphericState, TransposedStateCache using ..RTE using ..RTESolver using ..BCs @@ -60,8 +61,12 @@ Construct it with the `RRTMGPSolver` constructor and drive it with - `lws`: longwave RTE solver and its flux/scratch buffers. - `as`: the atmospheric state (solver inputs). - `lookups`: the [`LookupBundle`](@ref) from `lookup_tables` (for gray radiation only the band counts are carried). -- `clear_flux_lw`: clear-sky longwave fluxes, or `nothing`. -- `clear_flux_sw`: clear-sky shortwave fluxes, or `nothing`. +- `presented_flux_lw`, `presented_flux_sw`: host-facing `(nlev, ncol)` + [`FluxPresentation`](@ref RRTMGP.Fluxes.FluxPresentation) copies of the + longwave/shortwave fluxes, refreshed at the end of `update_fluxes!`; the flux + getters return views of these. +- `clear_flux_lw`: clear-sky longwave fluxes (`(nlev, ncol)` presentation layout), or `nothing`. +- `clear_flux_sw`: clear-sky shortwave fluxes (`(nlev, ncol)` presentation layout), or `nothing`. - `center_z`: layer-center altitudes [m], or `nothing`. - `face_z`: level (face) altitudes [m], or `nothing`. - `deep_atmosphere_inverse_scaling`: `(nlev, ncol)` factor multiplied into the fluxes for deep-atmosphere geometric scaling (the host supplies the multiplicative inverse of its metric scaling), or `nothing` (default) for the shallow-atmosphere approximation. @@ -93,6 +98,8 @@ struct RRTMGPSolver{ LWS, AS <: AbstractAtmosphericState, LU, + PFL <: Fluxes.FluxPresentation, + PFS <: Fluxes.FluxPresentation, CFLW, CFSW, ZC <: Union{AbstractArray, Nothing}, @@ -110,6 +117,8 @@ struct RRTMGPSolver{ lws::LWS as::AS lookups::LU + presented_flux_lw::PFL + presented_flux_sw::PFS clear_flux_lw::CFLW clear_flux_sw::CFSW center_z::ZC @@ -184,11 +193,19 @@ function RRTMGPSolver( flux_lw = Fluxes.FluxLW(grid_params) flux_sw = Fluxes.FluxSW(grid_params) - net_flux_buffer = similar(flux_lw.flux_net) + # Host-facing (nlev, ncol) presentation of the fluxes (see + # `Fluxes.FluxPresentation`): the flux getters return plain views of these + # arrays, filled by transposing copies at the end of `update_fluxes!`. + presented_flux_lw = Fluxes.FluxPresentation(grid_params; direct = false) + presented_flux_sw = Fluxes.FluxPresentation(grid_params; direct = true) + net_flux_buffer = similar(presented_flux_lw.flux_net) if radiation_method isa AllSkyRadiationWithClearSkyDiagnostics - clear_flux_lw = Fluxes.FluxLW(grid_params) - clear_flux_sw = Fluxes.FluxSW(grid_params) - clear_net_flux_buffer = similar(flux_lw.flux_net) + # The clear-sky snapshots are read only through the getters, so they + # are stored directly in the (nlev, ncol) presentation layout (the + # snapshot copy between the clear and all-sky solves transposes). + clear_flux_lw = Fluxes.FluxPresentation(grid_params; direct = false) + clear_flux_sw = Fluxes.FluxPresentation(grid_params; direct = true) + clear_net_flux_buffer = similar(presented_flux_lw.flux_net) else clear_flux_lw = nothing clear_flux_sw = nothing @@ -212,9 +229,16 @@ function RRTMGPSolver( band_flux_sw = nothing end + # Column-first copies of the hot state arrays for coalesced gas-optics + # reads; one cache shared by the longwave and shortwave workspaces + # (refreshed at each spectral solve). Gray radiation reads its analytic + # state directly and needs none. + state_cache = + radiation_method isa GrayRadiation ? nothing : + TransposedStateCache(grid_params) sws = op_sw isa OneScalar ? - RTE.NoScatSWRTE(context, op_sw, bcs_sw, fluxb_sw, flux_sw) : + RTE.NoScatSWRTE(context, op_sw, bcs_sw, fluxb_sw, flux_sw, state_cache) : RTE.TwoStreamSWRTE( context, op_sw, @@ -223,6 +247,7 @@ function RRTMGPSolver( fluxb_sw, # inferable, need radiation_method flux_sw, # views attached band_flux_sw, # optional per-band fluxes (or `nothing`) + state_cache, ) lws = op_lw isa OneScalar ? @@ -234,6 +259,7 @@ function RRTMGPSolver( fluxb_lw, # inferable, need radiation_method flux_lw, # views attached AngularDiscretizations.AngularDiscretization(grid_params, 1), + state_cache, ) : RTE.TwoStreamLWRTE( context, @@ -243,6 +269,7 @@ function RRTMGPSolver( fluxb_lw, # inferable, need radiation_method flux_lw, # views attached band_flux_lw, # optional per-band fluxes (or `nothing`) + state_cache, ) return RRTMGPSolver( grid_params, @@ -254,6 +281,8 @@ function RRTMGPSolver( lws, as, lookups, + presented_flux_lw, + presented_flux_sw, clear_flux_lw, clear_flux_sw, center_z, diff --git a/src/api/update_fluxes.jl b/src/api/update_fluxes.jl index 02b2554bf..6f2541bdd 100644 --- a/src/api/update_fluxes.jl +++ b/src/api/update_fluxes.jl @@ -5,9 +5,15 @@ """ update_lw_fluxes!(s::RRTMGPSolver) -Update the longwave fluxes. +Update the longwave fluxes, leaving the longwave flux getters consistent (the +`(ncol, nlev)` compute buffers are transposed into the `(nlev, ncol)` +presentation the getters expose). """ -update_lw_fluxes!(s::RRTMGPSolver) = update_lw_fluxes!(s, _radiation_method(s)) +function update_lw_fluxes!(s::RRTMGPSolver) + update_lw_fluxes!(s, _radiation_method(s)) + Fluxes.update_presentation!(s.presented_flux_lw, s.lws.flux) + return nothing +end update_lw_fluxes!(s::RRTMGPSolver, ::GrayRadiation) = RTESolver.solve_lw!( _longwave_solver(s), @@ -46,9 +52,8 @@ function update_lw_fluxes!( lookups.lookup_lw_aero, ms, ) - s.clear_flux_lw.flux_up .= s.lws.flux.flux_up - s.clear_flux_lw.flux_dn .= s.lws.flux.flux_dn - s.clear_flux_lw.flux_net .= s.lws.flux.flux_net + # snapshot the clear-sky fluxes into their (nlev, ncol) presentation + Fluxes.update_presentation!(s.clear_flux_lw, s.lws.flux) RTESolver.solve_lw!( lw_solver, as, @@ -62,9 +67,15 @@ end """ update_sw_fluxes!(s::RRTMGPSolver) -Update the shortwave fluxes. +Update the shortwave fluxes, leaving the shortwave flux getters consistent +(the `(ncol, nlev)` compute buffers are transposed into the `(nlev, ncol)` +presentation the getters expose). """ -update_sw_fluxes!(s::RRTMGPSolver) = update_sw_fluxes!(s, _radiation_method(s)) +function update_sw_fluxes!(s::RRTMGPSolver) + update_sw_fluxes!(s, _radiation_method(s)) + Fluxes.update_presentation!(s.presented_flux_sw, s.sws.flux) + return nothing +end update_sw_fluxes!(s::RRTMGPSolver, ::GrayRadiation) = RTESolver.solve_sw!( _shortwave_solver(s), @@ -103,10 +114,8 @@ function update_sw_fluxes!( lookups.lookup_sw_aero, ms, ) - s.clear_flux_sw.flux_up .= s.sws.flux.flux_up - s.clear_flux_sw.flux_dn .= s.sws.flux.flux_dn - s.clear_flux_sw.flux_dn_dir .= s.sws.flux.flux_dn_dir - s.clear_flux_sw.flux_net .= s.sws.flux.flux_net + # snapshot the clear-sky fluxes into their (nlev, ncol) presentation + Fluxes.update_presentation!(s.clear_flux_sw, s.sws.flux) RTESolver.solve_sw!( sw_solver, @@ -156,9 +165,13 @@ Combine the longwave and shortwave net fluxes into `net_flux(s)` (and, for update_net_fluxes!(s::RRTMGPSolver) = update_net_fluxes!(s, _radiation_method(s)) function update_net_fluxes!(s::RRTMGPSolver, ::AbstractRRTMGPMethod) - # Operate on the raw (boundary-extended) buffers, not `parent(getter(s))`: the - # getters allocate a `view` per call that would then just be stripped by `parent`. - s.net_flux_buffer .= s.lws.flux.flux_net .+ s.sws.flux.flux_net + # Operate on the raw (boundary-extended) buffers: the compute buffers are + # (ncol, nlev), the net-flux buffer is host-facing (nlev, ncol). + Fluxes.transpose_sum_into!( + s.net_flux_buffer, + s.lws.flux.flux_net, + s.sws.flux.flux_net, + ) _update_band_net_fluxes!(s.lws) _update_band_net_fluxes!(s.sws) return nothing @@ -167,7 +180,12 @@ function update_net_fluxes!( s::RRTMGPSolver, ::AllSkyRadiationWithClearSkyDiagnostics, ) - s.net_flux_buffer .= s.lws.flux.flux_net .+ s.sws.flux.flux_net + Fluxes.transpose_sum_into!( + s.net_flux_buffer, + s.lws.flux.flux_net, + s.sws.flux.flux_net, + ) + # the clear-sky snapshots are already stored (nlev, ncol) s.clear_net_flux_buffer .= s.clear_flux_lw.flux_net .+ s.clear_flux_sw.flux_net _update_band_net_fluxes!(s.lws) diff --git a/src/optics/AtmosphericStates.jl b/src/optics/AtmosphericStates.jl index d31bf3d63..dbb346f75 100644 --- a/src/optics/AtmosphericStates.jl +++ b/src/optics/AtmosphericStates.jl @@ -2,6 +2,7 @@ module AtmosphericStates import ClimaComms using Adapt import ..Parameters as RP +import ..RRTMGPGridParams using ..VolumeMixingRatios @@ -10,6 +11,7 @@ export AbstractAtmosphericState, GrayAtmosphericState, CloudState, AerosolState, + TransposedStateCache, setup_gray_as_pr_grid, MaxRandomOverlap, AbstractCloudMask, @@ -114,6 +116,106 @@ Adapt.@adapt_structure AtmosphericState @inline getview_rel_hum(as::AtmosphericState, gcol) = @inbounds view(as.layerdata, 4, :, gcol) +""" + TransposedStateCache{D2, D3} + +Column-first copies of the hot [`AtmosphericState`](@ref) fields, refreshed +once per solve by [`refresh_transposed_state!`](@ref). The gas-optics kernels +re-read the layer pressures, temperatures, dry-air column amounts, and level +temperatures for every g-point; in the caller-owned `AtmosphericState` those +arrays are vertical-first, so neighboring GPU threads (one per column) read +memory `nlay` or `4 nlay` elements apart. Reading the transposed copies +instead makes those loads consecutive (coalesced) across the threads of a +warp. + +The cache trades one `permutedims!` per solve plus `(4 nlay + nlev) ncol` +extra storage for coalesced reads in the g-point loop. Pass the same cache to +the longwave and shortwave workspaces to share the storage, or `nothing` to +opt out (the kernels then read the `AtmosphericState` directly, as before). + +# Fields +- `layerdata`: `(ncol, nlay, 4)` copy of `AtmosphericState.layerdata` + (dry-air column amount, layer pressure, layer temperature, relative + humidity). +- `t_lev`: `(ncol, nlay + 1)` copy of `AtmosphericState.t_lev`. +""" +struct TransposedStateCache{D2, D3} + layerdata::D3 + t_lev::D2 +end +Adapt.@adapt_structure TransposedStateCache + +function TransposedStateCache(grid_params::RRTMGPGridParams) + (; ncol, nlay) = grid_params + DA = ClimaComms.array_type(grid_params) + FT = eltype(grid_params) + return TransposedStateCache( + DA{FT, 3}(undef, ncol, nlay, 4), + DA{FT, 2}(undef, ncol, nlay + 1), + ) +end + +""" + refresh_transposed_state!(cache, as::AtmosphericState) + +Copy the current `AtmosphericState` layer data and level temperatures into +the column-first [`TransposedStateCache`](@ref) `cache`. Called at the start +of every spectral `solve_lw!`/`solve_sw!`; a no-op when `cache === nothing`. +""" +# Type-stable lazy permutations (the `PermutedDimsArray(x, perm)` constructor +# does not constant-fold `perm` before Julia 1.12); used by the refresh +# broadcasts and the no-cache fallbacks below. +@inline _lazy_permute(x::AbstractArray{T, 2}) where {T} = + PermutedDimsArray{T, 2, (2, 1), (2, 1), typeof(x)}(x) +@inline _lazy_permute(x::AbstractArray{T, 3}) where {T} = + PermutedDimsArray{T, 3, (3, 2, 1), (3, 2, 1), typeof(x)}(x) + +@inline refresh_transposed_state!(::Nothing, as::AtmosphericState) = nothing +function refresh_transposed_state!( + cache::TransposedStateCache, + as::AtmosphericState, +) + _transpose_into!(cache.layerdata, as.layerdata) + _transpose_into!(cache.t_lev, as.t_lev) + return nothing +end + +# Device arrays: the broadcast compiles to a single transpose kernel. +_transpose_into!(dest::AbstractArray, src) = (dest .= _lazy_permute(src); nothing) +# Host arrays: explicit loops (the broadcast through the permuted wrapper +# allocates a few bytes on Julia ≤ 1.11, tripping the zero-allocation tests). +function _transpose_into!(dest::Array{T, 3}, src::Array{T, 3}) where {T} + nfld, nlay, ncol = size(src) + @inbounds for f in 1:nfld, glay in 1:nlay, gcol in 1:ncol + dest[gcol, glay, f] = src[f, glay, gcol] + end + return nothing +end +function _transpose_into!(dest::Array{T, 2}, src::Array{T, 2}) where {T} + nlev, ncol = size(src) + @inbounds for glev in 1:nlev, gcol in 1:ncol + dest[gcol, glev] = src[glev, gcol] + end + return nothing +end + +# Kernel-facing accessors for the per-column layer data (`(nlay, 4)`: dry-air +# column amount, pressure, temperature, relative humidity — indexed +# `[glay, fld]`) and level temperatures (`(nlev,)`). With a cache, the +# returned views read the coalesced column-first copies; without one, they +# lazily present the `AtmosphericState` arrays in the same orientation (with +# the original memory access pattern, and — on Julia ≤ 1.11 — a small wrapper +# allocation per call on the CPU; the opt-out exists for memory-constrained +# configurations). +@inline layerdata_view(cache::TransposedStateCache, as, gcol) = + @inbounds view(cache.layerdata, gcol, :, :) +@inline layerdata_view(::Nothing, as::AtmosphericState, gcol) = + @inbounds view(_lazy_permute(as.layerdata), gcol, :, :) +@inline t_lev_view(cache::TransposedStateCache, as, gcol) = + @inbounds view(cache.t_lev, gcol, :) +@inline t_lev_view(::Nothing, as::AtmosphericState, gcol) = + @inbounds view(as.t_lev, :, gcol) + """ CloudState{CD, CF, CC, CM, CMT} diff --git a/src/optics/Fluxes.jl b/src/optics/Fluxes.jl index 9c3abf2a2..4a66855e6 100644 --- a/src/optics/Fluxes.jl +++ b/src/optics/Fluxes.jl @@ -8,12 +8,31 @@ export AbstractFlux, FluxLW, FluxSW, FluxBand, + FluxPresentation, + update_presentation!, + transpose_into!, + transpose_sum_into!, set_flux_to_zero!, set_band_flux_to_zero!, accumulate_band_flux!, apply_metric_scaling!, compute_net_flux! +""" + lazy_transpose(x) + +Lazily present the `(ncol, nlev[, n_bnd])` internal buffer `x` as +`(nlev, ncol[, n_bnd])` without copying. Equivalent to +`PermutedDimsArray(x, perm)`, but with the permutation in the type parameters: +before Julia 1.12, the `PermutedDimsArray(x, perm)` constructor does not +constant-fold `perm`, which would make every flux getter type-unstable (and +allocating) on Julia 1.9–1.11. +""" +lazy_transpose(x::AbstractArray{T, 2}) where {T} = + PermutedDimsArray{T, 2, (2, 1), (2, 1), typeof(x)}(x) +lazy_transpose(x::AbstractArray{T, 3}) where {T} = + PermutedDimsArray{T, 3, (2, 1, 3), (2, 1, 3), typeof(x)}(x) + """ AbstractFlux @@ -29,9 +48,9 @@ abstract type AbstractFlux{FT <: AbstractFloat, FTA2D <: AbstractArray{FT, 2}} e Upward, downward and net longwave fluxes at each level. # Fields -- `flux_up`: Upward flux [W/m²] `(nlev, ncol)`. -- `flux_dn`: Downward flux [W/m²] `(nlev, ncol)`. -- `flux_net`: Net flux [W/m²] `(nlev, ncol)`. +- `flux_up`: Upward flux [W/m²] `(ncol, nlev)`. +- `flux_dn`: Downward flux [W/m²] `(ncol, nlev)`. +- `flux_net`: Net flux [W/m²] `(ncol, nlev)`. """ struct FluxLW{FT <: AbstractFloat, FTA2D <: AbstractArray{FT, 2}} <: AbstractFlux{FT, FTA2D} @@ -95,14 +114,20 @@ end FluxBand{FT, FTA3D} Optional per-band upward, downward, and net radiative fluxes at each level, -`(ncol, nlev, n_bnd)`. Only allocated when spectrally-resolved fluxes are requested. -Band `b`'s slice `[:, :, b]` has the same `(ncol, nlev)` layout as the broadband -fluxes, and summing over the band dimension recovers the broadband fluxes. +`(nlev, ncol, n_bnd)`. Only allocated when spectrally-resolved fluxes are requested. +Summing over the band dimension recovers the broadband fluxes. + +Unlike the broadband compute buffers, the band buffers keep the host-facing +vertical-first layout: they are an opt-in diagnostic that the `spectral_*` +getters expose as plain views, and keeping them `(nlev, ncol, n_bnd)` avoids +doubling their (large) memory with separate presentation copies. The +accumulation writes are uncoalesced on the GPU, a cost paid only when +per-band fluxes are requested. # Fields -- `flux_up`: upward flux per band [W/m²], `(ncol, nlev, n_bnd)`. -- `flux_dn`: downward flux per band [W/m²], `(ncol, nlev, n_bnd)`. -- `flux_net`: net flux per band (`flux_up - flux_dn`) [W/m²], `(ncol, nlev, n_bnd)`. +- `flux_up`: upward flux per band [W/m²], `(nlev, ncol, n_bnd)`. +- `flux_dn`: downward flux per band [W/m²], `(nlev, ncol, n_bnd)`. +- `flux_net`: net flux per band (`flux_up - flux_dn`) [W/m²], `(nlev, ncol, n_bnd)`. """ struct FluxBand{FT <: AbstractFloat, FTA3D <: AbstractArray{FT, 3}} flux_up::FTA3D @@ -115,9 +140,9 @@ function FluxBand(grid_params::RRTMGPGridParams, n_bnd::Int) (; nlay, ncol) = grid_params FT = eltype(grid_params) DA = ClimaComms.array_type(grid_params) - flux_up = DA{FT}(undef, ncol, nlay + 1, n_bnd) - flux_dn = DA{FT}(undef, ncol, nlay + 1, n_bnd) - flux_net = DA{FT}(undef, ncol, nlay + 1, n_bnd) + flux_up = DA{FT}(undef, nlay + 1, ncol, n_bnd) + flux_dn = DA{FT}(undef, nlay + 1, ncol, n_bnd) + flux_net = DA{FT}(undef, nlay + 1, ncol, n_bnd) return FluxBand{FT, typeof(flux_up)}(flux_up, flux_dn, flux_net) end @@ -145,8 +170,8 @@ end ) bu, bd = band.flux_up, band.flux_dn @inbounds for ilev in 1:nlev - bu[gcol, ilev, ibnd] += flux_up[gcol, ilev] - bd[gcol, ilev, ibnd] += flux_dn[gcol, ilev] + bu[ilev, gcol, ibnd] += flux_up[gcol, ilev] + bd[ilev, gcol, ibnd] += flux_dn[gcol, ilev] end return nothing end @@ -233,11 +258,117 @@ function apply_metric_scaling!( flux::AbstractFlux, metric_scaling::AbstractArray, ) - sc = PermutedDimsArray(metric_scaling, (2, 1)) - flux.flux_up .= flux.flux_up .* sc - flux.flux_dn .= flux.flux_dn .* sc - flux.flux_net .= flux.flux_net .* sc - flux isa FluxSW && (flux.flux_dn_dir .= flux.flux_dn_dir .* sc) + _scale_by_transpose!(flux.flux_up, metric_scaling) + _scale_by_transpose!(flux.flux_dn, metric_scaling) + _scale_by_transpose!(flux.flux_net, metric_scaling) + flux isa FluxSW && _scale_by_transpose!(flux.flux_dn_dir, metric_scaling) + return nothing +end + +# Multiply the column-first `(ncol, nlev[, n_bnd])` array `a` by the +# vertical-first `(nlev, ncol)` scaling `sc`. Device arrays broadcast through +# the lazy transpose (a single kernel); host arrays use explicit loops, since +# broadcasts with a permuted-wrapper operand allocate a few bytes on +# Julia ≤ 1.11, tripping the zero-allocation tests. +_scale_by_transpose!(a::AbstractArray, sc) = + (a .= a .* lazy_transpose(sc); nothing) +function _scale_by_transpose!(a::Array{FT, 2}, sc::Array{FT, 2}) where {FT} + ncol, nlev = size(a) + @inbounds for ilev in 1:nlev, gcol in 1:ncol + a[gcol, ilev] *= sc[ilev, gcol] + end + return nothing +end +""" + FluxPresentation{FTA2D, FD} + +Host-facing `(nlev, ncol)` copies of a broadband flux set, filled by +[`update_presentation!`](@ref) from the column-first `(ncol, nlev)` compute +buffers at the end of every `update_lw_fluxes!`/`update_sw_fluxes!` (and hence +of [`update_fluxes!`](@ref RRTMGP.update_fluxes!)). The Layer-2 flux getters +return plain domain-masked views of these arrays, so the getter contract +(materializable with `Array`, broadcastable, reducible — also on the GPU, where +lazily transposed views of the compute buffers would fall outside the +wrapper types CUDA.jl dispatches on) holds without per-getter laziness. + +# Fields +- `flux_up`: upward flux [W/m²] `(nlev, ncol)`. +- `flux_dn`: downward flux [W/m²] `(nlev, ncol)`. +- `flux_net`: net flux [W/m²] `(nlev, ncol)`. +- `flux_dn_dir`: direct downward flux [W/m²] `(nlev, ncol)`, or `nothing` + (longwave). +""" +struct FluxPresentation{FTA2D <: AbstractArray, FD} + flux_up::FTA2D + flux_dn::FTA2D + flux_net::FTA2D + flux_dn_dir::FD +end +Adapt.@adapt_structure FluxPresentation + +function FluxPresentation(grid_params::RRTMGPGridParams; direct::Bool) + (; nlay, ncol) = grid_params + FT = eltype(grid_params) + DA = ClimaComms.array_type(grid_params) + alloc() = DA{FT}(undef, nlay + 1, ncol) + return FluxPresentation( + alloc(), + alloc(), + alloc(), + direct ? alloc() : nothing, + ) +end + +""" + update_presentation!(pres::FluxPresentation, flux::AbstractFlux) + +Fill the `(nlev, ncol)` presentation arrays from the `(ncol, nlev)` compute +buffers of `flux` (transposing copies; the direct beam only when present). +""" +function update_presentation!(pres::FluxPresentation, flux::AbstractFlux) + transpose_into!(pres.flux_up, flux.flux_up) + transpose_into!(pres.flux_dn, flux.flux_dn) + transpose_into!(pres.flux_net, flux.flux_net) + isnothing(pres.flux_dn_dir) || + transpose_into!(pres.flux_dn_dir, flux.flux_dn_dir) + return nothing +end + +""" + transpose_into!(dest, src) + +Copy the 2D array `src` into `dest` with the two dimensions swapped +(`dest[i, j] = src[j, i]`). Device arrays use a single broadcast kernel; host +arrays use explicit loops, since broadcasts with a permuted-wrapper operand +allocate a few bytes on Julia ≤ 1.11, tripping the zero-allocation tests. +""" +transpose_into!(dest::AbstractArray, src) = + (dest .= lazy_transpose(src); nothing) +function transpose_into!(dest::Array{FT, 2}, src::Array{FT, 2}) where {FT} + n1, n2 = size(dest) + @inbounds for j in 1:n2, i in 1:n1 + dest[i, j] = src[j, i] + end + return nothing +end + +""" + transpose_sum_into!(dest, a, b) + +Set `dest[i, j] = a[j, i] + b[j, i]` (transposing sum of two 2D arrays); the +device/host split follows [`transpose_into!`](@ref). +""" +transpose_sum_into!(dest::AbstractArray, a, b) = + (dest .= lazy_transpose(a) .+ lazy_transpose(b); nothing) +function transpose_sum_into!( + dest::Array{FT, 2}, + a::Array{FT, 2}, + b::Array{FT, 2}, +) where {FT} + n1, n2 = size(dest) + @inbounds for j in 1:n2, i in 1:n1 + dest[i, j] = a[j, i] + b[j, i] + end return nothing end function apply_metric_scaling!( @@ -253,9 +384,10 @@ end apply_metric_scaling!(::Nothing, metric_scaling) = nothing apply_metric_scaling!(band::FluxBand, metric_scaling::Nothing) = nothing function apply_metric_scaling!(band::FluxBand, metric_scaling::AbstractArray) - sc = PermutedDimsArray(metric_scaling, (2, 1)) - band.flux_up .= band.flux_up .* sc - band.flux_dn .= band.flux_dn .* sc + # The band buffers share `metric_scaling`'s vertical-first layout, so the + # `(nlev, ncol)` scaling broadcasts across the band dimension directly. + band.flux_up .= band.flux_up .* metric_scaling + band.flux_dn .= band.flux_dn .* metric_scaling return nothing end diff --git a/src/optics/Sources.jl b/src/optics/Sources.jl index ba55aafd1..8b969fbeb 100644 --- a/src/optics/Sources.jl +++ b/src/optics/Sources.jl @@ -24,8 +24,8 @@ and at the surface for no-scattering calculations. # Fields - `param_set`: Parameter set. - `sfc_source`: Surface source [W/m²] `(ncol)`. -- `lay_source`: Planck source at layer average temperature [W/m²] `(nlay, ncol)`. -- `lev_source`: Planck level source at layer edges [W/m²] `(nlay+1, ncol)`; includes +- `lay_source`: Planck source at layer average temperature [W/m²] `(ncol, nlay)`. +- `lev_source`: Planck level source at layer edges [W/m²] `(ncol, nlay+1)`; includes spectral weighting that accounts for the state-dependent frequency to g-space mapping. """ struct SourceLWNoScat{S, D, PS} <: AbstractSourceLW diff --git a/src/optics/compute_optical_props.jl b/src/optics/compute_optical_props.jl index b56c60d20..6cded6372 100644 --- a/src/optics/compute_optical_props.jl +++ b/src/optics/compute_optical_props.jl @@ -2,6 +2,7 @@ compute_optical_props!( op::AbstractOpticalProps, as::AtmosphericState{FT}, + state_cache::Union{TransposedStateCache, Nothing}, sf::AbstractSourceLW{FT}, gcol::Int, igpt::Int, @@ -10,11 +11,14 @@ lkp_aero::Union{LookUpAerosolMerra, Nothing} = nothing, ) where {FT<:AbstractFloat} -Compute optical properties for the longwave problem. +Compute optical properties for the longwave problem. The layer data and level +temperatures are read through `state_cache` (column-first copies for coalesced +GPU access) when one is provided, and from `as` directly otherwise. """ @inline function compute_optical_props!( op::OneScalar, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, sf::SourceLWNoScat, gcol::Int, igpt::Int, @@ -30,8 +34,8 @@ Compute optical properties for the longwave problem. t_sfc = as.t_sfc[gcol] ibnd = lkp.band_data.major_gpt2bnd[igpt] totplnk = view(lkp.planck.tot_planck, :, ibnd) - as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) - t_lev_col = view(as.t_lev, :, gcol) + as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) + t_lev_col = AtmosphericStates.t_lev_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) lev_src_inc_prev = zero(t_sfc) @@ -39,9 +43,9 @@ Compute optical properties for the longwave problem. t_lev_dec = t_lev_col[1] for glay in 1:nlay - col_dry, p_lay, t_lay = as_layerdata[1, glay], - as_layerdata[2, glay], - as_layerdata[3, glay] + col_dry, p_lay, t_lay = as_layerdata[glay, 1], + as_layerdata[glay, 2], + as_layerdata[glay, 3] # compute gas optics τ[glay], _, _, planckfrac = compute_gas_optics( lkp, @@ -102,7 +106,7 @@ Compute optical properties for the longwave problem. aero_mask = view(as.aerosol_state.aero_mask, :, gcol) aero_size = view(as.aerosol_state.aero_size, :, :, gcol) aero_mass = view(as.aerosol_state.aero_mass, :, :, gcol) - rel_hum = AtmosphericStates.getview_rel_hum(as, gcol) + rel_hum = view(as_layerdata, :, 4) add_aerosol_optics_1scalar!( τ, @@ -124,6 +128,7 @@ end @inline function compute_optical_props!( op::TwoStream, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, sf::SourceLW2Str, gcol::Int, igpt::Int, @@ -139,8 +144,8 @@ end t_sfc = as.t_sfc[gcol] ibnd = lkp.band_data.major_gpt2bnd[igpt] totplnk = view(lkp.planck.tot_planck, :, ibnd) - as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) - t_lev_col = view(as.t_lev, :, gcol) + as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) + t_lev_col = AtmosphericStates.t_lev_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) ssa = view(op.ssa, gcol, :) g = view(op.g, gcol, :) @@ -152,9 +157,9 @@ end @inbounds begin t_lev_dec = t_lev_col[1] for glay in 1:nlay - col_dry, p_lay, t_lay = as_layerdata[1, glay], - as_layerdata[2, glay], - as_layerdata[3, glay] + col_dry, p_lay, t_lay = as_layerdata[glay, 1], + as_layerdata[glay, 2], + as_layerdata[glay, 3] # compute gas optics τ[glay], ssa[glay], g[glay], planckfrac = compute_gas_optics( lkp, @@ -217,7 +222,7 @@ end aero_mask = view(as.aerosol_state.aero_mask, :, gcol) aero_size = view(as.aerosol_state.aero_size, :, :, gcol) aero_mass = view(as.aerosol_state.aero_mass, :, :, gcol) - rel_hum = AtmosphericStates.getview_rel_hum(as, gcol) + rel_hum = view(as_layerdata, :, 4) add_aerosol_optics_2stream!( τ, @@ -241,6 +246,7 @@ end compute_optical_props!( op::AbstractOpticalProps, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, gcol::Int, igpt::Int, lkp::LookUpSW, @@ -248,11 +254,14 @@ end lkp_aero::Union{LookUpAerosolMerra, Nothing} = nothing, ) -Compute optical properties for the shortwave problem. +Compute optical properties for the shortwave problem. The layer data are read +through `state_cache` (column-first copies for coalesced GPU access) when one +is provided, and from `as` directly otherwise. """ @inline function compute_optical_props!( op::OneScalar, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, gcol::Int, igpt::Int, lkp::LookUpSW, @@ -263,12 +272,12 @@ Compute optical properties for the shortwave problem. @inbounds begin ibnd = lkp.band_data.major_gpt2bnd[igpt] t_sfc = as.t_sfc[gcol] - as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) + as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) end @inbounds for glay in 1:nlay col_dry, p_lay, t_lay = - as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] + as_layerdata[glay, 1], as_layerdata[glay, 2], as_layerdata[glay, 3] # compute gas optics τ[glay], _, _ = compute_gas_optics( lkp, @@ -288,6 +297,7 @@ end @inline function compute_optical_props!( op::TwoStream, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, gcol::Int, igpt::Int, lkp::LookUpSW, @@ -299,14 +309,14 @@ end @inbounds begin ibnd = lkp.band_data.major_gpt2bnd[igpt] t_sfc = as.t_sfc[gcol] - as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) + as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) ssa = view(op.ssa, gcol, :) g = view(op.g, gcol, :) end @inbounds for glay in 1:nlay col_dry, p_lay, t_lay = - as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] + as_layerdata[glay, 1], as_layerdata[glay, 2], as_layerdata[glay, 3] # compute gas optics τ[glay], ssa[glay], g[glay] = compute_gas_optics( lkp, @@ -350,7 +360,7 @@ end aero_mask = view(as.aerosol_state.aero_mask, :, gcol) aero_size = view(as.aerosol_state.aero_size, :, :, gcol) aero_mass = view(as.aerosol_state.aero_mass, :, :, gcol) - rel_hum = AtmosphericStates.getview_rel_hum(as, gcol) + rel_hum = view(as_layerdata, :, 4) add_aerosol_optics_2stream!( τ, diff --git a/src/rte/RTE.jl b/src/rte/RTE.jl index 92aaa2066..92dc613c8 100644 --- a/src/rte/RTE.jl +++ b/src/rte/RTE.jl @@ -28,6 +28,9 @@ configurations for a non-scattering longwave simulation. - `fluxb`: Temporary storage for bandwise calculations. - `flux`: Longwave fluxes. - `angle_disc`: Angular discretization. +- `state_cache`: [`TransposedStateCache`](@ref + RRTMGP.AtmosphericStates.TransposedStateCache) with column-first copies of + the hot state arrays (refreshed at each spectral solve), or `nothing`. """ struct NoScatLWRTE{ C, @@ -37,6 +40,7 @@ struct NoScatLWRTE{ FXBL, FXL <: FluxLW, AD, + TSC, } context::C op::OP @@ -45,10 +49,17 @@ struct NoScatLWRTE{ fluxb::FXBL flux::FXL angle_disc::AD + state_cache::TSC end Adapt.@adapt_structure NoScatLWRTE -function NoScatLWRTE(grid_params::RRTMGPGridParams; params, sfc_emis, inc_flux) +function NoScatLWRTE( + grid_params::RRTMGPGridParams; + params, + sfc_emis, + inc_flux, + state_cache = TransposedStateCache(grid_params), +) (; context) = grid_params op = OneScalar(grid_params) src = SourceLWNoScat(grid_params; params) @@ -56,7 +67,7 @@ function NoScatLWRTE(grid_params::RRTMGPGridParams; params, sfc_emis, inc_flux) fluxb = FluxLW(grid_params) flux = FluxLW(grid_params) ad = AngularDiscretization(grid_params, 1) - return NoScatLWRTE(context, op, src, bcs, fluxb, flux, ad) + return NoScatLWRTE(context, op, src, bcs, fluxb, flux, ad, state_cache) end """ @@ -74,6 +85,9 @@ configurations for a `2-stream` longwave simulation. - `fluxb`: Temporary storage for bandwise calculations. - `flux`: Longwave fluxes. - `band_flux`: Optional per-band longwave fluxes `(nlev, ncol, n_bnd)`, or `nothing`. +- `state_cache`: [`TransposedStateCache`](@ref + RRTMGP.AtmosphericStates.TransposedStateCache) with column-first copies of + the hot state arrays (refreshed at each spectral solve), or `nothing`. """ struct TwoStreamLWRTE{ C, @@ -83,6 +97,7 @@ struct TwoStreamLWRTE{ FXBL, FXL <: FluxLW, FXBND, + TSC, } context::C op::OP @@ -91,17 +106,19 @@ struct TwoStreamLWRTE{ fluxb::FXBL flux::FXL band_flux::FXBND + state_cache::TSC end Adapt.@adapt_structure TwoStreamLWRTE # By default no per-band (spectrally-resolved) fluxes are retained. -TwoStreamLWRTE(context, op, src, bcs, fluxb, flux) = - TwoStreamLWRTE(context, op, src, bcs, fluxb, flux, nothing) +TwoStreamLWRTE(context, op, src, bcs, fluxb, flux, band_flux) = + TwoStreamLWRTE(context, op, src, bcs, fluxb, flux, band_flux, nothing) function TwoStreamLWRTE( grid_params::RRTMGPGridParams; params, sfc_emis, inc_flux, + state_cache = TransposedStateCache(grid_params), ) (; context) = grid_params op = TwoStream(grid_params) @@ -109,7 +126,16 @@ function TwoStreamLWRTE( bcs = LwBCs(sfc_emis, inc_flux) fluxb = FluxLW(grid_params) flux = FluxLW(grid_params) - return TwoStreamLWRTE(context, op, src, bcs, fluxb, flux) + return TwoStreamLWRTE( + context, + op, + src, + bcs, + fluxb, + flux, + nothing, + state_cache, + ) end """ @@ -125,13 +151,17 @@ configurations for a non-scattering shortwave simulation. - `bcs`: Shortwave boundary conditions. - `fluxb`: Temporary storage for bandwise calculations. - `flux`: Shortwave fluxes. +- `state_cache`: [`TransposedStateCache`](@ref + RRTMGP.AtmosphericStates.TransposedStateCache) with column-first copies of + the hot state arrays (refreshed at each spectral solve), or `nothing`. """ -struct NoScatSWRTE{C, OP <: OneScalar, BC <: SwBCs, FXBS, FXS <: FluxSW} +struct NoScatSWRTE{C, OP <: OneScalar, BC <: SwBCs, FXBS, FXS <: FluxSW, TSC} context::C op::OP bcs::BC fluxb::FXBS flux::FXS + state_cache::TSC end Adapt.@adapt_structure NoScatSWRTE @@ -142,6 +172,7 @@ function NoScatSWRTE( sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse, + state_cache = TransposedStateCache(grid_params), ) (; context) = grid_params op = OneScalar(grid_params) @@ -154,7 +185,7 @@ function NoScatSWRTE( ) fluxb = FluxSW(grid_params) flux = FluxSW(grid_params) - return NoScatSWRTE(context, op, bcs, fluxb, flux) + return NoScatSWRTE(context, op, bcs, fluxb, flux, state_cache) end """ @@ -172,6 +203,9 @@ configurations for a `2-stream` shortwave simulation. - `fluxb`: Temporary storage for bandwise calculations. - `flux`: Shortwave fluxes. - `band_flux`: Optional per-band shortwave fluxes `(nlev, ncol, n_bnd)`, or `nothing`. +- `state_cache`: [`TransposedStateCache`](@ref + RRTMGP.AtmosphericStates.TransposedStateCache) with column-first copies of + the hot state arrays (refreshed at each spectral solve), or `nothing`. """ struct TwoStreamSWRTE{ C, @@ -181,6 +215,7 @@ struct TwoStreamSWRTE{ FXBS, FXS <: FluxSW, FXBND, + TSC, } context::C op::OP @@ -189,11 +224,12 @@ struct TwoStreamSWRTE{ fluxb::FXBS flux::FXS band_flux::FXBND + state_cache::TSC end Adapt.@adapt_structure TwoStreamSWRTE # By default no per-band (spectrally-resolved) fluxes are retained. -TwoStreamSWRTE(context, op, src, bcs, fluxb, flux) = - TwoStreamSWRTE(context, op, src, bcs, fluxb, flux, nothing) +TwoStreamSWRTE(context, op, src, bcs, fluxb, flux, band_flux) = + TwoStreamSWRTE(context, op, src, bcs, fluxb, flux, band_flux, nothing) function TwoStreamSWRTE( grid_params::RRTMGPGridParams; @@ -202,6 +238,7 @@ function TwoStreamSWRTE( sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse, + state_cache = TransposedStateCache(grid_params), ) (; context) = grid_params op = TwoStream(grid_params) @@ -215,7 +252,16 @@ function TwoStreamSWRTE( ) fluxb = FluxSW(grid_params) flux = FluxSW(grid_params) - return TwoStreamSWRTE(context, op, src, bcs, fluxb, flux) + return TwoStreamSWRTE( + context, + op, + src, + bcs, + fluxb, + flux, + nothing, + state_cache, + ) end diff --git a/src/rte/RTESolver.jl b/src/rte/RTESolver.jl index 358622e30..7a88266f7 100644 --- a/src/rte/RTESolver.jl +++ b/src/rte/RTESolver.jl @@ -70,13 +70,14 @@ Additionally, takes an optional argument `metric_scaling` which scales the resul corresponding factor to account for column expansion in `deep-atmosphere` configurations. """ function solve_lw!( - (; context, fluxb, flux, src, bcs, op, angle_disc)::NoScatLWRTE, + (; context, fluxb, flux, src, bcs, op, angle_disc, state_cache)::NoScatLWRTE, as::AtmosphericState, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing} = nothing, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, metric_scaling::M = nothing, ) where {M} + AtmosphericStates.refresh_transposed_state!(state_cache, as) rte_lw_noscat_solve!( context.device, fluxb, @@ -86,6 +87,7 @@ function solve_lw!( op, angle_disc, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, @@ -108,13 +110,14 @@ Additionally, takes an optional argument `metric_scaling` which scales the resul corresponding factor to account for column expansion in `deep-atmosphere` configurations. """ function solve_lw!( - (; context, fluxb, flux, band_flux, src, bcs, op)::TwoStreamLWRTE, + (; context, fluxb, flux, band_flux, src, bcs, op, state_cache)::TwoStreamLWRTE, as::AtmosphericState, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing} = nothing, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, metric_scaling::M = nothing, ) where {M} + AtmosphericStates.refresh_transposed_state!(state_cache, as) rte_lw_2stream_solve!( context.device, fluxb, @@ -124,6 +127,7 @@ function solve_lw!( bcs, op, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, @@ -177,12 +181,22 @@ Additionally, takes an optional argument `metric_scaling` which scales the resul corresponding factor to account for column expansion in `deep-atmosphere` configurations. """ function solve_sw!( - (; context, fluxb, flux, bcs, op)::NoScatSWRTE, + (; context, fluxb, flux, bcs, op, state_cache)::NoScatSWRTE, as::AtmosphericState, lookup_sw::LookUpSW, metric_scaling::M = nothing, ) where {M} - rte_sw_noscat_solve!(context.device, fluxb, flux, op, bcs, as, lookup_sw) + AtmosphericStates.refresh_transposed_state!(state_cache, as) + rte_sw_noscat_solve!( + context.device, + fluxb, + flux, + op, + bcs, + as, + state_cache, + lookup_sw, + ) apply_metric_scaling!(flux, metric_scaling) end @@ -201,13 +215,14 @@ Additionally, takes an optional argument `metric_scaling` which scales the resul corresponding factor to account for column expansion in `deep-atmosphere` configurations. """ function solve_sw!( - (; context, fluxb, flux, band_flux, src, bcs, op)::TwoStreamSWRTE, + (; context, fluxb, flux, band_flux, src, bcs, op, state_cache)::TwoStreamSWRTE, as::AtmosphericState, lookup_sw::LookUpSW, lookup_sw_cld::Union{LookUpCld, Nothing} = nothing, lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, metric_scaling::M = nothing, ) where {M} + AtmosphericStates.refresh_transposed_state!(state_cache, as) rte_sw_2stream_solve!( context.device, fluxb, @@ -217,6 +232,7 @@ function solve_sw!( bcs, src, as, + state_cache, lookup_sw, lookup_sw_cld, lookup_sw_aero, diff --git a/src/rte/longwave_2stream.jl b/src/rte/longwave_2stream.jl index 846ababdd..d0dc5b62a 100644 --- a/src/rte/longwave_2stream.jl +++ b/src/rte/longwave_2stream.jl @@ -43,6 +43,7 @@ end bcs_lw, op, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, @@ -54,6 +55,7 @@ end compute_optical_props!( op, as, + state_cache, src_lw, gcol, igpt, @@ -77,6 +79,7 @@ function rte_lw_2stream_solve!( bcs_lw::LwBCs, op::TwoStream, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing} = nothing, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, @@ -111,6 +114,7 @@ function rte_lw_2stream_solve!( bcs_lw, op, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, diff --git a/src/rte/longwave_noscat.jl b/src/rte/longwave_noscat.jl index 24b8e6758..f687d6572 100644 --- a/src/rte/longwave_noscat.jl +++ b/src/rte/longwave_noscat.jl @@ -48,6 +48,7 @@ end Ds, w_μ, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, @@ -59,6 +60,7 @@ end compute_optical_props!( op, as, + state_cache, src_lw, gcol, igpt, @@ -92,6 +94,7 @@ function rte_lw_noscat_solve!( op::OneScalar, angle_disc::AngularDiscretization, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_lw::LookUpLW, lookup_lw_cld::Union{LookUpCld, Nothing} = nothing, lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, @@ -127,6 +130,7 @@ function rte_lw_noscat_solve!( Ds, w_μ, as, + state_cache, lookup_lw, lookup_lw_cld, lookup_lw_aero, diff --git a/src/rte/shortwave_2stream.jl b/src/rte/shortwave_2stream.jl index 233ef1ce4..435140bf5 100644 --- a/src/rte/shortwave_2stream.jl +++ b/src/rte/shortwave_2stream.jl @@ -54,6 +54,7 @@ end bcs_sw, src_sw, as, + state_cache, lookup_sw, lookup_sw_cld, lookup_sw_aero, @@ -66,6 +67,7 @@ end compute_optical_props!( op, as, + state_cache, gcol, igpt, lookup_sw, @@ -109,6 +111,7 @@ function rte_sw_2stream_solve!( bcs_sw::SwBCs, src_sw::SourceSW2Str, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_sw::LookUpSW, lookup_sw_cld::Union{LookUpCld, Nothing} = nothing, lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing, @@ -143,6 +146,7 @@ function rte_sw_2stream_solve!( bcs_sw, src_sw, as, + state_cache, lookup_sw, lookup_sw_cld, lookup_sw_aero, diff --git a/src/rte/shortwave_noscat.jl b/src/rte/shortwave_noscat.jl index bb1963c81..6436bbd7a 100644 --- a/src/rte/shortwave_noscat.jl +++ b/src/rte/shortwave_noscat.jl @@ -45,11 +45,12 @@ end op, bcs_sw, as, + state_cache, lookup_sw, n_gpt, nlev, ) - compute_optical_props!(op, as, gcol, igpt, lookup_sw, nothing) + compute_optical_props!(op, as, state_cache, gcol, igpt, lookup_sw, nothing) @inbounds solar_frac = lookup_sw.solar_src_scaled[igpt] rte_sw_noscat!(flux, op, bcs_sw, igpt, n_gpt, solar_frac, gcol, nlev) _accumulate_fluxes!(flux_sw, flux, gcol, nlev, igpt) @@ -63,6 +64,7 @@ function rte_sw_noscat_solve!( op::OneScalar, bcs_sw::SwBCs, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, lookup_sw::LookUpSW, ) nlay, ncol = AtmosphericStates.get_dims(as) @@ -81,6 +83,7 @@ function rte_sw_noscat_solve!( op, bcs_sw, as, + state_cache, lookup_sw, n_gpt, nlev, diff --git a/test/clear_sky_utils.jl b/test/clear_sky_utils.jl index 5d3d0d7ea..668ea9ceb 100644 --- a/test/clear_sky_utils.jl +++ b/test/clear_sky_utils.jl @@ -146,6 +146,19 @@ function clear_sky( )) ≤ 448 end + # The transposed state cache is a pure read-path optimization: a solver + # constructed without one must produce bitwise-identical fluxes. + slv_lw_nocache = SLVLW( + grid_params; + params = param_set, + sfc_emis, + inc_flux, + state_cache = nothing, + ) + solve_lw!(slv_lw_nocache, as, lookup_lw, nothing, nothing, metric_scaling) + @test Array(slv_lw_nocache.flux.flux_up) == Array(slv_lw.flux.flux_up) + @test Array(slv_lw_nocache.flux.flux_dn) == Array(slv_lw.flux.flux_dn) + # comparing longwave fluxes with data from RRTMGP FORTRAN code comp_flux_up_lw, comp_flux_dn_lw, comp_flux_up_sw, comp_flux_dn_sw = load_comparison_data(expt_no, bot_at_1, ncol) diff --git a/test/cloudy_sky_utils.jl b/test/cloudy_sky_utils.jl index 1e386b4e1..fda135945 100644 --- a/test/cloudy_sky_utils.jl +++ b/test/cloudy_sky_utils.jl @@ -27,6 +27,73 @@ using RRTMGP.ArtifactPaths include("reference_files.jl") include("read_cloudy_sky.jl") +# Function barrier for the type-stability and allocation checks: +# `setup_cloudy_sky_as` is not type-stable, so calling the solvers directly +# from `cloudy_sky` measures the call-site dispatch boxing of the test +# function itself (a few hundred bytes) rather than solver allocations. With +# concrete argument types, the solves must be allocation-free. +function _test_cloudy_zero_alloc( + slv_lw, + slv_sw, + as, + lookup_lw, + lookup_lw_cld, + lookup_sw, + lookup_sw_cld, + metric_scaling, +) + JET.@test_opt solve_lw!( + slv_lw, + as, + lookup_lw, + lookup_lw_cld, + nothing, + metric_scaling, + ) + # `@allocated`'s first measurement after JET includes one-time setup; discard it. + @allocated solve_lw!( + slv_lw, + as, + lookup_lw, + lookup_lw_cld, + nothing, + metric_scaling, + ) + @test (@allocated solve_lw!( + slv_lw, + as, + lookup_lw, + lookup_lw_cld, + nothing, + metric_scaling, + )) == 0 + JET.@test_opt solve_sw!( + slv_sw, + as, + lookup_sw, + lookup_sw_cld, + nothing, + metric_scaling, + ) + @allocated solve_sw!( + slv_sw, + as, + lookup_sw, + lookup_sw_cld, + nothing, + metric_scaling, + ) + @test (@allocated solve_sw!( + slv_sw, + as, + lookup_sw, + lookup_sw_cld, + nothing, + metric_scaling, + )) == 0 + return nothing +end + function cloudy_sky( context, ::Type{SLVLW}, @@ -109,45 +176,18 @@ function cloudy_sky( #------calling solvers metric_scaling = DA(one.(as.p_lev)) solve_lw!(slv_lw, as, lookup_lw, lookup_lw_cld, nothing, metric_scaling) + solve_sw!(slv_sw, as, lookup_sw, lookup_sw_cld, nothing, metric_scaling) if device isa ClimaComms.CPUSingleThreaded - JET.@test_opt solve_lw!( - slv_lw, - as, - lookup_lw, - lookup_lw_cld, - nothing, - metric_scaling, - ) - #@test (@allocated solve_lw!(slv_lw, as, lookup_lw, lookup_lw_cld, nothing, metric_scaling)) == 0 - @test (@allocated solve_lw!( + _test_cloudy_zero_alloc( slv_lw, + slv_sw, as, lookup_lw, lookup_lw_cld, - nothing, - metric_scaling, - )) ≤ 448 - end - - - solve_sw!(slv_sw, as, lookup_sw, lookup_sw_cld, nothing, metric_scaling) - if device isa ClimaComms.CPUSingleThreaded - JET.@test_opt solve_sw!( - slv_sw, - as, lookup_sw, lookup_sw_cld, - nothing, metric_scaling, ) - @test (@allocated solve_sw!( - slv_sw, - as, - lookup_sw, - lookup_sw_cld, - nothing, - metric_scaling, - )) ≤ 368 end #------------- # comparison diff --git a/test/scalar_indexing.jl b/test/scalar_indexing.jl index af985cf18..c11555317 100644 --- a/test/scalar_indexing.jl +++ b/test/scalar_indexing.jl @@ -4,12 +4,12 @@ import ClimaComms import RRTMGP import CUDA -# The `(nlev, ncol)` flux buffers are exposed through `_domain_view` (src/api/getters.jl) -# as plain domain-masked `SubArray` views of the device array. Bringing one to the CPU -# with `Array(getter(solver))` must not fall back to element-wise GPU scalar indexing -# (CUDA.jl is expected to convert these strided views natively). This test pins that -# contract for the flux getters; if it fails on GPU, an RRTMGP-owned `to_cpu` is the -# fix -- not a `Base.Array` overload on foreign types. +# The flux getters return plain domain-masked `SubArray` views of the +# `(nlev, ncol)` presentation arrays that `update_fluxes!` fills from the +# column-first compute buffers (src/api/getters.jl; `Fluxes.FluxPresentation`). +# Bringing one to the CPU with `Array(getter(solver))` must not fall back to +# element-wise GPU scalar indexing (CUDA.jl converts single-wrap strided views +# natively). This test pins that contract for the flux getters. # A gray two-stream lw+sw solver with an optional isothermal boundary layer. function gray_solver( diff --git a/test/standalone.jl b/test/standalone.jl index 4dc13a467..7a39a6bc7 100644 --- a/test/standalone.jl +++ b/test/standalone.jl @@ -242,8 +242,8 @@ end RRTMGP.update_fluxes!(solver) - # internal buffers are boundary-extended and coalesced (ncol, nlev); getters are domain-masked - @test size(solver.net_flux_buffer, 2) == nlay + 1 # full levels + # the net-flux buffer is boundary-extended, host-facing (nlev, ncol); getters are domain-masked + @test size(solver.net_flux_buffer, 1) == nlay + 1 # full levels @test size(RRTMGP.net_flux(solver), 1) == domain_nlay + 1 # domain levels @test size(RRTMGP.level_pressure(solver), 1) == domain_nlay + 1 @test size(RRTMGP.lw_flux_net(solver), 1) == domain_nlay + 1 From b1071b5e938db6549d3b1800623a0604d846b16b Mon Sep 17 00:00:00 2001 From: Tapio Schneider Date: Sun, 12 Jul 2026 14:42:39 +0000 Subject: [PATCH 3/6] Corrections and performance table updates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docs/src/howto/gpu.md | 14 +++--- docs/src/optics/AtmosphericStates.md | 1 + src/optics/AtmosphericStates.jl | 75 +++++++++++++--------------- src/optics/compute_optical_props.jl | 36 +++++++------ test/standalone.jl | 26 ++++++---- 5 files changed, 81 insertions(+), 71 deletions(-) diff --git a/docs/src/howto/gpu.md b/docs/src/howto/gpu.md index 4e1dcf411..be53ac49a 100644 --- a/docs/src/howto/gpu.md +++ b/docs/src/howto/gpu.md @@ -114,13 +114,13 @@ Every configuration evaluated by the benchmark (see [`perf/benchmark_ratchet.jl` ### Single precision (`Float32`) — Primary target -At `Float32` precision, memory traffic across lookup-table interpolations and solver evaluations is halved compared to `Float64`, enabling massive GPU speedups when combined with contiguous `(ncol, nlev)` internal memory coalescing: +At `Float32` precision, memory traffic across lookup-table interpolations and solver evaluations is halved compared to `Float64`, enabling large GPU speedups: | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup | |---|---|---|---| -| Clear-sky (two-stream) | 77.95 s | 0.95 s | 82.5× | -| All-sky, McICA clouds | 98.42 s | 1.39 s | 70.8× | -| All-sky with aerosols | 99.39 s | 1.52 s | 65.5× | +| Clear-sky (two-stream) | 77.95 s | 0.95 s | 81.7× | +| All-sky, McICA clouds | 98.42 s | 1.40 s | 70.4× | +| All-sky with aerosols | 99.39 s | 1.54 s | 64.5× | ### Double precision (`Float64`) @@ -128,6 +128,6 @@ For reference, running the identical uniform benchmark sweep in double precision | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup | |---|---|---|---| -| Clear-sky (two-stream) | 79.64 s | 1.32 s | 60.1× | -| All-sky, McICA clouds | 101.60 s | 1.93 s | 52.6× | -| All-sky with aerosols | 106.33 s | 2.07 s | 51.4× | +| Clear-sky (two-stream) | 79.64 s | 1.30 s | 61.1× | +| All-sky, McICA clouds | 101.60 s | 1.91 s | 53.3× | +| All-sky with aerosols | 106.33 s | 2.06 s | 51.6× | diff --git a/docs/src/optics/AtmosphericStates.md b/docs/src/optics/AtmosphericStates.md index fa5083d3c..5a0a6158c 100644 --- a/docs/src/optics/AtmosphericStates.md +++ b/docs/src/optics/AtmosphericStates.md @@ -10,6 +10,7 @@ AtmosphericState GrayAtmosphericState CloudState AerosolState +TransposedStateCache AbstractCloudMask MaxRandomOverlap AbstractGrayOpticalThickness diff --git a/src/optics/AtmosphericStates.jl b/src/optics/AtmosphericStates.jl index dbb346f75..6d5b70412 100644 --- a/src/optics/AtmosphericStates.jl +++ b/src/optics/AtmosphericStates.jl @@ -89,38 +89,27 @@ Adapt.@adapt_structure AtmosphericState @inline get_dims(as::AtmosphericState) = size(as.layerdata, 2), size(as.layerdata, 3) -@inline getview_layerdata(as::AtmosphericState, gcol) = - @inbounds view(as.layerdata, :, :, gcol) - +# Whole-column-set views of the packed layer data (each `(nlay, ncol)`); the +# per-column `(nlay,)` field slices used by the gas-optics kernels are the +# `col_dry_view`/`p_lay_view`/`t_lay_view`/`rel_hum_view` accessors below. # view of column amounts of dry air [molecules per cm^2 of dry air] @inline getview_col_dry(as::AtmosphericState) = @inbounds view(as.layerdata, 1, :, :) -@inline getview_col_dry(as::AtmosphericState, gcol) = - @inbounds view(as.layerdata, 1, :, gcol) - # view of layer pressures [Pa, mb] @inline getview_p_lay(as::AtmosphericState) = @inbounds view(as.layerdata, 2, :, :) -@inline getview_p_lay(as::AtmosphericState, gcol) = - @inbounds view(as.layerdata, 2, :, gcol) - # view of layer temperatures [K] @inline getview_t_lay(as::AtmosphericState) = @inbounds view(as.layerdata, 3, :, :) -@inline getview_t_lay(as::AtmosphericState, gcol) = - @inbounds view(as.layerdata, 3, :, gcol) - # view of relative humidity @inline getview_rel_hum(as::AtmosphericState) = @inbounds view(as.layerdata, 4, :, :) -@inline getview_rel_hum(as::AtmosphericState, gcol) = - @inbounds view(as.layerdata, 4, :, gcol) """ TransposedStateCache{D2, D3} Column-first copies of the hot [`AtmosphericState`](@ref) fields, refreshed -once per solve by [`refresh_transposed_state!`](@ref). The gas-optics kernels +once per solve (by `refresh_transposed_state!`). The gas-optics kernels re-read the layer pressures, temperatures, dry-air column amounts, and level temperatures for every g-point; in the caller-owned `AtmosphericState` those arrays are vertical-first, so neighboring GPU threads (one per column) read @@ -162,14 +151,6 @@ Copy the current `AtmosphericState` layer data and level temperatures into the column-first [`TransposedStateCache`](@ref) `cache`. Called at the start of every spectral `solve_lw!`/`solve_sw!`; a no-op when `cache === nothing`. """ -# Type-stable lazy permutations (the `PermutedDimsArray(x, perm)` constructor -# does not constant-fold `perm` before Julia 1.12); used by the refresh -# broadcasts and the no-cache fallbacks below. -@inline _lazy_permute(x::AbstractArray{T, 2}) where {T} = - PermutedDimsArray{T, 2, (2, 1), (2, 1), typeof(x)}(x) -@inline _lazy_permute(x::AbstractArray{T, 3}) where {T} = - PermutedDimsArray{T, 3, (3, 2, 1), (3, 2, 1), typeof(x)}(x) - @inline refresh_transposed_state!(::Nothing, as::AtmosphericState) = nothing function refresh_transposed_state!( cache::TransposedStateCache, @@ -180,10 +161,13 @@ function refresh_transposed_state!( return nothing end -# Device arrays: the broadcast compiles to a single transpose kernel. -_transpose_into!(dest::AbstractArray, src) = (dest .= _lazy_permute(src); nothing) -# Host arrays: explicit loops (the broadcast through the permuted wrapper -# allocates a few bytes on Julia ≤ 1.11, tripping the zero-allocation tests). +# Device arrays: `permutedims!` maps to a single native transpose kernel. +_transpose_into!(dest::AbstractArray{T, 3}, src::AbstractArray{T, 3}) where {T} = + (permutedims!(dest, src, (3, 2, 1)); nothing) +_transpose_into!(dest::AbstractArray{T, 2}, src::AbstractArray{T, 2}) where {T} = + (permutedims!(dest, src, (2, 1)); nothing) +# Host arrays: explicit loops (`permutedims!` allocates a few bytes on +# Julia ≤ 1.11, tripping the zero-allocation tests). function _transpose_into!(dest::Array{T, 3}, src::Array{T, 3}) where {T} nfld, nlay, ncol = size(src) @inbounds for f in 1:nfld, glay in 1:nlay, gcol in 1:ncol @@ -199,18 +183,31 @@ function _transpose_into!(dest::Array{T, 2}, src::Array{T, 2}) where {T} return nothing end -# Kernel-facing accessors for the per-column layer data (`(nlay, 4)`: dry-air -# column amount, pressure, temperature, relative humidity — indexed -# `[glay, fld]`) and level temperatures (`(nlev,)`). With a cache, the -# returned views read the coalesced column-first copies; without one, they -# lazily present the `AtmosphericState` arrays in the same orientation (with -# the original memory access pattern, and — on Julia ≤ 1.11 — a small wrapper -# allocation per call on the CPU; the opt-out exists for memory-constrained -# configurations). -@inline layerdata_view(cache::TransposedStateCache, as, gcol) = - @inbounds view(cache.layerdata, gcol, :, :) -@inline layerdata_view(::Nothing, as::AtmosphericState, gcol) = - @inbounds view(_lazy_permute(as.layerdata), gcol, :, :) +# Kernel-facing per-column accessors for the hot state fields (each a `(nlay,)` +# view of dry-air column amount / layer pressure / layer temperature / relative +# humidity, plus the `(nlev,)` level temperatures). With a cache, they read the +# coalesced column-first copies (`cache.layerdata[gcol, :, fld]`, stride-1 over +# columns); without one, they read the `AtmosphericState` arrays directly +# (uncoalesced — the opt-out for memory-constrained configurations). Both cases +# are plain `view`s, so they compile on the GPU (unlike an in-kernel +# `PermutedDimsArray`). The cache stores the fields in the `AtmosphericState` +# order: dry-air column amount 1, pressure 2, temperature 3, relative humidity 4. +@inline col_dry_view(cache::TransposedStateCache, as, gcol) = + @inbounds view(cache.layerdata, gcol, :, 1) +@inline col_dry_view(::Nothing, as::AtmosphericState, gcol) = + @inbounds view(as.layerdata, 1, :, gcol) +@inline p_lay_view(cache::TransposedStateCache, as, gcol) = + @inbounds view(cache.layerdata, gcol, :, 2) +@inline p_lay_view(::Nothing, as::AtmosphericState, gcol) = + @inbounds view(as.layerdata, 2, :, gcol) +@inline t_lay_view(cache::TransposedStateCache, as, gcol) = + @inbounds view(cache.layerdata, gcol, :, 3) +@inline t_lay_view(::Nothing, as::AtmosphericState, gcol) = + @inbounds view(as.layerdata, 3, :, gcol) +@inline rel_hum_view(cache::TransposedStateCache, as, gcol) = + @inbounds view(cache.layerdata, gcol, :, 4) +@inline rel_hum_view(::Nothing, as::AtmosphericState, gcol) = + @inbounds view(as.layerdata, 4, :, gcol) @inline t_lev_view(cache::TransposedStateCache, as, gcol) = @inbounds view(cache.t_lev, gcol, :) @inline t_lev_view(::Nothing, as::AtmosphericState, gcol) = diff --git a/src/optics/compute_optical_props.jl b/src/optics/compute_optical_props.jl index 6cded6372..3b879f43e 100644 --- a/src/optics/compute_optical_props.jl +++ b/src/optics/compute_optical_props.jl @@ -34,7 +34,9 @@ GPU access) when one is provided, and from `as` directly otherwise. t_sfc = as.t_sfc[gcol] ibnd = lkp.band_data.major_gpt2bnd[igpt] totplnk = view(lkp.planck.tot_planck, :, ibnd) - as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) + col_dry_col = AtmosphericStates.col_dry_view(state_cache, as, gcol) + p_lay_col = AtmosphericStates.p_lay_view(state_cache, as, gcol) + t_lay_col = AtmosphericStates.t_lay_view(state_cache, as, gcol) t_lev_col = AtmosphericStates.t_lev_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) @@ -43,9 +45,8 @@ GPU access) when one is provided, and from `as` directly otherwise. t_lev_dec = t_lev_col[1] for glay in 1:nlay - col_dry, p_lay, t_lay = as_layerdata[glay, 1], - as_layerdata[glay, 2], - as_layerdata[glay, 3] + col_dry, p_lay, t_lay = + col_dry_col[glay], p_lay_col[glay], t_lay_col[glay] # compute gas optics τ[glay], _, _, planckfrac = compute_gas_optics( lkp, @@ -106,7 +107,7 @@ GPU access) when one is provided, and from `as` directly otherwise. aero_mask = view(as.aerosol_state.aero_mask, :, gcol) aero_size = view(as.aerosol_state.aero_size, :, :, gcol) aero_mass = view(as.aerosol_state.aero_mass, :, :, gcol) - rel_hum = view(as_layerdata, :, 4) + rel_hum = AtmosphericStates.rel_hum_view(state_cache, as, gcol) add_aerosol_optics_1scalar!( τ, @@ -144,7 +145,9 @@ end t_sfc = as.t_sfc[gcol] ibnd = lkp.band_data.major_gpt2bnd[igpt] totplnk = view(lkp.planck.tot_planck, :, ibnd) - as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) + col_dry_col = AtmosphericStates.col_dry_view(state_cache, as, gcol) + p_lay_col = AtmosphericStates.p_lay_view(state_cache, as, gcol) + t_lay_col = AtmosphericStates.t_lay_view(state_cache, as, gcol) t_lev_col = AtmosphericStates.t_lev_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) ssa = view(op.ssa, gcol, :) @@ -157,9 +160,8 @@ end @inbounds begin t_lev_dec = t_lev_col[1] for glay in 1:nlay - col_dry, p_lay, t_lay = as_layerdata[glay, 1], - as_layerdata[glay, 2], - as_layerdata[glay, 3] + col_dry, p_lay, t_lay = + col_dry_col[glay], p_lay_col[glay], t_lay_col[glay] # compute gas optics τ[glay], ssa[glay], g[glay], planckfrac = compute_gas_optics( lkp, @@ -222,7 +224,7 @@ end aero_mask = view(as.aerosol_state.aero_mask, :, gcol) aero_size = view(as.aerosol_state.aero_size, :, :, gcol) aero_mass = view(as.aerosol_state.aero_mass, :, :, gcol) - rel_hum = view(as_layerdata, :, 4) + rel_hum = AtmosphericStates.rel_hum_view(state_cache, as, gcol) add_aerosol_optics_2stream!( τ, @@ -272,12 +274,14 @@ is provided, and from `as` directly otherwise. @inbounds begin ibnd = lkp.band_data.major_gpt2bnd[igpt] t_sfc = as.t_sfc[gcol] - as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) + col_dry_col = AtmosphericStates.col_dry_view(state_cache, as, gcol) + p_lay_col = AtmosphericStates.p_lay_view(state_cache, as, gcol) + t_lay_col = AtmosphericStates.t_lay_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) end @inbounds for glay in 1:nlay col_dry, p_lay, t_lay = - as_layerdata[glay, 1], as_layerdata[glay, 2], as_layerdata[glay, 3] + col_dry_col[glay], p_lay_col[glay], t_lay_col[glay] # compute gas optics τ[glay], _, _ = compute_gas_optics( lkp, @@ -309,14 +313,16 @@ end @inbounds begin ibnd = lkp.band_data.major_gpt2bnd[igpt] t_sfc = as.t_sfc[gcol] - as_layerdata = AtmosphericStates.layerdata_view(state_cache, as, gcol) + col_dry_col = AtmosphericStates.col_dry_view(state_cache, as, gcol) + p_lay_col = AtmosphericStates.p_lay_view(state_cache, as, gcol) + t_lay_col = AtmosphericStates.t_lay_view(state_cache, as, gcol) τ = view(op.τ, gcol, :) ssa = view(op.ssa, gcol, :) g = view(op.g, gcol, :) end @inbounds for glay in 1:nlay col_dry, p_lay, t_lay = - as_layerdata[glay, 1], as_layerdata[glay, 2], as_layerdata[glay, 3] + col_dry_col[glay], p_lay_col[glay], t_lay_col[glay] # compute gas optics τ[glay], ssa[glay], g[glay] = compute_gas_optics( lkp, @@ -360,7 +366,7 @@ end aero_mask = view(as.aerosol_state.aero_mask, :, gcol) aero_size = view(as.aerosol_state.aero_size, :, :, gcol) aero_mass = view(as.aerosol_state.aero_mass, :, :, gcol) - rel_hum = view(as_layerdata, :, 4) + rel_hum = AtmosphericStates.rel_hum_view(state_cache, as, gcol) add_aerosol_optics_2stream!( τ, diff --git a/test/standalone.jl b/test/standalone.jl index 7a39a6bc7..87dc6b2cb 100644 --- a/test/standalone.jl +++ b/test/standalone.jl @@ -265,6 +265,12 @@ end # overloads must preserve that topology across a device transfer (adapt the buffer once, # rebuild the views into it) so a checkpoint restore stays zero-allocation. @testset "Adapt preserves optics/sources view topology" begin + # Helper to verify that every scratch view aliases its single packed parent buffer. + # On GPU (CuArray), contiguous 3D slices unwrap into 2D CuArrays with a distinct + # parent identity, so we verify memory sharing via pointer bounds checking. + _is_view_into(v, buf) = v isa SubArray ? parent(v) === buf : + (UInt(pointer(buf)) <= UInt(pointer(v)) < UInt(pointer(buf)) + UInt(sizeof(buf))) + for FT in (Float32, Float64) nlay, ncol = 40, 4 solver = RRTMGP.solve_gray(FT; nlay, ncol).solver # two-stream lw + sw @@ -275,15 +281,15 @@ end # the adapted buffers are genuinely fresh copies, not the originals @test solver2.lws.op.layerdata !== solver.lws.op.layerdata # every scratch view aliases its own solver's single adapted buffer - @test parent(solver2.lws.op.τ) === solver2.lws.op.layerdata - @test parent(solver2.lws.op.ssa) === solver2.lws.op.layerdata - @test parent(solver2.lws.op.g) === solver2.lws.op.layerdata - @test parent(solver2.sws.op.τ) === solver2.sws.op.layerdata - @test parent(solver2.lws.src.lev_source) === solver2.lws.src.leveldata - @test parent(solver2.lws.src.albedo) === solver2.lws.src.leveldata - @test parent(solver2.lws.src.src) === solver2.lws.src.leveldata - @test parent(solver2.sws.src.albedo) === solver2.sws.src.leveldata - @test parent(solver2.sws.src.src) === solver2.sws.src.leveldata + @test _is_view_into(solver2.lws.op.τ, solver2.lws.op.layerdata) + @test _is_view_into(solver2.lws.op.ssa, solver2.lws.op.layerdata) + @test _is_view_into(solver2.lws.op.g, solver2.lws.op.layerdata) + @test _is_view_into(solver2.sws.op.τ, solver2.sws.op.layerdata) + @test _is_view_into(solver2.lws.src.lev_source, solver2.lws.src.leveldata) + @test _is_view_into(solver2.lws.src.albedo, solver2.lws.src.leveldata) + @test _is_view_into(solver2.lws.src.src, solver2.lws.src.leveldata) + @test _is_view_into(solver2.sws.src.albedo, solver2.sws.src.leveldata) + @test _is_view_into(solver2.sws.src.src, solver2.sws.src.leveldata) # functional: the adapted solver still solves and reproduces the fluxes RRTMGP.update_fluxes!(solver2) @@ -298,7 +304,7 @@ end ncol = 3, ) op1 = Adapt.adapt(CopyToArray(), RRTMGP.Optics.OneScalar(gp)) - @test parent(op1.τ) === op1.layerdata + @test _is_view_into(op1.τ, op1.layerdata) end # Opt-in input validation: `RRTMGP.check_values[] = true` makes `update_fluxes!` From 7aba4223c692ec04370ed3adcef70c6c69d7e3ea Mon Sep 17 00:00:00 2001 From: Tapio Schneider Date: Mon, 13 Jul 2026 22:33:55 +0000 Subject: [PATCH 4/6] Swap datalayouts for CPU --- docs/src/howto/gpu.md | 12 +++--- src/RRTMGP.jl | 2 +- src/api/solver.jl | 2 +- src/optics/Fluxes.jl | 82 ++++++++++++++++++++++++++++++++----- src/optics/Sources.jl | 9 ++-- src/optics/optical_props.jl | 6 ++- src/rte/RTE.jl | 16 ++++++-- 7 files changed, 100 insertions(+), 29 deletions(-) diff --git a/docs/src/howto/gpu.md b/docs/src/howto/gpu.md index be53ac49a..bdc2cf133 100644 --- a/docs/src/howto/gpu.md +++ b/docs/src/howto/gpu.md @@ -118,9 +118,9 @@ At `Float32` precision, memory traffic across lookup-table interpolations and so | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup | |---|---|---|---| -| Clear-sky (two-stream) | 77.95 s | 0.95 s | 81.7× | -| All-sky, McICA clouds | 98.42 s | 1.40 s | 70.4× | -| All-sky with aerosols | 99.39 s | 1.54 s | 64.5× | +| Clear-sky (two-stream) | 80.05 s | 0.95 s | 84.3× | +| All-sky, McICA clouds | 99.44 s | 1.40 s | 71.0× | +| All-sky with aerosols | 106.77 s | 1.56 s | 68.4× | ### Double precision (`Float64`) @@ -128,6 +128,6 @@ For reference, running the identical uniform benchmark sweep in double precision | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup | |---|---|---|---| -| Clear-sky (two-stream) | 79.64 s | 1.30 s | 61.1× | -| All-sky, McICA clouds | 101.60 s | 1.91 s | 53.3× | -| All-sky with aerosols | 106.33 s | 2.06 s | 51.6× | +| Clear-sky (two-stream) | 81.88 s | 1.30 s | 63.0× | +| All-sky, McICA clouds | 103.97 s | 1.92 s | 54.2× | +| All-sky with aerosols | 108.38 s | 2.07 s | 52.4× | diff --git a/src/RRTMGP.jl b/src/RRTMGP.jl index 539d4784a..333a577bf 100644 --- a/src/RRTMGP.jl +++ b/src/RRTMGP.jl @@ -26,12 +26,12 @@ include(joinpath("api", "aerosols.jl")) # ── Layer 0: core containers — states, sources, optics, fluxes, BCs, lookups include(joinpath("optics", "VolumeMixingRatios.jl")) const Vmrs = VolumeMixingRatios # compat alias (deprecated name; remove in 0.24) +include(joinpath("optics", "Fluxes.jl")) include(joinpath("optics", "LookUpTables.jl")) include(joinpath("optics", "AngularDiscretizations.jl")) include(joinpath("optics", "AtmosphericStates.jl")) include(joinpath("optics", "Sources.jl")) include(joinpath("optics", "Optics.jl")) -include(joinpath("optics", "Fluxes.jl")) include(joinpath("optics", "GrayAtmosphere.jl")) const GrayUtils = GrayAtmosphere # compat alias (deprecated name; remove in 0.24) include(joinpath("optics", "BCs.jl")) diff --git a/src/api/solver.jl b/src/api/solver.jl index 676ec734c..5308a1934 100644 --- a/src/api/solver.jl +++ b/src/api/solver.jl @@ -235,7 +235,7 @@ function RRTMGPSolver( # state directly and needs none. state_cache = radiation_method isa GrayRadiation ? nothing : - TransposedStateCache(grid_params) + RTE._default_state_cache(grid_params) sws = op_sw isa OneScalar ? RTE.NoScatSWRTE(context, op_sw, bcs_sw, fluxb_sw, flux_sw, state_cache) : diff --git a/src/optics/Fluxes.jl b/src/optics/Fluxes.jl index 4a66855e6..4354f00cf 100644 --- a/src/optics/Fluxes.jl +++ b/src/optics/Fluxes.jl @@ -33,6 +33,50 @@ lazy_transpose(x::AbstractArray{T, 2}) where {T} = lazy_transpose(x::AbstractArray{T, 3}) where {T} = PermutedDimsArray{T, 3, (2, 1, 3), (2, 1, 3), typeof(x)}(x) +# Unwrap a PermutedDimsArray whose permutation is the identity swap — avoids +# double-wrapping when the CPU uses transposed physical layout. +lazy_transpose(x::PermutedDimsArray{T, 2, (2, 1), (2, 1)}) where {T} = parent(x) +lazy_transpose(x::PermutedDimsArray{T, 3, (2, 1, 3), (2, 1, 3)}) where {T} = parent(x) + +""" + _coalesced_2d(DA, FT, ncol, nlev_or_nlay) + +Allocate a 2D buffer that kernels index as `(ncol, nlev)`. On GPU, the physical +layout is `(ncol, nlev)` (coalesced across columns). On CPU, the physical layout +is `(nlev, ncol)` (stride-1 vertical sweeps), wrapped in a `PermutedDimsArray` +so the kernel indexing convention is unchanged. +""" +@inline _coalesced_2d(DA, ::Type{FT}, d1, d2) where {FT} = + DA{FT}(undef, d1, d2) # GPU default: (ncol, nlev) +@inline _coalesced_2d(::Type{Array}, ::Type{FT}, d1, d2) where {FT} = + PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix{FT}}( + Matrix{FT}(undef, d2, d1), + ) + +""" + _coalesced_3d(DA, FT, ncol, nlev_or_nlay, nfld) + +3D variant of [`_coalesced_2d`](@ref): on GPU `(ncol, nlev, nfld)`, on CPU the +physical array is `(nlev, ncol, nfld)` wrapped so kernels see `(ncol, nlev, nfld)`. +""" +@inline _coalesced_3d(DA, ::Type{FT}, d1, d2, d3) where {FT} = + DA{FT}(undef, d1, d2, d3) +@inline _coalesced_3d(::Type{Array}, ::Type{FT}, d1, d2, d3) where {FT} = + PermutedDimsArray{FT, 3, (2, 1, 3), (2, 1, 3), Array{FT, 3}}( + Array{FT, 3}(undef, d2, d1, d3), + ) + +# Zeroed variant for optical properties (mirrors `zeros`). +@inline function _coalesced_3d_zeros(DA, ::Type{FT}, d1, d2, d3) where {FT} + buf = DA{FT}(zeros(d1, d2, d3)) + return buf +end +@inline function _coalesced_3d_zeros(::Type{Array}, ::Type{FT}, d1, d2, d3) where {FT} + phys = zeros(FT, d2, d1, d3) + return PermutedDimsArray{FT, 3, (2, 1, 3), (2, 1, 3), Array{FT, 3}}(phys) +end + + """ AbstractFlux @@ -66,9 +110,9 @@ function FluxLW(grid_params::RRTMGPGridParams) (; ncol, nlay) = grid_params DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) - flux_up = DA{FT}(undef, ncol, nlay + 1) - flux_dn = DA{FT}(undef, ncol, nlay + 1) - flux_net = DA{FT}(undef, ncol, nlay + 1) + flux_up = _coalesced_2d(DA, FT, ncol, nlay + 1) + flux_dn = _coalesced_2d(DA, FT, ncol, nlay + 1) + flux_net = _coalesced_2d(DA, FT, ncol, nlay + 1) return FluxLW{FT, typeof(flux_net)}(flux_up, flux_dn, flux_net) end @@ -103,10 +147,10 @@ function FluxSW(grid_params::RRTMGPGridParams) (; nlay, ncol) = grid_params FT = eltype(grid_params) DA = ClimaComms.array_type(grid_params) - flux_up = DA{FT}(undef, ncol, nlay + 1) - flux_dn = DA{FT}(undef, ncol, nlay + 1) - flux_net = DA{FT}(undef, ncol, nlay + 1) - flux_dn_dir = DA{FT}(undef, ncol, nlay + 1) + flux_up = _coalesced_2d(DA, FT, ncol, nlay + 1) + flux_dn = _coalesced_2d(DA, FT, ncol, nlay + 1) + flux_net = _coalesced_2d(DA, FT, ncol, nlay + 1) + flux_dn_dir = _coalesced_2d(DA, FT, ncol, nlay + 1) return FluxSW{FT, typeof(flux_net)}(flux_up, flux_dn, flux_net, flux_dn_dir) end @@ -279,6 +323,16 @@ function _scale_by_transpose!(a::Array{FT, 2}, sc::Array{FT, 2}) where {FT} end return nothing end +# CPU dual-layout: `a` is a PermutedDimsArray{...(2,1)...} wrapping (nlev, ncol); +# `sc` is the plain (nlev, ncol) metric_scaling. Access the parent directly. +function _scale_by_transpose!(a::PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix{FT}}, sc::Array{FT, 2}) where {FT} + pa = parent(a) + nlev, ncol = size(pa) + @inbounds for gcol in 1:ncol, ilev in 1:nlev + pa[ilev, gcol] *= sc[ilev, gcol] + end + return nothing +end """ FluxPresentation{FTA2D, FD} @@ -326,14 +380,20 @@ Fill the `(nlev, ncol)` presentation arrays from the `(ncol, nlev)` compute buffers of `flux` (transposing copies; the direct beam only when present). """ function update_presentation!(pres::FluxPresentation, flux::AbstractFlux) - transpose_into!(pres.flux_up, flux.flux_up) - transpose_into!(pres.flux_dn, flux.flux_dn) - transpose_into!(pres.flux_net, flux.flux_net) + _copy_transposed!(pres.flux_up, flux.flux_up) + _copy_transposed!(pres.flux_dn, flux.flux_dn) + _copy_transposed!(pres.flux_net, flux.flux_net) isnothing(pres.flux_dn_dir) || - transpose_into!(pres.flux_dn_dir, flux.flux_dn_dir) + _copy_transposed!(pres.flux_dn_dir, flux.flux_dn_dir) return nothing end +# GPU / generic path: transpose from (ncol, nlev) to (nlev, ncol). +_copy_transposed!(dest, src) = transpose_into!(dest, src) +# CPU dual-layout path: the PermutedDimsArray's parent is already (nlev, ncol). +_copy_transposed!(dest::Array{FT, 2}, src::PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix{FT}}) where {FT} = + (copyto!(dest, parent(src)); nothing) + """ transpose_into!(dest, src) diff --git a/src/optics/Sources.jl b/src/optics/Sources.jl index 8b969fbeb..8f4c1dcf2 100644 --- a/src/optics/Sources.jl +++ b/src/optics/Sources.jl @@ -5,6 +5,7 @@ using Adapt using ClimaComms import ..Parameters as RP import ..RRTMGPGridParams +import ..Fluxes: _coalesced_2d, _coalesced_3d export AbstractSourceLW, SourceLWNoScat, SourceLW2Str, SourceSW2Str @@ -41,8 +42,8 @@ function SourceLWNoScat(grid_params::RRTMGPGridParams; params::RP.ARP) DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) sfc_source = DA{FT, 1}(undef, ncol) - lay_source = DA{FT, 2}(undef, ncol, nlay) - lev_source = DA{FT, 2}(undef, ncol, nlay + 1) + lay_source = _coalesced_2d(DA, FT, ncol, nlay) + lev_source = _coalesced_2d(DA, FT, ncol, nlay + 1) return SourceLWNoScat{ typeof(sfc_source), @@ -106,7 +107,7 @@ function SourceLW2Str(grid_params::RRTMGPGridParams; params::RP.ARP) DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) sfc_source = DA{FT, 1}(undef, ncol) - leveldata = DA{FT, 3}(undef, ncol, nlay + 1, 3) + leveldata = _coalesced_3d(DA, FT, ncol, nlay + 1, 3) lev_source = view(leveldata, :, :, 1) albedo = view(leveldata, :, :, 2) src = view(leveldata, :, :, 3) @@ -164,7 +165,7 @@ function SourceSW2Str(grid_params::RRTMGPGridParams) DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) sfc_source = DA{FT, 1}(undef, ncol) - leveldata = DA{FT, 3}(undef, ncol, nlay + 1, 2) + leveldata = _coalesced_3d(DA, FT, ncol, nlay + 1, 2) albedo = view(leveldata, :, :, 1) src = view(leveldata, :, :, 2) diff --git a/src/optics/optical_props.jl b/src/optics/optical_props.jl index 854f5b021..a10feff97 100644 --- a/src/optics/optical_props.jl +++ b/src/optics/optical_props.jl @@ -1,3 +1,5 @@ +import ..Fluxes: _coalesced_3d, _coalesced_3d_zeros + """ OneScalar{D, V} <: AbstractOpticalProps @@ -23,7 +25,7 @@ function OneScalar(grid_params::RRTMGPGridParams) (; ncol, nlay) = grid_params DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) - layerdata = DA{FT, 3}(undef, ncol, nlay, 1) + layerdata = _coalesced_3d(DA, FT, ncol, nlay, 1) τ = view(layerdata, :, :, 1) V = typeof(τ) return OneScalar{typeof(layerdata), V}(layerdata, τ) @@ -62,7 +64,7 @@ function TwoStream(grid_params::RRTMGPGridParams) (; ncol, nlay) = grid_params DA = ClimaComms.array_type(grid_params) FT = eltype(grid_params) - layerdata = DA{FT, 3}(zeros(ncol, nlay, 3)) + layerdata = _coalesced_3d_zeros(DA, FT, ncol, nlay, 3) V = typeof(view(layerdata, :, :, 1)) return TwoStream{typeof(layerdata), V}( layerdata, diff --git a/src/rte/RTE.jl b/src/rte/RTE.jl index 92dc613c8..709135f66 100644 --- a/src/rte/RTE.jl +++ b/src/rte/RTE.jl @@ -13,6 +13,14 @@ import ..Parameters as RP export NoScatLWRTE, TwoStreamLWRTE, NoScatSWRTE, TwoStreamSWRTE +# On GPU, use the TransposedStateCache for coalesced reads across warps. +# On CPU, pass `nothing` to read the AtmosphericState directly with its +# cache-friendly stride-4 vertical layout. +_default_state_cache(gp::RRTMGPGridParams) = + _default_state_cache(ClimaComms.device(gp), gp) +_default_state_cache(::ClimaComms.AbstractCPUDevice, gp) = nothing +_default_state_cache(device, gp) = TransposedStateCache(gp) + """ NoScatLWRTE(grid_params::RRTMGPGridParams; params, sfc_emis, inc_flux) @@ -58,7 +66,7 @@ function NoScatLWRTE( params, sfc_emis, inc_flux, - state_cache = TransposedStateCache(grid_params), + state_cache = _default_state_cache(grid_params), ) (; context) = grid_params op = OneScalar(grid_params) @@ -118,7 +126,7 @@ function TwoStreamLWRTE( params, sfc_emis, inc_flux, - state_cache = TransposedStateCache(grid_params), + state_cache = _default_state_cache(grid_params), ) (; context) = grid_params op = TwoStream(grid_params) @@ -172,7 +180,7 @@ function NoScatSWRTE( sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse, - state_cache = TransposedStateCache(grid_params), + state_cache = _default_state_cache(grid_params), ) (; context) = grid_params op = OneScalar(grid_params) @@ -238,7 +246,7 @@ function TwoStreamSWRTE( sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse, - state_cache = TransposedStateCache(grid_params), + state_cache = _default_state_cache(grid_params), ) (; context) = grid_params op = TwoStream(grid_params) From f06670ec948f7acc65a82d8a581c92c187316fcf Mon Sep 17 00:00:00 2001 From: Tapio Schneider Date: Mon, 13 Jul 2026 16:58:47 -0700 Subject: [PATCH 5/6] Review polish for the dual-layout commit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- NEWS.md | 25 ++++++---- docs/src/functional_core.md | 13 ++++-- docs/src/getters.md | 10 ++-- docs/src/howto/gpu.md | 25 ++++++---- src/optics/AtmosphericStates.jl | 3 ++ src/optics/Fluxes.jl | 82 +++++++++++++++++++-------------- src/rte/RTE.jl | 15 ++++-- 7 files changed, 104 insertions(+), 69 deletions(-) diff --git a/NEWS.md b/NEWS.md index d4cc31bc0..e55f2e9b6 100644 --- a/NEWS.md +++ b/NEWS.md @@ -3,11 +3,15 @@ RRTMGP.jl Release Notes main ------ -- Internal array layouts transposed to column-first `(ncol, nlay/nlev)` for - coalesced GPU memory access: optical properties (`OneScalar`, `TwoStream`), - source functions (`SourceLWNoScat`, `SourceLW2Str`, `SourceSW2Str`), and the - broadband flux buffers (`FluxLW`, `FluxSW`). On an A100 at DYAMOND scale - this cuts longwave kernel times by 25-36% and shortwave all-sky by 14-20%. +- Internal compute buffers now use device-dependent physical layouts behind a + uniform column-first `(ncol, nlay/nlev)` indexing convention: optical + properties (`OneScalar`, `TwoStream`), source functions (`SourceLWNoScat`, + `SourceLW2Str`, `SourceSW2Str`), and the broadband flux buffers (`FluxLW`, + `FluxSW`). On the GPU the storage is column-first, giving coalesced access + (on an A100 at DYAMOND scale this cuts longwave kernel times by 25-36% and + shortwave all-sky by 14-20%); on the CPU the storage stays vertical-first + under a lazy `PermutedDimsArray` wrapper, preserving the stride-1 vertical + sweeps (and CPU performance) while running the same kernel code. The Layer-2 flux getters (`lw_flux_up`, `net_flux`, ...) still present plain `(nlev, ncol)` views: they read `Fluxes.FluxPresentation` arrays that `update_fluxes!` fills from the compute buffers with one fused transposing @@ -23,11 +27,12 @@ main `AtmosphericState` arrays (layer pressure/temperature/dry-air column amount/relative humidity, level temperatures), refreshed by one `permutedims!` per spectral solve, so the gas-optics g-point loop reads - coalesced memory on GPUs. On by default (the RTE workspace constructors - build one; the `RRTMGPSolver` shares a single cache between the longwave - and shortwave workspaces); pass `state_cache = nothing` to a workspace - constructor to opt out, or pass the same cache to several workspaces to - share the storage. + coalesced memory on GPUs. On by default on the GPU (the RTE workspace + constructors build one; the `RRTMGPSolver` shares a single cache between + the longwave and shortwave workspaces) and off (`nothing`) on the CPU, + where the state's vertical-first layout is already cache-friendly; pass + `state_cache = nothing` to opt out, or pass the same cache to several + workspaces to share the storage. - Documentation reorganized along tutorial / how-to / explanation / reference lines: two new executable (Literate.jl) tutorials — "A first radiation calculation" and "Radiative-convective equilibrium", which reproduces diff --git a/docs/src/functional_core.md b/docs/src/functional_core.md index 8c60750b6..4bc65b696 100644 --- a/docs/src/functional_core.md +++ b/docs/src/functional_core.md @@ -35,11 +35,14 @@ solve_lw!(slv_lw, state, lookup_lw) # or solve_lw!(slv_lw, state) for gray F = slv_lw.flux.flux_net # (ncol, nlev) [W/m²] ``` -Note the axis order: the raw flux buffers are stored column-first, `(ncol, -nlev)`, so that neighboring GPU threads (one per column) access consecutive -memory. Only the Layer-2 getters (`net_flux`, `lw_flux_up`, ...) present the -vertical-first `(nlev, ncol)` orientation, from presentation copies that -`update_fluxes!` refreshes at the end of each call. +Note the axis order: the raw flux buffers are indexed column-first, `(ncol, +nlev)`, on every device. On the GPU they are also stored that way, so that +neighboring threads (one per column) access consecutive memory; on the CPU +the storage is vertical-first under a lazy wrapper, which keeps the +per-column sweeps stride-1 without changing the indexing. Only the Layer-2 +getters (`net_flux`, `lw_flux_up`, ...) present the vertical-first +`(nlev, ncol)` orientation, from presentation copies that `update_fluxes!` +refreshes at the end of each call. This is what RRTMGP's own tests drive, and what [`RRTMGPSolver`](@ref RRTMGP.RRTMGPSolver) wraps behind [`update_fluxes!`](@ref RRTMGP.update_fluxes!) diff --git a/docs/src/getters.md b/docs/src/getters.md index 90dc3d37f..3c29af390 100644 --- a/docs/src/getters.md +++ b/docs/src/getters.md @@ -55,11 +55,11 @@ wraps an input view as a `Field` (writes flow into the solver's buffer), and `field2array(dst) .= getter(solver)` copies an output into a destination field. On the GPU, the returned views are plain `SubArray`s of device arrays; to bring one to the CPU, use `Array(getter(solver))` (or the element-wise `.=` -copy above). One subtlety: the solver's compute buffers are stored -column-first, `(ncol, nlev)`, for coalesced GPU access, and `update_fluxes!` -ends by transposing them into the `(nlev, ncol)` presentation arrays the flux -getters expose — so read fluxes after `update_fluxes!`, not between manual -Layer-1 solves. +copy above). One subtlety: the solver's compute buffers are indexed +column-first, `(ncol, nlev)` (with the physical layout chosen per device for +performance), and `update_fluxes!` ends by copying them into the +`(nlev, ncol)` presentation arrays the flux getters expose — so read fluxes +after `update_fluxes!`, not between manual Layer-1 solves. ## Responsibilities: what the host provides vs what RRTMGP derives diff --git a/docs/src/howto/gpu.md b/docs/src/howto/gpu.md index bdc2cf133..4d593f0c6 100644 --- a/docs/src/howto/gpu.md +++ b/docs/src/howto/gpu.md @@ -50,7 +50,11 @@ outside hot loops. `Adapt.adapt(Array, solver)` produces a host copy (e.g., for serialization) and `Adapt.adapt(CuArray, solver)` moves it back. RRTMGP's custom `adapt_structure` methods preserve the internal view topology (scratch views into packed -buffers), so an adapted solver keeps the zero-allocation property. +buffers), so an adapted solver keeps the zero-allocation property. Adapt +preserves each buffer's physical layout, so round-tripping a GPU solver +through the host is lossless; adapting a CPU-*constructed* solver to the GPU +runs correctly but keeps the CPU layout (uncoalesced) — for GPU work, +construct the solver on the GPU. ## Reproducibility caveat: McICA cloud sampling @@ -78,14 +82,17 @@ both backends: g-point's lookup-table slices hot in cache while sweeping the columns. With a single-threaded context, the same code runs serially. -The memory layout matches the thread mapping: the internal scratch and output -buffers are stored column-first, `(ncol, nlev)`, so consecutive GPU threads -write consecutive addresses (coalesced access), and the frequently re-read -state arrays (layer pressures and temperatures, level temperatures) are -copied into a column-first `TransposedStateCache` once per solve so the -g-point loop also reads coalesced memory. The caller-owned -`AtmosphericState` and the getter views keep their vertical-first -`(nlev, ncol)` presentation. +The memory layout matches the thread mapping on each device. On the GPU, the +internal scratch and output buffers are stored column-first, `(ncol, nlev)`, +so consecutive threads write consecutive addresses (coalesced access), and +the frequently re-read state arrays (layer pressures and temperatures, level +temperatures) are copied into a column-first `TransposedStateCache` once per +solve so the g-point loop also reads coalesced memory. On the CPU, the same +buffers are stored vertical-first under a lazy wrapper — each thread's +vertical sweep is then stride-1 — and the kernels read the +`AtmosphericState` directly, so both devices run identical kernel code on +their preferred layout. The caller-owned `AtmosphericState` and the getter +views keep their vertical-first `(nlev, ncol)` presentation throughout. ## Performance expectations diff --git a/src/optics/AtmosphericStates.jl b/src/optics/AtmosphericStates.jl index 6d5b70412..3ce243ff5 100644 --- a/src/optics/AtmosphericStates.jl +++ b/src/optics/AtmosphericStates.jl @@ -121,6 +121,9 @@ The cache trades one `permutedims!` per solve plus `(4 nlay + nlev) ncol` extra storage for coalesced reads in the g-point loop. Pass the same cache to the longwave and shortwave workspaces to share the storage, or `nothing` to opt out (the kernels then read the `AtmosphericState` directly, as before). +The workspace constructors default to a cache on the GPU and to `nothing` on +the CPU, where the state's vertical-first layout is already cache-friendly +for the per-column sweeps. # Fields - `layerdata`: `(ncol, nlay, 4)` copy of `AtmosphericState.layerdata` diff --git a/src/optics/Fluxes.jl b/src/optics/Fluxes.jl index 4354f00cf..a6b0f3999 100644 --- a/src/optics/Fluxes.jl +++ b/src/optics/Fluxes.jl @@ -41,14 +41,16 @@ lazy_transpose(x::PermutedDimsArray{T, 3, (2, 1, 3), (2, 1, 3)}) where {T} = par """ _coalesced_2d(DA, FT, ncol, nlev_or_nlay) -Allocate a 2D buffer that kernels index as `(ncol, nlev)`. On GPU, the physical -layout is `(ncol, nlev)` (coalesced across columns). On CPU, the physical layout -is `(nlev, ncol)` (stride-1 vertical sweeps), wrapped in a `PermutedDimsArray` -so the kernel indexing convention is unchanged. +Allocate a 2D compute buffer that the kernels index as `(ncol, nlev)`. The +physical layout is chosen per device (by dispatch on the array type `DA`): on +the GPU, the buffer is stored `(ncol, nlev)`, so the one-thread-per-column +kernels access consecutive addresses (coalesced); on the CPU (`DA === Array`), +it is stored `(nlev, ncol)` — stride-1 vertical sweeps, matching the +per-column loop order — and wrapped in a `PermutedDimsArray` so the kernel +indexing convention is the same on both devices. """ -@inline _coalesced_2d(DA, ::Type{FT}, d1, d2) where {FT} = - DA{FT}(undef, d1, d2) # GPU default: (ncol, nlev) -@inline _coalesced_2d(::Type{Array}, ::Type{FT}, d1, d2) where {FT} = +_coalesced_2d(DA, ::Type{FT}, d1, d2) where {FT} = DA{FT}(undef, d1, d2) +_coalesced_2d(::Type{Array}, ::Type{FT}, d1, d2) where {FT} = PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix{FT}}( Matrix{FT}(undef, d2, d1), ) @@ -56,26 +58,28 @@ so the kernel indexing convention is unchanged. """ _coalesced_3d(DA, FT, ncol, nlev_or_nlay, nfld) -3D variant of [`_coalesced_2d`](@ref): on GPU `(ncol, nlev, nfld)`, on CPU the -physical array is `(nlev, ncol, nfld)` wrapped so kernels see `(ncol, nlev, nfld)`. +3D variant of [`_coalesced_2d`](@ref) for field-packed buffers: indexed +`(ncol, nlev, nfld)` on both devices; stored `(nlev, ncol, nfld)` on the CPU. """ -@inline _coalesced_3d(DA, ::Type{FT}, d1, d2, d3) where {FT} = +_coalesced_3d(DA, ::Type{FT}, d1, d2, d3) where {FT} = DA{FT}(undef, d1, d2, d3) -@inline _coalesced_3d(::Type{Array}, ::Type{FT}, d1, d2, d3) where {FT} = +_coalesced_3d(::Type{Array}, ::Type{FT}, d1, d2, d3) where {FT} = PermutedDimsArray{FT, 3, (2, 1, 3), (2, 1, 3), Array{FT, 3}}( Array{FT, 3}(undef, d2, d1, d3), ) -# Zeroed variant for optical properties (mirrors `zeros`). -@inline function _coalesced_3d_zeros(DA, ::Type{FT}, d1, d2, d3) where {FT} - buf = DA{FT}(zeros(d1, d2, d3)) +# Zero-initialized variant of `_coalesced_3d` (the `TwoStream` optics start +# from zero). Fill the physical storage, not the wrapper: `fill!` on a +# `PermutedDimsArray` falls back to slow wrapped iteration. +function _coalesced_3d_zeros(DA, ::Type{FT}, d1, d2, d3) where {FT} + buf = _coalesced_3d(DA, FT, d1, d2, d3) + fill!(_physical(buf), FT(0)) return buf end -@inline function _coalesced_3d_zeros(::Type{Array}, ::Type{FT}, d1, d2, d3) where {FT} - phys = zeros(FT, d2, d1, d3) - return PermutedDimsArray{FT, 3, (2, 1, 3), (2, 1, 3), Array{FT, 3}}(phys) -end +# The physical storage underlying a (possibly layout-wrapped) compute buffer. +_physical(x::AbstractArray) = x +_physical(x::PermutedDimsArray) = parent(x) """ AbstractFlux @@ -309,23 +313,19 @@ function apply_metric_scaling!( return nothing end -# Multiply the column-first `(ncol, nlev[, n_bnd])` array `a` by the -# vertical-first `(nlev, ncol)` scaling `sc`. Device arrays broadcast through -# the lazy transpose (a single kernel); host arrays use explicit loops, since -# broadcasts with a permuted-wrapper operand allocate a few bytes on -# Julia ≤ 1.11, tripping the zero-allocation tests. +# Multiply the `(ncol, nlev)`-indexed compute buffer `a` by the vertical-first +# `(nlev, ncol)` scaling `sc`. Device arrays broadcast through the lazy +# transpose (a single kernel); host arrays use explicit loops, since broadcasts +# with a permuted-wrapper operand allocate a few bytes on Julia ≤ 1.11, +# tripping the zero-allocation tests. _scale_by_transpose!(a::AbstractArray, sc) = (a .= a .* lazy_transpose(sc); nothing) -function _scale_by_transpose!(a::Array{FT, 2}, sc::Array{FT, 2}) where {FT} - ncol, nlev = size(a) - @inbounds for ilev in 1:nlev, gcol in 1:ncol - a[gcol, ilev] *= sc[ilev, gcol] - end - return nothing -end -# CPU dual-layout: `a` is a PermutedDimsArray{...(2,1)...} wrapping (nlev, ncol); -# `sc` is the plain (nlev, ncol) metric_scaling. Access the parent directly. -function _scale_by_transpose!(a::PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix{FT}}, sc::Array{FT, 2}) where {FT} +# CPU dual-layout buffers: the wrapper's parent and `sc` share the +# (nlev, ncol) layout, so scale the parent directly, stride-1. +function _scale_by_transpose!( + a::PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix{FT}}, + sc::Array{FT, 2}, +) where {FT} pa = parent(a) nlev, ncol = size(pa) @inbounds for gcol in 1:ncol, ilev in 1:nlev @@ -333,6 +333,16 @@ function _scale_by_transpose!(a::PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix end return nothing end +# Defensive fallback for hand-constructed Layer-1 flux buffers that store a +# plain (ncol, nlev) host array (the internal CPU constructors always wrap). +function _scale_by_transpose!(a::Array{FT, 2}, sc::Array{FT, 2}) where {FT} + ncol, nlev = size(a) + @inbounds for ilev in 1:nlev, gcol in 1:ncol + a[gcol, ilev] *= sc[ilev, gcol] + end + return nothing +end + """ FluxPresentation{FTA2D, FD} @@ -376,8 +386,10 @@ end """ update_presentation!(pres::FluxPresentation, flux::AbstractFlux) -Fill the `(nlev, ncol)` presentation arrays from the `(ncol, nlev)` compute -buffers of `flux` (transposing copies; the direct beam only when present). +Fill the `(nlev, ncol)` presentation arrays from the `(ncol, nlev)`-indexed +compute buffers of `flux` (the direct beam only when present). On the GPU this +is a transposing copy; on the CPU, the compute buffers' physical parents +already have the presentation layout, so it is a plain `copyto!`. """ function update_presentation!(pres::FluxPresentation, flux::AbstractFlux) _copy_transposed!(pres.flux_up, flux.flux_up) diff --git a/src/rte/RTE.jl b/src/rte/RTE.jl index 709135f66..085f96197 100644 --- a/src/rte/RTE.jl +++ b/src/rte/RTE.jl @@ -13,13 +13,18 @@ import ..Parameters as RP export NoScatLWRTE, TwoStreamLWRTE, NoScatSWRTE, TwoStreamSWRTE -# On GPU, use the TransposedStateCache for coalesced reads across warps. -# On CPU, pass `nothing` to read the AtmosphericState directly with its -# cache-friendly stride-4 vertical layout. +# Device-dependent default for the workspaces' `state_cache` keyword. On the +# GPU, a `TransposedStateCache` gives the g-point loop coalesced state reads +# across the warp's columns. On the CPU, the kernels read the +# `AtmosphericState` directly (`nothing`): its vertical-first layout is +# already cache-friendly for the per-column sweeps, and a transposed copy +# would only add work. _default_state_cache(gp::RRTMGPGridParams) = _default_state_cache(ClimaComms.device(gp), gp) -_default_state_cache(::ClimaComms.AbstractCPUDevice, gp) = nothing -_default_state_cache(device, gp) = TransposedStateCache(gp) +_default_state_cache(::ClimaComms.AbstractCPUDevice, gp::RRTMGPGridParams) = + nothing +_default_state_cache(::ClimaComms.AbstractDevice, gp::RRTMGPGridParams) = + TransposedStateCache(gp) """ NoScatLWRTE(grid_params::RRTMGPGridParams; params, sfc_emis, inc_flux) From f77a5427a6faed40b58aa98ab1a1041352f22091 Mon Sep 17 00:00:00 2001 From: Tapio Schneider Date: Mon, 13 Jul 2026 19:10:09 -0700 Subject: [PATCH 6/6] Docs: fix missing_docs docbuild error on the dual-layout helpers `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 --- src/optics/AtmosphericStates.jl | 11 +++------ src/optics/Fluxes.jl | 44 +++++++++++++-------------------- 2 files changed, 21 insertions(+), 34 deletions(-) diff --git a/src/optics/AtmosphericStates.jl b/src/optics/AtmosphericStates.jl index 3ce243ff5..1637821f7 100644 --- a/src/optics/AtmosphericStates.jl +++ b/src/optics/AtmosphericStates.jl @@ -147,13 +147,10 @@ function TransposedStateCache(grid_params::RRTMGPGridParams) ) end -""" - refresh_transposed_state!(cache, as::AtmosphericState) - -Copy the current `AtmosphericState` layer data and level temperatures into -the column-first [`TransposedStateCache`](@ref) `cache`. Called at the start -of every spectral `solve_lw!`/`solve_sw!`; a no-op when `cache === nothing`. -""" +# refresh_transposed_state!(cache, as): copy the current `AtmosphericState` +# layer data and level temperatures into the column-first `TransposedStateCache` +# `cache`. Called at the start of every spectral `solve_lw!`/`solve_sw!`; a +# no-op when `cache === nothing`. @inline refresh_transposed_state!(::Nothing, as::AtmosphericState) = nothing function refresh_transposed_state!( cache::TransposedStateCache, diff --git a/src/optics/Fluxes.jl b/src/optics/Fluxes.jl index a6b0f3999..f168ba4fd 100644 --- a/src/optics/Fluxes.jl +++ b/src/optics/Fluxes.jl @@ -18,16 +18,12 @@ export AbstractFlux, apply_metric_scaling!, compute_net_flux! -""" - lazy_transpose(x) - -Lazily present the `(ncol, nlev[, n_bnd])` internal buffer `x` as -`(nlev, ncol[, n_bnd])` without copying. Equivalent to -`PermutedDimsArray(x, perm)`, but with the permutation in the type parameters: -before Julia 1.12, the `PermutedDimsArray(x, perm)` constructor does not -constant-fold `perm`, which would make every flux getter type-unstable (and -allocating) on Julia 1.9–1.11. -""" +# lazy_transpose(x): lazily present the `(ncol, nlev[, n_bnd])` internal buffer +# `x` as `(nlev, ncol[, n_bnd])` without copying. Equivalent to +# `PermutedDimsArray(x, perm)`, but with the permutation in the type parameters: +# before Julia 1.12, the `PermutedDimsArray(x, perm)` constructor does not +# constant-fold `perm`, which would make every flux getter type-unstable (and +# allocating) on Julia 1.9–1.11. lazy_transpose(x::AbstractArray{T, 2}) where {T} = PermutedDimsArray{T, 2, (2, 1), (2, 1), typeof(x)}(x) lazy_transpose(x::AbstractArray{T, 3}) where {T} = @@ -38,29 +34,23 @@ lazy_transpose(x::AbstractArray{T, 3}) where {T} = lazy_transpose(x::PermutedDimsArray{T, 2, (2, 1), (2, 1)}) where {T} = parent(x) lazy_transpose(x::PermutedDimsArray{T, 3, (2, 1, 3), (2, 1, 3)}) where {T} = parent(x) -""" - _coalesced_2d(DA, FT, ncol, nlev_or_nlay) - -Allocate a 2D compute buffer that the kernels index as `(ncol, nlev)`. The -physical layout is chosen per device (by dispatch on the array type `DA`): on -the GPU, the buffer is stored `(ncol, nlev)`, so the one-thread-per-column -kernels access consecutive addresses (coalesced); on the CPU (`DA === Array`), -it is stored `(nlev, ncol)` — stride-1 vertical sweeps, matching the -per-column loop order — and wrapped in a `PermutedDimsArray` so the kernel -indexing convention is the same on both devices. -""" +# _coalesced_2d(DA, FT, ncol, nlev_or_nlay): allocate a 2D compute buffer that +# the kernels index as `(ncol, nlev)`. The physical layout is chosen per device +# (by dispatch on the array type `DA`): on the GPU, the buffer is stored +# `(ncol, nlev)`, so the one-thread-per-column kernels access consecutive +# addresses (coalesced); on the CPU (`DA === Array`), it is stored +# `(nlev, ncol)` — stride-1 vertical sweeps, matching the per-column loop order +# — and wrapped in a `PermutedDimsArray` so the kernel indexing convention is +# the same on both devices. _coalesced_2d(DA, ::Type{FT}, d1, d2) where {FT} = DA{FT}(undef, d1, d2) _coalesced_2d(::Type{Array}, ::Type{FT}, d1, d2) where {FT} = PermutedDimsArray{FT, 2, (2, 1), (2, 1), Matrix{FT}}( Matrix{FT}(undef, d2, d1), ) -""" - _coalesced_3d(DA, FT, ncol, nlev_or_nlay, nfld) - -3D variant of [`_coalesced_2d`](@ref) for field-packed buffers: indexed -`(ncol, nlev, nfld)` on both devices; stored `(nlev, ncol, nfld)` on the CPU. -""" +# _coalesced_3d(DA, FT, ncol, nlev_or_nlay, nfld): 3D variant of `_coalesced_2d` +# for field-packed buffers: indexed `(ncol, nlev, nfld)` on both devices; stored +# `(nlev, ncol, nfld)` on the CPU. _coalesced_3d(DA, ::Type{FT}, d1, d2, d3) where {FT} = DA{FT}(undef, d1, d2, d3) _coalesced_3d(::Type{Array}, ::Type{FT}, d1, d2, d3) where {FT} =