Add some NVTX annotations

CliMA · Apr 3, 2024 · d3c1004 · d3c1004
1 parent ba3cf37
commit d3c1004
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 16 deletions.
diff --git 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,15 +14,15 @@ 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)
         compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     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
@@ -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
@@ -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
@@ -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)