Merge pull request #3858 from CliMA/ck/slim_cache5

Remove use of matching subfields

Author: charleskawczynski

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)

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

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

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.)

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?

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

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

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)