Merge #847

847: Unify exp vert adv r=charleskawczynski a=charleskawczynski This PR unifies the - `horizontal_advection_tendency` functions and - `explicit_vertical_advection_tendency!` functions Co-authored-by: Charles Kawczynski <[email protected]>
CliMA · Sep 16, 2022 · e3cfbd5 · e3cfbd5
2 parents 094ac0b + a586c0b
commit e3cfbd5
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 151 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -251,6 +251,8 @@ steps:
       - label: ":computer: aquaplanet (ρe_tot) equilmoist clearsky radiation monin_obukhov"
         command: "julia --color=yes --project=examples examples/hybrid/driver.jl --vert_diff true --surface_scheme monin_obukhov --moist equil --rad clearsky --microphy 0M --job_id sphere_aquaplanet_rhoe_equilmoist_clearsky_mo --dt 200secs --t_end 2.5days"
         artifact_paths: "sphere_aquaplanet_rhoe_equilmoist_clearsky_mo/*"
+        agents:
+          slurm_mem: 20GB
 
       - label: ":computer: aquaplanet (ρe_tot) equilmoist gray radiation"
         command: "julia --color=yes --project=examples examples/hybrid/driver.jl --vert_diff true --surface_scheme bulk --moist equil --rad gray --microphy 0M --job_id sphere_aquaplanet_rhoe_equilmoist_gray --dt 450secs --t_end 2.5days"

diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl
@@ -101,21 +101,7 @@ function additional_cache(Y, params, model_spec, dt; use_tempest_mode = false)
     (; microphysics_model, forcing_type, radiation_model, turbconv_model) =
         model_spec
 
-    default_remaining_tendencies = if model_spec.anelastic_dycore
-        nothing
-    else
-        if :ρe_tot in propertynames(Y.c) && enable_threading()
-            (;
-                horizontal_advection_tendency! = horizontal_advection_tendency_special!,
-                explicit_vertical_advection_tendency! = explicit_vertical_advection_tendency_special!,
-            )
-        else
-            (;
-                horizontal_advection_tendency! = horizontal_advection_tendency_generic!,
-                explicit_vertical_advection_tendency! = explicit_vertical_advection_tendency_generic!,
-            )
-        end
-    end
+    anelastic_dycore = model_spec.anelastic_dycore
 
     return merge(
         hyperdiffusion_cache(
@@ -172,7 +158,7 @@ function additional_cache(Y, params, model_spec, dt; use_tempest_mode = false)
             )
         ),
         (; Δt = dt),
-        (; default_remaining_tendencies),
+        (; anelastic_dycore),
         !isnothing(turbconv_model) ?
         (; edmf_cache = TCU.get_edmf_cache(Y, namelist, params, parsed_args)) :
         NamedTuple(),

diff --git a/examples/hybrid/staggered_nonhydrostatic_model.jl b/examples/hybrid/staggered_nonhydrostatic_model.jl
@@ -302,12 +302,12 @@ function _implicit_tendency!(Yₜ, Y, p, t)
 end
 
 function remaining_tendency!(Yₜ, Y, p, t)
-    default_tends = p.default_remaining_tendencies
+    anelastic_dycore = p.anelastic_dycore
     @nvtx "remaining tendency" color = colorant"yellow" begin
         Yₜ .= zero(eltype(Yₜ))
-        if !isnothing(default_tends)
-            default_tends.horizontal_advection_tendency!(Yₜ, Y, p, t)
-            default_tends.explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
+        if !anelastic_dycore
+            horizontal_advection_tendency!(Yₜ, Y, p, t)
+            explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         end
         @nvtx "additional_tendency!" color = colorant"orange" begin
             additional_tendency!(Yₜ, Y, p, t)
@@ -326,11 +326,11 @@ end
 
 function remaining_tendency_increment!(Y⁺, Y, p, t, dtγ)
     (; Yₜ, limiters) = p
-    default_tends = p.default_remaining_tendencies
+    anelastic_dycore = p.anelastic_dycore
     @nvtx "remaining tendency increment" color = colorant"yellow" begin
         Yₜ .= zero(eltype(Yₜ))
-        if !isnothing(default_tends)
-            default_tends.horizontal_advection_tendency!(Yₜ, Y, p, t)
+        if !anelastic_dycore
+            horizontal_advection_tendency!(Yₜ, Y, p, t)
             # Apply limiter
             if !isnothing(limiters)
                 @. Y⁺ += dtγ * Yₜ
@@ -343,7 +343,7 @@ function remaining_tendency_increment!(Y⁺, Y, p, t, dtγ)
                 end
                 Yₜ .= zero(eltype(Yₜ))
             end
-            default_tends.explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
+            explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         end
         @nvtx "additional_tendency! increment" color = colorant"orange" begin
             additional_tendency!(Yₜ, Y, p, t)
