debugging GPU compat

CliMA · Apr 13, 2024 · 8f897bd · 8f897bd
1 parent d043446
commit 8f897bd
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 63 deletions.
diff --git a/src/standalone/Soil/rre.jl b/src/standalone/Soil/rre.jl
@@ -427,70 +427,57 @@ Using this Jacobian with a backwards Euler timestepper is equivalent
 to using the modified Picard scheme of Celia et al. (1990).
 """
 function ClimaLand.make_update_jacobian(model::RichardsModel{FT}) where {FT}
+    return ((args...) -> update_jacobian!(model, args...))
+end
+
+function update_jacobian!(
+    model,
+    jacobian::ImplicitEquationJacobian,
+    Y,
+    p,
+    dtγ,
+    t,
+)
+    FT = eltype(Y)
     update_aux! = make_update_aux(model)
     update_boundary_fluxes! = make_update_boundary_fluxes(model)
-    function update_jacobian!(jacobian::ImplicitEquationJacobian, Y, p, dtγ, t)
-        (; matrix) = jacobian
-        (; ν, hydrology_cm, S_s, θ_r) = model.parameters
-        update_aux!(p, Y, t)
-        update_boundary_fluxes!(p, Y, t)
-        # Create divergence operator
-        divf2c_op = Operators.DivergenceF2C()
-        divf2c_matrix = MatrixFields.operator_matrix(divf2c_op)
-        # Create gradient operator, and set gradient at boundaries to 0
-        gradc2f_op = Operators.GradientC2F(
-            top = Operators.SetGradient(Geometry.WVector.(FT(0))),
-            bottom = Operators.SetGradient(Geometry.WVector.(FT(0))),
-        )
-        gradc2f_matrix = MatrixFields.operator_matrix(gradc2f_op)
-        # Create interpolation operator
-        interpc2f_op = Operators.InterpolateC2F(
-            bottom = Operators.Extrapolate(),
-            top = Operators.Extrapolate(),
-        )
 
-        # The derivative of the residual with respect to the prognostic variable
-        ∂ϑres∂ϑ = matrix[@name(soil.ϑ_l), @name(soil.ϑ_l)]
+    (; matrix) = jacobian
+    (; ν, hydrology_cm, S_s, θ_r) = model.parameters
+    update_aux!(p, Y, t)
+    update_boundary_fluxes!(p, Y, t)
+    # Create divergence operator
+    divf2c_op = Operators.DivergenceF2C()
+    divf2c_matrix = MatrixFields.operator_matrix(divf2c_op)
+    # Create gradient operator, and set gradient at boundaries to 0
+    gradc2f_op = Operators.GradientC2F(
+        top = Operators.SetGradient(Geometry.WVector.(FT(0))),
+        bottom = Operators.SetGradient(Geometry.WVector.(FT(0))),
+    )
+    gradc2f_matrix = MatrixFields.operator_matrix(gradc2f_op)
+    # Create interpolation operator
+    interpc2f_op = Operators.InterpolateC2F(
+        bottom = Operators.Extrapolate(),
+        top = Operators.Extrapolate(),
+    )
 
-        # If the top BC is a `MoistureStateBC`, add the term from the top BC
-        #  flux before applying divergence
-        if haskey(p.soil, :dfluxBCdY)
-            dfluxBCdY = p.soil.dfluxBCdY
-            topBC_op = Operators.SetBoundaryOperator(
-                top = Operators.SetValue(dfluxBCdY),
-                bottom = Operators.SetValue(
-                    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,
-                                ν,
-                                θ_r,
-                                S_s,
-                            ),
-                        ) + MatrixFields.LowerDiagonalMatrixRow(
-                            topBC_op(
-                                Geometry.Covariant3Vector(
-                                    zero(interpc2f_op(p.soil.K)),
-                                ),
-                            ),
-                        )
-                    )
-                ) - (I,)
-        else
-            @. ∂ϑres∂ϑ =
-                -dtγ * (
-                    divf2c_matrix() ⋅
+    # The derivative of the residual with respect to the prognostic variable
+    ∂ϑres∂ϑ = matrix[@name(soil.ϑ_l), @name(soil.ϑ_l)]
+
+    # If the top BC is a `MoistureStateBC`, add the term from the top BC
+    #  flux before applying divergence
+    if haskey(p.soil, :dfluxBCdY)
+        dfluxBCdY = p.soil.dfluxBCdY
+        topBC_op = Operators.SetBoundaryOperator(
+            top = Operators.SetValue(dfluxBCdY),
+            bottom = Operators.SetValue(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ϑ(
@@ -500,11 +487,41 @@ function ClimaLand.make_update_jacobian(model::RichardsModel{FT}) where {FT}
                             θ_r,
                             S_s,
                         ),
+                    ) + MatrixFields.LowerDiagonalMatrixRow(
+                        topBC_op(
+                            Geometry.Covariant3Vector(
+                                zero(interpc2f_op(p.soil.K)),
+                            ),
+                        ),
                     )
-                ) - (I,)
-        end
+                )
+            ) - (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, ν, θ_r, S_s),
+                )
+            ) - (I,)
     end
-    return update_jacobian!
+end
+
+import Adapt
+# function Adapt.adapt_structure(to, op::Operators.AbstractOperator)
+#     @info "adapt_structure"
+#     typeof(op).name.wrapper(map(bc -> Adapt.adapt_structure(to, bc), op.bcs))
+# end
+function Adapt.adapt_structure(to, op::Operators.SetBoundaryOperator)
+    @info "adapt_structure topBC"
+    bcs_adapted = NamedTuple{keys(op.bcs)}(
+        UnrolledUtilities.unrolled_map(
+            bc -> Adapt.adapt_structure(to, bc),
+            values(op.bcs),
+        ),
+    )
+    Operators.SetBoundaryOperator(bcs_adapted)
 end
 
 """

diff --git a/test/shared_utilities/implicit_timestepping/richards_model.jl b/test/shared_utilities/implicit_timestepping/richards_model.jl
@@ -1,6 +1,8 @@
 using Test
 using LinearAlgebra
 import ClimaCore: MatrixFields
+import ClimaCore: Operators, Geometry, Fields
+import ClimaCore.MatrixFields: @name, ⋅
 import ClimaParams as CP
 using ClimaLand
 using ClimaLand.Domains: Column, HybridBox