Preallocate bucket auxs

.buildkite/target/pipeline.yml

@@ -27,10 +27,9 @@ steps:
   - group: "Target Benchmark"
     steps:
       - label: ":bucket: Bucket"
-        command: "nsys launch julia --color=yes --project=.buildkite experiments/benchmarks/bucket.jl"
+        command: "julia --color=yes --project=.buildkite experiments/benchmarks/bucket.jl"
         artifact_paths:
           - "bucket_benchmark_gpu/*html"
-          - "bucket_benchmark_gpu/*nsys*"
         env:
           CLIMACOMMS_DEVICE: CUDA
         agents:

experiments/benchmarks/bucket.jl

@@ -206,11 +206,11 @@ ProfileCanvas.html_file(alloc_flame_file, profile)
 
 if ClimaComms.device() isa ClimaComms.CUDADevice
     import CUDA
-    CUDA.@profile external = true setup_and_solve_problem()
+    CUDA.@profile setup_and_solve_problem()
 end
 
 if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
-    PREVIOUS_BEST_TIME = 5.4
+    PREVIOUS_BEST_TIME = 3.6
     if average_timing_s > 1.1PREVIOUS_BEST_TIME
         @info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
     elseif average_timing_s < 0.8PREVIOUS_BEST_TIME

src/standalone/Bucket/Bucket.jl

@@ -317,11 +317,41 @@ auxiliary_types(::BucketModel{FT}) where {FT} = (
     FT,
     FT,
     FT,
+    FT,
+    FT,
+    NamedTuple{(:F_melt, :F_into_snow, :G_under_snow), Tuple{FT, FT, FT}},
+    FT,
+    FT,
+    FT,
+)
+auxiliary_vars(::BucketModel) = (
+    :q_sfc,
+    :turbulent_fluxes,
+    :R_n,
+    :T_sfc,
+    :α_sfc,
+    :ρ_sfc,
+    :snow_cover_fraction,
+    :F_sfc,
+    :partitioned_fluxes,
+    :G,
+    :snow_melt,
+    :infiltration,
+)
+auxiliary_domain_names(::BucketModel) = (
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
+    :surface,
 )
-auxiliary_vars(::BucketModel) =
-    (:q_sfc, :turbulent_fluxes, :R_n, :T_sfc, :α_sfc, :ρ_sfc)
-auxiliary_domain_names(::BucketModel) =
-    (:surface, :surface, :surface, :surface, :surface, :surface, :surface)
 
 
 """
@@ -331,80 +361,39 @@ Creates the compute_exp_tendency! function for the bucket model.
 """
 function make_compute_exp_tendency(model::BucketModel{FT}) where {FT}
     function compute_exp_tendency!(dY, Y, p, t)
-        (; κ_soil, ρc_soil, σS_c, W_f) = model.parameters
-
-        #Currently, the entire surface is assumed to be
-        # snow covered entirely or not at all.
-        snow_cover_fraction = heaviside.(Y.bucket.σS)
+        (; κ_soil, ρc_soil, σS_c) = model.parameters
 
-        # In this case, there is just one set of surface fluxes to compute.
-        # Since q_sfc itself has already been modified to account for
-        # snow covered regions, and since the surface temperature is
-        # assumed to be the same for snow and underlying land,
-        # the returned fluxes are computed correctly for the cell
-        # regardless of snow-coverage.
-
-        # The below must be adjusted to compute F_sfc over snow and over soil
-        # if we want the snow cover fraction to be intermediate between 0 and 1.
-        (; turbulent_fluxes, R_n) = p.bucket
-        F_sfc = @. (turbulent_fluxes.shf .+ turbulent_fluxes.lhf + R_n) # Eqn (21)
-        _T_freeze = LP.T_freeze(model.parameters.earth_param_set)
-        _LH_f0 = LP.LH_f0(model.parameters.earth_param_set)
-        _ρ_liq = LP.ρ_cloud_liq(model.parameters.earth_param_set)
-        _ρLH_f0 = _ρ_liq * _LH_f0 # Latent heat per unit volume.
-        # partition energy fluxes for snow covered area
-        partitioned_fluxes =
-            partition_snow_surface_fluxes.(
-                Y.bucket.σS,
-                p.bucket.T_sfc,
-                model.parameters.τc,
-                snow_cover_fraction,
-                turbulent_fluxes.vapor_flux,
-                F_sfc,
-                _ρLH_f0,
-                _T_freeze,
-            )
-        (; F_melt, F_into_snow, G_under_snow) = partitioned_fluxes
-        G = @. F_sfc * (1 - snow_cover_fraction) +
-           G_under_snow * snow_cover_fraction # Equation 22
         # Temperature profile of soil.
         gradc2f = ClimaCore.Operators.GradientC2F()
         divf2c = ClimaCore.Operators.DivergenceF2C(
-            top = ClimaCore.Operators.SetValue(ClimaCore.Geometry.WVector.(G)),
+            top = ClimaCore.Operators.SetValue(
+                ClimaCore.Geometry.WVector.(p.bucket.G),
+            ),
             bottom = ClimaCore.Operators.SetValue(
                 ClimaCore.Geometry.WVector.(FT(0.0)),
             ),
         )
 
         @. dY.bucket.T = -1 / ρc_soil * (divf2c(-κ_soil * gradc2f(Y.bucket.T))) # Simple heat equation, Eq 10
 
