diff --git a/NEWS.md b/NEWS.md index daf98159c..e55f2e9b6 100644 --- a/NEWS.md +++ b/NEWS.md @@ -3,6 +3,36 @@ RRTMGP.jl Release Notes main ------ +- 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 + 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 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 948ce3c2c..4bc65b696 100644 --- a/docs/src/functional_core.md +++ b/docs/src/functional_core.md @@ -32,9 +32,18 @@ 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 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!) and the named getters. @@ -81,10 +90,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..3c29af390 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 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 188fa6e32..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,6 +82,18 @@ 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 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 CI benchmarks the kernels on GPUs (see `perf/benchmark_ratchet.jl` @@ -105,13 +121,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 large GPU speedups: | 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) | 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`) @@ -119,6 +135,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) | 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/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/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/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/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/getters.jl b/src/api/getters.jl index f059e9b2d..8304ddef3 100644 --- a/src/api/getters.jl +++ b/src/api/getters.jl @@ -57,17 +57,17 @@ _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) +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) = _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) +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) diff --git a/src/api/solver.jl b/src/api/solver.jl index 3abe81d76..5308a1934 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 : + RTE._default_state_cache(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 d216fe819..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, ) - 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)) + # 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, ) - 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)) + # 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..1637821f7 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, @@ -87,32 +89,129 @@ 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) = + +""" + TransposedStateCache{D2, D3} + +Column-first copies of the hot [`AtmosphericState`](@ref) fields, refreshed +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 +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). +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` + (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): 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, + as::AtmosphericState, +) + _transpose_into!(cache.layerdata, as.layerdata) + _transpose_into!(cache.t_lev, as.t_lev) + return nothing +end + +# 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 + 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 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) = + @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 ea6e738bd..f168ba4fd 100644 --- a/src/optics/Fluxes.jl +++ b/src/optics/Fluxes.jl @@ -8,12 +8,69 @@ 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) + +# 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 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` +# 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} = + PermutedDimsArray{FT, 3, (2, 1, 3), (2, 1, 3), Array{FT, 3}}( + Array{FT, 3}(undef, d2, d1, 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 + +# The physical storage underlying a (possibly layout-wrapped) compute buffer. +_physical(x::AbstractArray) = x +_physical(x::PermutedDimsArray) = parent(x) + """ AbstractFlux @@ -29,9 +86,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} @@ -47,9 +104,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 = _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 @@ -59,10 +116,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 +141,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 = _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 @@ -96,8 +153,14 @@ end 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 -fluxes, and summing over the band dimension recovers the broadband fluxes. +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²], `(nlev, ncol, n_bnd)`. @@ -131,7 +194,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 +208,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[ilev, gcol, ibnd] += flux_up[gcol, ilev] + bd[ilev, gcol, ibnd] += flux_dn[gcol, ilev] end return nothing end @@ -163,14 +226,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 +248,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 +271,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 +296,141 @@ 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) + _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 `(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) +# 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 + pa[ilev, gcol] *= sc[ilev, gcol] + 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} + +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)`-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) + _copy_transposed!(pres.flux_dn, flux.flux_dn) + _copy_transposed!(pres.flux_net, flux.flux_net) + isnothing(pres.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) + +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!( @@ -252,6 +446,8 @@ 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) + # 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 diff --git a/src/optics/Sources.jl b/src/optics/Sources.jl index 6090af81f..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 @@ -24,8 +25,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 @@ -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, nlay, ncol) - lev_source = DA{FT, 2}(undef, nlay + 1, ncol) + lay_source = _coalesced_2d(DA, FT, ncol, nlay) + lev_source = _coalesced_2d(DA, FT, ncol, nlay + 1) return SourceLWNoScat{ typeof(sfc_source), @@ -66,10 +67,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 +84,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 +107,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 = _coalesced_3d(DA, FT, ncol, nlay + 1, 3) + lev_source = view(leveldata, :, :, 1) + albedo = view(leveldata, :, :, 2) + src = view(leveldata, :, :, 3) return SourceLW2Str{ typeof(sfc_source), @@ -135,9 +136,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 +150,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 +165,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 = _coalesced_3d(DA, FT, 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..3b879f43e 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,18 +34,19 @@ 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) - τ = view(op.τ, :, 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, :) lev_src_inc_prev = zero(t_sfc) lev_src_dec_prev = zero(t_sfc) 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 = + col_dry_col[glay], p_lay_col[glay], t_lay_col[glay] # compute gas optics τ[glay], _, _, planckfrac = compute_gas_optics( lkp, @@ -57,7 +62,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 +71,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) @@ -102,7 +107,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 = AtmosphericStates.rel_hum_view(state_cache, as, gcol) add_aerosol_optics_1scalar!( τ, @@ -124,6 +129,7 @@ end @inline function compute_optical_props!( op::TwoStream, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, sf::SourceLW2Str, gcol::Int, igpt::Int, @@ -139,11 +145,13 @@ 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) - τ = view(op.τ, :, gcol) - ssa = view(op.ssa, :, gcol) - g = view(op.g, :, 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, :) + g = view(op.g, gcol, :) end lev_src_inc_prev = zero(t_sfc) @@ -152,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[1, glay], - as_layerdata[2, glay], - as_layerdata[3, glay] + 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, @@ -177,15 +184,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 @@ -217,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 = AtmosphericStates.getview_rel_hum(as, gcol) + rel_hum = AtmosphericStates.rel_hum_view(state_cache, as, gcol) add_aerosol_optics_2stream!( τ, @@ -241,6 +248,7 @@ end compute_optical_props!( op::AbstractOpticalProps, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, gcol::Int, igpt::Int, lkp::LookUpSW, @@ -248,28 +256,32 @@ 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, ::Nothing, - ::Nothing, ) nlay = AtmosphericStates.get_nlay(as) (; vmr) = as @inbounds begin ibnd = lkp.band_data.major_gpt2bnd[igpt] t_sfc = as.t_sfc[gcol] - as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) - τ = view(op.τ, :, 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[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] + col_dry_col[glay], p_lay_col[glay], t_lay_col[glay] # compute gas optics τ[glay], _, _ = compute_gas_optics( lkp, @@ -289,6 +301,7 @@ end @inline function compute_optical_props!( op::TwoStream, as::AtmosphericState, + state_cache::Union{TransposedStateCache, Nothing}, gcol::Int, igpt::Int, lkp::LookUpSW, @@ -300,14 +313,16 @@ end @inbounds begin 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) + 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[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] + col_dry_col[glay], p_lay_col[glay], t_lay_col[glay] # compute gas optics τ[glay], ssa[glay], g[glay] = compute_gas_optics( lkp, @@ -351,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 = AtmosphericStates.getview_rel_hum(as, gcol) + rel_hum = AtmosphericStates.rel_hum_view(state_cache, as, gcol) add_aerosol_optics_2stream!( τ, 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..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 @@ -15,7 +17,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 +25,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 = _coalesced_3d(DA, FT, ncol, nlay, 1) + τ = view(layerdata, :, :, 1) V = typeof(τ) return OneScalar{typeof(layerdata), V}(layerdata, τ) end @@ -39,9 +41,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 +54,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 +64,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 = _coalesced_3d_zeros(DA, FT, 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/RTE.jl b/src/rte/RTE.jl index 92aaa2066..085f96197 100644 --- a/src/rte/RTE.jl +++ b/src/rte/RTE.jl @@ -13,6 +13,19 @@ import ..Parameters as RP export NoScatLWRTE, TwoStreamLWRTE, NoScatSWRTE, TwoStreamSWRTE +# 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::RRTMGPGridParams) = + nothing +_default_state_cache(::ClimaComms.AbstractDevice, gp::RRTMGPGridParams) = + TransposedStateCache(gp) + """ NoScatLWRTE(grid_params::RRTMGPGridParams; params, sfc_emis, inc_flux) @@ -28,6 +41,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 +53,7 @@ struct NoScatLWRTE{ FXBL, FXL <: FluxLW, AD, + TSC, } context::C op::OP @@ -45,10 +62,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 = _default_state_cache(grid_params), +) (; context) = grid_params op = OneScalar(grid_params) src = SourceLWNoScat(grid_params; params) @@ -56,7 +80,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 +98,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 +110,7 @@ struct TwoStreamLWRTE{ FXBL, FXL <: FluxLW, FXBND, + TSC, } context::C op::OP @@ -91,17 +119,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 = _default_state_cache(grid_params), ) (; context) = grid_params op = TwoStream(grid_params) @@ -109,7 +139,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 +164,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 +185,7 @@ function NoScatSWRTE( sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse, + state_cache = _default_state_cache(grid_params), ) (; context) = grid_params op = OneScalar(grid_params) @@ -154,7 +198,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 +216,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 +228,7 @@ struct TwoStreamSWRTE{ FXBS, FXS <: FluxSW, FXBND, + TSC, } context::C op::OP @@ -189,11 +237,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 +251,7 @@ function TwoStreamSWRTE( sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse, + state_cache = _default_state_cache(grid_params), ) (; context) = grid_params op = TwoStream(grid_params) @@ -215,7 +265,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/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..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, @@ -257,21 +261,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 +291,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 +323,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..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, @@ -240,26 +244,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 +271,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..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, @@ -315,27 +319,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 +348,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..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, @@ -129,17 +132,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..668ea9ceb 100644 --- a/test/clear_sky_utils.jl +++ b/test/clear_sky_utils.jl @@ -146,15 +146,28 @@ 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) 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 +199,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..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}, @@ -107,47 +174,20 @@ 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) + 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 @@ -158,9 +198,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 +229,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/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 d0c5051b8..87dc6b2cb 100644 --- a/test/standalone.jl +++ b/test/standalone.jl @@ -242,7 +242,7 @@ end RRTMGP.update_fluxes!(solver) - # internal buffers are boundary-extended; getters are domain-masked + # 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 @@ -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!`