@@ -361,70 +361,45 @@ function remaining_tendency_increment!(Y⁺, Y, p, t, dtγ)
     return Y⁺
 end
 
-function horizontal_advection_tendency_special!(Yₜ, Y, p, t)
-    ᶜρ = Y.c.ρ
-    ᶜuₕ = Y.c.uₕ
-    ᶠw = Y.f.w
-    (; ᶜuvw, ᶜK, ᶜΦ, ᶜts, ᶜp, ᶜω³, ᶠω¹², params) = p
-    point_type = eltype(Fields.local_geometry_field(axes(Y.c)).coordinates)
+function _precomputed_quantities!(Yₜ, Y, p, t)
+    Fields.bycolumn(axes(Y.c)) do colidx
+        ᶜuₕ = Y.c.uₕ
+        ᶠw = Y.f.w
+        (; ᶜuvw, ᶜK, ᶜts, ᶜp, params) = p
+
+        @. ᶜuvw[colidx] = C123(ᶜuₕ[colidx]) + C123(ᶜinterp(ᶠw[colidx]))
+        @. ᶜK[colidx] = norm_sqr(ᶜuvw[colidx]) / 2
+        thermo_state!(
+            ᶜts[colidx],
+            Y.c[colidx],
+            params,
+            ᶜinterp,
+            ᶜK[colidx],
+            Y.f.w[colidx],
+        )
+        thermo_params = CAP.thermodynamics_params(params)
+        @. ᶜp[colidx] = TD.air_pressure(thermo_params, ᶜts[colidx])
+        nothing
+    end
+    return nothing
+end
+
+function horizontal_advection_tendency!(Yₜ, Y, p, t)
     @nvtx "precomputed quantities" color = colorant"orange" begin
-        Fields.bycolumn(axes(Y.c)) do colidx
-            @. ᶜuvw[colidx] = C123(ᶜuₕ[colidx]) + C123(ᶜinterp(ᶠw[colidx]))
-            @. ᶜK[colidx] = norm_sqr(ᶜuvw[colidx]) / 2
-            thermo_state!(
-                ᶜts[colidx],
-                Y.c[colidx],
-                params,
-                ᶜinterp,
-                ᶜK[colidx],
-                Y.f.w[colidx],
-            )
-            thermo_params = CAP.thermodynamics_params(params)
-            @. ᶜp[colidx] = TD.air_pressure(thermo_params, ᶜts[colidx])
-            nothing
-        end
+        _precomputed_quantities!(Yₜ, Y, p, t)
     end
     @nvtx "horizontal" color = colorant"orange" begin
-        # Mass conservation
-        @. Yₜ.c.ρ -= divₕ(ᶜρ * ᶜuvw)
-
-        # Energy conservation
-        @. Yₜ.c.ρe_tot -= divₕ((Y.c.ρe_tot + ᶜp) * ᶜuvw)
-
-        # Momentum conservation
-        if point_type <: Geometry.Abstract3DPoint
-            @. ᶜω³ = curlₕ(ᶜuₕ)
-            @. ᶠω¹² = curlₕ(ᶠw)
-            @. Yₜ.c.uₕ -= gradₕ(ᶜp) / ᶜρ + gradₕ(ᶜK + ᶜΦ)
-        elseif point_type <: Geometry.Abstract2DPoint
-            ᶜω³ .= Ref(zero(eltype(ᶜω³)))
-            @. ᶠω¹² = Geometry.Contravariant12Vector(curlₕ(ᶠw))
-            @. Yₜ.c.uₕ -=
-                Geometry.Covariant12Vector(gradₕ(ᶜp) / ᶜρ + gradₕ(ᶜK + ᶜΦ))
-        end
-
-        # Tracer conservation
-        for ᶜρc_name in filter(is_tracer_var, propertynames(Y.c))
-            ᶜρc = getproperty(Y.c, ᶜρc_name)
-            ᶜρcₜ = getproperty(Yₜ.c, ᶜρc_name)
-            @. ᶜρcₜ -= divₕ(ᶜρc * ᶜuvw)
-        end
+        _horizontal_advection_tendency!(Yₜ, Y, p, t)
     end
+    return nothing
 end
 
-function horizontal_advection_tendency_generic!(Yₜ, Y, p, t)
+function _horizontal_advection_tendency!(Yₜ, Y, p, t)
     ᶜρ = Y.c.ρ
     ᶜuₕ = Y.c.uₕ
     ᶠw = Y.f.w
     (; ᶜuvw, ᶜK, ᶜΦ, ᶜts, ᶜp, ᶜω³, ᶠω¹², params) = p
     point_type = eltype(Fields.local_geometry_field(axes(Y.c)).coordinates)
