Fix adapt for new limiters, test on gpu

CliMA · Dec 23, 2024 · a1261e9 · a1261e9
1 parent 6040306
commit a1261e9
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 15 deletions.
diff --git a/.buildkite/pipeline.yml b/.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: "cpu_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:

diff --git a/src/Operators/finitedifference.jl b/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...)