Merge #2270 #2280

2270: Add AtmosNumerics r=charleskawczynski a=Sbozzolo

Group a few numeric choices into a single object.


2280: Use ClimaComms at-elapsed r=charleskawczynski a=charleskawczynski

This PR switches to using the `ClimaComms.`@elapsed`,` which behaves the same as the replaced code. This removes an explicit direct dependency on CUDA in the perf folder.

Co-authored-by: Gabriele Bozzola <gbozzola@caltech.edu>
Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

perf/benchmark_dump.jl

@@ -1,7 +1,6 @@
 import Random
 Random.seed!(1234)
 import ClimaAtmos as CA
-using CUDA
 import ClimaComms
 using Plots
 using PrettyTables
@@ -27,17 +26,9 @@ for h_elem in 8:8:40
     n_steps = 10
     comms_ctx = ClimaComms.context(integrator.u.c)
     device = ClimaComms.device(comms_ctx)
-    if device isa ClimaComms.CUDADevice
-        e = CUDA.@elapsed begin
-            s = CA.@timed_str begin
-                CA.benchmark_step!(integrator, Y₀, n_steps) # run
-            end
-        end
-    else
-        e = @elapsed begin
-            s = CA.@timed_str begin
-                CA.benchmark_step!(integrator, Y₀, n_steps) # run
-            end
+    e = ClimaComms.@elapsed device begin
+        s = CA.@timed_str begin
+            CA.benchmark_step!(integrator, Y₀, n_steps) # run
         end
     end
     @info "Ran step! $n_steps times in $s, ($(CA.prettytime(e/n_steps*1e9)) per step)"

perf/flame.jl

@@ -63,7 +63,7 @@ allocs_limit["flame_perf_target"] = 4656
 allocs_limit["flame_perf_target_tracers"] = 204288
 allocs_limit["flame_perf_target_edmfx"] = 253440
 allocs_limit["flame_perf_diagnostics"] = 3016136
-allocs_limit["flame_perf_target_diagnostic_edmfx"] = 919872
+allocs_limit["flame_perf_target_diagnostic_edmfx"] = 920512
 allocs_limit["flame_perf_target_edmf"] = 6386476224
 allocs_limit["flame_perf_target_threaded"] = 5857808
 allocs_limit["flame_perf_target_callbacks"] = 46407936

src/cache/cache.jl

@@ -26,30 +26,31 @@ import ClimaCore.Fields: ColumnField
 # This is a constant Coriolis frequency that is only used if space is flat
 function default_cache(
     Y,
-    parsed_args,
     params,
     atmos,
     spaces,
     numerics,
     simulation,
     surface_setup,
 )
+
     FT = eltype(params)
-    (; energy_upwinding, tracer_upwinding, density_upwinding, edmfx_upwinding) =
-        numerics
-    (; apply_limiter) = numerics
+
     ᶜcoord = Fields.local_geometry_field(Y.c).coordinates
-    R_d = FT(CAP.R_d(params))
-    MSLP = FT(CAP.MSLP(params))
     grav = FT(CAP.grav(params))
-    T_ref = FT(255)
-    ᶜΦ = CAP.grav(params) .* ᶜcoord.z
-    ᶜρ_ref = @. MSLP * exp(-grav * ᶜcoord.z / (R_d * T_ref)) / (R_d * T_ref)
-    ᶜp_ref = @. ᶜρ_ref * R_d * T_ref
-    if !parsed_args["use_reference_state"]
-        ᶜρ_ref .*= 0
-        ᶜp_ref .*= 0
+    ᶜΦ = grav .* ᶜcoord.z
+
+    if atmos.numerics.use_reference_state
+        R_d = FT(CAP.R_d(params))
+        MSLP = FT(CAP.MSLP(params))
+        T_ref = FT(255)
+        ᶜρ_ref = @. MSLP * exp(-grav * ᶜcoord.z / (R_d * T_ref)) / (R_d * T_ref)
+        ᶜp_ref = @. ᶜρ_ref * R_d * T_ref
+    else
+        ᶜρ_ref = zero(ᶜΦ)
+        ᶜp_ref = zero(ᶜΦ)
     end
