update tests for GPU

Author: simonbyrne

.buildkite/JuliaProject.toml

@@ -0,0 +1,14 @@
+[extras]
+HDF5_jll = "0234f1f7-429e-5d53-9886-15a909be8d59"
+MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
+
+[preferences.HDF5_jll]
+libhdf5_path = "libhdf5"
+libhdf5_hl_path = "libhdf5_hl"
+
+[preferences.MPIPreferences]
+_format = "1.0"
+abi = "OpenMPI"
+binary = "system"
+libmpi = "libmpi"
+mpiexec = "mpiexec"

.buildkite/pipeline.yml

@@ -1,19 +1,21 @@
+
+agents:
+  queue: central
+  slurm_mem: 8G
+  modules: julia/1.8.5 openmpi/4.1.1 cuda/11.3 hdf5/1.12.1-ompi411
+
 env:
-  JULIA_VERSION: "1.8.3"
-  CUDA_VERSION: "11.2"
-  OPENMPI_VERSION: "4.1.1"
+  JULIA_LOAD_PATH: "${JULIA_LOAD_PATH}:${BUILDKITE_BUILD_CHECKOUT_PATH}/.buildkite"
+  JULIA_CUDA_USE_BINARYBUILDER: false
   OPENBLAS_NUM_THREADS: 1
-  CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: "true"
 
 steps:
   - label: "init cpu env"
     key: "init_cpu_env"
     command:
-      - echo "--- Configure MPI"
-      - julia -e 'using Pkg; Pkg.add("MPIPreferences"); using MPIPreferences; use_system_binary()'
-
       - echo "--- Instantiate project"
       - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true)'"
+      - "julia --project -e 'using CUDA; CUDA.precompile_runtime()'"
 
       - echo "--- Instantiate test"
       - "julia --project=test -e 'using Pkg; Pkg.develop(path=\".\"); Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true)'"
@@ -23,85 +25,38 @@ steps:
 
       - echo "--- Package status"
       - "julia --project -e 'using Pkg; Pkg.status()'"
-
     agents:
-      config: cpu
-      queue: central
-      slurm_ntasks: 1
-
-  # - label: "init gpu env"
-  #   key: "init_gpu_env"
-  #   command:
-  #     - echo "--- Configure MPI"
-  #     - julia -e 'using Pkg; Pkg.add("MPIPreferences"); using MPIPreferences; use_system_binary()'
-
-  #     - echo "--- Instantiate project"
-  #     - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true)'"
-
-  #     - echo "--- Instantiate test"
-  #     - "julia --project=test -e 'using Pkg; Pkg.develop(path=\".\"); Pkg.instantiate(;verbose=true); Pkg.precompile()'"
-
-  #     - echo "--- Initialize CUDA runtime"
-  #     - "julia --project -e 'using CUDA; CUDA.precompile_runtime(); CUDA.versioninfo()'"
-
-  #     - echo "--- Package status"
-  #     - "julia --project -e 'using Pkg; Pkg.status()'"
-  #   agents:
-  #     config: gpu
-  #     queue: central
-  #     slurm_ntasks: 1
-  #     slurm_gres: "gpu:1"
+      slurm_cpus_per_task: 8
+    env:
+      JULIA_NUM_PRECOMPILE_TASKS: 8
 
   - wait
 
   - label: "CPU tests"
     command:
       - "julia --project=test --check-bounds=yes test/runtests.jl"
     artifact_paths: "output/*"
-    agents:
-      config: cpu
-      queue: central
-      slurm_ntasks: 1
-
-  # - label: "GPU tests"
-  #   command:
-  #     - "julia --project=test --check-bounds=yes test/runtests.jl CuArray"
-  #   artifact_paths: "output/*"
-  #   agents:
-  #     config: gpu
-  #     queue: central
-  #     slurm_ntasks: 1
-  #     slurm_gres: "gpu:1"
 
   - label: "Flame graph (1D diffusion)"
     command:
       - "julia --project=perf perf/flame.jl --job_id diffusion_1D"
     artifact_paths: "diffusion_1D/*"
-    agents:
-      config: cpu
-      queue: central
-      slurm_ntasks: 1
 
   - label: "Jet graph (ode fun)"
     command:
       - "julia --project=perf perf/jet.jl --problem ode_fun"
-    agents:
-      config: cpu
-      queue: central
-      slurm_ntasks: 1
 
   - label: "Jet (forward euler)"
     command:
       - "julia --project=perf perf/jet.jl --problem fe"
-    agents:
-      config: cpu
-      queue: central
-      slurm_ntasks: 1
 
   - 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:
-      config: cpu
-      queue: central
-      slurm_ntasks: 1
+      slurm_gpus: 1

test/problems.jl

@@ -21,16 +21,21 @@ 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 +101,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
 
 """

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