From d3c1004e9b34844c81e3c42867a224b9306fc355 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 3 Apr 2024 15:57:47 -0400 Subject: [PATCH] Add some NVTX annotations --- src/cache/cloud_fraction.jl | 7 +++-- .../diagnostic_edmf_precomputed_quantities.jl | 31 ++++++++++++++----- .../prognostic_edmf_precomputed_quantities.jl | 25 ++++++++++++--- src/prognostic_equations/gm_sgs_closures.jl | 3 +- 4 files changed, 50 insertions(+), 16 deletions(-) diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index 6c85ac85d74..bc6204d9ba1 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -1,3 +1,4 @@ +import NVTX import StaticArrays as SA import ClimaCore.RecursiveApply: rzero, ⊞, ⊠ @@ -13,7 +14,7 @@ end """ Compute the grid scale cloud fraction based on sub-grid scale properties """ -function set_cloud_fraction!(Y, p, ::DryModel, _) +NVTX.@annotate function set_cloud_fraction!(Y, p, ::DryModel, _) (; ᶜmixing_length) = p.precomputed (; turbconv_model) = p.atmos if isnothing(turbconv_model) @@ -21,7 +22,7 @@ function set_cloud_fraction!(Y, p, ::DryModel, _) end @. p.precomputed.ᶜcloud_fraction = 0 end -function set_cloud_fraction!( +NVTX.@annotate function set_cloud_fraction!( Y, p, ::Union{EquilMoistModel, NonEquilMoistModel}, @@ -36,7 +37,7 @@ function set_cloud_fraction!( end @. ᶜcloud_fraction = ifelse(TD.has_condensate(thermo_params, ᶜts), 1, 0) end -function set_cloud_fraction!( +NVTX.@annotate function set_cloud_fraction!( Y, p, ::Union{EquilMoistModel, NonEquilMoistModel}, diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index f6f6982a33d..d94c8671d0a 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -1,10 +1,11 @@ ##### ##### Precomputed quantities ##### +import NVTX import Thermodynamics as TD import ClimaCore: Spaces, Fields, RecursiveApply -function kinetic_energy!( +NVTX.@annotate function kinetic_energy!( K_level, uₕ_level, u³_halflevel, @@ -28,7 +29,7 @@ function kinetic_energy!( ) / 2 end -function set_diagnostic_edmfx_draft_quantities_level!( +NVTX.@annotate function set_diagnostic_edmfx_draft_quantities_level!( thermo_params, ts_level, ρ_level, @@ -50,7 +51,7 @@ function set_diagnostic_edmfx_draft_quantities_level!( return nothing end -function set_diagnostic_edmfx_env_quantities_level!( +NVTX.@annotate function set_diagnostic_edmfx_env_quantities_level!( ρ_level, ρaʲs_level, u³_halflevel, @@ -86,7 +87,11 @@ end Updates the bottom boundary conditions in precomputed quantities stored in `p` for diagnostic edmfx. """ -function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(Y, p, t) +NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( + Y, + p, + t, +) (; turbconv_model) = p.atmos FT = eltype(Y) n = n_mass_flux_subdomains(turbconv_model) @@ -207,7 +212,11 @@ function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(Y, p, t) return nothing end -function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t) +NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( + Y, + p, + t, +) (; turbconv_model, precip_model) = p.atmos FT = eltype(Y) n = n_mass_flux_subdomains(turbconv_model) @@ -693,7 +702,11 @@ end Updates the top boundary condition of precomputed quantities stored in `p` for diagnostic edmfx. """ -function set_diagnostic_edmf_precomputed_quantities_top_bc!(Y, p, t) +NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_top_bc!( + Y, + p, + t, +) n = n_mass_flux_subdomains(p.atmos.turbconv_model) (; ᶜentrʲs, ᶜdetrʲs, ᶜS_q_totʲs, ᶜS_e_totʲs_helper) = p.precomputed (; ᶠu³⁰, ᶠu³ʲs, ᶜuʲs, ᶠnh_pressure³ʲs) = p.precomputed @@ -741,7 +754,11 @@ end Updates the environment closures in precomputed quantities stored in `p` for diagnostic edmfx. """ -function set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t) +NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!( + Y, + p, + t, +) (; moisture_model, turbconv_model, precip_model) = p.atmos n = n_mass_flux_subdomains(turbconv_model) ᶜz = Fields.coordinate_field(Y.c).z diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index a00c7fa4d81..fcc17e0f376 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -1,6 +1,7 @@ ##### ##### Precomputed quantities ##### +import NVTX import Thermodynamics as TD import ClimaCore: Spaces, Fields @@ -9,7 +10,12 @@ import ClimaCore: Spaces, Fields Updates the edmf environment precomputed quantities stored in `p` for edmfx. """ -function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t) +NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( + Y, + p, + ᶠuₕ³, + t, +) @assert !(p.atmos.moisture_model isa DryModel) thermo_params = CAP.thermodynamics_params(p.params) @@ -49,7 +55,12 @@ end Updates the draft thermo state and boundary conditions precomputed quantities stored in `p` for edmfx. """ -function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t) +NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!( + Y, + p, + ᶠuₕ³, + t, +) (; moisture_model, turbconv_model) = p.atmos #EDMFX BCs only support total energy as state variable @assert !(moisture_model isa DryModel) @@ -161,7 +172,11 @@ end Updates the precomputed quantities stored in `p` for edmfx closures. """ -function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t) +NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!( + Y, + p, + t, +) (; moisture_model, turbconv_model) = p.atmos @assert !(moisture_model isa DryModel) @@ -339,7 +354,7 @@ function set_prognostic_edmf_precomputed_quantities_precipitation!( ) return nothing end -function set_prognostic_edmf_precomputed_quantities_precipitation!( +NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation!( Y, p, ::Microphysics0Moment, @@ -374,7 +389,7 @@ function set_prognostic_edmf_precomputed_quantities_precipitation!( ) return nothing end -function set_prognostic_edmf_precomputed_quantities_precipitation!( +NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation!( Y, p, ::Microphysics1Moment, diff --git a/src/prognostic_equations/gm_sgs_closures.jl b/src/prognostic_equations/gm_sgs_closures.jl index 07a752e7fb2..6635d2a4ce6 100644 --- a/src/prognostic_equations/gm_sgs_closures.jl +++ b/src/prognostic_equations/gm_sgs_closures.jl @@ -2,6 +2,7 @@ ##### Grid-mean SGS closures (mixing length) ##### +import NVTX import ClimaCore.Geometry as Geometry import ClimaCore.Fields as Fields @@ -22,7 +23,7 @@ function smagorinsky_lilly_length(c_smag, N_eff, dz, Pr, ϵ_st) c_smag * dz end -function compute_gm_mixing_length!(ᶜmixing_length, Y, p) +NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p) (; params) = p thermo_params = CAP.thermodynamics_params(params)