Define allowscalar

NEWS.md

@@ -1,7 +1,14 @@
 ClimaComms.jl Release Notes
 ========================
 
+v0.6.2
+-------
+
+- We added a device-agnostic `allowscalar(f, ::AbstractDevice, args...; kwargs...)` to further assist in making CUDA an extension.
+
 v0.6.1
+-------
+
 - Macros have been refactored to hopefully fix some code loading issues.
 
 v0.6.0

docs/src/index.md

@@ -27,6 +27,7 @@ ClimaComms.CUDADevice
 ClimaComms.device
 ClimaComms.device_functional
 ClimaComms.array_type
+ClimaComms.allowscalar
 ClimaComms.@threaded
 ClimaComms.@time
 ClimaComms.@elapsed

ext/ClimaCommsCUDAExt.jl

@@ -14,6 +14,8 @@ function ClimaComms.device_functional(::ClimaComms.CUDADevice)
 end
 
 ClimaComms.array_type(::ClimaComms.CUDADevice) = CUDA.CuArray
+ClimaComms.allowscalar(f, ::ClimaComms.CUDADevice, args...; kwargs...) =
+    CUDA.@allowscalar f(args...; kwargs...)
 
 # Extending ClimaComms methods that operate on expressions (cannot use dispatch here)
 ClimaComms.cuda_sync(expr) = CUDA.@sync expr

src/devices.jl

@@ -265,3 +265,27 @@ macro cuda_sync(device, expr)
         end
     end)
 end
+
+"""
+    allowscalar(f, ::AbstractDevice, args...; kwargs...)
+
+Device-flexible version of `CUDA.@allowscalar`.
+
+Lowers to
+```julia
+f(args...)
+```
+for CPU devices and
+```julia
+CUDA.@allowscalar f(args...)
+```
+for CUDA devices.
+
+This is usefully written with closures via
+```julia
+allowscalar(device) do
+    f()
+end
+```
+"""
+allowscalar(f, ::AbstractDevice, args...; kwargs...) = f(args...; kwargs...)

test/runtests.jl

@@ -217,3 +217,13 @@ end
     @test ClimaComms.bcast(context, AT(fill(Float64(pid), 3))) ==
           AT(fill(Float64(1), 3))
 end
+
+@testset "allowscalar" begin
+    a = AT(rand(3))
+    local x
+    ClimaComms.allowscalar(device) do
+        x = a[1]
+    end
+    device isa ClimaComms.CUDADevice && @test_throws ErrorException a[1]
+    @test x == Array(a)[1]
+end