-        # Partition water fluxes
-        liquid_precip = liquid_precipitation(model.atmos, p, t) # always negative
-        snow_precip = snow_precipitation(model.atmos, p, t) # always negative
-        # F_melt is negative as it is a downward welling flux warming the snow
-        snow_melt = @. F_melt / _ρLH_f0 # defined after Equation (22)
-
-        infiltration = @. infiltration_at_point(
-            Y.bucket.W,
-            snow_cover_fraction * snow_melt,
-            liquid_precip,
-            (1 - snow_cover_fraction) * turbulent_fluxes.vapor_flux,
-            W_f,
-        ) # Equation (2) of the text.
-
         # Positive infiltration -> net (negative) flux into soil
-        @. dY.bucket.W = -infiltration # Equation (2) of the text.
+        @. dY.bucket.W = -p.bucket.infiltration # Equation (2) of the text.
+
+        liquid_precip = liquid_precipitation(model.atmos, p, t) # always negative
+        snow_precip = snow_precipitation(model.atmos, p, t)
 
         dY.bucket.Ws = @. -(
             liquid_precip +
-            snow_cover_fraction * snow_melt +
-            (1 - snow_cover_fraction) * turbulent_fluxes.vapor_flux -
-            infiltration
+            p.bucket.snow_cover_fraction * p.bucket.snow_melt +
+            (1 - p.bucket.snow_cover_fraction) *
+            p.bucket.turbulent_fluxes.vapor_flux - p.bucket.infiltration
         ) # Equation (3) of the text
 
         dY.bucket.σS = @. -(
-            snow_precip + snow_cover_fraction * turbulent_fluxes.vapor_flux -
-            snow_cover_fraction * snow_melt
+            snow_precip +
+            p.bucket.snow_cover_fraction *
+            p.bucket.turbulent_fluxes.vapor_flux -
+            p.bucket.snow_cover_fraction * p.bucket.snow_melt
         ) # Equation (6)
     end
     return compute_exp_tendency!
@@ -435,7 +424,6 @@ function make_update_aux(model::BucketModel{FT}) where {FT}
                     ),
                 ),
             )
-
         # Compute turbulent surface fluxes
         p.bucket.turbulent_fluxes .=
             turbulent_fluxes(model.atmos, model, Y, p, t)
@@ -452,6 +440,59 @@ function make_update_aux(model::BucketModel{FT}) where {FT}
             p,
             t,
         )
