Skip to content

Commit

Permalink
Add more NVTX annotations
Browse files Browse the repository at this point in the history
  • Loading branch information
@charleskawczynski
charleskawczynski committed Feb 8, 2024
1 parent 8cc7f03 commit 537ec64
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 104 deletions.
101 changes: 53 additions & 48 deletions src/prognostic_equations/hyperdiffusion.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,43 @@ function hyperdiffusion_cache(Y, hyperdiff::ClimaHyperdiffusion, turbconv_model)
return (; quantities..., ᶜ∇²u, ᶜ∇²uʲs)
end

NVTX.@annotate function dss_hyperdiffusion_tendency!(Yₜ, Y, p, t)
buffer = p.hyperdiff.hyperdiffusion_ghost_buffer
(; hyperdiff, turbconv_model) = p.atmos
n = n_mass_flux_subdomains(turbconv_model)
diffuse_tke = use_prognostic_tke(turbconv_model)
if turbconv_model isa PrognosticEDMFX || turbconv_model isa DiagnosticEDMFX
(; ᶜ∇²tke⁰) = p.hyperdiff
(; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
end
(; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff
# 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
return nothing
end

NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
(; hyperdiff, turbconv_model) = p.atmos
isnothing(hyperdiff) && return nothing
Expand Down Expand Up @@ -83,9 +120,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 =
Expand All @@ -108,36 +142,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 =
Expand Down Expand Up @@ -171,6 +176,20 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
end
end

NVTX.@annotate function dss_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
(; hyperdiff, turbconv_model) = p.atmos
(; ᶜ∇²specific_tracers) = hyperdiff
buffer = 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
return nothing
end

NVTX.@annotate function tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
(; hyperdiff, turbconv_model) = p.atmos
isnothing(hyperdiff) && return nothing
Expand All @@ -187,9 +206,6 @@ NVTX.@annotate function tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
if turbconv_model isa PrognosticEDMFX
(; ᶜ∇²q_totʲs) = p.hyperdiff
end
if do_dss
buffer = p.hyperdiff.hyperdiffusion_ghost_buffer
end

for χ_name in propertynames(ᶜ∇²specific_tracers)
@. ᶜ∇²specific_tracers.:($$χ_name) = wdivₕ(gradₕ(ᶜspecific.:($$χ_name)))
Expand All @@ -202,18 +218,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
Expand Down
110 changes: 54 additions & 56 deletions src/prognostic_equations/implicit/implicit_solver.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -293,65 +293,63 @@ _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

function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
NVTX.@annotate function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
(; matrix, diffusion_flag, topography_flag) = A
(; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p
(; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, ᶠp_grad_matrix, ᶜadvection_matrix) = p
Expand Down

0 comments on commit 537ec64

Please sign in to comment.