Merge pull request #85 from CliMA/ck/bugfix

Define cuda_sync in ext, restrict cpu implementations
CliMA · Jun 21, 2024 · 5cbb27f · 5cbb27f · charleskawczynski · Jun 21, 2024
2 parents e62f126 + 3c44543
commit 5cbb27f
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 6 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,6 +1,11 @@
 ClimaComms.jl Release Notes
 ========================
 
+v0.6.3
+-------
+
+- Bugfix: `cuda_sync` was missing in the extension and, as a result, `ClimaComms.@cuda_sync` was not actually synchronizing. We've also removed the abstract fallback, so that we instead method-error if we pass a cuda device when the cuda extension does not exist.
+
 v0.6.2
 -------
 

diff --git a/Project.toml b/Project.toml
@@ -8,7 +8,7 @@ authors = [
     "Jake Bolewski <[email protected]>",
     "Gabriele Bozzola <[email protected]>",
 ]
-version = "0.6.2"
+version = "0.6.3"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

diff --git a/ext/ClimaCommsCUDAExt.jl b/ext/ClimaCommsCUDAExt.jl
@@ -21,6 +21,8 @@ ClimaComms.allowscalar(f, ::CUDADevice, args...; kwargs...) =
 # Extending ClimaComms methods that operate on expressions (cannot use dispatch here)
 ClimaComms.sync(f::F, ::CUDADevice, args...; kwargs...) where {F} =
     CUDA.@sync f(args...; kwargs...)
+ClimaComms.cuda_sync(f::F, ::CUDADevice, args...; kwargs...) where {F} =
+    CUDA.@sync f(args...; kwargs...)
 ClimaComms.time(f::F, ::CUDADevice, args...; kwargs...) where {F} =
     CUDA.@time f(args...; kwargs...)
 ClimaComms.elapsed(f::F, ::CUDADevice, args...; kwargs...) where {F} =

diff --git a/src/devices.jl b/src/devices.jl
@@ -144,7 +144,7 @@ CUDA.@time f(args...; kwargs...)
 ```
 for CUDA devices.
 """
-function time(f::F, device::AbstractDevice, args...; kwargs...) where {F}
+function time(f::F, device::AbstractCPUDevice, args...; kwargs...) where {F}
     Base.@time begin
         f(args...; kwargs...)
     end
@@ -165,7 +165,7 @@ CUDA.@elapsed f(args...; kwargs...)
 ```
 for CUDA devices.
 """
-function elapsed(f::F, device::AbstractDevice, args...; kwargs...) where {F}
+function elapsed(f::F, device::AbstractCPUDevice, args...; kwargs...) where {F}
     Base.@elapsed begin
         f(args...; kwargs...)
     end
@@ -205,7 +205,7 @@ If the CPU version of the above example does not leverage
 spawned tasks (which require using `Base.sync` or `Threads.wait`
 to synchronize), then you may want to simply use [`cuda_sync`](@ref).
 """
-function sync(f::F, ::AbstractDevice, args...; kwargs...) where {F}
+function sync(f::F, ::AbstractCPUDevice, args...; kwargs...) where {F}
     Base.@sync begin
         f(args...; kwargs...)
     end
@@ -226,7 +226,7 @@ CUDA.@sync f(args...; kwargs...)
 ```
 for CUDA devices.
 """
-function cuda_sync(f::F, ::AbstractDevice, args...; kwargs...) where {F}
+function cuda_sync(f::F, ::AbstractCPUDevice, args...; kwargs...) where {F}
     f(args...; kwargs...)
 end
 
@@ -252,7 +252,7 @@ allowscalar(device) do
 end
 ```
 """
-allowscalar(f, ::AbstractDevice, args...; kwargs...) = f(args...; kwargs...)
+allowscalar(f, ::AbstractCPUDevice, args...; kwargs...) = f(args...; kwargs...)
 
 """
     @time device expr