Merge pull request #1554 from CliMA/sb/point-space-adapt

Define Adapt for PointSpace

.buildkite/pipeline.yml

@@ -278,6 +278,14 @@ steps:
         agents:
           slurm_gpus: 1
 
+      - label: "Unit: cuda point spaces"
+        key: "point_space_cuda"
+        command: "julia --color=yes --check-bounds=yes --project=test test/Spaces/point_cuda.jl CUDA"
+        artifact_paths:
+          - output/point_cuda
+        agents:
+          slurm_gpus: 1
+
       - label: "Unit: cuda dss 2-process test"
         key: "gpu_ddss2_test"
         command:

src/Spaces/pointspace.jl

@@ -39,6 +39,10 @@ function PointSpace(
     return PointSpace(Adapt.adapt(ArrayType, local_geometry_data))
 end
 
+
+Adapt.adapt_structure(to, space::PointSpace) =
+    PointSpace(Adapt.adapt(to, space.local_geometry))
+
 function PointSpace(
     context::ClimaComms.AbstractCommsContext,
     coord::Geometry.Abstract1DPoint{FT},

test/Spaces/point_cuda.jl

@@ -0,0 +1,23 @@
+import ClimaCore
+import ClimaCore: Domains, Topologies, Meshes, Spaces, Geometry, column, Fields
+import ClimaComms
+using Test
+
+compare(cpu, gpu) = all(parent(cpu) .≈ Array(parent(gpu)))
+compare(cpu, gpu, sym) =
+    all(parent(getproperty(cpu, sym)) .≈ Array(parent(getproperty(gpu, sym))))
+
+@testset "CuArray-backed point spaces" begin
+    cpu_context =
+        ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())
+    gpu_context = ClimaComms.SingletonCommsContext(ClimaComms.CUDADevice())
+
+    point = Geometry.ZPoint(1.0)
+
+    cpuspace = Spaces.PointSpace(cpu_context, point)
+    gpuspace = Spaces.PointSpace(gpu_context, point)
+
+    # Test that all geometries match with CPU version:
+    @test compare(cpuspace, gpuspace, :local_geometry)
+
+end