Skip to content

Commit

Permalink
Define device-flexible @sync
Browse files Browse the repository at this point in the history
  • Loading branch information
charleskawczynski committed Oct 10, 2023
1 parent 8eba9c0 commit 7de96ef
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 0 deletions.
1 change: 1 addition & 0 deletions docs/src/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ ClimaComms.array_type
ClimaComms.@threaded
ClimaComms.@time
ClimaComms.@elapsed
ClimaComms.@sync
```

## Contexts
Expand Down
31 changes: 31 additions & 0 deletions src/devices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
4 changes: 4 additions & 0 deletions test/hygiene.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down

0 comments on commit 7de96ef

Please sign in to comment.