Merge pull request #2544 from CliMA/ck/buoyancy_grad_perf

Try improving buoyancy gradient kernel performance
CliMA · Jan 19, 2024 · 82fd759 · 82fd759
2 parents 3cf3f10 + 9570912
commit 82fd759
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 31 deletions.
diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -765,17 +765,13 @@ function set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t)
         EnvBuoyGradVars(
             thermo_params,
             ᶜts,
-            projected_vector_data(
+            projected_vector_buoy_grad_vars(
                 C3,
-                ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))),
+                ᶜ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
                 ᶜlg,
-            ),                                                                 # ∂θv∂z_unsat
-            projected_vector_data(C3, ᶜgradᵥ(ᶠinterp(q_tot)), ᶜlg),            # ∂qt∂z_sat
-            projected_vector_data(
-                C3,
-                ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))),
-                ᶜlg,
-            ),                                                                 # ∂θl∂z_sat
+            ),
         ),
     )
 

diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl
@@ -245,17 +245,13 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
         EnvBuoyGradVars(
             thermo_params,
             ᶜts⁰,
-            projected_vector_data(
+            projected_vector_buoy_grad_vars(
                 C3,
-                ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))),
+                ᶜ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
                 ᶜlg,
-            ),                                                                 # ∂θv∂z_unsat
-            projected_vector_data(C3, ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)), ᶜlg),          # ∂qt∂z_sat
-            projected_vector_data(
-                C3,
-                ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))),
-                ᶜlg,
-            ),                                                                 # ∂θl∂z_sat
+            ),
         ),
     )
 

diff --git a/src/prognostic_equations/gm_sgs_closures.jl b/src/prognostic_equations/gm_sgs_closures.jl
@@ -39,21 +39,13 @@ function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
         EnvBuoyGradVars(
             thermo_params,
             ᶜts,
-            projected_vector_data(
+            projected_vector_buoy_grad_vars(
                 C3,
-                ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))),
+                ᶜ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
                 ᶜlg,
-            ),                                                                 # ∂θv∂z_unsat
-            projected_vector_data(
-                C3,
-                ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts))),
-                ᶜlg,
-            ),            # ∂qt∂z_sat
-            projected_vector_data(
-                C3,
-                ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))),
-                ᶜlg,
-            ),                                                                 # ∂θl∂z_sat
+            ),
         ),
     )
 

diff --git a/src/solver/types.jl b/src/solver/types.jl
@@ -160,6 +160,14 @@ Base.@kwdef struct EnvBuoyGradVars{FT}
     ∂θl∂z_sat::FT
 end
 
+function EnvBuoyGradVars(
+    thermo_params,
+    ts::TD.ThermodynamicState,
+    ∂θv∂z_unsat_∂qt∂z_sat_∂θl∂z_sat,
+)
+    (; ∂θv∂z_unsat, ∂qt∂z_sat, ∂θl∂z_sat) = ∂θv∂z_unsat_∂qt∂z_sat_∂θl∂z_sat
+    return EnvBuoyGradVars(thermo_params, ts, ∂θv∂z_unsat, ∂qt∂z_sat, ∂θl∂z_sat)
+end
 
 function EnvBuoyGradVars(
     thermo_params,

diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl
@@ -201,6 +201,15 @@ The type should correspond to a vector with only one component, i.e., a basis ve
 projected_vector_data(::Type{V}, vector, local_geometry) where {V} =
     V(vector, local_geometry)[1] / unit_basis_vector_data(V, local_geometry)
 
+function projected_vector_buoy_grad_vars(::Type{V}, v1, v2, v3, lg) where {V}
+    ubvd = unit_basis_vector_data(V, lg)
+    return (;
+        ∂θv∂z_unsat = V(v1, lg)[1] / ubvd,
+        ∂qt∂z_sat = V(v2, lg)[1] / ubvd,
+        ∂θl∂z_sat = V(v3, lg)[1] / ubvd,
+    )
+end
+
 """
     get_physical_w(u, local_geometry)