Decorate with at-fused

CliMA · Apr 12, 2024 · a35f3af · a35f3af
1 parent e4906ea
commit a35f3af
Show file tree

Hide file tree

Showing 11 changed files with 129 additions and 100 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -23,11 +23,6 @@ steps:
     command:
       - "echo $$JULIA_DEPOT_PATH"
 
-      - echo "--- Instantiate project"
-      - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
-      - "julia --project -e 'using Pkg; Pkg.precompile()'"
-      - "julia --project -e 'using Pkg; Pkg.status()'"
-
       - echo "--- Instantiate examples"
       - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
       - "julia --project=examples -e 'using Pkg; Pkg.precompile()'"

diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
@@ -2,6 +2,7 @@ module ClimaAtmos
 
 using NVTX, Colors
 import Thermodynamics as TD
+import ClimaCore.DataLayouts: @fused
 
 include(joinpath("parameters", "Parameters.jl"))
 import .Parameters as CAP

diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl
@@ -21,22 +21,24 @@ function set_precipitation_precomputed_quantities!(Y, p, t)
 
     cmp = CAP.microphysics_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 begin
+        # compute the precipitation specific humidities
+        @. ᶜ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
diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -472,8 +472,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 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 =
@@ -484,9 +489,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

diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/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 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 begin
+        @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
+        @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
+    end
     return nothing
 end
 

diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl
@@ -283,17 +283,19 @@ function edmfx_entr_detr_tendency!(
 
     for j in 1:n
 
-        @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
-            Y.c.sgsʲs.:($$j).ρa[colidx] *
-            (ᶜentrʲs.:($$j)[colidx] - ᶜdetrʲs.:($$j)[colidx])
+        @fused begin
+            @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
+                Y.c.sgsʲs.:($$j).ρa[colidx] *
+                (ᶜentrʲs.:($$j)[colidx] - ᶜdetrʲs.:($$j)[colidx])
 
-        @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
-            ᶜentrʲs.:($$j)[colidx] *
-            (ᶜmse⁰[colidx] - Y.c.sgsʲs.:($$j).mse[colidx])
+            @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
+                ᶜentrʲs.:($$j)[colidx] *
+                (ᶜmse⁰[colidx] - Y.c.sgsʲs.:($$j).mse[colidx])
 
-        @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
-            ᶜentrʲs.:($$j)[colidx] *
-            (ᶜq_tot⁰[colidx] - Y.c.sgsʲs.:($$j).q_tot[colidx])
+            @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
+                ᶜentrʲs.:($$j)[colidx] *
+                (ᶜq_tot⁰[colidx] - Y.c.sgsʲs.:($$j).q_tot[colidx])
+        end
 
         @. Yₜ.f.sgsʲs.:($$j).u₃[colidx] +=
             ᶠinterp(ᶜentrʲs.:($$j)[colidx]) *

diff --git a/src/prognostic_equations/edmfx_precipitation.jl b/src/prognostic_equations/edmfx_precipitation.jl
@@ -28,20 +28,22 @@ function edmfx_precipitation_tendency!(
 
     for j in 1:n
 
-        @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
-            Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx]
+        @fused begin
+            @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
+                Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx]
 
-        @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
-            ᶜSqₜᵖʲs.:($$j)[colidx] * (
-                e_tot_0M_precipitation_sources_helper(
-                    thermo_params,
-                    ᶜtsʲs.:($$j)[colidx],
-                    ᶜΦ[colidx],
-                ) - TD.internal_energy(thermo_params, ᶜtsʲs.:($$j)[colidx])
-            )
+            @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
+                ᶜSqₜᵖʲs.:($$j)[colidx] * (
+                    e_tot_0M_precipitation_sources_helper(
+                        thermo_params,
+                        ᶜtsʲs.:($$j)[colidx],
+                        ᶜΦ[colidx],
+                    ) - TD.internal_energy(thermo_params, ᶜtsʲs.:($$j)[colidx])
+                )
 
-        @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
-            ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx])
+            @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
+                ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx])
+        end
     end
     return nothing
 end
@@ -62,16 +64,18 @@ function edmfx_precipitation_tendency!(
 
     for j in 1:n
 
-        @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
-            Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx]
+        @fused begin
+            @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
+                Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜSqₜᵖʲs.:($$j)[colidx]
 
-        @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
-            ᶜSeₜᵖʲs.:($$j)[colidx] -
-            ᶜSqₜᵖʲs.:($$j)[colidx] *
-            TD.internal_energy(thp, ᶜtsʲs.:($$j)[colidx])
+            @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
+                ᶜSeₜᵖʲs.:($$j)[colidx] -
+                ᶜSqₜᵖʲs.:($$j)[colidx] *
+                TD.internal_energy(thp, ᶜtsʲs.:($$j)[colidx])
 
-        @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
-            ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx])
+            @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
+                ᶜSqₜᵖʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx])
+        end
     end
     return nothing
 end
diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl
@@ -26,12 +26,15 @@ function edmfx_sgs_mass_flux_tendency!(
         ᶠu³_diff_colidx = p.scratch.ᶠtemp_CT3[colidx]
         ᶜa_scalar_colidx = p.scratch.ᶜtemp_scalar[colidx]
         for j in 1:n
-            @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
-            @. ᶜa_scalar_colidx =
-                (
-                    Y.c.sgsʲs.:($$j).mse[colidx] + ᶜKʲs.:($$j)[colidx] -
-                    ᶜh_tot[colidx]
-                ) * draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx])
+            @fused begin
+                @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
+                @. ᶜa_scalar_colidx =
+                    (
+                        Y.c.sgsʲs.:($$j).mse[colidx] + ᶜKʲs.:($$j)[colidx] -
+                        ᶜh_tot[colidx]
+                    ) *
+                    draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx])
+            end
             vertical_transport!(
                 Yₜ.c.ρe_tot[colidx],
                 ᶜJ[colidx],
@@ -42,10 +45,12 @@ function edmfx_sgs_mass_flux_tendency!(
                 edmfx_sgsflux_upwinding,
             )
         end
-        @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx]
-        @. ᶜa_scalar_colidx =
-            (ᶜmse⁰[colidx] + ᶜK⁰[colidx] - ᶜh_tot[colidx]) *
-            draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx])
+        @fused begin
+            @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx]
+            @. ᶜa_scalar_colidx =
+                (ᶜmse⁰[colidx] + ᶜK⁰[colidx] - ᶜh_tot[colidx]) *
+                draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx])
+        end
         vertical_transport!(
             Yₜ.c.ρe_tot[colidx],
             ᶜJ[colidx],
@@ -73,10 +78,12 @@ function edmfx_sgs_mass_flux_tendency!(
                     edmfx_sgsflux_upwinding,
                 )
             end
-            @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx]
-            @. ᶜa_scalar_colidx =
-                (ᶜq_tot⁰[colidx] - ᶜspecific.q_tot[colidx]) *
-                draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx])
+            @fused begin
+                @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx]
+                @. ᶜa_scalar_colidx =
+                    (ᶜq_tot⁰[colidx] - ᶜspecific.q_tot[colidx]) *
+                    draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx])
+            end
             vertical_transport!(
                 Yₜ.c.ρq_tot[colidx],
                 ᶜJ[colidx],
@@ -253,8 +260,10 @@ function edmfx_sgs_diffusive_flux_tendency!(
             )
             @. ᶜρχₜ_diffusion[colidx] =
                 ᶜdivᵥ_ρq_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜq_tot⁰[colidx])))
-            @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx]
-            @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx]
+            @fused begin
+                @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx]
+                @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx]
+            end
         end
 
         # momentum
@@ -339,8 +348,10 @@ function edmfx_sgs_diffusive_flux_tendency!(
             @. ᶜρχₜ_diffusion[colidx] = ᶜdivᵥ_ρq_tot(
                 -(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜspecific.q_tot[colidx])),
             )
-            @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx]
-            @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx]
+            @fused begin
+                @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx]
+                @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx]
+            end
         end
 
         # momentum

diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl
@@ -36,11 +36,13 @@ function edmfx_tke_tendency!(
 
     if use_prognostic_tke(turbconv_model)
         # shear production
-        @. Yₜ.c.sgs⁰.ρatke[colidx] +=
-            2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx]
-        # buoyancy production
-        @. Yₜ.c.sgs⁰.ρatke[colidx] -=
-            ᶜρa⁰[colidx] * ᶜK_h[colidx] * ᶜlinear_buoygrad[colidx]
+        @fused begin
+            @. Yₜ.c.sgs⁰.ρatke[colidx] +=
+                2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx]
+            # buoyancy production
+            @. Yₜ.c.sgs⁰.ρatke[colidx] -=
+                ᶜρa⁰[colidx] * ᶜK_h[colidx] * ᶜlinear_buoygrad[colidx]
+        end
         # entrainment and detraiment
         # using ᶜu⁰ and local geometry results in allocation
         for j in 1:n

diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl
@@ -92,8 +92,10 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)
                 C123(wgradₕ(divₕ(p.precomputed.ᶜuʲs.:($$j)))) -
                 C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜuʲs.:($$j)))))
             @. ᶜ∇²mseʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).mse))
-            @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j))
-            @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j))
+            @fused begin
+                @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j))
+                @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j))
+            end
         end
     end
 end

diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl
@@ -95,9 +95,13 @@ function vertical_diffusion_boundary_layer_tendency!(
                 ᶠgradᵥ(ᶜχ[colidx])
             ),
         )
-        @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx]
         if !(χ_name in (:q_rai, :q_sno))
-            @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx]
+            @fused begin
+                @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx]
+                @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx]
+            end
+        else
+            @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx]
         end
     end
 end