Merge pull request #4093 from CliMA/dr/gpu_perf/manual_sparse_jacobian_1

Improve performance of CUDA kernels generated from manual_sparse_jacobian

Author: daverumph

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-274
+275
 
 # **README**
 #
@@ -20,6 +20,11 @@
 
 
 #=
+275
+- Change order of GPU calculations for better performance, but it
+  results in slightly different floating point rounding.  Artifacts
+  all look the same.
+
 274
 - Remove unused calculation of TKE exchange in mixing length
 

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -626,11 +626,13 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
         FT = eltype(params)
         (; vertical_diffusion, smagorinsky_lilly) = p.atmos
         (; ᶜp) = p.precomputed
+        ᶜK_u = p.scratch.ᶜtemp_scalar_4
+        ᶜK_h = p.scratch.ᶜtemp_scalar_6
         if vertical_diffusion isa DecayWithHeightDiffusion
-            ᶜK_h = ᶜcompute_eddy_diffusivity_coefficient(Y.c.ρ, vertical_diffusion)
+            ᶜK_h .= ᶜcompute_eddy_diffusivity_coefficient(Y.c.ρ, vertical_diffusion)
             ᶜK_u = ᶜK_h
         elseif vertical_diffusion isa VerticalDiffusion
-            ᶜK_h = ᶜcompute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vertical_diffusion)
+            ᶜK_h .= ᶜcompute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vertical_diffusion)
             ᶜK_u = ᶜK_h
         elseif is_smagorinsky_vertical(smagorinsky_lilly)
             set_smagorinsky_lilly_precomputed_quantities!(Y, p, smagorinsky_lilly)
@@ -650,22 +652,18 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             @. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
         end
 
-
-        @. p.scratch.ᶜbidiagonal_adjoint_matrix_c3 =
-            ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h))
-        @. ᶜdiffusion_h_matrix =
-            p.scratch.ᶜbidiagonal_adjoint_matrix_c3 ⋅ ᶠgradᵥ_matrix()
-
+        ∂ᶠρχ_dif_flux_∂ᶜχ = ᶠp_grad_matrix
+        @. ∂ᶠρχ_dif_flux_∂ᶜχ =
+            DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h)) ⋅ ᶠgradᵥ_matrix()
+        @. ᶜdiffusion_h_matrix = ᶜadvdivᵥ_matrix() ⋅ ∂ᶠρχ_dif_flux_∂ᶜχ
         if (
             MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) ||
             !isnothing(p.atmos.turbconv_model) ||
             !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion)
         )
-            @. p.scratch.ᶜbidiagonal_adjoint_matrix_c3 =
-                ᶜadvdivᵥ_matrix() ⋅
-                DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u))
-            @. ᶜdiffusion_u_matrix =
-                p.scratch.ᶜbidiagonal_adjoint_matrix_c3 ⋅ ᶠgradᵥ_matrix()
+            @. ∂ᶠρχ_dif_flux_∂ᶜχ =
+                DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u)) ⋅ ᶠgradᵥ_matrix()
+            @. ᶜdiffusion_u_matrix = ᶜadvdivᵥ_matrix() ⋅ ∂ᶠρχ_dif_flux_∂ᶜχ
         end
 
         ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)]