Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Reduce memory footprint #397

Merged
merged 1 commit into from
Nov 23, 2023
Merged

Reduce memory footprint #397

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,5 +18,6 @@ CUDA = "4, 5"
ClimaComms = "0.5.1"
DocStringExtensions = "0.8, 0.9"
GaussQuadrature = "0.5"
Random = "1"
StaticArrays = "1.4"
julia = "1.7"
4 changes: 2 additions & 2 deletions examples/Manifest.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

julia_version = "1.9.3"
manifest_format = "2.0"
project_hash = "243c6060b2bc9f1d650bfee030845a40aea5b195"
project_hash = "ae0ac46db0377691aef20bea1543b1fb30104413"

[[deps.AbstractFFTs]]
deps = ["LinearAlgebra"]
Expand Down Expand Up @@ -657,7 +657,7 @@ uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
deps = ["Adapt", "CUDA", "ClimaComms", "DocStringExtensions", "GaussQuadrature", "Random", "StaticArrays"]
path = ".."
uuid = "a01a1ee8-cea4-48fc-987c-fc7878d79da1"
version = "0.9.1"
version = "0.9.3"

[[deps.Random]]
deps = ["SHA", "Serialization"]
Expand Down
8 changes: 4 additions & 4 deletions src/optics/AtmosphericStates.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@ struct AtmosphericState{
FTA1DN <: Union{AbstractArray{FT, 1}, Nothing},
FTA2D <: AbstractArray{FT, 2},
CLDP <: Union{AbstractArray{FT, 2}, Nothing},
CLDM <: Union{AbstractArray{Bool, 3}, Nothing},
RND <: Union{AbstractArray{FT, 3}, Nothing},
CLDM <: Union{AbstractArray{Bool, 2}, Nothing},
RND <: Union{AbstractArray{FT, 2}, Nothing},
CMASK <: Union{AbstractCloudMask, Nothing},
VMR <: AbstractVmr{FT},
} <: AbstractAtmosphericState{FT, FTA1D}
Expand Down Expand Up @@ -76,9 +76,9 @@ struct AtmosphericState{
cld_path_ice::CLDP
"cloud fraction"
cld_frac::CLDP
"random number storage for longwave bands `(ngpt, nlay, ncol)`"
"random number storage for longwave bands `(nlay, ncol)`"
random_lw::RND
"random number storage for shortwave bands `(ngpt, nlay, ncol)`"
"random number storage for shortwave bands `(nlay, ncol)`"
random_sw::RND
"cloud mask (longwave), = true if clouds are present"
cld_mask_lw::CLDM
Expand Down
48 changes: 21 additions & 27 deletions src/optics/CloudOptics.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ to the TwoStream longwave gas optics properties.
"""
function add_cloud_optics_2stream(op::TwoStream, as::AtmosphericState, lkp::LookUpLW, lkp_cld, glay, gcol, ibnd, igpt)
if as.cld_mask_lw isa AbstractArray
cld_mask = as.cld_mask_lw[igpt, glay, gcol]
cld_mask = as.cld_mask_lw[glay, gcol]
if cld_mask
τ_cl, τ_cl_ssa, τ_cl_ssag = compute_cld_props(lkp_cld, as, cld_mask, glay, gcol, ibnd, igpt)
increment!(op, τ_cl, τ_cl_ssa, τ_cl_ssag, glay, gcol, igpt)
Expand All @@ -40,7 +40,7 @@ to the TwoStream shortwave gase optics properties.
"""
function add_cloud_optics_2stream(op::TwoStream, as::AtmosphericState, lkp::LookUpSW, lkp_cld, glay, gcol, ibnd, igpt)
if as.cld_mask_sw isa AbstractArray
cld_mask = as.cld_mask_sw[igpt, glay, gcol]
cld_mask = as.cld_mask_sw[glay, gcol]
if cld_mask
τ_cl, τ_cl_ssa, τ_cl_ssag = compute_cld_props(lkp_cld, as, cld_mask, glay, gcol, ibnd, igpt)
τ_cl, τ_cl_ssa, τ_cl_ssag = delta_scale(τ_cl, τ_cl_ssa, τ_cl_ssag)
Expand Down Expand Up @@ -204,13 +204,13 @@ Builds McICA-sampled cloud mask from cloud fraction data for maximum-random over

Reference: https://github.com/AER-RC/RRTMG_SW/
"""
function build_cloud_mask!(cld_mask, cld_frac, random_arr, gcol, ::MaxRandomOverlap)
function build_cloud_mask!(cld_mask, cld_frac, random_arr, gcol, igpt, ::MaxRandomOverlap)

FT = eltype(cld_frac)
local_rand = view(random_arr, :, :, gcol)
local_rand = view(random_arr, :, gcol)
local_cld_frac = view(cld_frac, :, gcol)
local_cld_mask = view(cld_mask, :, :, gcol)
ngpt, nlay = size(local_rand)
local_cld_mask = view(cld_mask, :, gcol)
nlay = size(local_rand, 1)

Random.rand!(local_rand)
start = 0
Expand All @@ -222,8 +222,8 @@ function build_cloud_mask!(cld_mask, cld_frac, random_arr, gcol, ::MaxRandomOver
end

if start == 0 # no clouds in entire column
@inbounds for ilay in 1:nlay, igpt in 1:ngpt
local_cld_mask[igpt, ilay] = false
@inbounds for ilay in 1:nlay
local_cld_mask[ilay] = false
end
else
# finish = findlast(local_cld_frac .> FT(0)) # for GPU compatibility
Expand All @@ -235,33 +235,27 @@ function build_cloud_mask!(cld_mask, cld_frac, random_arr, gcol, ::MaxRandomOver
end
end
# set cloud mask for non-cloudy layers
@inbounds for ilay in 1:(start - 1), igpt in 1:ngpt
local_cld_mask[igpt, ilay] = false
@inbounds for ilay in 1:(start - 1)
local_cld_mask[ilay] = false
end
@inbounds for ilay in (finish + 1):nlay, igpt in 1:ngpt
local_cld_mask[igpt, ilay] = false
@inbounds for ilay in (finish + 1):nlay
local_cld_mask[ilay] = false
end
# set cloud mask for cloudy layers
# last layer
# RRTMG uses local_rand[igpt, finish] > (FT(1) - local_cld_frac[finish]), we change > to >= to address edge cases
@inbounds for igpt in 1:ngpt
local_cld_mask[igpt, finish] = local_rand[igpt, finish] >= (FT(1) - local_cld_frac[finish])
end
# RRTMG uses local_rand[finish] > (FT(1) - local_cld_frac[finish]), we change > to >= to address edge cases
@inbounds local_cld_mask[finish] = local_rand[finish] >= (FT(1) - local_cld_frac[finish])
@inbounds for ilay in (finish - 1):-1:start
if local_cld_frac[ilay] > FT(0)
for igpt in 1:ngpt
if local_cld_mask[igpt, ilay + 1]
local_rand[igpt, ilay] = local_rand[igpt, ilay + 1] # use same random number from the layer above if layer above is cloudy
else
local_rand[igpt, ilay] *= (FT(1) - local_cld_frac[ilay + 1]) # update random numbers if layer above is not cloudy
end
# RRTMG uses local_rand[igpt, ilay] > (FT(1) - local_cld_frac[ilay]), we change > to >= to address edge cases
local_cld_mask[igpt, ilay] = local_rand[igpt, ilay] >= (FT(1) - local_cld_frac[ilay])
if local_cld_mask[ilay + 1]
local_rand[ilay] = local_rand[ilay + 1] # use same random number from the layer above if layer above is cloudy
else
local_rand[ilay] *= (FT(1) - local_cld_frac[ilay + 1]) # update random numbers if layer above is not cloudy
end
# RRTMG uses local_rand[ilay] > (FT(1) - local_cld_frac[ilay]), we change > to >= to address edge cases
local_cld_mask[ilay] = local_rand[ilay] >= (FT(1) - local_cld_frac[ilay])
else
for igpt in 1:ngpt
local_cld_mask[igpt, ilay] = false
end
local_cld_mask[ilay] = false
end
end
end
Expand Down
42 changes: 24 additions & 18 deletions src/rte/RTESolver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,17 +186,19 @@ function rte_lw_solve!(
args = (flux, flux_lw, src_lw, bcs_lw, op, nlay, ncol, major_gpt2bnd, as, lookup_lw, lookup_lw_cld)
@cuda threads = (tx) blocks = (bx) rte_lw_2stream_solve_CUDA!(args...)
else
if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
@inbounds begin
ClimaComms.@threaded device for gcol in 1:ncol
Optics.build_cloud_mask!(as.cld_mask_lw, as.cld_frac, as.random_lw, gcol, as.cld_mask_type)
end
end
end
bld_cld_mask = as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
@inbounds begin
for igpt in 1:n_gpt
ClimaComms.@threaded device for gcol in 1:ncol
igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
bld_cld_mask && Optics.build_cloud_mask!(
as.cld_mask_lw,
as.cld_frac,
as.random_lw,
gcol,
igpt,
as.cld_mask_type,
)
compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld)
rte_lw_2stream_combine_sources!(src_lw, gcol, nlev, ncol)
rte_lw_2stream_source!(op, src_lw, gcol, nlay, ncol)
Expand Down Expand Up @@ -226,11 +228,11 @@ function rte_lw_2stream_solve_CUDA!(
nlev = nlay + 1
n_gpt = length(major_gpt2bnd)
if gcol ≤ ncol
if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
Optics.build_cloud_mask!(as.cld_mask_lw, as.cld_frac, as.random_lw, gcol, as.cld_mask_type)
end
@inbounds for igpt in 1:n_gpt
igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
Optics.build_cloud_mask!(as.cld_mask_lw, as.cld_frac, as.random_lw, gcol, igpt, as.cld_mask_type)
end
compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld)
rte_lw_2stream_combine_sources!(src_lw, gcol, nlev, ncol)
rte_lw_2stream_source!(op, src_lw, gcol, nlay, ncol)
Expand Down Expand Up @@ -443,15 +445,19 @@ function rte_sw_2stream_solve!(
@cuda threads = (tx) blocks = (bx) rte_sw_2stream_solve_CUDA!(args...)
else
@inbounds begin
if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
ClimaComms.@threaded device for gcol in 1:ncol
Optics.build_cloud_mask!(as.cld_mask_sw, as.cld_frac, as.random_sw, gcol, as.cld_mask_type)
end
end
bld_cld_mask = as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
# setting references for flux_sw
for igpt in 1:n_gpt
ClimaComms.@threaded device for gcol in 1:ncol
igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
bld_cld_mask && Optics.build_cloud_mask!(
as.cld_mask_sw,
as.cld_frac,
as.random_sw,
gcol,
igpt,
as.cld_mask_type,
)
compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
sw_two_stream!(op, src_sw, bcs_sw, gcol, nlay) # Cell properties: transmittance and reflectance for direct and diffuse radiation
sw_source_2str!(src_sw, bcs_sw, gcol, flux, igpt, solar_src_scaled, major_gpt2bnd, nlay) # Direct-beam and source for diffuse radiation
Expand Down Expand Up @@ -482,11 +488,11 @@ function rte_sw_2stream_solve_CUDA!(
nlev = nlay + 1
n_gpt = length(major_gpt2bnd)
if gcol ≤ ncol
if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
Optics.build_cloud_mask!(as.cld_mask_sw, as.cld_frac, as.random_sw, gcol, as.cld_mask_type)
end
@inbounds for igpt in 1:n_gpt
igpt == 1 && set_flux_to_zero!(flux_sw, gcol)
if as isa AtmosphericState && as.cld_mask_type isa AbstractCloudMask
Optics.build_cloud_mask!(as.cld_mask_sw, as.cld_frac, as.random_sw, gcol, igpt, as.cld_mask_type)
end
compute_optical_props!(op, as, gcol, igpt, lookup_sw, lookup_sw_cld)
sw_two_stream!(op, src_sw, bcs_sw, gcol, nlay) # Cell properties: transmittance and reflectance for direct and diffuse radiation
sw_source_2str!(src_sw, bcs_sw, gcol, flux, igpt, solar_src_scaled, major_gpt2bnd, nlay) # Direct-beam and source for diffuse radiation
Expand Down
3 changes: 1 addition & 2 deletions test/aqua.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,7 @@ end
Aqua.test_stale_deps(RRTMGP)
Aqua.test_deps_compat(RRTMGP)
Aqua.test_project_extras(RRTMGP)
Aqua.test_project_toml_formatting(RRTMGP)
Aqua.test_piracy(RRTMGP)
Aqua.test_piracies(RRTMGP)
end

nothing
8 changes: 4 additions & 4 deletions test/read_all_sky.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,10 +91,10 @@ function setup_allsky_as(
vmr_h2o = view(vmr.vmr, :, :, idx_gases["h2o"])

cld_frac = zeros(FT, nlay, ncol)
random_lw = DA{FT}(undef, ngpt_lw, nlay, ncol)
random_sw = DA{FT}(undef, ngpt_sw, nlay, ncol)
cld_mask_lw = zeros(Bool, ngpt_lw, nlay, ncol)
cld_mask_sw = zeros(Bool, ngpt_lw, nlay, ncol)
random_lw = DA{FT}(undef, nlay, ncol)
random_sw = DA{FT}(undef, nlay, ncol)
cld_mask_lw = zeros(Bool, nlay, ncol)
cld_mask_sw = zeros(Bool, nlay, ncol)
cld_r_eff_liq = zeros(FT, nlay, ncol)
cld_r_eff_ice = zeros(FT, nlay, ncol)
cld_path_liq = zeros(FT, nlay, ncol)
Expand Down
Loading