diff --git a/perf/benchmark.jl b/perf/benchmark.jl index e6a7eae513..01ac9c034a 100644 --- a/perf/benchmark.jl +++ b/perf/benchmark.jl @@ -49,6 +49,8 @@ are_boundschecks_forced = Base.JLOptions().check_bounds == 1 if haskey(trials, name) if trials[name].memory ≤ mem @warn "Allocation limits for $name can be reduced to $(trials[name].memory)." + else + @info "trials[$name].memory: $(trials[name].memory)" end return trials[name].memory ≤ mem else @@ -58,8 +60,8 @@ are_boundschecks_forced = Base.JLOptions().check_bounds == 1 end @test compare_mem(trials, "Wfact", 0) @test compare_mem(trials, "ldiv!", 0) - @test compare_mem(trials, "T_imp!", 1000000000000000000000) - @test compare_mem(trials, "T_exp_T_lim!", 10496) + @test compare_mem(trials, "T_imp!", 96) + @test compare_mem(trials, "T_exp_T_lim!", 190420) @test compare_mem(trials, "lim!", 0) @test compare_mem(trials, "dss!", 0) @test compare_mem(trials, "cache!", 120) diff --git a/perf/flame.jl b/perf/flame.jl index 72bd594b8f..61b56888b7 100644 --- a/perf/flame.jl +++ b/perf/flame.jl @@ -40,12 +40,12 @@ ProfileCanvas.html_file(joinpath(output_dir, "flame.html"), results) allocs_limit = Dict() allocs_limit["flame_baroclinic_wave_moist_gpu"] = 1_243_048 allocs_limit["flame_default"] = 99_928 -allocs_limit["flame_default_1m"] = 441_904 +allocs_limit["flame_default_1m"] = 527_712 allocs_limit["flame_diagnostics"] = 10_677_144 allocs_limit["flame_aquaplanet_diagedmf"] = 11_644_128 -allocs_limit["flame_aquaplanet_progedmf"] = 697_520 -allocs_limit["flame_aquaplanet_progedmf_dense_autodiff"] = 727_512 -allocs_limit["flame_diffusion"] = 100_360 +allocs_limit["flame_aquaplanet_progedmf"] = 774_712 +allocs_limit["flame_aquaplanet_progedmf_dense_autodiff"] = 774_712 +allocs_limit["flame_diffusion"] = 138_432 allocs_limit["flame_threaded"] = 2047_736 allocs_limit["flame_callbacks"] = 391_864 allocs_limit["flame_gravity_wave"] = 581_381_976 diff --git a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl index 1013a0cca2..1ac289a006 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -107,22 +107,23 @@ function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLill @. Yₜ.c.ρe_tot += wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜh_tot)) ## Tracer diffusion and associated mass changes - for (ᶜρχₜ, ᶜχ, χ_name) in CA.matching_subfields(Yₜ.c, ᶜspecific) - χ_name == :e_tot && continue - ᶜρχₜ_diffusion = - @. p.scratch.ᶜtemp_scalar = wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜχ)) + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) + ᶜρχₜ_diffusion = @. lazy(wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜχ))) @. ᶜρχₜ += ᶜρχₜ_diffusion # Rain and snow does not affect the mass - if χ_name ∉ (:q_rai, :q_sno) + if ρχ_name == @name(ρq_tot) @. Yₜ.c.ρ += ᶜρχₜ_diffusion end end end +import UnrolledUtilities as UU + function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) FT = eltype(Y) - (; sfc_temp_C3, ᶠtemp_scalar, ᶜtemp_scalar) = p.scratch + (; sfc_temp_C3, ᶠtemp_scalar) = p.scratch (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, ᶜspecific, ᶜh_tot, sfc_conditions) = p.precomputed (; ρ_flux_uₕ, ρ_flux_h_tot) = sfc_conditions @@ -155,18 +156,17 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜh_tot))) ## Tracer diffusion and associated mass changes - for (ᶜρχₜ, ᶜχ, χ_name) in CA.matching_subfields(Yₜ.c, ᶜspecific) - χ_name == :e_tot && continue - - ᶜdivᵥ_ρχ = Operators.DivergenceF2C(; - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(C3(FT(0))), - ) + ᶜdivᵥ_ρχ = Operators.DivergenceF2C(; + top = Operators.SetValue(C3(FT(0))), + bottom = Operators.SetValue(C3(FT(0))), + ) - ᶜ∇ᵥρD∇χₜ = @. ᶜtemp_scalar = ᶜdivᵥ_ρχ(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜχ))) + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + ᶜ∇ᵥρD∇χₜ = + @. lazy(ᶜdivᵥ_ρχ(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(specific(ᶜρχ, Y.c.ρ))))) @. ᶜρχₜ -= ᶜ∇ᵥρD∇χₜ # Rain and snow does not affect the mass - if χ_name == :q_tot + if ρχ_name == @name(ρq_tot) @. Yₜ.c.ρ -= ᶜ∇ᵥρD∇χₜ end end diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index daf1e27617..a932682049 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -213,10 +213,12 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) ᶜρ = Y.c.ρ # Full vertical advection of passive tracers (like liq, rai, etc) ... - for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific) - χ_name in (:e_tot, :q_tot) && continue - vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding) - @. ᶜρχₜ += vtt + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + if !(ρχ_name in (@name(ρe_tot), @name(ρq_tot))) + ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding) + @. ᶜρχₜ += vtt + end end # ... and upwinding correction of energy and total water. # (The central advection of energy and total water is done implicitly.) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 7ebd4b7fc5..48cc759f68 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -166,8 +166,8 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) @. Yₜ.c.ρe_tot += vst_ρe_tot # TODO: can we write this out explicitly? - for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific) - χ_name == :e_tot && continue + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge) @. ᶜρχₜ += vst_tracer @. Yₜ.c.ρ += vst_tracer # TODO: This doesn't look right for all tracers here. Remove? diff --git a/src/prognostic_equations/surface_flux.jl b/src/prognostic_equations/surface_flux.jl index 376d0d9bd3..da2d58c0fb 100644 --- a/src/prognostic_equations/surface_flux.jl +++ b/src/prognostic_equations/surface_flux.jl @@ -115,16 +115,16 @@ function surface_flux_tendency!(Yₜ, Y, p, t) btt = boundary_tendency_scalar(ᶜh_tot, sfc_conditions.ρ_flux_h_tot) @. Yₜ.c.ρe_tot -= btt ρ_flux_χ = p.scratch.sfc_temp_C3 - for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific) - χ_name == :e_tot && continue - if χ_name == :q_tot + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) + if ρχ_name == @name(ρq_tot) @. ρ_flux_χ = sfc_conditions.ρ_flux_q_tot else @. ρ_flux_χ = C3(FT(0)) end btt = boundary_tendency_scalar(ᶜχ, ρ_flux_χ) @. ᶜρχₜ -= btt - if χ_name == :q_tot + if ρχ_name == @name(ρq_tot) @. Yₜ.c.ρ -= btt end end diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index 1bce2989b5..569e65cfc3 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -46,7 +46,7 @@ This function is dispatched based on the type of the vertical diffusion model for scalars. Zero-flux boundary conditions are explicitly applied. - **Note on mass conservation for `q_tot` diffusion**: The current implementation also modifies the tendency of total moist air density `Yₜ.c.ρ` based on the - diffusion tendency of total specific humidity `ρq_tot`: + diffusion tendency of total specific humidity `ρq_tot`: `Yₜ.c.ρ -= ᶜρχₜ_diffusion_for_q_tot`. Arguments for all methods: @@ -58,7 +58,7 @@ Arguments for all methods: - `t`: Current simulation time (not directly used in diffusion calculations). - `vert_diff_model` (for dispatched methods): The specific vertical diffusion model instance. -Modifies components of tendency vector `Yₜ.c` (e.g., `Yₜ.c.uₕ`, `Yₜ.c.ρe_tot`, `Yₜ.c.ρ`, and +Modifies components of tendency vector `Yₜ.c` (e.g., `Yₜ.c.uₕ`, `Yₜ.c.ρe_tot`, `Yₜ.c.ρ`, and various tracer fields such as `Yₜ.c.ρq_tot`). """ @@ -96,23 +96,27 @@ function vertical_diffusion_boundary_layer_tendency!( ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar ᶜK_h_scaled = p.scratch.ᶜtemp_scalar_2 - for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific) - χ_name == :e_tot && continue - if χ_name in (:q_rai, :q_sno, :n_rai) + ᶜdivᵥ_ρχ = Operators.DivergenceF2C( + top = Operators.SetValue(C3(0)), + bottom = Operators.SetValue(C3(0)), + ) + foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + if ρχ_name in (@name(ρq_rai), @name(ρq_sno), @name(ρn_rai)) @. ᶜK_h_scaled = α_vert_diff_tracer * ᶜK_h else @. ᶜK_h_scaled = ᶜK_h end - ᶜdivᵥ_ρχ = Operators.DivergenceF2C( - top = Operators.SetValue(C3(0)), - bottom = Operators.SetValue(C3(0)), + @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρχ( + -( + ᶠinterp(Y.c.ρ) * + ᶠinterp(ᶜK_h_scaled) * + ᶠgradᵥ(specific(ᶜρχ, Y.c.ρ)) + ), ) - @. ᶜρχₜ_diffusion = - ᶜdivᵥ_ρχ(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h_scaled) * ᶠgradᵥ(ᶜχ))) @. ᶜρχₜ -= ᶜρχₜ_diffusion - # Only add contribution from total water diffusion to mass tendency - # (exclude contributions from diffusion of condensate, precipitation) - if χ_name == :q_tot + # Only add contribution from total water diffusion to mass tendency + # (exclude contributions from diffusion of condensate, precipitation) + if ρχ_name == @name(ρq_tot) @. Yₜ.c.ρ -= ᶜρχₜ_diffusion end end diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 5da8c6e32c..021c0c79dc 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -1,3 +1,5 @@ +import ClimaCore.MatrixFields: @name + """ specific(ρχ, ρ) specific(ρaχ, ρa, ρχ, ρ, turbconv_model) @@ -51,6 +53,66 @@ Arguments: return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ end +""" + tracer_names(field) + +Filters and returns the names of the variables from a given state +vector component, excluding `ρ`, `ρe_tot`, and `uₕ` and SGS fields. + +Arguments: + +- `field`: A component of the state vector `Y.c`. + +Returns: + +- A `Tuple` of `ClimaCore.MatrixFields.FieldName`s corresponding to the tracers. +""" +tracer_names(field) = + unrolled_filter(MatrixFields.top_level_names(field)) do name + !( + name in + (@name(ρ), @name(ρe_tot), @name(uₕ), @name(sgs⁰), @name(sgsʲs)) + ) + end + +""" + foreach_gs_tracer(f::F, Yₜ, Y) where {F} + +Applies a given function `f` to each grid-scale scalar variable (except `ρ` and `ρe_tot`) +in the state `Y` and its corresponding tendency `Yₜ`. +This utility abstracts the process of iterating over all scalars. It uses +`tracer_names` to identify the relevant variables and `unrolled_foreach` to +ensure a performant loop. For each tracer, it calls the provided function `f` +with the tendency field, the state field, and a boolean flag indicating if +the current tracer is `ρq_tot` (to allow for special handling). + +Arguments: + +- `f`: A function to apply to each grid-scale scalar. It must have the signature `f + (ᶜρχₜ, ᶜρχ, ρχ_name)`, where `ᶜρχₜ` is the tendency field, `ᶜρχ` + is the state field, and `ρχ_name` is a `MatrixFields.@name` object. +- `Yₜ`: The tendency state vector. +- `Y`: The current state vector. + +# Example + +```julia +foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name + # Apply some operation, e.g., a sponge layer + @. ᶜρχₜ += some_sponge_function(ᶜρχ) + if ρχ_name == @name(ρq_tot) + # Perform an additional operation only for ρq_tot + end +end +``` +""" +foreach_gs_tracer(f::F, Yₜ, Y) where {F} = + unrolled_foreach(tracer_names(Y.c)) do scalar_name + ᶜρχₜ = MatrixFields.get_field(Yₜ.c, scalar_name) + ᶜρχ = MatrixFields.get_field(Y.c, scalar_name) + f(ᶜρχₜ, ᶜρχ, scalar_name) + end + """ sgs_weight_function(a, a_half)