Define device-agnostic fill

NEWS.md

@@ -5,6 +5,7 @@ v0.6.2
 -------
 
 - We added a device-agnostic `allowscalar(f, ::AbstractDevice, args...; kwargs...)` to further assist in making CUDA an extension.
+- We added a device-agnostic `fill(value, dims...)` to further assist in making CUDA an extension.
 
 v0.6.1
 -------

Project.toml

@@ -8,7 +8,7 @@ authors = [
     "Jake Bolewski <[email protected]>",
     "Gabriele Bozzola <[email protected]>",
 ]
-version = "0.6.1"
+version = "0.6.2"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

docs/src/index.md

@@ -28,6 +28,7 @@ ClimaComms.device
 ClimaComms.device_functional
 ClimaComms.array_type
 ClimaComms.allowscalar
+ClimaComms.fill
 ClimaComms.@threaded
 ClimaComms.@time
 ClimaComms.@elapsed

ext/ClimaCommsCUDAExt.jl

@@ -16,6 +16,7 @@ end
 ClimaComms.array_type(::ClimaComms.CUDADevice) = CUDA.CuArray
 ClimaComms.allowscalar(f, ::ClimaComms.CUDADevice, args...; kwargs...) =
     CUDA.@allowscalar f(args...; kwargs...)
+ClimaComms.fill(::CUDADevice, value, dims...) = CUDA.fill(value, dims...)
 
 # Extending ClimaComms methods that operate on expressions (cannot use dispatch here)
 ClimaComms.cuda_sync(expr) = CUDA.@sync expr

src/devices.jl

@@ -289,3 +289,12 @@ end
 ```
 """
 allowscalar(f, ::AbstractDevice, args...; kwargs...) = f(args...; kwargs...)
+
+"""
+    fill(::AbstractDevice, value, dims...)
+
+Device-flexible version of `fill`.
+
+Calls `Base.fill` on `AbstractCPUDevice`s and `CUDA.fill` on `CUDADevice`s.
+"""
+ClimaComms.fill(::AbstractDevice, value, dims...) = Base.fill(value, dims...)

test/runtests.jl

@@ -227,3 +227,7 @@ end
     device isa ClimaComms.CUDADevice && @test_throws ErrorException a[1]
     @test x == Array(a)[1]
 end
+
+@testset "fill" begin
+    @test ClimaComms.fill(device, 3, (2, 2)) isa AT
+end