Obtain device array from ClimaComms context. Add shallow-water CUDA

driver template.

Author: sriharshakandala

examples/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.8.3"
 manifest_format = "2.0"
-project_hash = "1a511843299a40488fdeb0ecc8eb50466fa341cb"
+project_hash = "6c93fbc837ba5dffbd7162f915124e75b23659a5"
 
 [[deps.AMD]]
 deps = ["Libdl", "LinearAlgebra", "SparseArrays", "Test"]

examples/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCommsMPI = "5f86816e-8b66-43b2-912e-75384f99de49"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"

examples/sphere/shallow_water_cuda.jl

@@ -0,0 +1,29 @@
+using CUDA
+using ClimaComms
+
+import ClimaCore:
+    Device,
+    Domains,
+    Fields,
+    Geometry,
+    Meshes,
+    Operators,
+    Spaces,
+    Topologies,
+    DataLayouts
+
+function shallow_water_driver_cuda(ARGS, ::Type{FT}) where {FT}
+    device = Device.device()
+    context = ClimaComms.SingletonCommsContext(device)
+    println("running serial simulation on $device device")
+    # Test case specifications
+    test_name = get(ARGS, 1, "steady_state") # default test case to run
+    test_angle_name = get(ARGS, 2, "alpha0") # default test case to run
+    α = parse(FT, replace(test_angle_name, "alpha" => ""))
+
+    println("Test name: $test_name, α = $(α)⁰")
+
+    return nothing
+end
+
+shallow_water_driver_cuda(ARGS, Float64)

src/ClimaCore.jl

@@ -5,6 +5,7 @@ const VERSION = PkgVersion.@Version
 
 function comm_context end
 
+include("Device.jl")
 include("interface.jl")
 include("Utilities/Utilities.jl")
 include("RecursiveApply/RecursiveApply.jl")

src/Device.jl

@@ -0,0 +1,11 @@
+module Device
+using ClimaComms
+using CUDA
+
+device(; disablegpu = false) =
+    CUDA.has_cuda_gpu() && !disablegpu ? ClimaComms.CUDA() : ClimaComms.CPU()
+
+device_array_type(::ClimaComms.CPU) = Array
+device_array_type(::ClimaComms.CUDA) = CUDA.CuArray
+
+end

test/gpu/device.jl

@@ -0,0 +1,22 @@
+using Test
+using CUDA
+using ClimaComms
+import ClimaCore: Device
+
+@testset "check device detection on GPU" begin
+    device = Device.device()
+    cuda_context = ClimaComms.SingletonCommsContext(device)
+    DA = Device.device_array_type(cuda_context.device)
+
+    @test device isa ClimaComms.CUDA
+    @test cuda_context.device == ClimaComms.CUDA()
+    @test DA == CuArray
+
+    override_device = Device.device(disablegpu = true)
+    override_cuda_context = ClimaComms.SingletonCommsContext(override_device)
+    DA = Device.device_array_type(override_cuda_context.device)
+
+    @test override_device isa ClimaComms.CPU
+    @test override_cuda_context.device == ClimaComms.CPU()
+    @test DA == Array
+end

test/runtests.jl

@@ -96,6 +96,7 @@ end
 if "CUDA" in ARGS
     @safetestset "GPU - cuda" begin @time include("gpu/cuda.jl") end
     @safetestset "GPU - data" begin @time include("gpu/data.jl") end
+    @safetestset "GPU - device" begin @time include("gpu/device.jl") end
 end
 
 #! format: on