Merge pull request #2662 from CliMA/ck/nvtx_annot

Add NVTX annotations
CliMA · Feb 9, 2024 · 095dbec · 095dbec
2 parents e73d5fc + 8d80ba2
commit 095dbec
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 102 deletions.
diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
@@ -35,7 +35,7 @@ function flux_accumulation!(integrator)
     return nothing
 end
 
-function cloud_fraction_model_callback!(integrator)
+NVTX.@annotate function cloud_fraction_model_callback!(integrator)
     Y = integrator.u
     p = integrator.p
     set_cloud_fraction!(Y, p, p.atmos.moisture_model, p.atmos.cloud_model)
@@ -183,7 +183,7 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
     return nothing
 end
 
-function save_state_to_disk_func(integrator, output_dir)
+NVTX.@annotate function save_state_to_disk_func(integrator, output_dir)
     (; t, u, p) = integrator
     Y = u
 

diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl
@@ -55,6 +55,44 @@ function hyperdiffusion_cache(Y, hyperdiff::ClimaHyperdiffusion, turbconv_model)
     return (; quantities..., ᶜ∇²u, ᶜ∇²uʲs)
 end
 
+NVTX.@annotate function dss_hyperdiffusion_tendency!(Yₜ, Y, p, t)
+    (; hyperdiff, turbconv_model) = p.atmos
+    buffer = p.hyperdiff.hyperdiffusion_ghost_buffer
+    (; ᶜp, ᶜspecific) = p.precomputed
+    (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff
+    diffuse_tke = use_prognostic_tke(turbconv_model)
+    n = n_mass_flux_subdomains(turbconv_model)
+    if turbconv_model isa PrognosticEDMFX
+        (; ᶜ∇²tke⁰, ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
+    elseif turbconv_model isa DiagnosticEDMFX
+        (; ᶜ∇²tke⁰) = p.hyperdiff
+    end
+    # DSS on Grid scale quantities
+    # Need to split the DSS computation here, because our DSS
+    # operations do not accept Covariant123Vector types
+    Spaces.weighted_dss!(
+        ᶜ∇²u => buffer.ᶜ∇²u,
+        ᶜ∇²specific_energy => buffer.ᶜ∇²specific_energy,
+        (diffuse_tke ? (ᶜ∇²tke⁰ => buffer.ᶜ∇²tke⁰,) : ())...,
+    )
+    if turbconv_model isa PrognosticEDMFX
+        # Need to split the DSS computation here, because our DSS
+        # operations do not accept Covariant123Vector types
+        for j in 1:n
+            @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j))
+            @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j))
+        end
+        Spaces.weighted_dss!(
+            ᶜ∇²uₕʲs => buffer.ᶜ∇²uₕʲs,
+            ᶜ∇²uᵥʲs => buffer.ᶜ∇²uᵥʲs,
+            ᶜ∇²mseʲs => buffer.ᶜ∇²mseʲs,
+        )
+        for j in 1:n
+            @. ᶜ∇²uʲs.:($$j) = C123(ᶜ∇²uₕʲs.:($$j)) + C123(ᶜ∇²uᵥʲs.:($$j))
+        end
+    end
+end
+
 NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
     (; hyperdiff, turbconv_model) = p.atmos
     isnothing(hyperdiff) && return nothing
@@ -83,9 +121,6 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
         (; ᶜtke⁰) = p.precomputed
         (; ᶜ∇²tke⁰) = p.hyperdiff
     end
-    if do_dss
-        buffer = p.hyperdiff.hyperdiffusion_ghost_buffer
-    end
 
     # Grid scale hyperdiffusion
     @. ᶜ∇²u =
@@ -108,36 +143,7 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
         end
     end
 
