Hoist and precompute common variables

perf/flame.jl

@@ -38,16 +38,16 @@ ProfileCanvas.html_file(joinpath(output_dir, "flame.html"), results)
 #####
 
 allocs_limit = Dict()
-allocs_limit["flame_perf_target"] = 275_800
-allocs_limit["flame_perf_target_tracers"] = 305_776
+allocs_limit["flame_perf_target"] = 278360
+allocs_limit["flame_perf_target_tracers"] = 308336
 allocs_limit["flame_perf_target_edmfx"] = 7_005_552
-allocs_limit["flame_perf_diagnostics"] = 108_877_816
-allocs_limit["flame_perf_target_diagnostic_edmfx"] = 401_944
+allocs_limit["flame_perf_diagnostics"] = 108880760
+allocs_limit["flame_perf_target_diagnostic_edmfx"] = 412056
 allocs_limit["flame_sphere_baroclinic_wave_rhoe_equilmoist_expvdiff"] =
     4_018_252_656
-allocs_limit["flame_perf_target_frierson"] = 4_015_545_328
+allocs_limit["flame_perf_target_frierson"] = 4015547056
 allocs_limit["flame_perf_target_threaded"] = 1_276_864
-allocs_limit["flame_perf_target_callbacks"] = 395048
+allocs_limit["flame_perf_target_callbacks"] = 398984
 allocs_limit["flame_perf_gw"] = 3_268_961_856
 allocs_limit["flame_perf_target_prognostic_edmfx_aquaplanet"] = 299_616
 allocs_limit["flame_gpu_implicit_barowave_moist"] = 4300000
@@ -88,7 +88,7 @@ using Test
 # Threaded allocations are not deterministic, so let's add a buffer
 # TODO: remove buffer, and threaded tests, when
 #       threaded/unthreaded functions are unified
-buffer = occursin("threaded", job_id) ? 1.4 : 1
+buffer = occursin("threaded", job_id) ? 1.4 : 1.1
 
 
 ## old allocation profiler (TODO: remove this)

src/cache/cloud_fraction.jl

@@ -55,29 +55,22 @@ function set_cloud_fraction!(
     coeff = FT(2.1) # TODO - move to parameters
     @. ᶜcloud_fraction = quad_loop(
         SG_quad,
-        ᶜp,
-        TD.total_specific_humidity(thermo_params, ᶜts),
-        TD.liquid_ice_pottemp(thermo_params, ᶜts),
-        Geometry.WVector(
-            ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts))),
-        ),
-        Geometry.WVector(
-            ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))),
-        ),
+        ᶜts,
+        Geometry.WVector(p.precomputed.ᶜgradᵥ_q_tot),
+        Geometry.WVector(p.precomputed.ᶜgradᵥ_θ_liq_ice),
         coeff,
         ᶜmixing_length,
         thermo_params,
     )
 end
 
 """
-    function quad_loop(SG_quad, p_c, q_mean, θ_mean, ᶜ∇q, ᶜ∇θ,
+    function quad_loop(SG_quad, ts, ᶜ∇q, ᶜ∇θ,
                        coeff, ᶜlength_scale, thermo_params)
 
 where:
   - SG_quad is a struct containing information about quadrature type and order
-  - p_c is the atmospheric pressure
-  - q_mean, θ_mean is the grid mean q_tot and liquid ice potential temperature
+  - ts is the thermodynamic state
   - ᶜ∇q, ᶜ∇θ are the gradients of q_tot and liquid ice potential temperature
   - coeff - a free parameter (to be moved into params)
   - ᶜlength_scale - mixing length for simulations with EDMF and Smagorinsky
@@ -89,15 +82,16 @@ returns cloud fraction computed as a sum over quadrature points.
 """
 function quad_loop(
     SG_quad::SGSQuadrature,
-    p_c,
-    q_mean,
-    θ_mean,
+    ts,
     ᶜ∇q,
     ᶜ∇θ,
     coeff,
     ᶜlength_scale,
     thermo_params,
 )
+    p_c = TD.air_pressure(thermo_params, ts)
+    q_mean = TD.total_specific_humidity(thermo_params, ts)
+    θ_mean = TD.liquid_ice_pottemp(thermo_params, ts)
     # Returns the physical values based on quadrature sampling points
     # and limited covarainces
     function get_x_hat(χ1, χ2)
