Define device-flexible @sync

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaComms"
 uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 authors = ["CliMA Contributors <[email protected]>"]
-version = "0.5.4"
+version = "0.5.5"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

docs/src/index.md

@@ -21,6 +21,7 @@ ClimaComms.array_type
 ClimaComms.@threaded
 ClimaComms.@time
 ClimaComms.@elapsed
+ClimaComms.@sync
 ```
 
 ## Contexts

src/devices.jl

@@ -163,3 +163,34 @@ macro elapsed(device, expr)
         end
     end
 end
+
+"""
+    @sync device expr
+
+Device-flexible `@sync`.
+
+Lowers to
+```julia
+@sync expr
+```
+for CPU devices and
+```julia
+CUDA.@sync expr
+```
+for CUDA devices.
+"""
+macro sync(device, expr)
+    # https://github.com/JuliaLang/julia/issues/28979#issuecomment-1756145207
+    return esc(quote
+        if $(device) isa $CUDADevice
+            $CUDA.@sync begin
+                $(expr)
+            end
+        else
+            @assert $(device) isa $AbstractDevice
+            $Base.@sync begin
+                $(expr)
+            end
+        end
+    end)
+end

test/hygiene.jl

@@ -14,6 +14,10 @@ function test_macro_hyhiene(dev)
     CC.@elapsed dev for i in 1:n
         sin.(rand(10))
     end
+
+    CC.@sync dev for i in 1:n
+        sin.(rand(10))
+    end
 end
 dev = CC.device()