diff --git a/examples/hybrid/hyperdiffusion.jl b/examples/hybrid/hyperdiffusion.jl index 392e23e133..ce89d804ba 100644 --- a/examples/hybrid/hyperdiffusion.jl +++ b/examples/hybrid/hyperdiffusion.jl @@ -30,154 +30,152 @@ function hyperdiffusion_cache( end function hyperdiffusion_tendency!(Yₜ, Y, P, t) - if P.use_tempest_mode - hyperdiffusion_tendency_tempest!(Yₜ, Y, P, t) - else - hyperdiffusion_tendency_clima!(Yₜ, Y, P, t) + @nvtx "hyperdiffusion tendency" color = colorant"yellow" begin + if P.use_tempest_mode + hyperdiffusion_tendency_tempest!(Yₜ, Y, P, t) + else + hyperdiffusion_tendency_clima!(Yₜ, Y, P, t) + end end end function hyperdiffusion_tendency_clima!(Yₜ, Y, p, t) - @nvtx "hyperdiffusion tendency" color = colorant"yellow" begin - ᶜρ = Y.c.ρ - ᶜuₕ = Y.c.uₕ - (; ᶜp, ᶜχ, ᶜχuₕ) = p # assume ᶜp has been updated - (; - ghost_buffer, - κ₄, - divergence_damping_factor, - use_tempest_mode, - disable_qt_hyperdiffusion, - ) = p - point_type = eltype(Fields.local_geometry_field(axes(Y.c)).coordinates) - is_ρq_tot = :ρq_tot in propertynames(Y.c) && !disable_qt_hyperdiffusion - is_2d_pt = point_type <: Geometry.Abstract2DPoint - is_3d_pt = !is_2d_pt - - ᶜρs = - :ρe_tot in propertynames(Y.c) ? Y.c.ρe_tot : - (:ρe_int in propertynames(Y.c) ? Y.c.ρe_int : Y.c.ρθ) - ᵗρs = - :ρe_tot in propertynames(Y.c) ? Yₜ.c.ρe_tot : - (:ρe_int in propertynames(Y.c) ? Yₜ.c.ρe_int : Yₜ.c.ρθ) - - (:ρθ in propertynames(Y.c)) && (@. ᶜχ = wdivₕ(gradₕ(ᶜρs / ᶜρ))) - !(:ρθ in propertynames(Y.c)) && (@. ᶜχ = wdivₕ(gradₕ((ᶜρs + ᶜp) / ᶜρ))) - - if is_ρq_tot - (; ᶜχρq_tot) = p - @. ᶜχρq_tot = wdivₕ(gradₕ(Y.c.ρq_tot / ᶜρ)) - end + ᶜρ = Y.c.ρ + ᶜuₕ = Y.c.uₕ + (; ᶜp, ᶜχ, ᶜχuₕ) = p # assume ᶜp has been updated + (; + ghost_buffer, + κ₄, + divergence_damping_factor, + use_tempest_mode, + disable_qt_hyperdiffusion, + ) = p + point_type = eltype(Fields.local_geometry_field(axes(Y.c)).coordinates) + is_ρq_tot = :ρq_tot in propertynames(Y.c) && !disable_qt_hyperdiffusion + is_2d_pt = point_type <: Geometry.Abstract2DPoint + is_3d_pt = !is_2d_pt + + ᶜρs = + :ρe_tot in propertynames(Y.c) ? Y.c.ρe_tot : + (:ρe_int in propertynames(Y.c) ? Y.c.ρe_int : Y.c.ρθ) + ᵗρs = + :ρe_tot in propertynames(Y.c) ? Yₜ.c.ρe_tot : + (:ρe_int in propertynames(Y.c) ? Yₜ.c.ρe_int : Yₜ.c.ρθ) + + (:ρθ in propertynames(Y.c)) && (@. ᶜχ = wdivₕ(gradₕ(ᶜρs / ᶜρ))) + !(:ρθ in propertynames(Y.c)) && (@. ᶜχ = wdivₕ(gradₕ((ᶜρs + ᶜp) / ᶜρ))) + + if is_ρq_tot + (; ᶜχρq_tot) = p + @. ᶜχρq_tot = wdivₕ(gradₕ(Y.c.ρq_tot / ᶜρ)) + end - is_3d_pt && (@. ᶜχuₕ = - wgradₕ(divₕ(ᶜuₕ)) - Geometry.project( - Geometry.Covariant12Axis(), - wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))), - )) - is_2d_pt && (@. ᶜχuₕ = - Geometry.project(Geometry.Covariant12Axis(), wgradₕ(divₕ(ᶜuₕ)))) - - @nvtx "dss_hyperdiffusion_tendency" color = colorant"green" begin - Spaces.weighted_dss_start!(ᶜχ, ghost_buffer.χ) - is_ρq_tot && - (Spaces.weighted_dss_start!(ᶜχρq_tot, ghost_buffer.ᶜχρq_tot)) - Spaces.weighted_dss_start!(ᶜχuₕ, ghost_buffer.χuₕ) - - Spaces.weighted_dss_internal!(ᶜχ, ghost_buffer.χ) - is_ρq_tot && - (Spaces.weighted_dss_internal!(ᶜχρq_tot, ghost_buffer.ᶜχρq_tot)) - Spaces.weighted_dss_internal!(ᶜχuₕ, ghost_buffer.χuₕ) - - Spaces.weighted_dss_ghost!(ᶜχ, ghost_buffer.χ) - is_ρq_tot && - (Spaces.weighted_dss_ghost!(ᶜχρq_tot, ghost_buffer.ᶜχρq_tot)) - Spaces.weighted_dss_ghost!(ᶜχuₕ, ghost_buffer.χuₕ) - end + is_3d_pt && (@. ᶜχuₕ = + wgradₕ(divₕ(ᶜuₕ)) - Geometry.project( + Geometry.Covariant12Axis(), + wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))), + )) + is_2d_pt && (@. ᶜχuₕ = + Geometry.project(Geometry.Covariant12Axis(), wgradₕ(divₕ(ᶜuₕ)))) + + @nvtx "dss_hyperdiffusion_tendency" color = colorant"green" begin + Spaces.weighted_dss_start!(ᶜχ, ghost_buffer.χ) + is_ρq_tot && + (Spaces.weighted_dss_start!(ᶜχρq_tot, ghost_buffer.ᶜχρq_tot)) + Spaces.weighted_dss_start!(ᶜχuₕ, ghost_buffer.χuₕ) + + Spaces.weighted_dss_internal!(ᶜχ, ghost_buffer.χ) + is_ρq_tot && + (Spaces.weighted_dss_internal!(ᶜχρq_tot, ghost_buffer.ᶜχρq_tot)) + Spaces.weighted_dss_internal!(ᶜχuₕ, ghost_buffer.χuₕ) + + Spaces.weighted_dss_ghost!(ᶜχ, ghost_buffer.χ) + is_ρq_tot && + (Spaces.weighted_dss_ghost!(ᶜχρq_tot, ghost_buffer.ᶜχρq_tot)) + Spaces.weighted_dss_ghost!(ᶜχuₕ, ghost_buffer.χuₕ) + end - @. ᵗρs -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχ)) - if is_ρq_tot - @. Yₜ.c.ρq_tot -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) - @. Yₜ.c.ρ -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) - end - if is_3d_pt - @. Yₜ.c.uₕ -= - κ₄ * ( - divergence_damping_factor * wgradₕ(divₕ(ᶜχuₕ)) - - Geometry.project( - Geometry.Covariant12Axis(), - wcurlₕ( - Geometry.project( - Geometry.Covariant3Axis(), - curlₕ(ᶜχuₕ), - ), + @. ᵗρs -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχ)) + if is_ρq_tot + @. Yₜ.c.ρq_tot -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) + @. Yₜ.c.ρ -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) + end + if is_3d_pt + @. Yₜ.c.uₕ -= + κ₄ * ( + divergence_damping_factor * wgradₕ(divₕ(ᶜχuₕ)) - + Geometry.project( + Geometry.Covariant12Axis(), + wcurlₕ( + Geometry.project( + Geometry.Covariant3Axis(), + curlₕ(ᶜχuₕ), ), - ) + ), ) - elseif is_2d_pt - @. Yₜ.c.uₕ -= - κ₄ * - divergence_damping_factor * - Geometry.Covariant12Vector(wgradₕ(divₕ(ᶜχuₕ))) - end + ) + elseif is_2d_pt + @. Yₜ.c.uₕ -= + κ₄ * + divergence_damping_factor * + Geometry.Covariant12Vector(wgradₕ(divₕ(ᶜχuₕ))) end return nothing end function hyperdiffusion_tendency_tempest!(Yₜ, Y, p, t) - @nvtx "hyperdiffusion tendency" color = colorant"yellow" begin - !(:ρθ in propertynames(Y.c)) && - (error("use_tempest_mode must be false when not using ρθ")) - ᶜρ = Y.c.ρ - ᶜuₕ = Y.c.uₕ - (; ᶜp, ᶜχ, ᶜχuₕ) = p # assume ᶜp has been updated - (; ghost_buffer, κ₄, divergence_damping_factor, use_tempest_mode) = p - point_type = eltype(Fields.local_geometry_field(axes(Y.c)).coordinates) - is_ρq_tot = :ρq_tot in propertynames(Y.c) - is_2d_pt = point_type <: Geometry.Abstract2DPoint - is_3d_pt = !is_2d_pt - - @. ᶜχ = wdivₕ(gradₕ(ᶜρ)) # ᶜχρ - Spaces.weighted_dss!(ᶜχ, ghost_buffer.χ) - @. Yₜ.c.ρ -= κ₄ * wdivₕ(gradₕ(ᶜχ)) - - @. ᶜχ = wdivₕ(gradₕ(Y.c.ρθ)) # ᶜχρθ - Spaces.weighted_dss!(ᶜχ, ghost_buffer.χ) - @. Yₜ.c.ρθ -= κ₄ * wdivₕ(gradₕ(ᶜχ)) - - (; ᶠχw_data) = p - @. ᶠχw_data = wdivₕ(gradₕ(Y.f.w.components.data.:1)) - Spaces.weighted_dss!(ᶠχw_data, ghost_buffer.χ) - @. Yₜ.f.w.components.data.:1 -= κ₄ * wdivₕ(gradₕ(ᶠχw_data)) - - if is_ρq_tot - (; ᶜχρq_tot) = p - @. ᶜχρq_tot = wdivₕ(gradₕ(Y.c.ρq_tot / ᶜρ)) - Spaces.weighted_dss!(ᶜχρq_tot, ghost_buffer.χ) - @. Yₜ.c.ρq_tot -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) - @. Yₜ.c.ρ -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) - end + !(:ρθ in propertynames(Y.c)) && + (error("use_tempest_mode must be false when not using ρθ")) + ᶜρ = Y.c.ρ + ᶜuₕ = Y.c.uₕ + (; ᶜp, ᶜχ, ᶜχuₕ) = p # assume ᶜp has been updated + (; ghost_buffer, κ₄, divergence_damping_factor, use_tempest_mode) = p + point_type = eltype(Fields.local_geometry_field(axes(Y.c)).coordinates) + is_ρq_tot = :ρq_tot in propertynames(Y.c) + is_2d_pt = point_type <: Geometry.Abstract2DPoint + is_3d_pt = !is_2d_pt + + @. ᶜχ = wdivₕ(gradₕ(ᶜρ)) # ᶜχρ + Spaces.weighted_dss!(ᶜχ, ghost_buffer.χ) + @. Yₜ.c.ρ -= κ₄ * wdivₕ(gradₕ(ᶜχ)) + + @. ᶜχ = wdivₕ(gradₕ(Y.c.ρθ)) # ᶜχρθ + Spaces.weighted_dss!(ᶜχ, ghost_buffer.χ) + @. Yₜ.c.ρθ -= κ₄ * wdivₕ(gradₕ(ᶜχ)) + + (; ᶠχw_data) = p + @. ᶠχw_data = wdivₕ(gradₕ(Y.f.w.components.data.:1)) + Spaces.weighted_dss!(ᶠχw_data, ghost_buffer.χ) + @. Yₜ.f.w.components.data.:1 -= κ₄ * wdivₕ(gradₕ(ᶠχw_data)) + + if is_ρq_tot + (; ᶜχρq_tot) = p + @. ᶜχρq_tot = wdivₕ(gradₕ(Y.c.ρq_tot / ᶜρ)) + Spaces.weighted_dss!(ᶜχρq_tot, ghost_buffer.χ) + @. Yₜ.c.ρq_tot -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) + @. Yₜ.c.ρ -= κ₄ * wdivₕ(ᶜρ * gradₕ(ᶜχρq_tot)) + end - if is_3d_pt - @. ᶜχuₕ = - wgradₕ(divₕ(ᶜuₕ)) - Geometry.Covariant12Vector( - wcurlₕ(Geometry.Covariant3Vector(curlₕ(ᶜuₕ))), + if is_3d_pt + @. ᶜχuₕ = + wgradₕ(divₕ(ᶜuₕ)) - Geometry.Covariant12Vector( + wcurlₕ(Geometry.Covariant3Vector(curlₕ(ᶜuₕ))), + ) + Spaces.weighted_dss!(ᶜχuₕ, ghost_buffer.χuₕ) + @. Yₜ.c.uₕ -= + κ₄ * ( + divergence_damping_factor * wgradₕ(divₕ(ᶜχuₕ)) - + Geometry.Covariant12Vector( + wcurlₕ(Geometry.Covariant3Vector(curlₕ(ᶜχuₕ))), ) - Spaces.weighted_dss!(ᶜχuₕ, ghost_buffer.χuₕ) - @. Yₜ.c.uₕ -= - κ₄ * ( - divergence_damping_factor * wgradₕ(divₕ(ᶜχuₕ)) - - Geometry.Covariant12Vector( - wcurlₕ(Geometry.Covariant3Vector(curlₕ(ᶜχuₕ))), - ) - ) - elseif is_2d_pt - @. ᶜχuₕ = Geometry.Covariant12Vector(wgradₕ(divₕ(ᶜuₕ))) - Spaces.weighted_dss!(ᶜχuₕ, ghost_buffer.χuₕ) - @. Yₜ.c.uₕ -= - κ₄ * - divergence_damping_factor * - Geometry.Covariant12Vector(wgradₕ(divₕ(ᶜχuₕ))) - end + ) + elseif is_2d_pt + @. ᶜχuₕ = Geometry.Covariant12Vector(wgradₕ(divₕ(ᶜuₕ))) + Spaces.weighted_dss!(ᶜχuₕ, ghost_buffer.χuₕ) + @. Yₜ.c.uₕ -= + κ₄ * + divergence_damping_factor * + Geometry.Covariant12Vector(wgradₕ(divₕ(ᶜχuₕ))) end return nothing end diff --git a/perf/allocs.jl b/perf/allocs.jl index a7365b0ab8..03eddf04d4 100644 --- a/perf/allocs.jl +++ b/perf/allocs.jl @@ -2,13 +2,19 @@ example_dir = joinpath(dirname(@__DIR__), "examples") include(joinpath(example_dir, "hybrid", "cli_options.jl")); +import ClimaAtmos import ClimaCore import SciMLBase import DiffEqBase import OrdinaryDiffEq +import ClimaTimeSteppers +import Thermodynamics +import SurfaceFluxes +import CloudMicrophysics import DiffEqOperators dirs_to_monitor = [ + pkgdir(ClimaAtmos), example_dir, joinpath(example_dir, "hybrid"), joinpath(example_dir, "hybrid", "sphere"), @@ -16,12 +22,19 @@ dirs_to_monitor = [ pkgdir(SciMLBase), pkgdir(DiffEqBase), pkgdir(OrdinaryDiffEq), + pkgdir(ClimaTimeSteppers), + pkgdir(Thermodynamics), + pkgdir(SurfaceFluxes), + pkgdir(CloudMicrophysics), pkgdir(DiffEqOperators), ] -@info "`dirs_to_monitor` = $dirs_to_monitor" -filter!(x -> x isa String, dirs_to_monitor) -@info "`dirs_to_monitor` (post filter) = $dirs_to_monitor" -dirs_to_monitor = String.(dirs_to_monitor) +@info "`dirs_to_monitor` (Pre) = $dirs_to_monitor" +dirs_to_monitor_filtered = filter(x -> x isa String, dirs_to_monitor) +@info "`dirs_to_monitor` (Post) = $dirs_to_monitor_filtered" +dirs_to_monitor_filtered = String.(dirs_to_monitor_filtered) +if length(dirs_to_monitor_filtered) ≠ length(dirs_to_monitor) + @warn "Some packages' directories not found." +end #! format: off @@ -47,7 +60,7 @@ for clio in cli_options ReportMetrics.report_allocs(; job_name = string(job_id), run_cmd = `$(Base.julia_cmd()) --project=perf/ --track-allocation=all perf/allocs_per_case.jl $clio_in`, - dirs_to_monitor = dirs_to_monitor, + dirs_to_monitor = dirs_to_monitor_filtered, n_unique_allocs = 20, process_filename = function process_fn(fn) fn = "ClimaAtmos.jl/" * last(split(fn, "climaatmos-ci/"))