From 8d80ba29e862c500298ff4e2f64764492bad3ea3 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 8 Feb 2024 21:23:04 -0500 Subject: [PATCH] Add NVTX annotations --- src/callbacks/callbacks.jl | 4 +- src/prognostic_equations/hyperdiffusion.jl | 97 ++++++++-------- .../implicit/implicit_solver.jl | 108 +++++++++--------- 3 files changed, 107 insertions(+), 102 deletions(-) diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl index 3963489b2e..f4b71196ec 100644 --- 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 index 2c54b167d4..c198018661 100644 --- 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 index e9ec942777..65bba0bd0c 100644 --- 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