diff --git a/perf/benchmark_dump.jl b/perf/benchmark_dump.jl index 4bce06a25c..bb5a456e6e 100644 --- a/perf/benchmark_dump.jl +++ b/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)" diff --git a/perf/flame.jl b/perf/flame.jl index 3359a597b0..35d5cdd89a 100644 --- a/perf/flame.jl +++ b/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 diff --git a/src/cache/cache.jl b/src/cache/cache.jl index c75265333b..6439132353 100644 --- a/src/cache/cache.jl +++ b/src/cache/cache.jl @@ -26,7 +26,6 @@ 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, @@ -34,22 +33,24 @@ function default_cache( 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, diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index aafede7dd1..992d9db88c 100644 --- a/src/prognostic_equations/advection.jl +++ b/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 diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index d8c574acd3..c2e861fe9e 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/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 diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index c85d9f335e..151acf1432 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/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))' diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 55b182300d..6c73ef4f08 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/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) diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 87309fe5e8..9eff208e40 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -80,6 +80,7 @@ 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))" @@ -87,16 +88,30 @@ function get_atmos(config::AtmosConfig, params) 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, diff --git a/src/solver/types.jl b/src/solver/types.jl index ed08bef070..419fdc9f1f 100644 --- a/src/solver/types.jl +++ b/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) diff --git a/src/utils/debug_utils.jl b/src/utils/debug_utils.jl index aaf85615c3..44b0e74a80 100644 --- a/src/utils/debug_utils.jl +++ b/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)