Try ClimaCore MultiBroadcastFusion

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

examples/Manifest.toml

@@ -319,8 +319,10 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 version = "0.5.8"
 
 [[deps.ClimaCore]]
-deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "Unrolled"]
-git-tree-sha1 = "3c8ba2a248d30aaa7f0e72a3516b142521447377"
+deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "Unrolled"]
+git-tree-sha1 = "65cd39bb2fe20bd30421aea4d84af6628675d7b5"
+repo-rev = "ck/mbf"
+repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 version = "0.13.3"
 weakdeps = ["Krylov"]
@@ -1582,6 +1584,11 @@ git-tree-sha1 = "cac9cc5499c25554cba55cd3c30543cff5ca4fab"
 uuid = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 version = "0.2.4"
 
+[[deps.MultiBroadcastFusion]]
+git-tree-sha1 = "66f29657d5efca986e4eef5eee9960a4dd222e2b"
+uuid = "c3c07f87-98de-43f2-a76f-835b330b2cbb"
+version = "0.1.0"
+
 [[deps.Multisets]]
 git-tree-sha1 = "8d852646862c96e226367ad10c8af56099b4047e"
 uuid = "3b2b4ff1-bcff-5658-a3ee-dbcf1ce5ac09"

perf/Manifest.toml

@@ -334,8 +334,10 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 version = "0.5.8"
 
 [[deps.ClimaCore]]
-deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "Unrolled"]
-git-tree-sha1 = "3c8ba2a248d30aaa7f0e72a3516b142521447377"
+deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "Unrolled"]
+git-tree-sha1 = "65cd39bb2fe20bd30421aea4d84af6628675d7b5"
+repo-rev = "ck/mbf"
+repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 version = "0.13.3"
 weakdeps = ["Krylov"]
@@ -1644,6 +1646,11 @@ git-tree-sha1 = "cac9cc5499c25554cba55cd3c30543cff5ca4fab"
 uuid = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 version = "0.2.4"
 
+[[deps.MultiBroadcastFusion]]
+git-tree-sha1 = "66f29657d5efca986e4eef5eee9960a4dd222e2b"
+uuid = "c3c07f87-98de-43f2-a76f-835b330b2cbb"
+version = "0.1.0"
+
 [[deps.Multisets]]
 git-tree-sha1 = "8d852646862c96e226367ad10c8af56099b4047e"
 uuid = "3b2b4ff1-bcff-5658-a3ee-dbcf1ce5ac09"

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

@@ -21,22 +21,24 @@ 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 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

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

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
 

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]) *

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

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -253,8 +253,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 +341,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

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

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

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