-    if do_dss
-        NVTX.@range "dss_hyperdiffusion_tendency" color = colorant"green" begin
-            # DSS on Grid scale quantities
-            # Need to split the DSS computation here, because our DSS
-            # operations do not accept Covariant123Vector types
-            Spaces.weighted_dss!(
-                ᶜ∇²u => buffer.ᶜ∇²u,
-                ᶜ∇²specific_energy => buffer.ᶜ∇²specific_energy,
-                (diffuse_tke ? (ᶜ∇²tke⁰ => buffer.ᶜ∇²tke⁰,) : ())...,
-            )
-            if turbconv_model isa PrognosticEDMFX
-                # Need to split the DSS computation here, because our DSS
-                # operations do not accept Covariant123Vector types
-                for j in 1:n
-                    @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j))
-                    @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j))
-                end
-                Spaces.weighted_dss!(
-                    ᶜ∇²uₕʲs => buffer.ᶜ∇²uₕʲs,
-                    ᶜ∇²uᵥʲs => buffer.ᶜ∇²uᵥʲs,
-                    ᶜ∇²mseʲs => buffer.ᶜ∇²mseʲs,
-                )
-                for j in 1:n
-                    @. ᶜ∇²uʲs.:($$j) =
-                        C123(ᶜ∇²uₕʲs.:($$j)) + C123(ᶜ∇²uᵥʲs.:($$j))
-                end
-            end
-
-        end
-    end
+    do_dss && dss_hyperdiffusion_tendency!(Yₜ, Y, p, t)
 
     # re-use to store the curl-curl part
     @. ᶜ∇²u =
@@ -170,6 +176,18 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
         @. Yₜ.c.sgs⁰.ρatke -= ν₄_vorticity * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²tke⁰))
     end
 end
+NVTX.@annotate function dss_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
+    (; turbconv_model) = p.atmos
+    (; ᶜ∇²specific_tracers) = p.hyperdiff
+    buffer = p.hyperdiff.hyperdiffusion_ghost_buffer
+    if !isempty(propertynames(ᶜ∇²specific_tracers))
+        Spaces.weighted_dss!(ᶜ∇²specific_tracers => buffer.ᶜ∇²specific_tracers)
+    end
+    if turbconv_model isa PrognosticEDMFX
+        (; ᶜ∇²q_totʲs) = p.hyperdiff
+        Spaces.weighted_dss!(ᶜ∇²q_totʲs => buffer.ᶜ∇²q_totʲs)
+    end
+end
 
 NVTX.@annotate function tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
     (; hyperdiff, turbconv_model) = p.atmos
@@ -202,18 +220,7 @@ NVTX.@annotate function tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
         end
     end
 
-    if do_dss
-        NVTX.@range "dss_hyperdiffusion_tendency" color = colorant"green" begin
-            if !isempty(propertynames(ᶜ∇²specific_tracers))
-                Spaces.weighted_dss!(
-                    ᶜ∇²specific_tracers => buffer.ᶜ∇²specific_tracers,
-                )
-            end
-            if turbconv_model isa PrognosticEDMFX
-                Spaces.weighted_dss!(ᶜ∇²q_totʲs => buffer.ᶜ∇²q_totʲs)
-            end
-        end
-    end
+    do_dss && dss_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
 
     # TODO: Since we are not applying the limiter to density (or area-weighted
     # density), the mass redistributed by hyperdiffusion will not be conserved

diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl
@@ -263,7 +263,7 @@ end
 Base.similar(A::ImplicitEquationJacobian) = A
 
 # This method specifies how to solve the equation E'(Y) * ΔY = E(Y) for ΔY.
