Skip to content

Commit 1822266

Browse files
committed
Reduce allocations
This commit reduces allocations by changing turbulent_fluxes and net_radiation to be in place. In addition to this, the workaround introduced to enable GPU-compatibility in energy_hydrology resulted in lots of additional allocations, so I modified that too. There are still allocations. One of the easy ones to tackle is in surface_specificy_humidity for snow, which accounts for a large fraction of the total allocations for the snowy_land benchmark
1 parent 13303d5 commit 1822266

File tree

24 files changed

+434
-366
lines changed

24 files changed

+434
-366
lines changed

.buildkite/Manifest-v1.11.toml

Lines changed: 83 additions & 84 deletions
Large diffs are not rendered by default.

.buildkite/Manifest.toml

Lines changed: 83 additions & 84 deletions
Large diffs are not rendered by default.

docs/src/APIs/shared_utilities.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ ClimaLand.CoupledAtmosphere
7575
ClimaLand.CoupledRadiativeFluxes
7676
ClimaLand.AbstractAtmosphericDrivers
7777
ClimaLand.AbstractRadiativeDrivers
78-
ClimaLand.turbulent_fluxes
78+
ClimaLand.turbulent_fluxes!
7979
ClimaLand.turbulent_fluxes_at_a_point
8080
ClimaLand.set_atmos_ts!
8181
ClimaLand.surface_air_density

docs/tutorials/standalone/Bucket/coupled_bucket.jl

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@
6161

6262
# These are stored in the [BucketModel](https://clima.github.io/ClimaLand.jl/dev/APIs/Bucket/#ClimaLand.Bucket.BucketModel) object,
6363
# along with [BucketParameters](https://clima.github.io/ClimaLand.jl/dev/APIs/Bucket/#ClimaLand.Bucket.BucketParameters).
64-
# In order to compute turbulent surface fluxes, we call [turbulent_fluxes](https://clima.github.io/ClimaLand.jl/dev/APIs/shared_utilities/#ClimaLand.turbulent_fluxes),
64+
# In order to compute turbulent surface fluxes, we call [turbulent_fluxes!](https://clima.github.io/ClimaLand.jl/dev/APIs/shared_utilities/#ClimaLand.turbulent_fluxes!),
6565
# with arguments including `prescribed_atmos`. Since this argument is of the type `PrescribedAtmosphere`, the method of `turbulent_fluxes` which is executed is one which computes the turbulent surface fluxes
6666
# using MOST. We have a similar function for [net_radiation](https://clima.github.io/ClimaLand.jl/dev/APIs/shared_utilities/#ClimaLand.net_radiation) and which computes the net radiation based on the prescribed downwelling radiative fluxes, stored in an argument
6767
# `prescribed_radiation`, which is of type `PrescribedRadiation`.
@@ -74,38 +74,40 @@
7474
# struct CoupledRadiativeFluxes{FT} <: AbstractRadiativeDrivers{FT} end
7575
# ```
7676

77-
# Then, we have defined a new method for `turbulent_fluxes` and `net_radiation` which dispatch for these types,
78-
# and simply return the fluxes that the coupler has updated `p.bucket.turbulent_fluxes` and `p.bucket.R_n` with.
77+
# Then, we have defined a new method for `turbulent_fluxes!` and `net_radiation!` which dispatch for these types,
78+
# and updates dest with the fluxes that the coupler has updated `p.bucket.turbulent_fluxes` and `p.bucket.R_n` with.
7979
# In pseudo code:
8080
#
8181
# ```julia
82-
# function ClimaLand.turbulent_fluxes(
82+
# function ClimaLand.turbulent_fluxes!(
83+
# dest,
8384
# atmos::CoupledAtmosphere,
8485
# model::BucketModel,
8586
# p)
86-
# return (
87+
# dest .= (
8788
# lhf = p.bucket.turbulent_fluxes.lhf,
8889
# shf = p.bucket.turbulent_fluxes.shf,
8990
# vapor_flux = p.bucket.turbulent_fluxes.vapor_flux,
9091
# )
9192
# end
9293
# ```
9394

94-
# similarily:
95+
# similarily:
9596

9697
# ```julia
97-
# function ClimaLand.net_radiation(
98+
# function ClimaLand.net_radiation!(
99+
# dest,
98100
# radiation::CoupledRadiativeFluxes{FT},
99101
# model::BucketModel{FT},
100102
# p)
101-
# return p.bucket.R_n
103+
# dest .= p.bucket.R_n
102104
# end
103105
# ```
104106

105-
# These methods simply returns the values stored in the auxiliary state `p`. Importantly, these functions are
107+
# These methods update the values stored in the auxiliary state `p`. Importantly, these functions are
106108
# called by the bucket model
107109
# each time step **after** the coupler has already computed these values
108-
# (or extracted them from another model) and modifed `p`!
110+
# (or extracted them from another model) and modified `p`!
109111

110112
# # Surface air density
111113
# Within the right hand side/ODE function calls for the bucket model, we need both the

experiments/benchmarks/bucket.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice
244244
end
245245

246246
if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
247-
PREVIOUS_BEST_TIME = 1.1
247+
PREVIOUS_BEST_TIME = 0.477
248248
if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
249249
@info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
250250
elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

experiments/benchmarks/land.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -428,7 +428,7 @@ ProfileCanvas.html_file(flame_file, results)
428428
@info "Saved compute flame to $flame_file"
429429