-    thermo_params = CAP.thermodynamics_params(params)
-
-    # Precomputed quantities
-    @. ᶜuvw = C123(ᶜuₕ) + C123(ᶜinterp(ᶠw))
-    @. ᶜK = norm_sqr(ᶜuvw) / 2
-    thermo_state!(ᶜts, Y, params, ᶜinterp, ᶜK)
-    @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
 
     # Mass conservation
     @. Yₜ.c.ρ -= divₕ(ᶜρ * ᶜuvw)
@@ -464,86 +439,60 @@ function horizontal_advection_tendency_generic!(Yₜ, Y, p, t)
         ᶜρcₜ = getproperty(Yₜ.c, ᶜρc_name)
         @. ᶜρcₜ -= divₕ(ᶜρc * ᶜuvw)
     end
+    return nothing
 end
 
-function explicit_vertical_advection_tendency_special!(Yₜ, Y, p, t)
-    ᶜρ = Y.c.ρ
-    ᶜuₕ = Y.c.uₕ
-    ᶠw = Y.f.w
-    (; ᶜuvw, ᶜK, ᶜp, ᶜω³, ᶠω¹², ᶠu¹², ᶠu³, ᶜf) = p
+function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     @nvtx "vertical" color = colorant"orange" begin
-        Fields.bycolumn(axes(Y.c)) do colidx
-
-            # Mass conservation
-            @. Yₜ.c.ρ[colidx] -= ᶜdivᵥ(ᶠinterp(ᶜρ[colidx] * ᶜuₕ[colidx]))
+        _explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
+    end
+    return nothing
+end
+function _explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
+    Fields.bycolumn(axes(Y.c)) do colidx
+        ᶜρ = Y.c.ρ
+        ᶜuₕ = Y.c.uₕ
+        ᶠw = Y.f.w
+        (; ᶜuvw, ᶜK, ᶜp, ᶜω³, ᶠω¹², ᶠu¹², ᶠu³, ᶜf) = p
+        # Mass conservation
+        @. Yₜ.c.ρ[colidx] -= ᶜdivᵥ(ᶠinterp(ᶜρ[colidx] * ᶜuₕ[colidx]))
 
-            # Energy conservation
+        # Energy conservation
+        if :ρθ in propertynames(Y.c)
+            @. Yₜ.c.ρθ[colidx] -= ᶜdivᵥ(ᶠinterp(Y.c.ρθ[colidx] * ᶜuₕ[colidx]))
+        elseif :ρe_tot in propertynames(Y.c)
             @. Yₜ.c.ρe_tot[colidx] -=
                 ᶜdivᵥ(ᶠinterp((Y.c.ρe_tot[colidx] + ᶜp[colidx]) * ᶜuₕ[colidx]))
-
-            # Momentum conservation
-            @. ᶠω¹²[colidx] += ᶠcurlᵥ(ᶜuₕ[colidx])
-            @. ᶠu¹²[colidx] = Geometry.project(
-                Geometry.Contravariant12Axis(),
-                ᶠinterp(ᶜuvw[colidx]),
-            )
-            @. ᶠu³[colidx] = Geometry.project(
-                Geometry.Contravariant3Axis(),
-                C123(ᶠinterp(ᶜuₕ[colidx])) + C123(ᶠw[colidx]),
-            )
-            @. Yₜ.c.uₕ[colidx] -=
-                ᶜinterp(ᶠω¹²[colidx] × ᶠu³[colidx]) +
-                (ᶜf[colidx] + ᶜω³[colidx]) ×
-                (Geometry.project(Geometry.Contravariant12Axis(), ᶜuvw[colidx]))
-            @. Yₜ.f.w[colidx] -=
-                ᶠω¹²[colidx] × ᶠu¹²[colidx] + ᶠgradᵥ(ᶜK[colidx])
-
-            # Tracer conservation
-            for ᶜρc_name in filter(is_tracer_var, propertynames(Y.c))
-                ᶜρc = getproperty(Y.c, ᶜρc_name)
-                ᶜρcₜ = getproperty(Yₜ.c, ᶜρc_name)
-                @. ᶜρcₜ[colidx] -= ᶜdivᵥ(ᶠinterp(ᶜρc[colidx] * ᶜuₕ[colidx]))
-            end
+        elseif :ρe_int in propertynames(Y.c)
+            @. Yₜ.c.ρe_int[colidx] -=
+                ᶜdivᵥ(ᶠinterp((Y.c.ρe_int[colidx] + ᶜp[colidx]) * ᶜuₕ[colidx]))
         end
-    end
-end
 
