From 0bef09060058d6326756c9dabe06b1abdc62e24d Mon Sep 17 00:00:00 2001 From: Simon Byrne Date: Thu, 14 Sep 2023 14:42:06 -0700 Subject: [PATCH] switch to using NVTX annotate instead of range --- src/cache/precomputed_quantities.jl | 2 +- src/callbacks/callbacks.jl | 22 ++++--- src/prognostic_equations/advection.jl | 6 +- src/prognostic_equations/dss.jl | 2 +- src/prognostic_equations/hyperdiffusion.jl | 4 +- .../implicit/implicit_tendency.jl | 35 ++++------- .../implicit/schur_complement_W.jl | 58 +++++++++---------- src/prognostic_equations/implicit/wfact.jl | 12 ++-- .../limited_tendencies.jl | 8 +-- .../remaining_tendency.jl | 32 ++++------ 10 files changed, 75 insertions(+), 106 deletions(-) diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 1a993c663b..a87ee32c6a 100644 --- 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 index 1a851275d5..c73da82cbb 100644 --- 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 index d511fe388a..aff688bd18 100644 --- 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 index e865eedc5e..4e0455b074 100644 --- 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 index 11051675df..27d396fd1a 100644 --- 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 index 2d9ae45170..de6dadd30a 100644 --- 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 index e2b5dcb27e..9ae95ed5bc 100644 --- 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 index 4c029351be..8144801093 100644 --- 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 index f9ec8c735a..9d10faafc7 100644 --- 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 index a4bc049319..a41f21a10a 100644 --- 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)