Refactor cubed sphere dss test

charleskawczynski · charleskawczynski · commit 9066a7345780 · 2024-08-29T14:37:30.000-04:00
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -278,17 +278,21 @@ steps:
         key: unit_ddss1
         command:
           - "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
-          - "julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/ddss1.jl CUDA"
+          - "julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/ddss1.jl"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:
           slurm_gpus: 1
 
+      - label: "Unit: ddss1 cs"
+        key: unit_cuda_ddss1_cs
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/ddss1_cs.jl"
+
       - label: "Unit: ddss1 cs"
         key: unit_ddss1_cs
         command:
           - "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
-          - "julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/ddss1_cs.jl CUDA"
+          - "julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/ddss1_cs.jl"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:
diff --git a/test/Spaces/ddss1.jl b/test/Spaces/ddss1.jl
@@ -1,5 +1,5 @@
 #=
-julia --project=test
+julia --project
 using Revise; include(joinpath("test", "Spaces", "ddss1.jl"))
 =#
 using Logging
diff --git a/test/Spaces/ddss1_cs.jl b/test/Spaces/ddss1_cs.jl
@@ -1,3 +1,7 @@
+#=
+julia --project
+using Revise; include(joinpath("test", "Spaces", "ddss1_cs.jl"))
+=#
 using Test
 using ClimaComms
 ClimaComms.@import_required_backends
@@ -13,29 +17,68 @@ import ClimaCore:
     Topologies,
     DataLayouts
 
-@testset "DSS on Equiangular Cubed Sphere mesh (ne = 3, serial run)" begin
-    device = ClimaComms.device() #ClimaComms.CUDADevice()
-    context = ClimaComms.SingletonCommsContext(device)
-    context_cpu =
-        ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())
-
-    println("running test on $device device")
-
-    domain = Domains.SphereDomain(300.0)
+function get_space_cs(::Type{FT}; context, R = 300.0) where {FT}
+    domain = Domains.SphereDomain{FT}(300.0)
     mesh = Meshes.EquiangularCubedSphere(domain, 3)
     topology = Topologies.Topology2D(context, mesh)
-    topology_cpu = Topologies.Topology2D(context_cpu, mesh)
     quad = Quadratures.GLL{4}()
     space = Spaces.SpectralElementSpace2D(topology, quad)
-    space_cpu = Spaces.SpectralElementSpace2D(topology_cpu, quad)
+    return space
+end
 
-    x = ones(space)
-    x_cpu = ones(space_cpu)
+function get_space_and_buffers(::Type{FT}; context) where {FT}
+    init_state_covariant12(local_geometry, p) =
+        Geometry.Covariant12Vector(1.0, -1.0)
+    init_state_covariant123(local_geometry, p) =
+        Geometry.Covariant123Vector(1.0, -1.0, 1.0)
+
+    R = FT(6.371229e6)
+    npoly = 2
+    z_max = FT(30e3)
+    z_elem = 3
+    h_elem = 2
+    device = ClimaComms.device(context)
+    @info "running dss-Covariant123Vector test on $(device)" h_elem z_elem npoly R z_max FT
+    # horizontal space
+    domain = Domains.SphereDomain{FT}(R)
+    horizontal_mesh = Meshes.EquiangularCubedSphere(domain, h_elem)
+    horizontal_topology = Topologies.Topology2D(
+        context,
+        horizontal_mesh,
+        Topologies.spacefillingcurve(horizontal_mesh),
+    )
+    quad = Quadratures.GLL{npoly + 1}()
+    h_space = Spaces.SpectralElementSpace2D(horizontal_topology, quad)
+    # vertical space
+    z_domain = Domains.IntervalDomain(
+        Geometry.ZPoint{FT}(zero(z_max)),
+        Geometry.ZPoint{FT}(z_max);
+        boundary_names = (:bottom, :top),
+    )
+    z_mesh = Meshes.IntervalMesh(z_domain, nelems = z_elem)
+    z_topology = Topologies.IntervalTopology(context, z_mesh)
+    z_center_space = Spaces.CenterFiniteDifferenceSpace(z_topology)
+    space = Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_center_space)
+    args = (Fields.local_geometry_field(space), Ref(nothing))
+    y12 = init_state_covariant12.(args...)
+    y123 = init_state_covariant123.(args...)
+    dss_buffer12 = Spaces.create_dss_buffer(y12)
+    dss_buffer123 = Spaces.create_dss_buffer(y123)
+    return (; space, y12, y123, dss_buffer12, dss_buffer123)
+end
 
+@testset "DSS on Equiangular Cubed Sphere mesh (ne = 3, serial run)" begin
+    FT = Float64
+    device = ClimaComms.device()
+    context = ClimaComms.SingletonCommsContext(device)
+    println("running test on $device device")
+    space = get_space_cs(FT; context)
+    space_cpu = get_space_cs(FT; context)
+    x = ones(space)
     Spaces.weighted_dss!(x)
-    Spaces.weighted_dss!(x_cpu)
 
