split computation in Jacobian (#1121)

Author: kmdeck

experiments/benchmarks/bucket.jl

@@ -23,7 +23,8 @@ import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput
 
 import ClimaTimeSteppers as CTS
 import NCDatasets
-using ClimaCore
+import ClimaCore
+@show pkgversion(ClimaCore)
 import ClimaLand
 using ClimaParams
 using ClimaLand.Bucket:
@@ -66,13 +67,10 @@ function setup_prob(t0, tf, Δt; nelements = (200, 7))
     # We set up the problem in a function so that we can make multiple copies (for profiling)
 
     # Set up simulation domain
-    soil_depth = FT(3.5)
-    bucket_domain = ClimaLand.Domains.SphericalShell(;
-        radius = FT(6.3781e6),
-        depth = soil_depth,
-        nelements = nelements,
-        dz_tuple = FT.((1.0, 0.05)),
-    )
+    dz_tuple = FT.((1.0, 0.05))
+    depth = FT(3.5)
+    bucket_domain =
+        ClimaLand.global_domain(FT; nelements, dz_tuple, depth = depth)
     start_date = DateTime(2005)
 
     # Initialize parameters
@@ -259,7 +257,7 @@ if profiler == "flamegraph"
     end
 
     if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
-        PREVIOUS_BEST_TIME = 1.69
+        PREVIOUS_BEST_TIME = 0.533
         if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
             @info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
         elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

experiments/benchmarks/richards.jl

@@ -36,7 +36,8 @@ ClimaComms.@import_required_backends
 import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput
 import ClimaUtilities.Regridders: InterpolationsRegridder
 import ClimaTimeSteppers as CTS
-using ClimaCore
+import ClimaCore
+@show pkgversion(ClimaCore)
 using ClimaParams
 
 import ClimaLand
@@ -305,7 +306,7 @@ if profiler == "flamegraph"
     end
 
     if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
-        PREVIOUS_BEST_TIME = 1.49
+        PREVIOUS_BEST_TIME = 0.6
         if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
             @info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
         elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

experiments/benchmarks/snowy_land.jl

@@ -30,7 +30,8 @@ import SciMLBase
 import ClimaComms
 ClimaComms.@import_required_backends
 import ClimaTimeSteppers as CTS
-using ClimaCore
+import ClimaCore
+@show pkgversion(ClimaCore)
 using ClimaUtilities.ClimaArtifacts
 
 import ClimaUtilities.TimeVaryingInputs:
@@ -490,7 +491,7 @@ if profiler == "flamegraph"
         @info "Saved allocation flame to $alloc_flame_file"
     end
     if get(ENV, "BUILDKITE_PIPELINE_SLUG", nothing) == "climaland-benchmark"
-        PREVIOUS_BEST_TIME = 3.29
+        PREVIOUS_BEST_TIME = 0.74
         if average_timing_s > PREVIOUS_BEST_TIME + std_timing_s
             @info "Possible performance regression, previous average time was $(PREVIOUS_BEST_TIME)"
         elseif average_timing_s < PREVIOUS_BEST_TIME - std_timing_s

experiments/long_runs/bucket.jl

@@ -14,7 +14,8 @@ import SciMLBase
 import ClimaComms
 ClimaComms.@import_required_backends
 import ClimaTimeSteppers as CTS
-using ClimaCore
+import ClimaCore
+@show pkgversion(ClimaCore)
 using ClimaUtilities.ClimaArtifacts
 import Interpolations
 import ClimaUtilities.TimeVaryingInputs:

experiments/long_runs/land_region.jl

@@ -19,7 +19,8 @@ import SciMLBase
 import ClimaComms
 ClimaComms.@import_required_backends
 import ClimaTimeSteppers as CTS
-using ClimaCore
+import ClimaCore
+@show pkgversion(ClimaCore)
 using ClimaUtilities.ClimaArtifacts
 
 import ClimaDiagnostics

experiments/long_runs/snowy_land.jl

@@ -19,7 +19,8 @@ import SciMLBase
 import ClimaComms
 ClimaComms.@import_required_backends
 import ClimaTimeSteppers as CTS
-using ClimaCore
+import ClimaCore
+@show pkgversion(ClimaCore)
 using ClimaUtilities.ClimaArtifacts
 import ClimaUtilities.OnlineLogging: WallTimeInfo, report_walltime
 import ClimaUtilities.TimeManager: ITime, date

experiments/long_runs/soil.jl

@@ -18,7 +18,8 @@ import SciMLBase
 import ClimaComms
 ClimaComms.@import_required_backends
 import ClimaTimeSteppers as CTS
-using ClimaCore
+import ClimaCore
+@show pkgversion(ClimaCore)
 using ClimaUtilities.ClimaArtifacts
 
 using ClimaDiagnostics

src/shared_utilities/Domains.jl

@@ -169,7 +169,12 @@ function Column(;
     subsurface_space =
         ClimaCore.Spaces.CenterFiniteDifferenceSpace(device, mesh)
     surface_space = obtain_surface_space(subsurface_space)
-    space = (; surface = surface_space, subsurface = subsurface_space)
+    subsurface_face_space = obtain_face_space(subsurface_space)
+    space = (;
+        surface = surface_space,
+        subsurface = subsurface_space,
+        subsurface_face = subsurface_face_space,
+    )
     fields = get_additional_coordinate_field_data(subsurface_space)
     return Column{FT}(
         zlim,
@@ -470,7 +475,12 @@ function HybridBox(;
     )
 
     surface_space = obtain_surface_space(subsurface_space)
-    space = (; surface = surface_space, subsurface = subsurface_space)
+    subsurface_face_space = obtain_face_space(subsurface_space)
+    space = (;
+        surface = surface_space,
+        subsurface = subsurface_space,
+        subsurface_face = subsurface_face_space,
+    )
     fields = get_additional_coordinate_field_data(subsurface_space)
     return HybridBox{FT}(
         horzdomain.xlim,
@@ -593,7 +603,12 @@ function SphericalShell(;
         vert_center_space,
     )
     surface_space = obtain_surface_space(subsurface_space)
-    space = (; surface = surface_space, subsurface = subsurface_space)
+    subsurface_face_space = obtain_face_space(subsurface_space)
+    space = (;
+        surface = surface_space,
+        subsurface = subsurface_space,
+        subsurface_face = subsurface_face_space,
+    )
     fields = get_additional_coordinate_field_data(subsurface_space)
     return SphericalShell{FT}(
         radius,

src/standalone/Soil/boundary_conditions.jl

@@ -883,7 +883,7 @@ top boundary.
 These variables are updated in place in `boundary_flux!`.
 """
 boundary_vars(bc::MoistureStateBC, ::ClimaLand.TopBoundary) =
-    (:top_bc, :top_bc_wvec, :dfluxBCdY)
+    (:top_bc, :top_bc_wvec, :dfluxBCdY, :topBC_scratch)
 
 """
     boundary_var_domain_names(::MoistureStateBC, ::ClimaLand.TopBoundary)
@@ -892,7 +892,7 @@ An extension of the `boundary_var_domain_names` method for MoistureStateBC at th
 top boundary.
 """
 boundary_var_domain_names(bc::MoistureStateBC, ::ClimaLand.TopBoundary) =
-    (:surface, :surface, :surface)
+    (:surface, :surface, :surface, :subsurface_face)
 """
     boundary_var_types(::RichardsModel{FT},
                         ::MoistureStateBC,
@@ -910,6 +910,7 @@ boundary_var_types(
     FT,
     ClimaCore.Geometry.WVector{FT},
     ClimaCore.Geometry.Covariant3Vector{FT},
+    ClimaCore.Geometry.Covariant3Vector{FT},
 )
 
 

src/standalone/Soil/energy_hydrology.jl

@@ -343,6 +343,29 @@ function ClimaLand.make_compute_jacobian(model::EnergyHydrology{FT}) where {FT}
         ∂ϑres∂ϑ = matrix[@name(soil.ϑ_l), @name(soil.ϑ_l)]
         ∂ρeres∂ρe = matrix[@name(soil.ρe_int), @name(soil.ρe_int)]
         ∂ρeres∂ϑ = matrix[@name(soil.ρe_int), @name(soil.ϑ_l)]
+
+        # Derivatives with respect to ϑ:
+
+        # Precompute intermediate quantities to improve performance
+        # due to fusing of broadcasted expressions involving matrices
+        # First, the gradient of ∂ψ∂ϑ
+        # This term is used again below, so we do not alter it once we have made it
+        @. p.soil.bidiag_matrix_scratch =
+            gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
+                ClimaLand.Soil.dψdϑ(
+                    hydrology_cm,
+                    Y.soil.ϑ_l,
+                    ν - Y.soil.θ_i, #ν_eff
+                    θ_r,
+                    S_s,
+                ),
+            )
+        # Now the full Darcy flux term. This term is the one that gets altered with the BC
+        # contribution in place, below
+        @. p.soil.full_bidiag_matrix_scratch =
+            MatrixFields.DiagonalMatrixRow(interpc2f_op(-p.soil.K)) ⋅
+            p.soil.bidiag_matrix_scratch
+
         # If the top BC is a `MoistureStateBC`, add the term from the top BC
         #  flux before applying divergence
         if haskey(p.soil, :dfluxBCdY)
@@ -353,80 +376,47 @@ function ClimaLand.make_compute_jacobian(model::EnergyHydrology{FT}) where {FT}
                     Geometry.Covariant3Vector(zero(FT)),
                 ),
             )
-            # Add term from top boundary condition before applying divergence
             # Note: need to pass 3D field on faces to `topBC_op`. Interpolating `K` to faces
             #  for this is inefficient - we should find a better solution.
-            @. ∂ϑres∂ϑ =
-                -dtγ * (
-                    divf2c_matrix() ⋅ (
-                        MatrixFields.DiagonalMatrixRow(
-                            interpc2f_op(-p.soil.K),
-                        ) ⋅ gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
-                            ClimaLand.Soil.dψdϑ(
-                                hydrology_cm,
-                                Y.soil.ϑ_l,
-                                ν - Y.soil.θ_i, #ν_eff
-                                θ_r,
-                                S_s,
-                            ),
-                        ) + MatrixFields.LowerDiagonalMatrixRow(
-                            topBC_op(
-                                Geometry.Covariant3Vector(
-                                    zero(interpc2f_op(p.soil.K)),
-                                ),
-                            ),
-                        )
-                    )
-                ) - (I,)
-        else
-            @. ∂ϑres∂ϑ =
-                -dtγ * (
-                    divf2c_matrix() ⋅
-                    MatrixFields.DiagonalMatrixRow(interpc2f_op(-p.soil.K)) ⋅
-                    gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
-                        ClimaLand.Soil.dψdϑ(
-                            hydrology_cm,
-                            Y.soil.ϑ_l,
-                            ν - Y.soil.θ_i, #ν_eff
-                            θ_r,
-                            S_s,
-                        ),
-                    )
-                ) - (I,)
+            @. p.soil.topBC_scratch = topBC_op(
+                Geometry.Covariant3Vector(zero(interpc2f_op(p.soil.K))),
+            )
+            @. p.soil.full_bidiag_matrix_scratch +=
+                MatrixFields.LowerDiagonalMatrixRow(p.soil.topBC_scratch)
         end
+
+        @. ∂ϑres∂ϑ =
+            -dtγ * (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
+
+        # Now create the flux term for ∂ρe∂ϑ using bidiag_matrix_scratch
+        # This overwrites full_bidiag_matrix_scratch
+        @. p.soil.full_bidiag_matrix_scratch =
+            MatrixFields.DiagonalMatrixRow(
+                -interpc2f_op(
+                    volumetric_internal_energy_liq(
+                        p.soil.T,
+                        model.parameters.earth_param_set,
+                    ) * p.soil.K,
+                ),
+            ) ⋅ p.soil.bidiag_matrix_scratch
         @. ∂ρeres∂ϑ =
-            -dtγ * (
-                divf2c_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
-                    -interpc2f_op(
-                        volumetric_internal_energy_liq(
-                            p.soil.T,
-                            model.parameters.earth_param_set,
-                        ) * p.soil.K,
-                    ),
-                ) ⋅ gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
-                    ClimaLand.Soil.dψdϑ(
-                        hydrology_cm,
-                        Y.soil.ϑ_l,
-                        ν - Y.soil.θ_i, #ν_eff
-                        θ_r,
-                        S_s,
-                    ),
-                )
-            ) - (I,)
+            -dtγ * (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
 
+        # Now overwrite bidiag_matrix_scratch and full_bidiag scratch for the ρe ρe bidiagonal
+        @. p.soil.bidiag_matrix_scratch =
+            gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
+                1 / ClimaLand.Soil.volumetric_heat_capacity(
+                    p.soil.θ_l,
+                    Y.soil.θ_i,
+                    ρc_ds,
+                    earth_param_set,
+                ),
+            )
+        @. p.soil.full_bidiag_matrix_scratch =
+            MatrixFields.DiagonalMatrixRow(interpc2f_op(-p.soil.κ)) ⋅
+            p.soil.bidiag_matrix_scratch
         @. ∂ρeres∂ρe =
