Merge pull request #2926 from CliMA/ck/gpu_prog_edmf

Hoist BidiagonalMatrixRow expressions for prog edmf

Author: charleskawczynski

.buildkite/pipeline.yml

@@ -762,7 +762,6 @@ steps:
         agents:
           slurm_gpus: 1
           slurm_mem: 20G
-        soft_fail: true
 
   - group: "GPU Performance"
     steps:

src/cache/precomputed_quantities.jl

@@ -93,6 +93,9 @@ function precomputed_quantities(Y, atmos)
             ᶜKʲs = similar(Y.c, NTuple{n, FT}),
             ᶠKᵥʲs = similar(Y.f, NTuple{n, FT}),
             ᶜtsʲs = similar(Y.c, NTuple{n, TST}),
+            bdmr_l = similar(Y.c, BidiagonalMatrixRow{FT}),
+            bdmr_r = similar(Y.c, BidiagonalMatrixRow{FT}),
+            bdmr = similar(Y.c, BidiagonalMatrixRow{FT}),
             ᶜρʲs = similar(Y.c, NTuple{n, FT}),
             ᶜentrʲs = similar(Y.c, NTuple{n, FT}),
             ᶜdetrʲs = similar(Y.c, NTuple{n, FT}),

src/prognostic_equations/implicit/implicit_solver.jl

@@ -429,6 +429,9 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
                 p.precomputed.ᶜρʲs,
                 p.precomputed.ᶠu³ʲs,
                 p.precomputed.ᶜtsʲs,
+                p.precomputed.bdmr_l,
+                p.precomputed.bdmr_r,
+                p.precomputed.bdmr,
             ) : (;)
         )...,
         p.core.ᶜΦ,
@@ -724,6 +727,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     if p.atmos.turbconv_model isa PrognosticEDMFX
         if use_derivative(sgs_advection_flag)
             (; ᶜgradᵥ_ᶠΦ, ᶜρʲs, ᶠu³ʲs, ᶜtsʲs) = p
+            (; bdmr_l, bdmr_r, bdmr) = p
             is_third_order = edmfx_upwinding == Val(:third_order)
             ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1
             ᶠset_upwind_bcs = Operators.SetBoundaryOperator(;
@@ -884,12 +888,19 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
                 matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)]
             ᶜu₃ʲ = ᶜtemp_C3
             @. ᶜu₃ʲ[colidx] = ᶜinterp(Y.f.sgsʲs.:(1).u₃[colidx])
+
+            @. bdmr_l[colidx] =
+                convert(BidiagonalMatrixRow{FT}, ᶜleft_bias_matrix())
+            @. bdmr_r[colidx] =
+                convert(BidiagonalMatrixRow{FT}, ᶜright_bias_matrix())
+            @. bdmr[colidx] = ifelse(
+                ᶜu₃ʲ[colidx].components.data.:1 > 0,
+                bdmr_l[colidx],
+                bdmr_r[colidx],
+            )
+
             @. ᶠtridiagonal_matrix_c3[colidx] =
-                -(ᶠgradᵥ_matrix()) ⋅ ifelse(
-                    ᶜu₃ʲ[colidx].components.data.:1 > 0,
-                    convert(BidiagonalMatrixRow{FT}, ᶜleft_bias_matrix()),
-                    convert(BidiagonalMatrixRow{FT}, ᶜright_bias_matrix()),
-                )
+                -(ᶠgradᵥ_matrix()) ⋅ bdmr[colidx]
             if p.atmos.rayleigh_sponge isa RayleighSponge
                 @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ[colidx] =
                     dtγ * (