-    @test parent(x_cpu) ≈ Array(parent(x))
+    @test Array(parent(x)) ≈ ones(size(parent(x))) # TODO: improve the quality of this test
+
     wrong_field = map(Fields.coordinate_field(space)) do cf
         (; a = Float64(0))
     end
@@ -60,113 +103,57 @@ end
 
 @testset "DSS of Covariant12Vector & Covariant123Vector on extruded Cubed Sphere mesh (ne = 3, serial run)" begin
     FT = Float64
-    context = ClimaComms.SingletonCommsContext(ClimaComms.CUDADevice())
-    context_cpu =
-        ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded()) # CPU context for comparison
-    R = FT(6.371229e6)
-
-    npoly = 4
-    z_max = FT(30e3)
-    z_elem = 10
-    h_elem = 4
-    println(
-        "running dss-Covariant123Vector test on $(context.device); h_elem = $h_elem; z_elem = $z_elem; npoly = $npoly; R = $R; z_max = $z_max; FT = $FT",
-    )
-    # horizontal space
-    domain = Domains.SphereDomain(R)
-    horizontal_mesh = Meshes.EquiangularCubedSphere(domain, h_elem)
-    horizontal_topology = Topologies.Topology2D(
-        context,
-        horizontal_mesh,
-        Topologies.spacefillingcurve(horizontal_mesh),
-    )
-    horizontal_topology_cpu = Topologies.Topology2D(
-        context_cpu,
-        horizontal_mesh,
-        Topologies.spacefillingcurve(horizontal_mesh),
-    )
-    quad = Quadratures.GLL{npoly + 1}()
-    h_space = Spaces.SpectralElementSpace2D(horizontal_topology, quad)
-    h_space_cpu = Spaces.SpectralElementSpace2D(horizontal_topology_cpu, quad)
-
-    # vertical space
-    z_domain = Domains.IntervalDomain(
-        Geometry.ZPoint(zero(z_max)),
-        Geometry.ZPoint(z_max);
-        boundary_names = (:bottom, :top),
-    )
-    z_mesh = Meshes.IntervalMesh(z_domain, nelems = z_elem)
-    z_topology = Topologies.IntervalTopology(context, z_mesh)
-    z_topology_cpu = Topologies.IntervalTopology(context_cpu, z_mesh)
-
-    z_center_space = Spaces.CenterFiniteDifferenceSpace(z_topology)
-    z_center_space_cpu = Spaces.CenterFiniteDifferenceSpace(z_topology_cpu)
-
-    hv_center_space =
-        Spaces.ExtrudedFiniteDifferenceSpace(h_space, z_center_space)
-
-    hv_center_space_cpu =
-        Spaces.ExtrudedFiniteDifferenceSpace(h_space_cpu, z_center_space_cpu)
+    device = ClimaComms.device()
+    nt = get_space_and_buffers(FT; context = ClimaComms.context(device))
 
     # test DSS for a Covariant12Vector
-    init_state_covariant12(local_geometry, p) =
-        Geometry.Covariant12Vector(1.0, -1.0)
-
-    y12 =
-        init_state_covariant12.(
-            Fields.local_geometry_field(hv_center_space),
-            Ref(nothing),
-        )
-    y12_cpu =
-        init_state_covariant12.(
-            Fields.local_geometry_field(hv_center_space_cpu),
-            Ref(nothing),
-        )
-
-    dss_buffer12 = Spaces.create_dss_buffer(y12)
-    dss_buffer12_cpu = Spaces.create_dss_buffer(y12_cpu)
     # ensure physical velocity is continous across SE boundary for initial state
-    Spaces.weighted_dss!(y12 => dss_buffer12)
-    Spaces.weighted_dss!(y12_cpu => dss_buffer12_cpu)
+    Spaces.weighted_dss!(nt.y12 => nt.dss_buffer12)
+    init = copy(nt.y12)
+    Spaces.weighted_dss!(nt.y12, nt.dss_buffer12)
+    @test init ≈ nt.y12
+    # ensure physical velocity is continous across SE boundary for initial state
+    Spaces.weighted_dss!(nt.y123, nt.dss_buffer123)
+    init = copy(nt.y123)
+    Spaces.weighted_dss!(nt.y123, nt.dss_buffer123)
+    @test init ≈ nt.y123
+end
 
-    yinit12 = copy(y12)
-    yinit12_cpu = copy(y12_cpu)
+# TODO: remove once the quality of the above test is improved
+(ClimaComms.device() isa ClimaComms.CUDADevice) &&
+    @testset "GPU-vs-CPU test: DSS of Covariant12Vector & Covariant123Vector on extruded Cubed Sphere mesh (ne = 3, serial run)" begin
+        FT = Float64
+        cpu_device = ClimaComms.CPUSingleThreaded()
+        gpu_device = ClimaComms.CUDADevice()
+        gpu =
+            get_space_and_buffers(FT; context = ClimaComms.context(gpu_device))
+        cpu =
+            get_space_and_buffers(FT; context = ClimaComms.context(cpu_device))
 
