Decorate with at-fused, try ClimaCore

.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()'"

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

src/cache/precipitation_precomputed_quantities.jl

@@ -17,17 +17,19 @@ function set_precipitation_precomputed_quantities!(Y, p, t)
     cmp = CAP.microphysics_params(p.params)
 
     # 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.ρ),
-    )
+    @fused begin
+        @. ᶜ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

@@ -454,8 +454,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 =
@@ -466,8 +471,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)

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -19,27 +19,31 @@ function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³
     (; ᶜ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
 
@@ -77,9 +81,6 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ
         ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot
 
         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ʲ)
 
         # EDMFX boundary condition:
 
@@ -109,49 +110,60 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ
         ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1))
         ᶜK_int_val = Fields.field_values(Fields.level(ᶜK, 1))
         ᶜmseʲ_int_val = Fields.field_values(Fields.level(ᶜmseʲ, 1))
-        @. ᶜmseʲ_int_val = sgs_scalar_first_interior_bc(
-            ᶜz_int_val - z_sfc_val,
-            ᶜρ_int_val,
-            ᶜh_tot_int_val - ᶜK_int_val,
-            buoyancy_flux_val,
-            ρ_flux_h_tot_val,
-            ustar_val,
-            obukhov_length_val,
-            sfc_local_geometry_val,
-        )
-
         # ... and the first interior point for EDMFX ᶜq_totʲ.
         ᶜq_tot_int_val = Fields.field_values(Fields.level(ᶜspecific.q_tot, 1))
         ᶜq_totʲ_int_val = Fields.field_values(Fields.level(ᶜq_totʲ, 1))
-        @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc(
-            ᶜz_int_val - z_sfc_val,
-            ᶜρ_int_val,
-            ᶜq_tot_int_val,
-            buoyancy_flux_val,
-            ρ_flux_q_tot_val,
-            ustar_val,
-            obukhov_length_val,
-            sfc_local_geometry_val,
-        )
 
-        # Then overwrite the prognostic variables at first inetrior point.
         ᶜΦ_int_val = Fields.field_values(Fields.level(ᶜΦ, 1))
         ᶜtsʲ_int_val = Fields.field_values(Fields.level(ᶜtsʲ, 1))
-        @. ᶜtsʲ_int_val = TD.PhaseEquil_phq(
-            thermo_params,
-            ᶜp_int_val,
-            ᶜmseʲ_int_val - ᶜΦ_int_val,
-            ᶜq_totʲ_int_val,
-        )
+
         sgsʲs_ρ_int_val = Fields.field_values(Fields.level(ᶜρʲs.:($j), 1))
         sgsʲs_ρa_int_val =
             Fields.field_values(Fields.level(Y.c.sgsʲs.:($j).ρa, 1))
 
         turbconv_params = CAP.turbconv_params(params)
-        @. sgsʲs_ρ_int_val = TD.air_density(thermo_params, ᶜtsʲ_int_val)
-        @. sgsʲs_ρa_int_val =
-            $(FT(turbconv_params.surface_area)) *
-            TD.air_density(thermo_params, ᶜtsʲ_int_val)
+
+        @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
+        @fused begin
+            @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
+            @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
+        end
+
+        @fused begin
+            @. ᶜmseʲ_int_val = sgs_scalar_first_interior_bc(
+                ᶜz_int_val - z_sfc_val,
+                ᶜρ_int_val,
+                ᶜh_tot_int_val - ᶜK_int_val,
+                buoyancy_flux_val,
+                ρ_flux_h_tot_val,
+                ustar_val,
+                obukhov_length_val,
+                sfc_local_geometry_val,
+            )
+
+            @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc(
+                ᶜz_int_val - z_sfc_val,
+                ᶜρ_int_val,
+                ᶜq_tot_int_val,
+                buoyancy_flux_val,
+                ρ_flux_q_tot_val,
+                ustar_val,
+                obukhov_length_val,
+                sfc_local_geometry_val,
+            )
+
+            # Then overwrite the prognostic variables at first inetrior point.
+            @. ᶜtsʲ_int_val = TD.PhaseEquil_phq(
+                thermo_params,
+                ᶜp_int_val,
+                ᶜmseʲ_int_val - ᶜΦ_int_val,
+                ᶜq_totʲ_int_val,
+            )
+            @. sgsʲs_ρ_int_val = TD.air_density(thermo_params, ᶜtsʲ_int_val)
+            @. sgsʲs_ρa_int_val =
+                $(FT(turbconv_params.surface_area)) *
+                TD.air_density(thermo_params, ᶜtsʲ_int_val)
+        end
     end
     return nothing
 end

src/prognostic_equations/edmfx_entr_detr.jl

@@ -293,17 +293,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]) *

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] * ᶜS_q_totʲs.:($$j)[colidx]
+        @fused begin
+            @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
+                Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜS_q_totʲs.:($$j)[colidx]
 
-        @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
-            ᶜS_q_totʲ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] +=
+                ᶜS_q_totʲ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] +=
-            ᶜS_q_totʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx])
+            @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
+                ᶜS_q_totʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx])
+        end
     end
     return nothing
 end

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],
@@ -59,10 +64,17 @@ function edmfx_sgs_mass_flux_tendency!(
         if !(p.atmos.moisture_model isa DryModel)
             # specific humidity
             for j in 1:n
-                @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
-                @. ᶜa_scalar_colidx =
-                    (Y.c.sgsʲs.:($$j).q_tot[colidx] - ᶜspecific.q_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).q_tot[colidx] -
+                            ᶜspecific.q_tot[colidx]
+                        ) * draft_area(
+                            Y.c.sgsʲs.:($$j).ρa[colidx],
+                            ᶜρʲs.:($$j)[colidx],
+                        )
+                end
                 vertical_transport!(
                     Yₜ.c.ρq_tot[colidx],
                     ᶜJ[colidx],
@@ -73,10 +85,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 +267,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