430430
prob, ode_algo, Δt, cb = setup_simulation()
431-
Profile.Allocs.@profile sample_rate = 0.005 SciMLBase.solve(
431+
Profile.Allocs.@profile sample_rate = 0.0025 SciMLBase.solve(
432432
prob,
433433
ode_algo;
434434
dt = Δt,
@@ -465,7 +465,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice
465465
end
466466

467467
if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
468-
PREVIOUS_BEST_TIME = 4.9
468+
PREVIOUS_BEST_TIME = 6.1
469469
if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
470470
@info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
471471
elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

experiments/benchmarks/richards.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice
295295
end
296296

297297
if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
298-
PREVIOUS_BEST_TIME = 5.9
298+
PREVIOUS_BEST_TIME = 5.8
299299
if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
300300
@info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
301301
elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

experiments/benchmarks/snowy_land.jl

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,12 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 15))
6969
start_date = DateTime(2008)
7070
# Forcing data
7171
era5_artifact_path =
72-
ClimaLand.Artifacts.era5_land_forcing_data2008_folder_path(; context)
73-
era5_ncdata_path = joinpath(era5_artifact_path, "era5_2008_1.0x1.0.nc")
72+
ClimaLand.Artifacts.era5_land_forcing_data2008_folder_path(;
73+
context,
74+
lowres = true,
75+
)
76+
era5_ncdata_path =
77+
joinpath(era5_artifact_path, "era5_2008_1.0x1.0_lowres.nc")
7478
atmos, radiation = ClimaLand.prescribed_forcing_era5(
7579
era5_ncdata_path,
7680
surface_space,
@@ -366,7 +370,7 @@ function setup_simulation(; greet = false)
366370
t0 = 0.0
367371
tf = 60 * 60.0 * 6
368372
Δt = 450.0
369-
nelements = (101, 15)
373+
nelements = (11, 15)
370374
if greet
371375
@info "Run: Global Soil-Canopy-Snow Model"
372376
@info "Resolution: $nelements"
@@ -434,7 +438,7 @@ ProfileCanvas.html_file(flame_file, results)
434438
@info "Saved compute flame to $flame_file"
435439

436440
prob, ode_algo, Δt, cb = setup_simulation()
437-
Profile.Allocs.@profile sample_rate = 0.005 SciMLBase.solve(
441+
Profile.Allocs.@profile sample_rate = 0.0025 SciMLBase.solve(
438442
prob,
439443
ode_algo;
440444
dt = Δt,
@@ -471,7 +475,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice
471475
end
472476

473477
if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
474-
PREVIOUS_BEST_TIME = 6.1
478+
PREVIOUS_BEST_TIME = 6.5
475479
if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
476480
@info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
477481
elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

src/integrated/land.jl

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -480,7 +480,7 @@ function soil_boundary_fluxes!(
480480
t,
481481
) where {FT}
482482
bc = soil.boundary_conditions.top
483-
p.soil.turbulent_fluxes .= turbulent_fluxes(bc.atmos, soil, Y, p, t)
483+
turbulent_fluxes!(p.soil.turbulent_fluxes, bc.atmos, soil, Y, p, t)
484484
# influx = maximum possible rate of infiltration given precip, snowmelt, evaporation/condensation
485485
# but if this exceeds infiltration capacity of the soil, runoff will
486486
# be generated.
@@ -505,6 +505,7 @@ function soil_boundary_fluxes!(
505505
) +
506506
p.excess_heat_flux +
507507
p.snow.snow_cover_fraction * p.ground_heat_flux
508+
return nothing
508509
end
509510

510511
function snow_boundary_fluxes!(
@@ -515,7 +516,7 @@ function snow_boundary_fluxes!(
515516
p,
516517
t,
517518
) where {FT}
518-
p.snow.turbulent_fluxes .= turbulent_fluxes(bc.atmos, model, Y, p, t)
519+
turbulent_fluxes!(p.snow.turbulent_fluxes, bc.atmos, model, Y, p, t)
519520
# How does rain affect the below?
520521
P_snow = p.drivers.P_snow
521522

@@ -539,6 +540,7 @@ function snow_boundary_fluxes!(
539540
p.snow.turbulent_fluxes.shf +
540541
p.snow.R_n - p.snow.energy_runoff - p.ground_heat_flux
541542
) * p.snow.snow_cover_fraction
543+
return nothing
542544
end
543545

544546

src/integrated/soil_canopy_model.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -369,7 +369,7 @@ function soil_boundary_fluxes!(
369369
t,
370370
) where {FT}
371371
bc = soil.boundary_conditions.top
372-
p.soil.turbulent_fluxes .= turbulent_fluxes(bc.atmos, soil, Y, p, t)
372+
turbulent_fluxes!(p.soil.turbulent_fluxes, bc.atmos, soil, Y, p, t)
373373
# influx = maximum possible rate of infiltration given precip, snowmelt, evaporation/condensation
374374
# but if this exceeds infiltration capacity of the soil, runoff will
375375
# be generated.
@@ -383,6 +383,7 @@ function soil_boundary_fluxes!(
383383
@. p.soil.top_bc.water = p.soil.infiltration
384384
@. p.soil.top_bc.heat =
385385
-p.soil.R_n + p.soil.turbulent_fluxes.lhf + p.soil.turbulent_fluxes.shf
386+
return nothing
386387
end
387388

388389
"""

0 commit comments

Comments
 (0)