From a1239f28eb1f3f2bd59b7497fabda9e26a2fef14 Mon Sep 17 00:00:00 2001
From: Simon Byrne <simonbyrne@gmail.com>
Date: Sun, 8 Oct 2023 18:14:44 -0700
Subject: [PATCH 1/2] add support for spatially dependent bonudary conditions
 on GPU

---
 src/Operators/finitedifference.jl | 16 +++++++++++++++-
 test/Operators/hybrid/3d.jl       | 32 +++++++++++++++++++++++--------
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/src/Operators/finitedifference.jl b/src/Operators/finitedifference.jl
index 31d3d6edc1..d2c8d44b9d 100644
--- a/src/Operators/finitedifference.jl
+++ b/src/Operators/finitedifference.jl
@@ -240,11 +240,16 @@ StencilBroadcasted{Style}(
 
 Adapt.adapt_structure(to, sbc::StencilBroadcasted{Style}) where {Style} =
     StencilBroadcasted{Style}(
-        sbc.op,
+        Adapt.adapt(to, sbc.op),
         Adapt.adapt(to, sbc.args),
         Adapt.adapt(to, sbc.axes),
     )
 
+function Adapt.adapt_structure(to, op::FiniteDifferenceOperator)
+    op
+end
+
+
 
 function Base.Broadcast.instantiate(sbc::StencilBroadcasted)
     op = sbc.op
@@ -2690,6 +2695,15 @@ Base.@propagate_inbounds function stencil_right_boundary(
     stencil_interior(op, loc, space, idx - 1, hidx, arg)
 end
 
+function Adapt.adapt_structure(to, op::DivergenceF2C)
+    DivergenceF2C(map(bc -> Adapt.adapt_structure(to, bc), op.bcs))
+end
+
+function Adapt.adapt_structure(to, bc::SetValue)
+    SetValue(Adapt.adapt_structure(to, bc.val))
+end
+
+
 """
     D = DivergenceC2F(;boundaryname=boundarycondition...)
     D.(v)
diff --git a/test/Operators/hybrid/3d.jl b/test/Operators/hybrid/3d.jl
index 5dc04f29b6..d7afe2bcb7 100644
--- a/test/Operators/hybrid/3d.jl
+++ b/test/Operators/hybrid/3d.jl
@@ -18,6 +18,8 @@ import ClimaCore.Geometry: WVector
 import ClimaCore.Utilities: half
 import ClimaCore.DataLayouts: level
 
+device = ClimaComms.device()
+
 @testset "sphere divergence" begin
     FT = Float64
     vertdomain = Domains.IntervalDomain(
@@ -26,12 +28,16 @@ import ClimaCore.DataLayouts: level
         boundary_tags = (:bottom, :top),
     )
     vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 10)
-    vert_center_space = Spaces.CenterFiniteDifferenceSpace(vertmesh)
+    verttopology = Topologies.IntervalTopology(
+        ClimaComms.SingletonCommsContext(device),
+        vertmesh,
+    )
+    vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopology)
 
     horzdomain = Domains.SphereDomain(30.0)
     horzmesh = Meshes.EquiangularCubedSphere(horzdomain, 4)
     horztopology = Topologies.Topology2D(
-        ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded()),
+        ClimaComms.SingletonCommsContext(device),
         horzmesh,
     )
     quad = Spaces.Quadratures.GLL{3 + 1}()
@@ -53,7 +59,8 @@ function hvspace_3D(
     xelem = 4,
     yelem = 4,
     zelem = 16,
-    npoly = 7,
+    npoly = 7;
+    device = ClimaComms.device(),
 )
     FT = Float64
     vertdomain = Domains.IntervalDomain(
@@ -62,7 +69,11 @@ function hvspace_3D(
         boundary_tags = (:bottom, :top),
     )
     vertmesh = Meshes.IntervalMesh(vertdomain, nelems = zelem)
-    vert_center_space = Spaces.CenterFiniteDifferenceSpace(vertmesh)
+    verttopology = Topologies.IntervalTopology(
+        ClimaComms.SingletonCommsContext(device),
+        vertmesh,
+    )
+    vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopology)
 
     horzdomain = Domains.RectangleDomain(
         Geometry.XPoint{FT}(xlim[1]) .. Geometry.XPoint{FT}(xlim[2]),
@@ -72,7 +83,7 @@ function hvspace_3D(
     )
     horzmesh = Meshes.RectilinearMesh(horzdomain, xelem, yelem)
     horztopology = Topologies.Topology2D(
-        ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded()),
+        ClimaComms.SingletonCommsContext(device),
         horzmesh,
     )
 
@@ -105,7 +116,8 @@ end
 
 @testset "2D SE, 1D FV Extruded Domain ∇ ODE Solve vertical" begin
 
-    hv_center_space, hv_face_space = hvspace_3D()
+    hv_center_space, hv_face_space =
+        hvspace_3D(; device = ClimaComms.CPUSingleThreaded())
     V =
         Geometry.UVWVector.(
             zeros(Float64, hv_face_space),
@@ -181,8 +193,12 @@ end
     #
     # NOTE: the equation setup is only correct for Cartesian domains!
 
-    hv_center_space, hv_face_space =
-        hvspace_3D((-1.0, 1.0), (-1.0, 1.0), (-1.0, 1.0))
+    hv_center_space, hv_face_space = hvspace_3D(
+        (-1.0, 1.0),
+        (-1.0, 1.0),
+        (-1.0, 1.0);
+        device = ClimaComms.CPUSingleThreaded(),
+    )
 
     abstract type BCtag end
     struct ZeroFlux <: BCtag end

From 9437aa5c17ca8f88e219cf44730e35c78a77e43a Mon Sep 17 00:00:00 2001
From: Simon Byrne <simonbyrne@gmail.com>
Date: Mon, 9 Oct 2023 10:07:00 -0700
Subject: [PATCH 2/2] add test

---
 .buildkite/pipeline.yml | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index dbb5b4865b..fd0a6ed54c 100755
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -435,6 +435,12 @@ steps:
         key: unit_hyb_ops_3d
         command: "julia --color=yes --check-bounds=yes --project=test test/Operators/hybrid/3d.jl"
 
+      - label: "Unit: hyb ops 3d CUDA"
+        key: unit_hyb_ops_3d_cuda
+        command: "julia --color=yes --check-bounds=yes --project=test test/Operators/hybrid/3d.jl"
+        agents:
+          slurm_gpus: 1
+
       - label: "Unit: remapping"
         key: unit_remapping
         command: "julia --color=yes --check-bounds=yes --project=test test/Operators/remapping.jl"