Improve tracer performance in jacobian update

Split up a high register pressure broadcast.

It appears that just the broadcast is nearly seven times as fast after the
change.

SYPD for 1m with 30 he improves by ~11%

Author: imreddyTeja

src/cache/temporary_quantities.jl

@@ -1,5 +1,6 @@
 using ClimaCore: Fields
 using ClimaCore.Utilities: half
+using ClimaCore.MatrixFields
 
 # Fields used to store variables that only need to be used in a single function
 # but cannot be computed on the fly. Unlike the precomputed quantities, these
@@ -62,6 +63,11 @@ function temporary_quantities(Y, atmos)
             NTuple{n_mass_flux_subdomains(atmos.turbconv_model), CT12{FT}},
             face_space,
         ), # ᶠω¹²ʲs
+        # TODO: better names for theis scratch fields
+        ᶜbidiagonal_adjoint_matrix_c3 = Fields.Field(
+            BidiagonalMatrixRow{Adjoint{FT, C3{FT}}},
+            center_space,
+        ),
         ᶠtemp_C123 = Fields.Field(C123{FT}, face_space), # χ₁₂₃
         ᶜtemp_UVW = Fields.Field(typeof(uvw_vec), center_space), # UVW(ᶜu)
         ᶠtemp_UVW = Fields.Field(typeof(uvw_vec), face_space), # UVW(ᶠu³)
@@ -76,6 +82,15 @@ function temporary_quantities(Y, atmos)
         ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}),
         ᶠp_grad_matrix = similar(Y.f, BidiagonalMatrixRow{C3{FT}}),
         ᶠbidiagonal_matrix_ct3 = similar(Y.f, BidiagonalMatrixRow{CT3{FT}}),
+        # TODO: better names for theis scratch fields
+        ᶠband_matrix_wvec = similar(
+            Y.f,
+            ClimaCore.MatrixFields.BandMatrixRow{
+                ClimaCore.Utilities.PlusHalf{Int64}(0),
+                1,
+                ClimaCore.Geometry.WVector{FT},
+            },
+        ),
         ᶠbidiagonal_matrix_ct3_2 = similar(Y.f, BidiagonalMatrixRow{CT3{FT}}),
         ᶜadvection_matrix = similar(
             Y.c,

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -565,16 +565,21 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
         #            p.precomputed.ᶜwₜqₜ / ᶜq_tot,
         #        ),
         #    ) - (I,)
-
         MatrixFields.unrolled_foreach(tracer_info) do (ρχₚ_name, wₚ_name)
             MatrixFields.has_field(Y, ρχₚ_name) || return
             ∂ᶜρχₚ_err_∂ᶜρχₚ = matrix[ρχₚ_name, ρχₚ_name]
             ᶜwₚ = MatrixFields.get_field(p.precomputed, wₚ_name)
+            # TODO: Maybe just reinterpret an existing scratch field to avoid adding 2 new ones
+            # TODO: come up with read-able names for the intermediate computations...
+            @. p.scratch.ᶜbidiagonal_adjoint_matrix_c3 =
+                -(ClimaAtmos.ᶜprecipdivᵥ_matrix()) ⋅
+                DiagonalMatrixRow(ClimaAtmos.ᶠinterp(ᶜρ * ᶜJ) / ᶠJ)
+            @. p.scratch.ᶠband_matrix_wvec =
+                ClimaAtmos.ᶠright_bias_matrix() ⋅
+                DiagonalMatrixRow(ClimaCore.Geometry.WVector(-(ᶜwₚ) / ᶜρ))
             @. ∂ᶜρχₚ_err_∂ᶜρχₚ =
-                dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
-                DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅
-                ᶠright_bias_matrix() ⋅
-                DiagonalMatrixRow(-Geometry.WVector(ᶜwₚ) / ᶜρ) - (I,)
+                dtγ * p.scratch.ᶜbidiagonal_adjoint_matrix_c3 ⋅
+                p.scratch.ᶠband_matrix_wvec - (I,)
         end
 
     end