@@ -139,8 +133,8 @@ function quad_loop(
         FT = eltype(x1_hat)
         @assert(x1_hat >= FT(0))
         @assert(x2_hat >= FT(0))
-        ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
-        hc = TD.has_condensate(thermo_params, ts)
+        _ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
+        hc = TD.has_condensate(thermo_params, _ts)
         return (;
             cf = hc ? FT(1) : FT(0), # cloud fraction
             q_tot_sat = hc ? x2_hat : FT(0), # cloudy/dry for buoyancy in TKE

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -767,9 +767,9 @@ function set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t)
             ᶜts,
             projected_vector_buoy_grad_vars(
                 C3,
-                ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))),    # ∂θv∂z_unsat
-                ᶜgradᵥ(ᶠinterp(q_tot)),                                     # ∂qt∂z_sat
-                ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))), # ∂θl∂z_sat
+                p.precomputed.ᶜgradᵥ_θ_virt,    # ∂θv∂z_unsat
+                p.precomputed.ᶜgradᵥ_q_tot,     # ∂qt∂z_sat
+                p.precomputed.ᶜgradᵥ_θ_liq_ice, # ∂θl∂z_sat
                 ᶜlg,
             ),
         ),

src/cache/precomputed_quantities.jl

@@ -41,6 +41,8 @@ function precomputed_quantities(Y, atmos)
     @assert isnothing(atmos.turbconv_model) || isnothing(atmos.vert_diff)
     TST = thermo_state_type(atmos.moisture_model, FT)
     SCT = SurfaceConditions.surface_conditions_type(atmos, FT)
+    cspace = axes(Y.c)
+    fspace = axes(Y.f)
     n = n_mass_flux_subdomains(atmos.turbconv_model)
     gs_quantities = (;
         ᶜspecific = specific_gs.(Y.c),
@@ -83,6 +85,12 @@ function precomputed_quantities(Y, atmos)
             ᶜS_q_totʲs = similar(Y.c, NTuple{n, FT}),
             ᶜS_q_tot⁰ = similar(Y.c, FT),
         ) : (;)
+    sgs_quantities = (;
+        ᶜgradᵥ_θ_virt = Fields.Field(C3{FT}, cspace),
+        ᶜgradᵥ_q_tot = Fields.Field(C3{FT}, cspace),
+        ᶜgradᵥ_θ_liq_ice = Fields.Field(C3{FT}, cspace),
+    )
+
     diagnostic_sgs_quantities =
         atmos.turbconv_model isa DiagnosticEDMFX ?
         (;
@@ -124,6 +132,7 @@ function precomputed_quantities(Y, atmos)
         (; ᶜwᵣ = similar(Y.c, FT), ᶜwₛ = similar(Y.c, FT)) : (;)
     return (;
         gs_quantities...,
+        sgs_quantities...,
         advective_sgs_quantities...,
         diagnostic_sgs_quantities...,
         vert_diff_quantities...,
@@ -447,6 +456,15 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ)
     @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
 
+    if turbconv_model isa AbstractEDMF
+        @. p.precomputed.ᶜgradᵥ_θ_virt =
+            ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts)))
+        @. p.precomputed.ᶜgradᵥ_q_tot =
+            ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts)))
+        @. p.precomputed.ᶜgradᵥ_θ_liq_ice =
+            ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
+    end
+
     (; ᶜh_tot) = p.precomputed
     @. ᶜh_tot = TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot)
 

src/callbacks/callbacks.jl

@@ -38,6 +38,16 @@ end
 NVTX.@annotate function cloud_fraction_model_callback!(integrator)
     Y = integrator.u
     p = integrator.p
+    (; ᶜts, ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
+    thermo_params = CAP.thermodynamics_params(p.params)
+    if isnothing(p.atmos.turbconv_model)
+        @. ᶜgradᵥ_θ_virt =
+            ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts)))
+        @. ᶜgradᵥ_q_tot =
+            ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts)))
+        @. ᶜgradᵥ_θ_liq_ice =
+            ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
+    end
     set_cloud_fraction!(Y, p, p.atmos.moisture_model, p.atmos.cloud_model)
 end
 

src/prognostic_equations/gm_sgs_closures.jl

@@ -41,9 +41,9 @@ function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
             ᶜts,
             projected_vector_buoy_grad_vars(
                 C3,
-                ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))),         # ∂θv∂z_unsat
-                ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts))), # ∂qt∂z_sat
-                ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))),      # ∂θl∂z_sat
+                p.precomputed.ᶜgradᵥ_θ_virt,    # ∂θv∂z_unsat
+                p.precomputed.ᶜgradᵥ_q_tot,     # ∂qt∂z_sat
+                p.precomputed.ᶜgradᵥ_θ_liq_ice, # ∂θl∂z_sat
                 ᶜlg,
             ),
         ),