From 44d4ac3c878d8a0b890d5c1442d31cac42bf69fd Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Thu, 6 Jun 2024 13:04:44 -0700 Subject: [PATCH 1/2] Preallocate aux variables for bucket --- experiments/benchmarks/bucket.jl | 2 +- src/standalone/Bucket/Bucket.jl | 163 ++++++++++++-------- test/standalone/Bucket/soil_bucket_tests.jl | 38 +++++ 3 files changed, 141 insertions(+), 62 deletions(-) diff --git a/experiments/benchmarks/bucket.jl b/experiments/benchmarks/bucket.jl index ed9c1be2e1..be87fe7b78 100644 --- a/experiments/benchmarks/bucket.jl +++ b/experiments/benchmarks/bucket.jl @@ -210,7 +210,7 @@ if ClimaComms.device() isa ClimaComms.CUDADevice 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 diff --git a/src/standalone/Bucket/Bucket.jl b/src/standalone/Bucket/Bucket.jl index 43dd218514..b65b6c6a66 100644 --- a/src/standalone/Bucket/Bucket.jl +++ b/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,46 +361,14 @@ 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)), ), @@ -378,33 +376,24 @@ function make_compute_exp_tendency(model::BucketModel{FT}) where {FT} @. 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 diff --git a/test/standalone/Bucket/soil_bucket_tests.jl b/test/standalone/Bucket/soil_bucket_tests.jl index bc808e9c73..401fb75d4e 100644 --- a/test/standalone/Bucket/soil_bucket_tests.jl +++ b/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 From cdcfa5f1b099f1a6b359ed970133c3d15b1dee38 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Wed, 12 Jun 2024 08:48:49 -0700 Subject: [PATCH 2/2] Remove nsys from benchmark --- .buildkite/target/pipeline.yml | 3 +-- experiments/benchmarks/bucket.jl | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/.buildkite/target/pipeline.yml b/.buildkite/target/pipeline.yml index a2491e029f..05952a7765 100644 --- a/.buildkite/target/pipeline.yml +++ b/.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: diff --git a/experiments/benchmarks/bucket.jl b/experiments/benchmarks/bucket.jl index be87fe7b78..d6f1d5120b 100644 --- a/experiments/benchmarks/bucket.jl +++ b/experiments/benchmarks/bucket.jl @@ -206,7 +206,7 @@ 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"