-            -dtγ * (
-                divf2c_matrix() ⋅
-                MatrixFields.DiagonalMatrixRow(interpc2f_op(-p.soil.κ)) ⋅
-                gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
-                    1 / ClimaLand.Soil.volumetric_heat_capacity(
-                        p.soil.θ_l,
-                        Y.soil.θ_i,
-                        ρc_ds,
-                        earth_param_set,
-                    ),
-                )
-            ) - (I,)
+            -dtγ * (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
     end
     return compute_jacobian!
 end
@@ -505,6 +495,8 @@ ClimaLand.auxiliary_vars(soil::EnergyHydrology) = (
     :θ_l,
     :T,
     :κ,
+    :bidiag_matrix_scratch,
+    :full_bidiag_matrix_scratch,
     boundary_vars(soil.boundary_conditions.top, ClimaLand.TopBoundary())...,
     boundary_vars(
         soil.boundary_conditions.bottom,
@@ -526,6 +518,8 @@ ClimaLand.auxiliary_types(soil::EnergyHydrology{FT}) where {FT} = (
     FT,
     FT,
     FT,
+    MatrixFields.BidiagonalMatrixRow{Geometry.Covariant3Vector{FT}},
+    MatrixFields.BidiagonalMatrixRow{Geometry.Covariant3Vector{FT}},
     boundary_var_types(
         soil,
         soil.boundary_conditions.top,
@@ -546,6 +540,8 @@ ClimaLand.auxiliary_domain_names(soil::EnergyHydrology) = (
     :subsurface,
     :subsurface,
     :subsurface,
+    :subsurface_face,
+    :subsurface_face,
     boundary_var_domain_names(
         soil.boundary_conditions.top,
         ClimaLand.TopBoundary(),