Fuse some broadcast expressions

Author: charleskawczynski

src/ClimaAtmos.jl

@@ -2,6 +2,7 @@ module ClimaAtmos
 
 using NVTX, Colors
 import Thermodynamics as TD
+import ClimaCore.Fields: @fused_direct
 
 include(joinpath("parameters", "Parameters.jl"))
 import .Parameters as CAP

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -39,15 +39,17 @@ NVTX.@annotate function set_diagnostic_edmfx_draft_quantities_level!(
     Φ_level,
 )
     FT = eltype(thermo_params)
-    @. ts_level = TD.PhaseEquil_phq(
-        thermo_params,
-        p_level,
-        mse_level - Φ_level,
-        q_tot_level,
-        8,
-        FT(0.0003),
-    )
-    @. ρ_level = TD.air_density(thermo_params, ts_level)
+    @fused_direct begin
+        @. ts_level = TD.PhaseEquil_phq(
+            thermo_params,
+            p_level,
+            mse_level - Φ_level,
+            q_tot_level,
+            8,
+            FT(0.0003),
+        )
+        @. ρ_level = TD.air_density(thermo_params, ts_level)
+    end
     return nothing
 end
 

src/cache/precipitation_precomputed_quantities.jl

@@ -22,21 +22,23 @@ function set_precipitation_precomputed_quantities!(Y, p, t)
     cmp = CAP.microphysics_precipitation_params(p.params)
 
     # compute the precipitation specific humidities
-    @. ᶜqᵣ = qₚ(Y.c.ρq_rai, Y.c.ρ)
-    @. ᶜqₛ = qₚ(Y.c.ρq_sno, Y.c.ρ)
+    @fused_direct begin
+        @. ᶜqᵣ = qₚ(Y.c.ρq_rai, Y.c.ρ)
+        @. ᶜqₛ = qₚ(Y.c.ρq_sno, Y.c.ρ)
 
-    # compute the precipitation terminal velocity [m/s]
-    @. ᶜwᵣ = CM1.terminal_velocity(
-        cmp.pr,
-        cmp.tv.rain,
-        Y.c.ρ,
-        abs(Y.c.ρq_rai / Y.c.ρ),
-    )
-    @. ᶜwₛ = CM1.terminal_velocity(
-        cmp.ps,
-        cmp.tv.snow,
-        Y.c.ρ,
-        abs(Y.c.ρq_sno / Y.c.ρ),
-    )
+        # compute the precipitation terminal velocity [m/s]
+        @. ᶜwᵣ = CM1.terminal_velocity(
+            cmp.pr,
+            cmp.tv.rain,
+            Y.c.ρ,
+            abs(Y.c.ρq_rai / Y.c.ρ),
+        )
+        @. ᶜwₛ = CM1.terminal_velocity(
+            cmp.ps,
+            cmp.tv.snow,
+            Y.c.ρ,
+            abs(Y.c.ρq_sno / Y.c.ρ),
+        )
+    end
     return nothing
 end

src/cache/precomputed_quantities.jl

@@ -479,8 +479,13 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
         # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ
         # TODO: We should think more about these increments before we use them.
     end
-    @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ)
-    @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
+    (; ᶜh_tot) = p.precomputed
+    @fused_direct begin
+        @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ)
+        @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
+        @. ᶜh_tot =
+            TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot)
+    end
 
     if turbconv_model isa AbstractEDMF
         @. p.precomputed.ᶜgradᵥ_θ_virt =
@@ -491,9 +496,6 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
             ᶜ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)
-
     if !isnothing(p.sfc_setup)
         SurfaceConditions.update_surface_conditions!(Y, p, t)
     end

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -25,27 +25,31 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
     (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
         p.precomputed
 
-    @. ᶜρa⁰ = ρa⁰(Y.c)
-    @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model)
-    @. ᶜmse⁰ = divide_by_ρa(
-        Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs),
-        ᶜρa⁰,
-        Y.c.ρ * (ᶜh_tot - ᶜK),
-        Y.c.ρ,
-        turbconv_model,
-    )
-    @. ᶜq_tot⁰ = divide_by_ρa(
-        Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs),
-        ᶜρa⁰,
-        Y.c.ρq_tot,
-        Y.c.ρ,
-        turbconv_model,
-    )
+    @fused_direct begin
+        @. ᶜρa⁰ = ρa⁰(Y.c)
+        @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model)
+        @. ᶜmse⁰ = divide_by_ρa(
+            Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs),
+            ᶜρa⁰,
+            Y.c.ρ * (ᶜh_tot - ᶜK),
+            Y.c.ρ,
+            turbconv_model,
+        )
+        @. ᶜq_tot⁰ = divide_by_ρa(
+            Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs),
+            ᶜρa⁰,
+            Y.c.ρq_tot,
+            Y.c.ρ,
+            turbconv_model,
+        )
+    end
     set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
     set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
     # @. ᶜK⁰ += ᶜtke⁰
-    @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
-    @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
+    @fused_direct begin
+        @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
+        @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
+    end
     return nothing
 end
 
@@ -90,8 +94,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
 
         set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³)
         @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
-        @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
-        @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
+        @fused_direct begin
+            @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
+            @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
+        end
 
         # EDMFX boundary condition:
 
@@ -324,12 +330,13 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
         ᶜtke_exch,
     )
 
-    @. ᶜmixing_length = ᶜmixing_length_tuple.master
-
     turbconv_params = CAP.turbconv_params(params)
     c_m = CAP.tke_ed_coeff(turbconv_params)
-    @. ᶜK_u = c_m * ᶜmixing_length * sqrt(max(ᶜtke⁰, 0))
-    @. ᶜK_h = ᶜK_u / ᶜprandtl_nvec
+    @fused_direct begin
+        @. ᶜmixing_length = ᶜmixing_length_tuple.master
+        @. ᶜK_u = c_m * ᶜmixing_length * sqrt(max(ᶜtke⁰, 0))
+        @. ᶜK_h = ᶜK_u / ᶜprandtl_nvec
+    end
 
     ρatke_flux_values = Fields.field_values(ρatke_flux)
     ρ_int_values = Fields.field_values(Fields.level(ᶜρa⁰, 1))

src/parameterized_tendencies/microphysics/cloud_condensate.jl

@@ -15,6 +15,9 @@ function cloud_condensate_tendency!(Yₜ, p, colidx, ::NonEquilMoistModel)
     thp = CAP.thermodynamics_params(params)
     cmc = CAP.microphysics_cloud_params(params)
 
-    @. Yₜ.c.ρq_liq[colidx] += cloud_sources(cmc.liquid, thp, ᶜts[colidx], dt)
-    @. Yₜ.c.ρq_ice[colidx] += cloud_sources(cmc.ice, thp, ᶜts[colidx], dt)
+    @fused_direct begin
+        @. Yₜ.c.ρq_liq[colidx] +=
+            cloud_sources(cmc.liquid, thp, ᶜts[colidx], dt)
+        @. Yₜ.c.ρq_ice[colidx] += cloud_sources(cmc.ice, thp, ᶜts[colidx], dt)
+    end
 end