-function explicit_vertical_advection_tendency_generic!(Yₜ, Y, p, t)
-    ᶜρ = Y.c.ρ
-    ᶜuₕ = Y.c.uₕ
-    ᶠw = Y.f.w
-    (; ᶜuvw, ᶜK, ᶜp, ᶜω³, ᶠω¹², ᶠu¹², ᶠu³, ᶜf) = p
-
-    # Mass conservation
-    @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠinterp(ᶜρ * ᶜuₕ))
-
-    # Energy conservation
-    if :ρθ in propertynames(Y.c)
-        @. Yₜ.c.ρθ -= ᶜdivᵥ(ᶠinterp(Y.c.ρθ * ᶜuₕ))
-    elseif :ρe_tot in propertynames(Y.c)
-        @. Yₜ.c.ρe_tot -= ᶜdivᵥ(ᶠinterp((Y.c.ρe_tot + ᶜp) * ᶜuₕ))
-    elseif :ρe_int in propertynames(Y.c)
-        @. Yₜ.c.ρe_int -= ᶜdivᵥ(ᶠinterp((Y.c.ρe_int + ᶜp) * ᶜuₕ))
-    end
-
-    # Momentum conservation
-    @. ᶠω¹² += ᶠcurlᵥ(ᶜuₕ)
-    @. ᶠu¹² = Geometry.project(Geometry.Contravariant12Axis(), ᶠinterp(ᶜuvw))
-    @. ᶠu³ = Geometry.project(
-        Geometry.Contravariant3Axis(),
-        C123(ᶠinterp(ᶜuₕ)) + C123(ᶠw),
-    )
-    @. Yₜ.c.uₕ -=
-        ᶜinterp(ᶠω¹² × ᶠu³) +
-        (ᶜf + ᶜω³) × (Geometry.project(Geometry.Contravariant12Axis(), ᶜuvw))
-    @. Yₜ.f.w -= ᶠω¹² × ᶠu¹² + ᶠgradᵥ(ᶜK)
+        # Momentum conservation
+        @. ᶠω¹²[colidx] += ᶠcurlᵥ(ᶜuₕ[colidx])
+        @. ᶠu¹²[colidx] = Geometry.project(
+            Geometry.Contravariant12Axis(),
+            ᶠinterp(ᶜuvw[colidx]),
+        )
+        @. ᶠu³[colidx] = Geometry.project(
+            Geometry.Contravariant3Axis(),
+            C123(ᶠinterp(ᶜuₕ[colidx])) + C123(ᶠw[colidx]),
+        )
+        @. Yₜ.c.uₕ[colidx] -=
+            ᶜinterp(ᶠω¹²[colidx] × ᶠu³[colidx]) +
+            (ᶜf[colidx] + ᶜω³[colidx]) ×
+            (Geometry.project(Geometry.Contravariant12Axis(), ᶜuvw[colidx]))
+        @. Yₜ.f.w[colidx] -= ᶠω¹²[colidx] × ᶠu¹²[colidx] + ᶠgradᵥ(ᶜK[colidx])
 
-    # Tracer conservation
-    for ᶜρc_name in filter(is_tracer_var, propertynames(Y.c))
-        ᶜρc = getproperty(Y.c, ᶜρc_name)
-        ᶜρcₜ = getproperty(Yₜ.c, ᶜρc_name)
-        @. ᶜρcₜ -= ᶜdivᵥ(ᶠinterp(ᶜρc * ᶜuₕ))
+        # Tracer conservation
+        for ᶜρc_name in filter(is_tracer_var, propertynames(Y.c))
+            ᶜρc = getproperty(Y.c, ᶜρc_name)
+            ᶜρcₜ = getproperty(Yₜ.c, ᶜρc_name)
+            @. ᶜρcₜ[colidx] -= ᶜdivᵥ(ᶠinterp(ᶜρc[colidx] * ᶜuₕ[colidx]))
+        end
+        nothing
     end
+    return nothing
 end
 
 # Allow one() to be called on vectors.

diff --git a/perf/flame.jl b/perf/flame.jl
@@ -67,10 +67,10 @@ allocs = @allocated OrdinaryDiffEq.step!(integrator)
 @info "`allocs ($job_id)`: $(allocs)"
 
 allocs_limit = Dict()
-allocs_limit["flame_perf_target_rhoe"] = 19867936
-allocs_limit["flame_perf_target_rhoe_threaded"] = 27249159
-allocs_limit["flame_perf_target_rhoe_callbacks"] = 31034008
-allocs_limit["flame_perf_target_rhoe_ARS343"] = 61185840
+allocs_limit["flame_perf_target_rhoe"] = 26060640
+allocs_limit["flame_perf_target_rhoe_threaded"] = 27560432
+allocs_limit["flame_perf_target_rhoe_callbacks"] = 37226712
+allocs_limit["flame_perf_target_rhoe_ARS343"] = 73571248
 
 if allocs < allocs_limit[job_id] * buffer
     @info "TODO: lower `allocs_limit[$job_id]` to: $(allocs)"