From aad6f727ceea68004573237c53faf5f86ac848dc Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Tue, 9 Apr 2024 14:39:40 -0400 Subject: [PATCH] Decorate with at-fused --- .buildkite/pipeline.yml | 5 --- src/ClimaAtmos.jl | 1 + .../precipitation_precomputed_quantities.jl | 34 +++++++------- src/cache/precomputed_quantities.jl | 12 ++--- .../prognostic_edmf_precomputed_quantities.jl | 40 +++++++++-------- src/prognostic_equations/edmfx_entr_detr.jl | 20 +++++---- .../edmfx_precipitation.jl | 44 ++++++++++--------- src/prognostic_equations/edmfx_sgs_flux.jl | 12 +++-- src/prognostic_equations/edmfx_tke.jl | 12 ++--- src/prognostic_equations/hyperdiffusion.jl | 6 ++- .../vertical_diffusion_boundary_layer.jl | 8 +++- 11 files changed, 108 insertions(+), 86 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 29ac02e5f6..0dc191d3dd 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -23,11 +23,6 @@ steps: command: - "echo $$JULIA_DEPOT_PATH" - - echo "--- Instantiate project" - - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project -e 'using Pkg; Pkg.precompile()'" - - "julia --project -e 'using Pkg; Pkg.status()'" - - echo "--- Instantiate examples" - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - "julia --project=examples -e 'using Pkg; Pkg.precompile()'" diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 49a43c0b92..3e8da1779e 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -2,6 +2,7 @@ module ClimaAtmos using NVTX, Colors import Thermodynamics as TD +import ClimaCore.DataLayouts: @fused include(joinpath("parameters", "Parameters.jl")) import .Parameters as CAP diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index f4b25dbbe4..002b4f90f0 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -21,22 +21,24 @@ function set_precipitation_precomputed_quantities!(Y, p, t) cmp = CAP.microphysics_precipitation_params(p.params) - # compute the precipitation specific humidities - @. ᶜqᵣ = qₚ(Y.c.ρq_rai, Y.c.ρ) - @. ᶜqₛ = qₚ(Y.c.ρq_sno, Y.c.ρ) + @fused begin + # compute the precipitation specific humidities + @. ᶜqᵣ = qₚ(Y.c.ρq_rai, Y.c.ρ) + @. ᶜqₛ = qₚ(Y.c.ρq_sno, Y.c.ρ) - # compute the precipitation terminal velocity [m/s] - @. ᶜwᵣ = CM1.terminal_velocity( - cmp.pr, - cmp.tv.rain, - Y.c.ρ, - abs(Y.c.ρq_rai / Y.c.ρ), - ) - @. ᶜwₛ = CM1.terminal_velocity( - cmp.ps, - cmp.tv.snow, - Y.c.ρ, - abs(Y.c.ρq_sno / Y.c.ρ), - ) + # compute the precipitation terminal velocity [m/s] + @. ᶜwᵣ = CM1.terminal_velocity( + cmp.pr, + cmp.tv.rain, + Y.c.ρ, + abs(Y.c.ρq_rai / Y.c.ρ), + ) + @. ᶜwₛ = CM1.terminal_velocity( + cmp.ps, + cmp.tv.snow, + Y.c.ρ, + abs(Y.c.ρq_sno / Y.c.ρ), + ) + end return nothing end diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 3376146682..d62e42ac67 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -472,8 +472,13 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ # TODO: We should think more about these increments before we use them. end - @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ) - @. ᶜp = TD.air_pressure(thermo_params, ᶜts) + (; ᶜh_tot) = p.precomputed + @fused begin + @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ) + @. ᶜp = TD.air_pressure(thermo_params, ᶜts) + @. ᶜh_tot = + TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) + end if turbconv_model isa AbstractEDMF @. p.precomputed.ᶜgradᵥ_θ_virt = @@ -484,9 +489,6 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))) end - (; ᶜh_tot) = p.precomputed - @. ᶜh_tot = TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) - if !isnothing(p.sfc_setup) SurfaceConditions.update_surface_conditions!(Y, p, t) end diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index c63e764cd8..eadf4be788 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -25,27 +25,31 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed - @. ᶜρa⁰ = ρa⁰(Y.c) - @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) - @. ᶜmse⁰ = divide_by_ρa( - Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρ * (ᶜh_tot - ᶜK), - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_tot⁰ = divide_by_ρa( - Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_tot, - Y.c.ρ, - turbconv_model, - ) + @fused begin + @. ᶜρa⁰ = ρa⁰(Y.c) + @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) + @. ᶜmse⁰ = divide_by_ρa( + Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρ * (ᶜh_tot - ᶜK), + Y.c.ρ, + turbconv_model, + ) + @. ᶜq_tot⁰ = divide_by_ρa( + Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρq_tot, + Y.c.ρ, + turbconv_model, + ) + end set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) # @. ᶜK⁰ += ᶜtke⁰ - @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) - @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) + @fused begin + @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) + @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) + end return nothing end diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index 2da439de5a..c7c8fa0f52 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -283,17 +283,19 @@ function edmfx_entr_detr_tendency!( for j in 1:n - @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += - Y.c.sgsʲs.:($$j).ρa[colidx] * - (ᶜentrʲs.:($$j)[colidx] - ᶜdetrʲs.:($$j)[colidx]) + @fused begin + @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += + Y.c.sgsʲs.:($$j).ρa[colidx] * + (ᶜentrʲs.:($$j)[colidx] - ᶜdetrʲs.:($$j)[colidx]) - @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += - ᶜentrʲs.:($$j)[colidx] * - (ᶜmse⁰[colidx] - Y.c.sgsʲs.:($$j).mse[colidx]) + @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += + ᶜentrʲs.:($$j)[colidx] * + (ᶜmse⁰[colidx] - Y.c.sgsʲs.:($$j).mse[colidx]) - @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += - ᶜentrʲs.:($$j)[colidx] * - (ᶜq_tot⁰[colidx] - Y.c.sgsʲs.:($$j).q_tot[colidx]) + @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += + ᶜentrʲs.:($$j)[colidx] * + (ᶜq_tot⁰[colidx] - Y.c.sgsʲs.:($$j).q_tot[colidx]) + end @. Yₜ.f.sgsʲs.:($$j).u₃[colidx] += ᶠinterp(ᶜentrʲs.:($$j)[colidx]) * diff --git a/src/prognostic_equations/edmfx_precipitation.jl b/src/prognostic_equations/edmfx_precipitation.jl index 0d68d6f18d..d424d39732 100644 --- a/src/prognostic_equations/edmfx_precipitation.jl +++ b/src/prognostic_equations/edmfx_precipitation.jl @@ -28,20 +28,22 @@ function edmfx_precipitation_tendency!( for j in 1:n - @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += - Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx] + @fused begin + @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += + Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx] - @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += - ᶜSqₜᵖʲs.:($$j)[colidx] * ( - e_tot_0M_precipitation_sources_helper( - thermo_params, - ᶜtsʲs.:($$j)[colidx], - ᶜΦ[colidx], - ) - TD.internal_energy(thermo_params, ᶜtsʲs.:($$j)[colidx]) - ) + @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += + ᶜSqₜᵖʲs.:($$j)[colidx] * ( + e_tot_0M_precipitation_sources_helper( + thermo_params, + ᶜtsʲs.:($$j)[colidx], + ᶜΦ[colidx], + ) - TD.internal_energy(thermo_params, ᶜtsʲs.:($$j)[colidx]) + ) - @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += - ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx]) + @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += + ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx]) + end end return nothing end @@ -62,16 +64,18 @@ function edmfx_precipitation_tendency!( for j in 1:n - @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += - Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx] + @fused begin + @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += + Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx] - @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += - ᶜSeₜᵖʲs.:($$j)[colidx] - - ᶜSqₜᵖʲs.:($$j)[colidx] * - TD.internal_energy(thp, ᶜtsʲs.:($$j)[colidx]) + @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += + ᶜSeₜᵖʲs.:($$j)[colidx] - + ᶜSqₜᵖʲs.:($$j)[colidx] * + TD.internal_energy(thp, ᶜtsʲs.:($$j)[colidx]) - @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += - ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx]) + @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += + ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx]) + end end return nothing end diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 10dba82f65..8a63b0e49e 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -253,8 +253,10 @@ function edmfx_sgs_diffusive_flux_tendency!( ) @. ᶜρχₜ_diffusion[colidx] = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜq_tot⁰[colidx]))) - @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx] - @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + @fused begin + @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx] + @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + end end # momentum @@ -339,8 +341,10 @@ function edmfx_sgs_diffusive_flux_tendency!( @. ᶜρχₜ_diffusion[colidx] = ᶜdivᵥ_ρq_tot( -(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜspecific.q_tot[colidx])), ) - @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx] - @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + @fused begin + @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx] + @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + end end # momentum diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index 477a2a36ee..8d1e1f6e3c 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -36,11 +36,13 @@ function edmfx_tke_tendency!( if use_prognostic_tke(turbconv_model) # shear production - @. Yₜ.c.sgs⁰.ρatke[colidx] += - 2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx] - # buoyancy production - @. Yₜ.c.sgs⁰.ρatke[colidx] -= - ᶜρa⁰[colidx] * ᶜK_h[colidx] * ᶜlinear_buoygrad[colidx] + @fused begin + @. Yₜ.c.sgs⁰.ρatke[colidx] += + 2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx] + # buoyancy production + @. Yₜ.c.sgs⁰.ρatke[colidx] -= + ᶜρa⁰[colidx] * ᶜK_h[colidx] * ᶜlinear_buoygrad[colidx] + end # entrainment and detraiment # using ᶜu⁰ and local geometry results in allocation for j in 1:n diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index cf046c05d9..877962787d 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -92,8 +92,10 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) C123(wgradₕ(divₕ(p.precomputed.ᶜuʲs.:($$j)))) - C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜuʲs.:($$j))))) @. ᶜ∇²mseʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).mse)) - @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j)) - @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j)) + @fused begin + @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j)) + @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j)) + end end end end diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index 39f03bb031..1f85816fe1 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -95,9 +95,13 @@ function vertical_diffusion_boundary_layer_tendency!( ᶠgradᵥ(ᶜχ[colidx]) ), ) - @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx] if !(χ_name in (:q_rai, :q_sno)) - @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + @fused begin + @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx] + @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + end + else + @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx] end end end