update tests for GPU

Author: simonbyrne

.buildkite/pipeline.yml

@@ -37,13 +37,6 @@ steps:
       - "julia --project=test --check-bounds=yes test/runtests.jl"
     artifact_paths: "output/*"
 
-  - label: "GPU tests"
-    command:
-      - "julia --project=test --check-bounds=yes test/runtests.jl CuArray"
-    artifact_paths: "output/*"
-    agents:
-      slurm_gpus: 1
-
   - label: "Flame graph (1D diffusion)"
     command:
       - "julia --project=perf perf/flame.jl --job_id diffusion_1D"
@@ -60,3 +53,10 @@ steps:
   - label: "Benchmark"
     command:
       - "julia --project=perf perf/benchmark.jl"
+
+  - label: "GPU tests"
+    command:
+      - julia --project=test --check-bounds=yes test/simple_gpu.jl
+    artifact_paths: "output/*"
+    agents:
+      slurm_gpus: 1

test/problems.jl

@@ -1,5 +1,6 @@
 using DiffEqBase, ClimaTimeSteppers, LinearAlgebra, StaticArrays
 using ClimaCore
+import ClimaCore.Device as Device
 import ClimaCore.Domains as Domains
 import ClimaCore.Geometry as Geometry
 import ClimaCore.Meshes as Meshes
@@ -21,16 +22,16 @@ u(t) = u_0 e^{αt}
 
 This is an in-place variant of the one from DiffEqProblemLibrary.jl.
 """
-function linear_prob()
+function linear_prob(::Type{ArrayType}=Array) where {ArrayType}
     ODEProblem(
         IncrementingODEFunction{true}((du, u, p, t, α = true, β = false) -> (du .= α .* p .* u .+ β .* du)),
-        [1 / 2],
+        ArrayType([1 / 2]),
         (0.0, 1.0),
         1.01,
     )
 end
-function linear_prob_fe()
-    ODEProblem(ForwardEulerODEFunction((un, u, p, t, dt) -> (un .= u .+ dt .* p .* u)), [1.0], (0.0, 1.0), -0.2)
+function linear_prob_fe(::Type{ArrayType}=Array) where {ArrayType}
+    ODEProblem(ForwardEulerODEFunction((un, u, p, t, dt) -> (un .= u .+ dt .* p .* u)), ArrayType([1.0]), (0.0, 1.0), -0.2)
 end
 
 function linear_prob_wfactt()
@@ -96,27 +97,29 @@ with initial condition ``u_0=[0,1]``, parameter ``α=2``, and solution
 u(t) = [cos(αt) sin(αt); -sin(αt) cos(αt) ] u_0
 ```
 """
-function sincos_prob()
+function sincos_prob(::Type{ArrayType}=Array) where {ArrayType}
     ODEProblem(
         IncrementingODEFunction{true}((du, u, p, t, α = true, β = false) -> (du[1] = α * p * u[2] + β * du[1];
         du[2] = -α * p * u[1] + β * du[2])),
-        [0.0, 1.0],
+        ArrayType([0.0, 1.0]),
         (0.0, 1.0),
         2.0,
     )
 end
-function sincos_prob_fe()
+function sincos_prob_fe(::Type{ArrayType}=Array) where {ArrayType}
     ODEProblem(
         ForwardEulerODEFunction((un, u, p, t, dt) -> (un[1] = u[1] + dt * p * u[2]; un[2] = u[2] - dt * p * u[1])),
-        [0.0, 1.0],
+        ArrayType([0.0, 1.0]),
         (0.0, 1.0),
         2.0,
     )
 end
 
 function sincos_sol(u0, p, t)
     s, c = sincos(p * t)
-    [c s; -s c] * u0
+    SC = similar(u0, (2,2))
+    copyto!(SC, [c s; -s c])
+    return SC * u0
 end
 
 """
@@ -428,9 +431,16 @@ end
 
 2D diffusion test problem. See [`2D diffusion problem`](@ref) for more details.
 """
-function climacore_2Dheat_test_cts(::Type{FT}) where {FT}
+function climacore_2Dheat_test_cts(::Type{FT}; print_arr_type = false) where {FT}
     dss_tendency = true
 
+    device = Device.device()
+    context = ClimaComms.SingletonCommsContext(device)
+
+    if print_arr_type
+        @info "Array type: $(Device.device_array_type(device))"
+    end
+
     n_elem_x = 2
     n_elem_y = 2
     n_poly = 2
@@ -445,13 +455,17 @@ function climacore_2Dheat_test_cts(::Type{FT}) where {FT}
         Domains.IntervalDomain(Geometry.YPoint(FT(0)), Geometry.YPoint(FT(1)), periodic = true),
     )
     mesh = Meshes.RectilinearMesh(domain, n_elem_x, n_elem_y)
-    topology = Topologies.Topology2D(mesh)
+    topology = Topologies.Topology2D(context, mesh)
     quadrature = Spaces.Quadratures.GLL{n_poly + 1}()
     space = Spaces.SpectralElementSpace2D(topology, quadrature)
     (; x, y) = Fields.coordinate_field(space)
 
     λ = (2 * FT(π))^2 * (n_x^2 + n_y^2)
-    φ_sin_sin = @. sin(2 * FT(π) * n_x * x) * sin(2 * FT(π) * n_y * y)
+
+    # Revert once https://github.com/CliMA/ClimaCore.jl/issues/1097
+    # is fixed
+    # φ_sin_sin = @. sin(2 * FT(π) * n_x * x) * sin(2 * FT(π) * n_y * y)
+    φ_sin_sin = @. sin(2 * π * n_x * x) * sin(2 * π * n_y * y)
 
     init_state = Fields.FieldVector(; u = φ_sin_sin)
 

test/simple_gpu.jl

@@ -0,0 +1,26 @@
+using ClimaTimeSteppers, LinearAlgebra, Test, CUDA
+
+include(joinpath(@__DIR__, "convergence_orders.jl"))
+include(joinpath(@__DIR__, "convergence_utils.jl"))
+include(joinpath(@__DIR__, "utils.jl"))
+include(joinpath(@__DIR__, "problems.jl"))
+
+
+@testset "LSRK and SSP convergence" begin
+    dts = 0.5 .^ (4:7)
+
+    for (prob, sol, tscale) in [
+        (linear_prob(CuArray), linear_sol, 1)
+        (sincos_prob(CuArray), sincos_sol, 1)
+    ]
+
+        @test convergence_order(prob, sol, LSRKEulerMethod(), dts .* tscale) ≈ 1 atol = 0.1
+        @test convergence_order(prob, sol, LSRK54CarpenterKennedy(), dts .* tscale) ≈ 4 atol = 0.05
+        @test convergence_order(prob, sol, LSRK144NiegemannDiehlBusch(), dts .* tscale) ≈ 4 atol = 0.05
+
+        @test convergence_order(prob, sol, SSPRK22Heuns(), dts .* tscale) ≈ 2 atol = 0.05
+        @test convergence_order(prob, sol, SSPRK22Ralstons(), dts .* tscale) ≈ 2 atol = 0.05
+        @test convergence_order(prob, sol, SSPRK33ShuOsher(), dts .* tscale) ≈ 3 atol = 0.05
+        @test convergence_order(prob, sol, SSPRK34SpiteriRuuth(), dts .* tscale) ≈ 3 atol = 0.05
+    end
+end