+
     if eltype(ᶜcoord) <: Geometry.LatLongZPoint
         Ω = CAP.Omega(params)
         ᶜf = @. 2 * Ω * sind(ᶜcoord.lat)
@@ -66,23 +67,17 @@ function default_cache(
         (; c = Spaces.create_dss_buffer(Y.c), f = Spaces.create_dss_buffer(Y.f))
 
     limiter =
-        apply_limiter ? Limiters.QuasiMonotoneLimiter(similar(Y.c, FT)) :
-        nothing
+        isnothing(numerics.limiter) ? nothing :
+        numerics.limiter(similar(Y.c, FT))
 
     net_energy_flux_toa = [Geometry.WVector(FT(0))]
     net_energy_flux_sfc = [Geometry.WVector(FT(0))]
-    test = if parsed_args["test_dycore_consistency"]
-        TestDycoreConsistency()
-    else
-        nothing
-    end
 
     default_cache = (;
         is_init = Ref(true),
         simulation,
         atmos,
         sfc_setup = surface_setup(params),
-        test,
         limiter,
         ᶜΦ,
         ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(ᶜΦ),
@@ -92,10 +87,6 @@ function default_cache(
         ᶜf,
         ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}),
         params,
-        energy_upwinding,
-        tracer_upwinding,
-        density_upwinding,
-        edmfx_upwinding,
         do_dss,
         ghost_buffer,
         net_energy_flux_toa,

src/prognostic_equations/advection.jl

@@ -73,12 +73,11 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     (; turbconv_model) = p.atmos
     n = n_prognostic_mass_flux_subdomains(turbconv_model)
     advect_tke = use_prognostic_tke(turbconv_model)
-    is_total_energy = p.atmos.energy_form isa TotalEnergy
     point_type = eltype(Fields.coordinate_field(Y.c))
     (; dt) = p.simulation
     ᶜJ = Fields.local_geometry_field(Y.c).J
     (; ᶜu, ᶠu³, ᶜK, ᶜf) = p
-    (; edmfx_upwinding) = n > 0 || advect_tke ? p : all_nothing
+    (; edmfx_upwinding) = n > 0 || advect_tke ? p.atmos.numerics : all_nothing
     (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜρʲs) = n > 0 ? p : all_nothing
     (; ᶜp, ᶜp_ref, ᶜρ_ref, ᶠgradᵥ_ᶜΦ) = n > 0 ? p : all_nothing
     (; ᶠu³⁰) = advect_tke ? p : all_nothing

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -15,7 +15,7 @@ function edmfx_sgs_mass_flux_tendency!(
 
     FT = Spaces.undertype(axes(Y.c))
     n = n_mass_flux_subdomains(turbconv_model)
-    (; edmfx_upwinding) = p
+    (; edmfx_upwinding) = p.atmos.numerics
     (; ᶠu³, ᶜh_tot, ᶜspecific) = p
     (; ᶠu³ʲs) = p
     (; ᶜρa⁰, ᶠu³⁰, ᶜu⁰, ᶜh_tot⁰, ᶜq_tot⁰) = p
@@ -98,7 +98,8 @@ function edmfx_sgs_mass_flux_tendency!(
 
     FT = Spaces.undertype(axes(Y.c))
     n = n_mass_flux_subdomains(turbconv_model)
-    (; edmfx_upwinding, sfc_conditions) = p
+    (; edmfx_upwinding) = p.atmos.numerics
+    (; sfc_conditions) = p
     (; ᶠu³, ᶜu, ᶜh_tot, ᶜspecific) = p
     (; ᶜρaʲs, ᶠu³ʲs, ᶜh_totʲs, ᶜq_totʲs) = p
     (; ᶜK_u, ᶜK_h) = p
@@ -223,7 +224,8 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     FT = Spaces.undertype(axes(Y.c))
     n = n_mass_flux_subdomains(turbconv_model)
-    (; edmfx_upwinding, sfc_conditions) = p
+    (; edmfx_upwinding) = p.atmos.numerics
+    (; sfc_conditions) = p
     (; ᶠu³, ᶜu, ᶜh_tot, ᶜspecific) = p
     (; ᶜρaʲs, ᶠu³ʲs, ᶜh_totʲs, ᶜq_totʲs) = p
     (; ᶜK_u, ᶜK_h) = p

src/prognostic_equations/implicit/implicit_solver.jl

@@ -219,14 +219,14 @@ function Wfact!(A, Y, p, dtγ, t)
             p.ᶜp_ref,
             p.ᶜtemp_scalar,
             p.params,
-            p.energy_upwinding,
-            p.tracer_upwinding,
-            p.density_upwinding,
             p.atmos,
             (energy_form isa TotalEnergy ? (; p.ᶜh_tot) : (;))...,
             (rayleigh_sponge isa RayleighSponge ? (; p.ᶠβ_rayleigh_w) : (;))...,
         )
 
+        (; energy_upwinding, tracer_upwinding, density_upwinding) =
+            p.atmos.numerics
+
         # Convert dtγ from a Float64 to an FT.
         FT = Spaces.undertype(axes(Y.c))
         dtγ′ = FT(dtγ)
@@ -242,7 +242,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     (; matrix, enthalpy_flag) = A
     (; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK_∂ᶠu₃) = p
     (; ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, params) = p
-    (; energy_upwinding, tracer_upwinding, density_upwinding) = p
+    (; energy_upwinding, tracer_upwinding, density_upwinding) = p.atmos.numerics
 
     FT = Spaces.undertype(axes(Y.c))
     one_ATC3 = CT3(FT(1))'

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -60,8 +60,7 @@ vertical_advection!(ᶜρχₜ, ᶠu³, ᶜχ, ::Val{:third_order}) =
     @. ᶜρχₜ -= ᶜadvdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³)
 
 function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
-    (; energy_upwinding, tracer_upwinding, density_upwinding, edmfx_upwinding) =
-        p
+    (; energy_upwinding, tracer_upwinding, density_upwinding) = p.atmos.numerics
     (; turbconv_model, rayleigh_sponge) = p.atmos
     (; dt) = p.simulation
     n = n_mass_flux_subdomains(turbconv_model)

src/solver/type_getters.jl

@@ -80,23 +80,38 @@ function get_atmos(config::AtmosConfig, params)
         rayleigh_sponge = get_rayleigh_sponge_model(parsed_args, params, FT),
         sfc_temperature = get_sfc_temperature_form(parsed_args),
         surface_model = get_surface_model(parsed_args),
+        numerics = get_numerics(parsed_args),
     )
 
     @info "AtmosModel: \n$(summary(atmos))"
     return atmos
 end
 
 function get_numerics(parsed_args)
