From a1a91457cee469808bddcac8c3710168790c61fe Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 11 Jan 2023 15:59:44 -0800 Subject: [PATCH] Add simple GPU test to test suite --- .buildkite/pipeline.yml | 12 +++++-- test/problems.jl | 79 +++++++++++++++++++++++++---------------- test/simple_gpu.jl | 20 +++++++++++ 3 files changed, 79 insertions(+), 32 deletions(-) create mode 100644 test/simple_gpu.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 260400945..f40f78728 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -117,9 +117,17 @@ steps: queue: central slurm_ntasks: 1 - - label: "GPU" + # - 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/runtests.jl CuArray" + - "julia --project=test test/simple_gpu.jl" agents: slurm_gres: "gpu:1" queue: central diff --git a/test/problems.jl b/test/problems.jl index a93dbc856..8142dfc1b 100644 --- a/test/problems.jl +++ b/test/problems.jl @@ -1,5 +1,13 @@ 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 +import ClimaCore.Topologies as Topologies +import ClimaCore.Spaces as Spaces +import ClimaCore.Fields as Fields +import ClimaCore.Operators as Operators """ Single variable linear ODE @@ -421,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 @@ -433,38 +448,42 @@ function climacore_2Dheat_test_cts(::Type{FT}) where {FT} Δλ = FT(1) # denoted by Δλ̂ above t_end = FT(0.05) # denoted by t̂ above - domain = ClimaCore.Domains.RectangleDomain( - ClimaCore.Domains.IntervalDomain( - ClimaCore.Geometry.XPoint(FT(0)), - ClimaCore.Geometry.XPoint(FT(1)), + domain = Domains.RectangleDomain( + Domains.IntervalDomain( + Geometry.XPoint(FT(0)), + Geometry.XPoint(FT(1)), periodic = true, ), - ClimaCore.Domains.IntervalDomain( - ClimaCore.Geometry.YPoint(FT(0)), - ClimaCore.Geometry.YPoint(FT(1)), + Domains.IntervalDomain( + Geometry.YPoint(FT(0)), + Geometry.YPoint(FT(1)), periodic = true, ), ) - mesh = ClimaCore.Meshes.RectilinearMesh(domain, n_elem_x, n_elem_y) - topology = ClimaCore.Topologies.Topology2D(mesh) - quadrature = ClimaCore.Spaces.Quadratures.GLL{n_poly + 1}() - space = ClimaCore.Spaces.SpectralElementSpace2D(topology, quadrature) - (; x, y) = ClimaCore.Fields.coordinate_field(space) + mesh = Meshes.RectilinearMesh(domain, n_elem_x, n_elem_y) + 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) - init_state = ClimaCore.Fields.FieldVector(; u = φ_sin_sin) + # 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) - wdiv = ClimaCore.Operators.WeakDivergence() - grad = ClimaCore.Operators.Gradient() + wdiv = Operators.WeakDivergence() + grad = Operators.Gradient() function T_exp!(tendency, state, _, t) @. tendency.u = wdiv(grad(state.u)) + f_0 * exp(-(λ + Δλ) * t) * φ_sin_sin - dss_tendency && ClimaCore.Spaces.weighted_dss!(tendency.u) + dss_tendency && Spaces.weighted_dss!(tendency.u) end function dss!(state, _, t) - dss_tendency || ClimaCore.Spaces.weighted_dss!(state.u) + dss_tendency || Spaces.weighted_dss!(state.u) end function analytic_sol(t) @@ -492,25 +511,25 @@ function climacore_1Dheat_test_cts(::Type{FT}) where {FT} Δλ = FT(1) # denoted by Δλ̂ above t_end = FT(0.1) # denoted by t̂ above - domain = ClimaCore.Domains.IntervalDomain( - ClimaCore.Geometry.ZPoint(FT(0)), - ClimaCore.Geometry.ZPoint(FT(1)), + domain = Domains.IntervalDomain( + Geometry.ZPoint(FT(0)), + Geometry.ZPoint(FT(1)), boundary_names = (:bottom, :top), ) - mesh = ClimaCore.Meshes.IntervalMesh(domain, nelems = n_elem_z) - space = ClimaCore.Spaces.FaceFiniteDifferenceSpace(mesh) - (; z) = ClimaCore.Fields.coordinate_field(space) + mesh = Meshes.IntervalMesh(domain, nelems = n_elem_z) + space = Spaces.FaceFiniteDifferenceSpace(mesh) + (; z) = Fields.coordinate_field(space) λ = (2 * FT(π) * n_z)^2 φ_sin = @. sin(2 * FT(π) * n_z * z) - init_state = ClimaCore.Fields.FieldVector(; u = φ_sin) + init_state = Fields.FieldVector(; u = φ_sin) - div = ClimaCore.Operators.DivergenceC2F(; - bottom = ClimaCore.Operators.SetDivergence(FT(0)), - top = ClimaCore.Operators.SetDivergence(FT(0)), + div = Operators.DivergenceC2F(; + bottom = Operators.SetDivergence(FT(0)), + top = Operators.SetDivergence(FT(0)), ) - grad = ClimaCore.Operators.GradientF2C() + grad = Operators.GradientF2C() function T_exp!(tendency, state, _, t) @. tendency.u = div(grad(state.u)) + f_0 * exp(-(λ + Δλ) * t) * φ_sin end diff --git a/test/simple_gpu.jl b/test/simple_gpu.jl new file mode 100644 index 000000000..562fe730e --- /dev/null +++ b/test/simple_gpu.jl @@ -0,0 +1,20 @@ +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 +