-function ldiv!(
+NVTX.@annotate function ldiv!(
     x::Fields.FieldVector,
     A::ImplicitEquationJacobian,
     b::Fields.FieldVector,
@@ -277,7 +277,7 @@ end
 # This method for ldiv! is called by Krylov.jl from inside ClimaTimeSteppers.jl.
 # See https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a
 # related issue that requires the same workaround.
-function ldiv!(
+NVTX.@annotate function ldiv!(
     x::AbstractVector,
     A::ImplicitEquationJacobian,
     b::AbstractVector,
@@ -293,61 +293,59 @@ _linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b)
 
 # This method specifies how to compute E'(Y), which is referred to as "Wfact" in
 # DiffEqBase.jl.
-function Wfact!(A, Y, p, dtγ, t)
-    NVTX.@range "Wfact!" color = colorant"green" begin
-        # Remove unnecessary values from p to avoid allocations in bycolumn.
-        p′ = (;
-            p.precomputed.ᶜspecific,
-            p.precomputed.ᶜK,
-            p.precomputed.ᶜts,
-            p.precomputed.ᶜp,
-            (
-                p.atmos.precip_model isa Microphysics1Moment ?
-                (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;)
-            )...,
-            p.precomputed.ᶜh_tot,
-            (
-                use_derivative(A.diffusion_flag) ?
-                (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;)
-            )...,
-            (
-                use_derivative(A.diffusion_flag) &&
-                p.atmos.turbconv_model isa AbstractEDMF ?
-                (; p.precomputed.ᶜtke⁰, p.precomputed.ᶜmixing_length) : (;)
-            )...,
-            (
-                use_derivative(A.diffusion_flag) &&
-                p.atmos.turbconv_model isa PrognosticEDMFX ?
-                (; p.precomputed.ᶜρa⁰) : (;)
-            )...,
-            p.core.ᶜΦ,
-            p.core.ᶠgradᵥ_ᶜΦ,
-            p.core.ᶜρ_ref,
-            p.core.ᶜp_ref,
-            p.scratch.ᶜtemp_scalar,
-            p.scratch.∂ᶜK_∂ᶜuₕ,
-            p.scratch.∂ᶜK_∂ᶠu₃,
-            p.scratch.ᶠp_grad_matrix,
-            p.scratch.ᶜadvection_matrix,
-            p.scratch.ᶜdiffusion_h_matrix,
-            p.scratch.ᶜdiffusion_u_matrix,
-            p.dt,
-            p.params,
-            p.atmos,
-            (
-                p.atmos.rayleigh_sponge isa RayleighSponge ?
-                (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;)
-            )...,
-        )
+NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
+    # Remove unnecessary values from p to avoid allocations in bycolumn.
+    p′ = (;
+        p.precomputed.ᶜspecific,
+        p.precomputed.ᶜK,
+        p.precomputed.ᶜts,
+        p.precomputed.ᶜp,
+        (
+            p.atmos.precip_model isa Microphysics1Moment ?
+            (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;)
+        )...,
+        p.precomputed.ᶜh_tot,
+        (
+            use_derivative(A.diffusion_flag) ?
+            (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;)
+        )...,
+        (
+            use_derivative(A.diffusion_flag) &&
+            p.atmos.turbconv_model isa AbstractEDMF ?
+            (; p.precomputed.ᶜtke⁰, p.precomputed.ᶜmixing_length) : (;)
+        )...,
+        (
+            use_derivative(A.diffusion_flag) &&
+            p.atmos.turbconv_model isa PrognosticEDMFX ?
+            (; p.precomputed.ᶜρa⁰) : (;)
+        )...,
+        p.core.ᶜΦ,
+        p.core.ᶠgradᵥ_ᶜΦ,
+        p.core.ᶜρ_ref,
+        p.core.ᶜp_ref,
+        p.scratch.ᶜtemp_scalar,
+        p.scratch.∂ᶜK_∂ᶜuₕ,
+        p.scratch.∂ᶜK_∂ᶠu₃,
+        p.scratch.ᶠp_grad_matrix,
+        p.scratch.ᶜadvection_matrix,
+        p.scratch.ᶜdiffusion_h_matrix,
+        p.scratch.ᶜdiffusion_u_matrix,
+        p.dt,
+        p.params,
+        p.atmos,
+        (
+            p.atmos.rayleigh_sponge isa RayleighSponge ?
+            (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;)
+        )...,
+    )
 
-        # Convert dtγ from a Float64 to an FT.
-        FT = Spaces.undertype(axes(Y.c))
-        dtγ′ = FT(dtγ)
+    # Convert dtγ from a Float64 to an FT.
+    FT = Spaces.undertype(axes(Y.c))
+    dtγ′ = FT(dtγ)
 
-        A.dtγ_ref[] = dtγ′
-        Fields.bycolumn(axes(Y.c)) do colidx
-            update_implicit_equation_jacobian!(A, Y, p′, dtγ′, colidx)
-        end
+    A.dtγ_ref[] = dtγ′
+    Fields.bycolumn(axes(Y.c)) do colidx
+        update_implicit_equation_jacobian!(A, Y, p′, dtγ′, colidx)
     end
 end