Add GPU tests

.buildkite/pipeline.yml

@@ -35,32 +35,33 @@ steps:
       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()'
+  - label: "init gpu env"
+    key: "init_gpu_env"
+    command:
+      - "echo $$JULIA_DEPOT_PATH"
+      - 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)'"
-  #     - "julia --project -e 'using Pkg; Pkg.precompile(;strict=true)'"
+      - echo "--- Instantiate project"
+      - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
+      - "julia --project -e 'using Pkg; Pkg.precompile(;strict=true)'"
 
-  #     - echo "--- Instantiate test"
-  #     - "julia --project=test -e 'using Pkg; Pkg.develop(path=\".\")'"
-  #     - "julia --project=test -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
-  #     - "julia --project=test -e 'using Pkg; Pkg.precompile()'"
+      - echo "--- Instantiate test"
+      - "julia --project=test -e 'using Pkg; Pkg.develop(path=\".\")'"
+      - "julia --project=test -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
+      - "julia --project=test -e 'using Pkg; Pkg.precompile(;strict=true)'"
 
-  #     - echo "--- Initialize CUDA runtime"
-  #     - "julia --project -e 'using CUDA; CUDA.precompile_runtime()'"
-  #     - "julia --project -e 'using CUDA; CUDA.versioninfo()'"
+      - echo "--- Initialize CUDA runtime"
+      - "julia --project -e 'using CUDA; CUDA.precompile_runtime()'"
+      - "julia --project -e 'using CUDA; CUDA.versioninfo()'"
 
-  #     - echo "--- Package status"
-  #     - "julia --project -e 'using Pkg; Pkg.status()'"
-  #   agents:
-  #     config: gpu
-  #     queue: central
-  #     slurm_ntasks: 1
-  #     slurm_gres: "gpu:1"
+      - echo "--- Package status"
+      - "julia --project -e 'using Pkg; Pkg.status()'"
+    agents:
+      config: gpu
+      queue: central
+      slurm_ntasks: 1
+      slurm_gres: "gpu:1"
 
   - wait
 
@@ -115,3 +116,19 @@ steps:
       config: cpu
       queue: central
       slurm_ntasks: 1
+
+  # - label: "GPU"
+  #   command:
+  #     - "julia --project=test test/runtests.jl CuArray"
+  #   agents:
+  #     slurm_gres: "gpu:1"
+  #     queue: central
+  #     slurm_ntasks: 1
+
+  - label: "Simple GPU"
+    command:
+      - "julia --project=test test/simple_gpu.jl"
+    agents:
+      slurm_gres: "gpu:1"
+      queue: central
+      slurm_ntasks: 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
@@ -428,9 +429,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 +453,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,19 @@
+using ClimaTimeSteppers
+import ClimaTimeSteppers as CTS
+import OrdinaryDiffEq as ODE
+
+include(joinpath(@__DIR__, "problems.jl"))
+
+function main(::Type{FT}) where {FT}
+    alg_name = ARS343()
+    test_case = climacore_2Dheat_test_cts(FT; print_arr_type = true)
+    prob = test_case.split_prob
+    alg = CTS.IMEXAlgorithm(alg_name, NewtonsMethod(; max_iters = 2))
+    integrator = ODE.init(prob, alg; dt = FT(0.01))
+    sol = ODE.solve!(integrator)
+    @info "Done!"
+    return integrator
+end
+
+integrator = main(Float64)
+nothing