From fe506bdbe21a7b1de8549de5b442c578cb355c8d Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Thu, 22 Feb 2024 12:36:33 -0800 Subject: [PATCH] Restructure datalayout in `AtmosphericStates` to enable coalesced memory access in optics calculations --- src/optics/AtmosphericStates.jl | 30 ++++++++++++------ src/optics/Optics.jl | 54 ++++++++++++--------------------- test/read_all_sky.jl | 7 ++++- test/read_rfmip_clear_sky.jl | 7 ++++- 4 files changed, 52 insertions(+), 46 deletions(-) diff --git a/src/optics/AtmosphericStates.jl b/src/optics/AtmosphericStates.jl index 42a6df220..d9dfc3983 100644 --- a/src/optics/AtmosphericStates.jl +++ b/src/optics/AtmosphericStates.jl @@ -38,23 +38,19 @@ Atmospheric conditions, used to compute optical properties. # Fields $(DocStringExtensions.FIELDS) """ -struct AtmosphericState{FTA1D, FTA1DN, FTA2D, VMR, CLD} <: AbstractAtmosphericState +struct AtmosphericState{FTA1D, FTA1DN, FTA2D, D, VMR, CLD} <: AbstractAtmosphericState "longitude, in degrees (`ncol`), optional" lon::FTA1DN "latitude, in degrees (`ncol`), optional" lat::FTA1DN - "Layer pressures `[Pa, mb]`; `(nlay,ncol)`" - p_lay::FTA2D + "storage for col_dry [`molecules per cm^2 of dry air`], play `[Pa, mb]`, tlay `[K]`; `(3, nlay, ncol)`" + layerdata::D "Level pressures `[Pa, mb]`; `(nlay+1,ncol)`" p_lev::FTA2D - "Layer temperatures `[K]`; `(nlay,ncol)`" - t_lay::FTA2D "Level temperatures `[K]`; `(nlay+1,ncol)`" t_lev::FTA2D "Surface temperatures `[K]`; `(ncol)`" t_sfc::FTA1D - "Number of molecules per cm^2 of dry air `(nlay, ncol)`" - col_dry::FTA2D "volume mixing ratio of all relevant gases" vmr::VMR "cloud state" @@ -63,11 +59,25 @@ end Adapt.@adapt_structure AtmosphericState # Number of layers -@inline get_nlay(as::AtmosphericState) = size(as.p_lay, 1) +@inline get_nlay(as::AtmosphericState) = size(as.layerdata, 2) # Number of columns -@inline get_ncol(as::AtmosphericState) = size(as.p_lay, 2) +@inline get_ncol(as::AtmosphericState) = size(as.layerdata, 3) # Number of layers and columns -@inline get_dims(as::AtmosphericState) = size(as.p_lay) +@inline get_dims(as::AtmosphericState) = size(as.layerdata, 2), size(as.layerdata, 3) + +@inline getview_layerdata(as::AtmosphericState, gcol) = @inbounds view(as.layerdata, :, :, gcol) + +# 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) """ CloudState{CD, CF, CM, CMT} diff --git a/src/optics/Optics.jl b/src/optics/Optics.jl index c6e8dcfc0..61cf021ba 100644 --- a/src/optics/Optics.jl +++ b/src/optics/Optics.jl @@ -169,17 +169,13 @@ Computes 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) - col_dry_col = view(as.col_dry, :, gcol) - p_lay_col = view(as.p_lay, :, gcol) - t_lay_col = view(as.t_lay, :, gcol) + as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) t_lev_col = view(as.t_lev, :, gcol) τ = view(op.τ, :, gcol) t_lev_dec = t_lev_col[1] for glay in 1:nlay - col_dry = col_dry_col[glay] - p_lay = p_lay_col[glay] - t_lay = t_lay_col[glay] + col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] # compute gas optics τ[glay], _, _, planckfrac = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol) # compute longwave source terms @@ -214,9 +210,7 @@ end t_sfc = as.t_sfc[gcol] ibnd = lkp.band_data.major_gpt2bnd[igpt] totplnk = view(lkp.planck.tot_planck, :, ibnd) - col_dry_col = view(as.col_dry, :, gcol) - p_lay_col = view(as.p_lay, :, gcol) - t_lay_col = view(as.t_lay, :, gcol) + as_layerdata = AtmosphericStates.getview_layerdata(as, gcol) t_lev_col = view(as.t_lev, :, gcol) τ = view(op.τ, :, gcol) ssa = view(op.ssa, :, gcol) @@ -229,9 +223,7 @@ end @inbounds begin t_lev_dec = t_lev_col[1] for glay in 1:nlay - col_dry = col_dry_col[glay] - p_lay = p_lay_col[glay] - t_lay = t_lay_col[glay] + col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] # compute gas optics τ[glay], ssa[glay], g[glay], planckfrac = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol) @@ -300,17 +292,14 @@ Computes optical properties for the shortwave problem. ) nlay = AtmosphericStates.get_nlay(as) (; vmr) = as - @inbounds ibnd = lkp.band_data.major_gpt2bnd[igpt] - @inbounds t_sfc = as.t_sfc[gcol] - col_dry_col = view(as.col_dry, :, gcol) - p_lay_col = view(as.p_lay, :, gcol) - t_lay_col = view(as.t_lay, :, gcol) - τ = view(op.τ, :, gcol) - + @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) + end @inbounds for glay in 1:nlay - col_dry = col_dry_col[glay] - p_lay = p_lay_col[glay] - t_lay = t_lay_col[glay] + col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] # compute gas optics τ[glay], _, _ = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol) end @@ -327,19 +316,16 @@ end ) nlay = AtmosphericStates.get_nlay(as) (; vmr) = as - @inbounds ibnd = lkp.band_data.major_gpt2bnd[igpt] - @inbounds t_sfc = as.t_sfc[gcol] - col_dry_col = view(as.col_dry, :, gcol) - p_lay_col = view(as.p_lay, :, gcol) - t_lay_col = view(as.t_lay, :, gcol) - τ = view(op.τ, :, gcol) - ssa = view(op.ssa, :, gcol) - g = view(op.g, :, gcol) - + @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) + end @inbounds for glay in 1:nlay - col_dry = col_dry_col[glay] - p_lay = p_lay_col[glay] - t_lay = t_lay_col[glay] + col_dry, p_lay, t_lay = as_layerdata[1, glay], as_layerdata[2, glay], as_layerdata[3, glay] # compute gas optics τ[glay], ssa[glay], g[glay] = compute_gas_optics(lkp, vmr, col_dry, igpt, ibnd, p_lay, t_lay, glay, gcol) end diff --git a/test/read_all_sky.jl b/test/read_all_sky.jl index 068dd0f3a..0eda70927 100644 --- a/test/read_all_sky.jl +++ b/test/read_all_sky.jl @@ -133,6 +133,11 @@ function setup_allsky_as( compute_col_gas!(context, p_lev, col_dry, param_set, vmr_h2o, lat) # the example skips lat based gravity calculation + layerdata = similar(p_lay, 3, nlay, ncol) + layerdata[1, :, :] .= col_dry + layerdata[2, :, :] .= p_lay + layerdata[3, :, :] .= t_lay + t_sfc = DA(t_sfc) cld_frac = DA(cld_frac) @@ -155,7 +160,7 @@ function setup_allsky_as( ice_rgh, ) return ( - AtmosphericState(lon, lat, p_lay, p_lev, t_lay, t_lev, t_sfc, col_dry, vmr, cloud_state), + AtmosphericState(lon, lat, layerdata, p_lev, t_lev, t_sfc, vmr, cloud_state), sfc_emis, sfc_alb_direct, sfc_alb_diffuse, diff --git a/test/read_rfmip_clear_sky.jl b/test/read_rfmip_clear_sky.jl index 92ecb3b4b..14228aea6 100644 --- a/test/read_rfmip_clear_sky.jl +++ b/test/read_rfmip_clear_sky.jl @@ -112,10 +112,15 @@ function setup_rfmip_as( # FORTRAN RRTMGP test case. compute_col_gas!(context, p_lev, col_dry, param_set, vmr_h2o, lat) # the example skips lat based gravity calculation + layerdata = similar(p_lay, 3, nlay, ncol) + layerdata[1, :, :] .= col_dry + layerdata[2, :, :] .= p_lay + layerdata[3, :, :] .= t_lay + vmr = VMR(vmr_h2o, vmr_o3, FTA1D(vmrat)) #------------------ return ( - AtmosphericState(lon, lat, p_lay, p_lev, t_lay, t_lev, t_sfc, col_dry, vmr, nothing), + AtmosphericState(lon, lat, layerdata, p_lev, t_lev, t_sfc, vmr, nothing), sfc_emis, sfc_alb, cos_zenith,