+    test_dycore =
+        parsed_args["test_dycore_consistency"] ? TestDycoreConsistency() :
+        nothing
+
+    energy_upwinding = Val(Symbol(parsed_args["energy_upwinding"]))
+    tracer_upwinding = Val(Symbol(parsed_args["tracer_upwinding"]))
+    density_upwinding = Val(Symbol(parsed_args["density_upwinding"]))
+    edmfx_upwinding = Val(Symbol(parsed_args["edmfx_upwinding"]))
+
+    limiter =
+        parsed_args["apply_limiter"] ? Limiters.QuasiMonotoneLimiter : nothing
+
+
     # wrap each upwinding mode in a Val for dispatch
-    numerics = (;
-        energy_upwinding = Val(Symbol(parsed_args["energy_upwinding"])),
-        tracer_upwinding = Val(Symbol(parsed_args["tracer_upwinding"])),
-        density_upwinding = Val(Symbol(parsed_args["density_upwinding"])),
-        edmfx_upwinding = Val(Symbol(parsed_args["edmfx_upwinding"])),
-        apply_limiter = parsed_args["apply_limiter"],
-        bubble = parsed_args["bubble"],
+    numerics = AtmosNumerics(;
+        energy_upwinding,
+        tracer_upwinding,
+        density_upwinding,
+        edmfx_upwinding,
+        limiter,
+        test_dycore_consistency = test_dycore,
+        use_reference_state = parsed_args["use_reference_state"],
     )
-    @info "numerics" numerics...
+    @info "numerics $(summary(numerics))"
 
     return numerics
 end
@@ -544,7 +559,6 @@ function get_cache(
 )
     _default_cache = default_cache(
         Y,
-        parsed_args,
         params,
         atmos,
         spaces,
@@ -581,7 +595,6 @@ function get_simulation(config::AtmosConfig)
     mkpath(output_dir)
 
     sim = (;
-        is_debugging_tc = parsed_args["debugging_tc"],
         output_dir,
         restart = haskey(ENV, "RESTART_FILE"),
         job_id,

src/solver/types.jl

@@ -268,6 +268,49 @@ struct TestDycoreConsistency end
 
 Base.broadcastable(x::AbstractPerformanceMode) = tuple(x)
 
+Base.@kwdef struct AtmosNumerics{EN_UP, TR_UP, DE_UP, ED_UP, DYCORE, LIM}
+
+    """Enable specific upwinding schemes for specific equations"""
+    energy_upwinding::EN_UP
+    tracer_upwinding::TR_UP
+    density_upwinding::DE_UP
+    edmfx_upwinding::ED_UP
+
+    """Add NaNs to certain equations to track down problems"""
+    test_dycore_consistency::DYCORE
+
+    limiter::LIM
+
+    """Whether to subtract the reference state from the evolved state"""
+    use_reference_state::Bool
+end
+Base.broadcastable(x::AtmosNumerics) = tuple(x)
+
+function Base.summary(io::IO, numerics::AtmosNumerics)
+    pns = string.(propertynames(numerics))
+    buf = maximum(length.(pns))
+    keys = propertynames(numerics)
+    vals = repeat.(" ", map(s -> buf - length(s) + 2, pns))
+    bufs = (; zip(keys, vals)...)
+    print(io, '\n')
+    for pn in propertynames(numerics)
+        prop = getproperty(numerics, pn)
+        s = string(
+            "  ", # needed for some reason
+            getproperty(bufs, pn),
+            '`',
+            string(pn),
+            '`',
+            "::",
+            '`',
+            typeof(prop),
+            '`',
+            '\n',
+        )
+        print(io, s)
+    end
+end
+
 Base.@kwdef struct AtmosModel{
     MC,
     PEM,
@@ -294,6 +337,7 @@ Base.@kwdef struct AtmosModel{
     RS,
     ST,
     SM,
+    NUM,
 }
     model_config::MC = nothing
     perf_mode::PEM = nothing
@@ -320,6 +364,7 @@ Base.@kwdef struct AtmosModel{
     rayleigh_sponge::RS = nothing
     sfc_temperature::ST = nothing
     surface_model::SM = nothing
+    numerics::NUM = nothing
 end
 
 Base.broadcastable(x::AtmosModel) = tuple(x)

src/utils/debug_utils.jl

@@ -195,6 +195,7 @@ end
 
 Fill a data structure's `Field`s / `FieldVector`s with NaNs.
 """
-fill_with_nans!(p) = fill_with_nans!(p, p.test)
+fill_with_nans!(p) =
+    fill_with_nans!(p, p.atmos.numerics.test_dycore_consistency)
 fill_with_nans!(p, ::Nothing) = nothing
 fill_with_nans!(p, ::TestDycoreConsistency) = fill_with_nans_generic!(p)