Merge pull request #2112 from CliMA/ck/adapt_limiter_fixes

Fix adapt for new limiters, test on gpu

Author: charleskawczynski

.buildkite/pipeline.yml

@@ -1454,21 +1454,32 @@ steps:
           - "julia --color=yes --project=.buildkite examples/column/fct_advection.jl"
         artifact_paths:
           - "examples/column/output/fct_advection/*"
-      
+
       - label: ":computer: Column TVD Slope-limited Advection Eq"
         key: "cpu_tvd_column_advect"
         command:
           - "julia --color=yes --project=.buildkite examples/column/tvd_advection.jl"
         artifact_paths:
           - "examples/column/output/tvd_advection/*"
-      
+
       - label: ":computer: Column Lin vanLeer Limiter Advection Eq"
         key: "cpu_lvl_column_advect"
         command:
           - "julia --color=yes --project=.buildkite examples/column/vanleer_advection.jl"
         artifact_paths:
           - "examples/column/output/vanleer_advection/*"
 
+      - label: ":computer: Column Lin vanLeer Limiter Advection Eq cuda"
+        key: "gpu_lvl_column_advect"
+        command:
+          - "julia --color=yes --project=.buildkite examples/column/vanleer_advection.jl"
+        artifact_paths:
+          - "examples/column/output/vanleer_advection/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
       - label: ":computer: Column BB FCT Advection Eq"
         key: "cpu_bb_fct_column_advect"
         command:

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaCore"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 authors = ["CliMA Contributors <clima-software@caltech.edu>"]
-version = "0.14.21"
+version = "0.14.23"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

examples/column/vanleer_advection.jl

@@ -1,3 +1,8 @@
+#=
+julia --project=.buildkite
+ENV["CLIMACOMMS_DEVICE"] = "CUDA"
+using Revise; include("examples/column/vanleer_advection.jl")
+=#
 using Test
 using LinearAlgebra
 import ClimaComms
@@ -51,13 +56,13 @@ end
 
 # Define a pulse wave or square wave
 
-FT = Float64
-t₀ = FT(0.0)
-t₁ = FT(6)
-z₀ = FT(0.0)
-zₕ = FT(2π)
-z₁ = FT(1.0)
-speed = FT(-1.0)
+const FT = Float64
+const t₀ = FT(0.0)
+const t₁ = FT(6)
+const z₀ = FT(0.0)
+const zₕ = FT(2π)
+const z₁ = FT(1.0)
+const speed = FT(-1.0)
 pulse(z, t, z₀, zₕ, z₁) = abs(z - speed * t) ≤ zₕ ? z₁ : z₀
 
 n = 2 .^ 8
@@ -135,6 +140,14 @@ for (i, stretch_fn) in enumerate(stretch_fns)
         err = norm(q_final .- q_analytic)
         rel_mass_err = norm((sum(q_final) - sum(q_init)) / sum(q_init))
 
+        @test err ≤ 0.25
+        @test rel_mass_err ≤ 10eps()
+
+        device = ClimaComms.device(q_init)
+        if device isa ClimaComms.CUDADevice
+            continue
+        end
+
         if j == 1
             fig = Plots.plot(q_analytic; label = "Exact", color = :red)
         end
@@ -162,7 +175,5 @@ for (i, stretch_fn) in enumerate(stretch_fns)
                 ),
             )
         end
-        @test err ≤ 0.25
-        @test rel_mass_err ≤ 10eps()
     end
 end

src/Operators/finitedifference.jl

@@ -3061,19 +3061,24 @@ function Adapt.adapt_structure(to, bc::AbstractBoundaryCondition)
 end
 
 # Extend `adapt_structure` for all operator types with boundary conditions.
-function Adapt.adapt_structure(to, op::FiniteDifferenceOperator)
-    if hasfield(typeof(op), :bcs)
-        bcs_adapted = NamedTuple{keys(op.bcs)}(
-            UnrolledFunctions.unrolled_map(
-                bc -> Adapt.adapt_structure(to, bc),
-                values(op.bcs),
-            ),
-        )
-        return unionall_type(typeof(op))(bcs_adapted)
-    else
-        return op
-    end
-end
+Adapt.adapt_structure(to, op::FiniteDifferenceOperator) =
+    hasfield(typeof(op), :bcs) ? adapt_fd_operator(to, op, op.bcs) : op
+
+@inline adapt_fd_operator(to, op::LinVanLeerC2F, bcs) =
+    LinVanLeerC2F(adapt_bcs(to, bcs), Adapt.adapt_structure(to, op.constraint))
+
+@inline adapt_fd_operator(to, op::TVDLimitedFluxC2F, bcs) =
+    TVDLimitedFluxC2F(adapt_bcs(to, bcs), Adapt.adapt_structure(to, op.method))
+
+@inline adapt_fd_operator(to, op, bcs) =
+    unionall_type(typeof(op))(adapt_bcs(to, bcs))
+
+@inline adapt_bcs(to, bcs) = NamedTuple{keys(bcs)}(
+    UnrolledFunctions.unrolled_map(
+        bc -> Adapt.adapt_structure(to, bc),
+        values(bcs),
+    ),
+)
 
 """
     D = DivergenceC2F(;boundaryname=boundarycondition...)