Merge #894

894: Widen allocation monitoring r=charleskawczynski a=charleskawczynski This PR - widens the allocation monitoring to include additional packages - Moves some ``@nvtx`` macro call to outer functions to help unindent some code Co-authored-by: Charles Kawczynski <[email protected]>
CliMA · Oct 5, 2022 · 5f70352 · 5f70352
2 parents f338619 + 25e7a6c
commit 5f70352
Show file tree

Hide file tree

Showing 2 changed files with 149 additions and 138 deletions.
diff --git 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
@@ -2,26 +2,39 @@
 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"),
     pkgdir(ClimaCore),
     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/"))