+
+        #Currently, the entire surface is assumed to be
+        # snow covered entirely or not at all.
+        p.bucket.snow_cover_fraction .= heaviside.(Y.bucket.σS)
+
+        # In this case, there is just one set of surface fluxes to compute.
+        # Since q_sfc itself has already been modified to account for
+        # snow covered regions, and since the surface temperature is
+        # assumed to be the same for snow and underlying land,
+        # the returned fluxes are computed correctly for the cell
+        # regardless of snow-coverage.
+
+        # The below must be adjusted to compute F_sfc over snow and over soil
+        # if we want the snow cover fraction to be intermediate between 0 and 1.
+        @. p.bucket.F_sfc = (
+            p.bucket.turbulent_fluxes.shf .+ p.bucket.turbulent_fluxes.lhf +
+            p.bucket.R_n
+        ) # Eqn (21)
+        _T_freeze = LP.T_freeze(model.parameters.earth_param_set)
+        _LH_f0 = LP.LH_f0(model.parameters.earth_param_set)
+        _ρ_liq = LP.ρ_cloud_liq(model.parameters.earth_param_set)
+        _ρLH_f0 = _ρ_liq * _LH_f0 # Latent heat per unit volume.
+        # partition energy fluxes for snow covered area
+        p.bucket.partitioned_fluxes .=
+            partition_snow_surface_fluxes.(
+                Y.bucket.σS,
+                p.bucket.T_sfc,
+                model.parameters.τc,
+                p.bucket.snow_cover_fraction,
+                p.bucket.turbulent_fluxes.vapor_flux,
+                p.bucket.F_sfc,
+                _ρLH_f0,
+                _T_freeze,
+            )
+        @. p.bucket.G =
+            p.bucket.F_sfc * (1 - p.bucket.snow_cover_fraction) +
+            p.bucket.partitioned_fluxes.G_under_snow *
+            p.bucket.snow_cover_fraction # Equation 22
+
+        # Partition water fluxes
+        liquid_precip = liquid_precipitation(model.atmos, p, t) # always negative
+        # F_melt is negative as it is a downward welling flux warming the snow
+        @.p.bucket.snow_melt = p.bucket.partitioned_fluxes.F_melt / _ρLH_f0 # defined after Equation (22)
+
+        @. p.bucket.infiltration = infiltration_at_point(
+            Y.bucket.W,
+            p.bucket.snow_cover_fraction * p.bucket.snow_melt,
+            liquid_precip,
+            (1 - p.bucket.snow_cover_fraction) *
+            p.bucket.turbulent_fluxes.vapor_flux,
+            model.parameters.W_f,
+        ) # Equation (2) of the text.
+
     end
     return update_aux!
 end

test/standalone/Bucket/soil_bucket_tests.jl

@@ -218,6 +218,44 @@ for FT in (Float32, Float64)
                 p.bucket.turbulent_fluxes.shf .+ p.bucket.R_n
             surface_space = model.domain.space.surface
             A_sfc = sum(ones(surface_space))
+
+            _T_freeze = LP.T_freeze(earth_param_set)
+            _LH_f0 = LP.LH_f0(earth_param_set)
+            _ρ_liq = LP.ρ_cloud_liq(earth_param_set)
+            _ρLH_f0 = _ρ_liq * _LH_f0
+
+            @test p.bucket.snow_cover_fraction ==
+                  ClimaLand.heaviside.(Y.bucket.σS)
+            @test p.bucket.F_sfc == F_sfc
+            @test p.bucket.partitioned_fluxes ==
+                  ClimaLand.Bucket.partition_snow_surface_fluxes.(
+                Y.bucket.σS,
+                p.bucket.T_sfc,
+                model.parameters.τc,
+                p.bucket.snow_cover_fraction,
+                p.bucket.turbulent_fluxes.vapor_flux,
+                p.bucket.F_sfc,
+                _ρLH_f0,
+                _T_freeze,
+            )
+            @test p.bucket.G ==
+                  F_sfc .* (1 .- p.bucket.snow_cover_fraction) .+
+                  p.bucket.partitioned_fluxes.G_under_snow .*
+                  p.bucket.snow_cover_fraction
+            @test p.bucket.snow_melt ==
+                  p.bucket.partitioned_fluxes.F_melt ./ _ρLH_f0
+            liquid_precip = ClimaLand.liquid_precipitation(model.atmos, p, t0)
+
+            @test p.bucket.infiltration ==
+                  ClimaLand.Bucket.infiltration_at_point.(
+                Y.bucket.W,
+                p.bucket.snow_cover_fraction .* p.bucket.snow_melt,
+                liquid_precip,
+                (1 .- p.bucket.snow_cover_fraction) .*
+                p.bucket.turbulent_fluxes.vapor_flux,
+                model.parameters.W_f,
+            )
+
             # For the point space, we actually want the flux itself, since our variables are per unit area.
             # Divide by A_sfc
             if typeof(model.domain.space.surface) <: ClimaCore.Spaces.PointSpace