Merge pull request #967 from CliMA/gb/allocations

Reduce allocations

Author: Sbozzolo

.buildkite/Manifest-v1.11.toml

@@ -1,7 +1,11 @@
 agents:
   queue: clima
   slurm_mem: 8G
-  modules: climacommon/2024_10_09
+  modules: climacommon/2024_12_16
+
+env:
+  JULIA_CUDA_MEMORY_POOL: "cuda"
+
 
 steps:
   - label: "init :GPU:"

.buildkite/Manifest.toml

@@ -1,13 +1,15 @@
 agents:
   queue: new-central
   slurm_mem: 12G
-  modules: climacommon/2024_10_09
+  modules: climacommon/2024_12_16
 
 env:
   JULIA_NVTX_CALLBACKS: gc
   OPENBLAS_NUM_THREADS: 1
   SLURM_KILL_BAD_EXIT: 1
   JULIA_DEPOT_PATH: "${BUILDKITE_BUILD_PATH}/${BUILDKITE_PIPELINE_SLUG}/depot/default"
+  JULIA_CUDA_MEMORY_POOL: "cuda"
+
 
 steps:
   - label: "init environment :computer:"

.buildkite/longruns_gpu/pipeline.yml

@@ -1,12 +1,13 @@
 agents:
   queue: clima
-  modules: climacommon/2024_10_09
+  modules: climacommon/2024_12_16
 
 env:
   JULIA_DEPOT_PATH: "${BUILDKITE_BUILD_PATH}/${BUILDKITE_PIPELINE_SLUG}/depot/default"
   JULIA_NVTX_CALLBACKS: gc
   OPENBLAS_NUM_THREADS: 1
   SLURM_KILL_BAD_EXIT: 1
+  JULIA_CUDA_MEMORY_POOL: "cuda"
 
 steps:
   - label: "init environment :computer:"

.buildkite/pipeline.yml

@@ -75,7 +75,7 @@ ClimaLand.CoupledAtmosphere
 ClimaLand.CoupledRadiativeFluxes
 ClimaLand.AbstractAtmosphericDrivers
 ClimaLand.AbstractRadiativeDrivers
-ClimaLand.turbulent_fluxes
+ClimaLand.turbulent_fluxes!
 ClimaLand.turbulent_fluxes_at_a_point
 ClimaLand.set_atmos_ts!
 ClimaLand.surface_air_density

.buildkite/target/pipeline.yml

@@ -61,7 +61,7 @@
 
 # These are stored in the [BucketModel](https://clima.github.io/ClimaLand.jl/dev/APIs/Bucket/#ClimaLand.Bucket.BucketModel) object,
 # along with [BucketParameters](https://clima.github.io/ClimaLand.jl/dev/APIs/Bucket/#ClimaLand.Bucket.BucketParameters).
-# In order to compute turbulent surface fluxes, we call [turbulent_fluxes](https://clima.github.io/ClimaLand.jl/dev/APIs/shared_utilities/#ClimaLand.turbulent_fluxes),
+# In order to compute turbulent surface fluxes, we call [turbulent_fluxes!](https://clima.github.io/ClimaLand.jl/dev/APIs/shared_utilities/#ClimaLand.turbulent_fluxes!),
 # 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
 # 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
 # `prescribed_radiation`, which is of type `PrescribedRadiation`.
@@ -74,38 +74,40 @@
 # struct CoupledRadiativeFluxes{FT} <: AbstractRadiativeDrivers{FT} end
 # ```
 
-# Then, we have defined a new method for `turbulent_fluxes` and `net_radiation` which dispatch for these types,
-# and simply return the fluxes that the coupler has updated `p.bucket.turbulent_fluxes` and `p.bucket.R_n` with.
+# Then, we have defined a new method for `turbulent_fluxes!` and `net_radiation!` which dispatch for these types,
+# and updates dest with the fluxes that the coupler has updated `p.bucket.turbulent_fluxes` and `p.bucket.R_n` with.
 # In pseudo code:
 #
 # ```julia
-# function ClimaLand.turbulent_fluxes(
+# function ClimaLand.turbulent_fluxes!(
+#    dest,
 #    atmos::CoupledAtmosphere,
 #    model::BucketModel,
 #    p)
-#    return (
+#    dest .= (
 #         lhf = p.bucket.turbulent_fluxes.lhf,
 #         shf = p.bucket.turbulent_fluxes.shf,
 #         vapor_flux = p.bucket.turbulent_fluxes.vapor_flux,
 #     )
 # end
 # ```
 
-# similarily: 
+# similarily:
 
 # ```julia
-# function ClimaLand.net_radiation(
+# function ClimaLand.net_radiation!(
+#     dest,
 #     radiation::CoupledRadiativeFluxes{FT},
 #     model::BucketModel{FT},
 #     p)
-#     return p.bucket.R_n
+#     dest .= p.bucket.R_n
 # end
 # ```
 
-# These methods simply returns the values stored in the auxiliary state `p`. Importantly, these functions are
+# These methods update the values stored in the auxiliary state `p`. Importantly, these functions are
 # called by the bucket model
 # each time step **after** the coupler has already computed these values
-# (or extracted them from another model) and modifed `p`!
+# (or extracted them from another model) and modified `p`!
 
 # # Surface air density
 # Within the right hand side/ODE function calls for the bucket model, we need both the 

docs/src/APIs/shared_utilities.md

@@ -9,6 +9,7 @@
 #
 # When run with buildkite on clima, this code also compares with the previous best time
 # saved at the bottom of this file
+delete!(ENV, "JULIA_CUDA_MEMORY_POOL")
 
 import SciMLBase
 using Dates
@@ -243,7 +244,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice
 end
 
 if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
-    PREVIOUS_BEST_TIME = 1.1
+    PREVIOUS_BEST_TIME = 0.477
     if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
         @info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
     elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

docs/tutorials/standalone/Bucket/coupled_bucket.jl

@@ -20,6 +20,8 @@
 # Fixed number of iterations: 3
 # Jacobian update: Every Newton iteration
 # Atmos forcing update: every 3 hours
+delete!(ENV, "JULIA_CUDA_MEMORY_POOL")
+
 import SciMLBase
 import ClimaComms
 ClimaComms.@import_required_backends
@@ -426,7 +428,7 @@ ProfileCanvas.html_file(flame_file, results)
 @info "Saved compute flame to $flame_file"
 
 prob, ode_algo, Δt, cb = setup_simulation()
-Profile.Allocs.@profile sample_rate = 0.005 SciMLBase.solve(
+Profile.Allocs.@profile sample_rate = 0.0025 SciMLBase.solve(
     prob,
     ode_algo;
     dt = Δt,
@@ -463,7 +465,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice
 end
 
 if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
-    PREVIOUS_BEST_TIME = 4.9
+    PREVIOUS_BEST_TIME = 6.1
     if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
         @info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
     elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

experiments/benchmarks/bucket.jl

@@ -21,6 +21,7 @@
 # Fixed number of iterations: 2
 # Jacobian update: Every Newton iteration
 # Precipitation data update: every timestep
+delete!(ENV, "JULIA_CUDA_MEMORY_POOL")
 
 import SciMLBase
 using Dates
@@ -294,7 +295,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice
 end
 
 if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
-    PREVIOUS_BEST_TIME = 5.9
+    PREVIOUS_BEST_TIME = 5.8
     if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
         @info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
     elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s