Merge #1489

1489: add support for spatially dependent boundary conditions on GPU r=simonbyrne a=simonbyrne Adapt spatially-varying boundary conditions for the GPU. Fixes CliMA/ClimaAtmos.jl#2192 - [x] Code follows the [style guidelines](https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/) OR N/A. - [x] Unit tests are included OR N/A. - [x] Code is exercised in an integration test OR N/A. - [x] Documentation has been added/updated OR N/A. Co-authored-by: Simon Byrne <[email protected]>
CliMA · Oct 12, 2023 · 32be46e · 32be46e · simonbyrne · Oct 12, 2023
2 parents 68d31f4 + 9437aa5
commit 32be46e
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 9 deletions.
diff --git 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"

diff --git 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
@@ -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