-    Spaces.weighted_dss!(y12, dss_buffer12)
-    Spaces.weighted_dss!(y12_cpu, dss_buffer12_cpu)
-    @test yinit12 ≈ y12
-    @test yinit12_cpu ≈ y12_cpu
-    @test parent(y12_cpu) ≈ Array(parent(y12))
+        # test DSS for a Covariant12Vector
+        # ensure physical velocity is continous across SE boundary for initial state
+        Spaces.weighted_dss!(cpu.y12 => cpu.dss_buffer12)
+        Spaces.weighted_dss!(gpu.y12 => gpu.dss_buffer12)
 
-    # test DSS for a Covariant123Vector
-    init_state_covariant123(local_geometry, p) =
-        Geometry.Covariant123Vector(1.0, -1.0, 1.0)
+        inity12 = (; cpu = copy(cpu.y12), gpu = copy(gpu.y12))
 
-    y123 =
-        init_state_covariant123.(
-            Fields.local_geometry_field(hv_center_space),
-            Ref(nothing),
-        )
-    y123_cpu =
-        init_state_covariant123.(
-            Fields.local_geometry_field(hv_center_space_cpu),
-            Ref(nothing),
-        )
+        Spaces.weighted_dss!(cpu.y12, cpu.dss_buffer12)
+        Spaces.weighted_dss!(gpu.y12, gpu.dss_buffer12)
 
-    dss_buffer123 = Spaces.create_dss_buffer(y123)
-    dss_buffer123_cpu = Spaces.create_dss_buffer(y123_cpu)
+        @test inity12.cpu ≈ cpu.y12
+        @test inity12.gpu ≈ gpu.y12
+        @test parent(cpu.y12) ≈ Array(parent(gpu.y12))
 
-    # ensure physical velocity is continous across SE boundary for initial state
-    Spaces.weighted_dss!(y123, dss_buffer123)
-    Spaces.weighted_dss!(y123_cpu, dss_buffer123_cpu)
+        # ensure physical velocity is continous across SE boundary for initial state
+        Spaces.weighted_dss!(cpu.y123, cpu.dss_buffer123)
+        Spaces.weighted_dss!(gpu.y123, gpu.dss_buffer123)
 
-    yinit123 = copy(y123)
-    yinit123_cpu = copy(y123_cpu)
+        inity123 = (; cpu = copy(cpu.y123), gpu = copy(gpu.y123))
 
-    Spaces.weighted_dss!(y123, dss_buffer123)
-    Spaces.weighted_dss!(y123_cpu, dss_buffer123_cpu)
+        Spaces.weighted_dss!(cpu.y123, cpu.dss_buffer123)
+        Spaces.weighted_dss!(gpu.y123, gpu.dss_buffer123)
 
-    @test yinit123 ≈ y123
-    @test yinit123_cpu ≈ y123_cpu
-    @test parent(y123_cpu) ≈ Array(parent(y123))
-end
+        @test inity123.cpu ≈ cpu.y123
+        @test inity123.gpu ≈ gpu.y123
+        @test parent(cpu.y123) ≈ Array(parent(gpu.y123))
+    end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -33,6 +33,7 @@ UnitTest("Cubedsphere surface topology"            ,"Topologies/cubedsphere_sfc.
 UnitTest("Quadratures"                             ,"Quadratures/Quadratures.jl"),
 UnitTest("Spaces"                                  ,"Spaces/unit_spaces.jl"),
 UnitTest("Spaces - serial CPU DSS"                 ,"Spaces/ddss1.jl"),
+UnitTest("Spaces - DSS cubed sphere"               ,"Spaces/ddss1_cs.jl"),
 UnitTest("Sphere spaces"                           ,"Spaces/sphere.jl"),
 # UnitTest("Terrain warp"                            ,"Spaces/terrain_warp.jl"), # appears to hang on GHA
 UnitTest("Fields"                                  ,"Fields/unit_field.jl"), # has benchmarks
@@ -107,7 +108,6 @@ UnitTest("Aqua"                                    ,"aqua.jl"),
 UnitTest("Deprecations"                            ,"deprecations.jl"),
 UnitTest("GPU - cuda"                              ,"gpu/cuda.jl";meta=:gpu_only),
 UnitTest("GPU - data"                              ,"DataLayouts/cuda.jl";meta=:gpu_only),
-UnitTest("Spaces - serial CUDA DSS on CubedSphere" ,"Spaces/ddss1_cs.jl";meta=:gpu_only),
 UnitTest("Operators - spectral element CUDA"       ,"Operators/spectralelement/rectilinear_cuda.jl";meta=:gpu_only),
 UnitTest("Operators - finite difference CUDA"      ,"Operators/hybrid/unit_cuda.jl";meta=:gpu_only),
 UnitTest("Operators - extruded sphere space ops"   ,"Operators/hybrid/extruded_sphere_cuda.jl";meta=:gpu_only),
