switch to using NVTX annotate instead of range

CliMA · Sep 14, 2023 · 0bef090 · 0bef090
1 parent f751e69
commit 0bef090
Show file tree

Hide file tree

Showing 10 changed files with 75 additions and 106 deletions.
diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -275,7 +275,7 @@ Note: If you need to use any of the precomputed quantities, please call this
 function instead of recomputing the value yourself. Otherwise, it will be
 difficult to ensure that the duplicated computations are consistent.
 """
-function set_precomputed_quantities!(Y, p, t)
+NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     (; energy_form, moisture_model, turbconv_model) = p.atmos
     thermo_params = CAP.thermodynamics_params(p.params)
     n = n_mass_flux_subdomains(turbconv_model)

diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
@@ -16,18 +16,16 @@ import ClimaCore.Fields: ColumnField
 
 include("callback_helpers.jl")
 
-function dss_callback!(integrator)
+NVTX.@annotate function dss_callback!(integrator)
     Y = integrator.u
     ghost_buffer = integrator.p.ghost_buffer
     if integrator.p.do_dss
-        NVTX.@range "dss callback" color = colorant"yellow" begin
-            Spaces.weighted_dss_start2!(Y.c, ghost_buffer.c)
-            Spaces.weighted_dss_start2!(Y.f, ghost_buffer.f)
-            Spaces.weighted_dss_internal2!(Y.c, ghost_buffer.c)
-            Spaces.weighted_dss_internal2!(Y.f, ghost_buffer.f)
-            Spaces.weighted_dss_ghost2!(Y.c, ghost_buffer.c)
-            Spaces.weighted_dss_ghost2!(Y.f, ghost_buffer.f)
-        end
+        Spaces.weighted_dss_start2!(Y.c, ghost_buffer.c)
+        Spaces.weighted_dss_start2!(Y.f, ghost_buffer.f)
+        Spaces.weighted_dss_internal2!(Y.c, ghost_buffer.c)
+        Spaces.weighted_dss_internal2!(Y.f, ghost_buffer.f)
+        Spaces.weighted_dss_ghost2!(Y.c, ghost_buffer.c)
+        Spaces.weighted_dss_ghost2!(Y.f, ghost_buffer.f)
     end
     return nothing
 end
@@ -51,7 +49,7 @@ function flux_accumulation!(integrator)
     return nothing
 end
 
-function turb_conv_affect_filter!(integrator)
+NVTX.@annotate function turb_conv_affect_filter!(integrator)
     p = integrator.p
     (; edmf_cache) = p
     (; edmf, param_set) = edmf_cache
@@ -73,7 +71,7 @@ function turb_conv_affect_filter!(integrator)
     SciMLBase.u_modified!(integrator, false)
 end
 
-function rrtmgp_model_callback!(integrator)
+NVTX.@annotate function rrtmgp_model_callback!(integrator)
     Y = integrator.u
     p = integrator.p
     t = integrator.t
@@ -210,7 +208,7 @@ function rrtmgp_model_callback!(integrator)
     return nothing
 end
 
-function compute_diagnostics(integrator)
+NVTX.@annotate function compute_diagnostics(integrator)
     (; t, u, p) = integrator
     Y = u
     (; params, env_thermo_quad) = p

diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl
@@ -6,7 +6,7 @@ using LinearAlgebra: ×, dot
 import ClimaCore.Fields as Fields
 import ClimaCore.Geometry as Geometry
 
-function horizontal_advection_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function horizontal_advection_tendency!(Yₜ, Y, p, t)
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
     (; ᶜu, ᶜK, ᶜp, ᶜΦ, ᶜp_ref) = p
     if p.atmos.turbconv_model isa AbstractEDMF
@@ -51,7 +51,7 @@ function horizontal_advection_tendency!(Yₜ, Y, p, t)
     return nothing
 end
 
-function horizontal_tracer_advection_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function horizontal_tracer_advection_tendency!(Yₜ, Y, p, t)
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
     (; ᶜu) = p
     if p.atmos.turbconv_model isa EDMFX
@@ -74,7 +74,7 @@ function horizontal_tracer_advection_tendency!(Yₜ, Y, p, t)
     return nothing
 end
 
-function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     (; turbconv_model) = p.atmos
     n = n_prognostic_mass_flux_subdomains(turbconv_model)
     advect_tke = use_prognostic_tke(turbconv_model)

diff --git a/src/prognostic_equations/dss.jl b/src/prognostic_equations/dss.jl
@@ -12,7 +12,7 @@ using ClimaCore.Utilities: half
 
 import ClimaCore.Fields: ColumnField
 
-function dss!(Y, p, t)
+NVTX.@annotate function dss!(Y, p, t)
     if p.do_dss
         Spaces.weighted_dss_start2!(Y.c, p.ghost_buffer.c)
         Spaces.weighted_dss_start2!(Y.f, p.ghost_buffer.f)

diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl
@@ -49,7 +49,7 @@ function hyperdiffusion_cache(Y, atmos, do_dss)
     return (; quantities..., ᶜ∇²u, ᶜ∇²uʲs)
 end
 
-function hyperdiffusion_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
     (; hyperdiff, turbconv_model) = p.atmos
     isnothing(hyperdiff) && return nothing
 
@@ -187,7 +187,7 @@ function hyperdiffusion_tendency!(Yₜ, Y, p, t)
     end
 end
 
-function tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
     (; hyperdiff, turbconv_model) = p.atmos
     isnothing(hyperdiff) && return nothing
 

diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -4,32 +4,21 @@
 
 import ClimaCore: Fields, Geometry
 
-function implicit_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
     fill_with_nans!(p)
-    NVTX.@range "implicit tendency" color = colorant"yellow" begin
-        Yₜ .= zero(eltype(Yₜ))
-        NVTX.@range "precomputed quantities" color = colorant"orange" begin
-            set_precomputed_quantities!(Y, p, t)
+    Yₜ .= zero(eltype(Yₜ))
+    set_precomputed_quantities!(Y, p, t)
+    Fields.bycolumn(axes(Y.c)) do colidx
+        implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
+        if p.turbconv_model isa TurbulenceConvection.EDMFModel
+            implicit_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
         end
-        Fields.bycolumn(axes(Y.c)) do colidx
-            implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
-            if p.turbconv_model isa TurbulenceConvection.EDMFModel
-                implicit_sgs_flux_tendency!(
-                    Yₜ,
-                    Y,
-                    p,
-                    t,
-                    colidx,
-                    p.turbconv_model,
-                )
-            end
-            # NOTE: All ρa tendencies should be applied before calling this function
-            pressure_work_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model)
+        # NOTE: All ρa tendencies should be applied before calling this function
+        pressure_work_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model)
 
-            # NOTE: This will zero out all monmentum tendencies in the edmfx advection test
-            # please DO NOT add additional velocity tendencies after this function
-            zero_velocity_tendency!(Yₜ, Y, p, t, colidx)
-        end
+        # NOTE: This will zero out all monmentum tendencies in the edmfx advection test
+        # please DO NOT add additional velocity tendencies after this function
+        zero_velocity_tendency!(Yₜ, Y, p, t, colidx)
     end
     return nothing
 end

diff --git a/src/prognostic_equations/implicit/schur_complement_W.jl b/src/prognostic_equations/implicit/schur_complement_W.jl
@@ -240,7 +240,7 @@ function LinearAlgebra.ldiv!(x, A::SchurComplementW, b)
     x .= A.temp2
 end
 
-function LinearAlgebra.ldiv!(
+NVTX.@annotate function LinearAlgebra.ldiv!(
     x::Fields.FieldVector,
     A::SchurComplementW,
     b::Fields.FieldVector,
@@ -257,35 +257,33 @@ function LinearAlgebra.ldiv!(
             changing the ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode or the energy variable."
         @warn str maxlog = 1
     end
-    NVTX.@range "linsolve" color = colorant"lime" begin
-        # Initialize x to -b, which correctly sets all the components of x that
-        # correspond to variables without implicit tendencies.
-        @. x = -b
-        # TODO: Figure out why moving this into _ldiv_serial! results in a lot
-        # of allocations for EDMFX.
-
-        # Compute Schur complement
-        Fields.bycolumn(axes(x.c)) do colidx
-            _ldiv_serial!(
-                A,
-                x.c[colidx],
-                x.f[colidx],
-                b.c[colidx],
-                b.f[colidx],
-                dtγ,
-                transform,
-                cond,
-                ∂ᶜρₜ∂ᶠ𝕄[colidx],
-                ∂ᶜ𝔼ₜ∂ᶠ𝕄[colidx],
-                ∂ᶠ𝕄ₜ∂ᶜ𝔼[colidx],
-                ∂ᶠ𝕄ₜ∂ᶜρ[colidx],
-                ∂ᶠ𝕄ₜ∂ᶠ𝕄[colidx],
-                ∂ᶜ𝕋ₜ∂ᶠ𝕄_field[colidx],
-                isnothing(∂ᶜTCₜ∂ᶜTC) ? nothing : ∂ᶜTCₜ∂ᶜTC[colidx],
-                isnothing(∂ᶠTCₜ∂ᶠTC) ? nothing : ∂ᶠTCₜ∂ᶠTC[colidx],
-                S[colidx],
-            )
-        end
+    # Initialize x to -b, which correctly sets all the components of x that
+    # correspond to variables without implicit tendencies.
+    @. x = -b
+    # TODO: Figure out why moving this into _ldiv_serial! results in a lot
+    # of allocations for EDMFX.
+
+    # Compute Schur complement
+    Fields.bycolumn(axes(x.c)) do colidx
+        _ldiv_serial!(
+            A,
+            x.c[colidx],
+            x.f[colidx],
+            b.c[colidx],
+            b.f[colidx],
+            dtγ,
+            transform,
+            cond,
+            ∂ᶜρₜ∂ᶠ𝕄[colidx],
+            ∂ᶜ𝔼ₜ∂ᶠ𝕄[colidx],
+            ∂ᶠ𝕄ₜ∂ᶜ𝔼[colidx],
+            ∂ᶠ𝕄ₜ∂ᶜρ[colidx],
+            ∂ᶠ𝕄ₜ∂ᶠ𝕄[colidx],
+            ∂ᶜ𝕋ₜ∂ᶠ𝕄_field[colidx],
+            isnothing(∂ᶜTCₜ∂ᶜTC) ? nothing : ∂ᶜTCₜ∂ᶜTC[colidx],
+            isnothing(∂ᶠTCₜ∂ᶠTC) ? nothing : ∂ᶠTCₜ∂ᶠTC[colidx],
+            S[colidx],
+        )
     end
 end
 

diff --git a/src/prognostic_equations/implicit/wfact.jl b/src/prognostic_equations/implicit/wfact.jl
@@ -58,13 +58,11 @@ function validate_flags!(Y, flags, energy_upwinding)
         )
 end
 
-function Wfact!(W, Y, p, dtγ, t)
-    NVTX.@range "Wfact!" color = colorant"green" begin
-        fill_with_nans!(p)
-        set_precomputed_quantities!(Y, p, t)
-        Fields.bycolumn(axes(Y.c)) do colidx
-            Wfact!(W, Y, p, dtγ, t, colidx)
-        end
+NVTX.@annotate function Wfact!(W, Y, p, dtγ, t)
+    fill_with_nans!(p)
+    set_precomputed_quantities!(Y, p, t)
+    Fields.bycolumn(axes(Y.c)) do colidx
+        Wfact!(W, Y, p, dtγ, t, colidx)
     end
 end
 

diff --git a/src/prognostic_equations/limited_tendencies.jl b/src/prognostic_equations/limited_tendencies.jl
@@ -12,17 +12,15 @@ using ClimaCore.Utilities: half
 
 import ClimaCore.Fields: ColumnField
 
-function limited_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function limited_tendency!(Yₜ, Y, p, t)
     Yₜ .= zero(eltype(Yₜ))
     set_precomputed_quantities!(Y, p, t)
     horizontal_tracer_advection_tendency!(Yₜ, Y, p, t)
-    NVTX.@range "tracer hyperdiffusion tendency" color = colorant"yellow" begin
-        tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
-    end
+    tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
     return nothing
 end
 
-function limiters_func!(Y, p, t, ref_Y)
+NVTX.@annotate function limiters_func!(Y, p, t, ref_Y)
     (; limiter) = p
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
     if !isnothing(limiter)

diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl
@@ -1,31 +1,19 @@
 
-function remaining_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function remaining_tendency!(Yₜ, Y, p, t)
     fill_with_nans!(p)
-    NVTX.@range "remaining tendency" color = colorant"yellow" begin
-        Yₜ .= zero(eltype(Yₜ))
-        NVTX.@range "precomputed quantities" color = colorant"orange" begin
-            set_precomputed_quantities!(Y, p, t)
-        end
-        NVTX.@range "horizontal" color = colorant"orange" begin
-            horizontal_advection_tendency!(Yₜ, Y, p, t)
-            NVTX.@range "hyperdiffusion tendency" color = colorant"yellow" begin
-                hyperdiffusion_tendency!(Yₜ, Y, p, t)
-            end
-        end
-        NVTX.@range "vertical" color = colorant"orange" begin
-            explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
-        end
-        NVTX.@range "additional_tendency!" color = colorant"orange" begin
-            additional_tendency!(Yₜ, Y, p, t)
-        end
-        NVTX.@range "dss_remaining_tendency" color = colorant"blue" begin
-            dss!(Yₜ, p, t)
-        end
+    Yₜ .= zero(eltype(Yₜ))
+    set_precomputed_quantities!(Y, p, t)
+    NVTX.@range "horizontal" color = colorant"orange" begin
+        horizontal_advection_tendency!(Yₜ, Y, p, t)
+        hyperdiffusion_tendency!(Yₜ, Y, p, t)
     end
+    explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
+    additional_tendency!(Yₜ, Y, p, t)
+    dss!(Yₜ, p, t)
     return Yₜ
 end
 
-function additional_tendency!(Yₜ, Y, p, t)
+NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     viscous_sponge_tendency!(Yₜ, Y, p, t, p.atmos.viscous_sponge)
     surface_temp_tendency!(Yₜ, Y, p, t, p.atmos.surface_model)