Skip to content

Commit 7bf3248

Browse files
committed
implement wait on GPU event via hostcall
1 parent c9eae1f commit 7bf3248

File tree

1 file changed

+9
-7
lines changed

1 file changed

+9
-7
lines changed

lib/CUDAKernels/src/CUDAKernels.jl

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -75,14 +75,16 @@ import Base: wait
7575

7676
wait(ev::CudaEvent, progress=yield) = wait(CPU(), ev, progress)
7777

78-
function wait(::CPU, ev::CudaEvent, progress=yield)
79-
if progress === nothing
80-
CUDA.synchronize(ev.event)
81-
else
82-
while !isdone(ev)
83-
progress()
84-
end
78+
function wait(::CPU, ev::CudaEvent, progress=nothing)
79+
isdone(ev) && return nothing
80+
81+
event = Base.Threads.Event()
82+
stream = next_stream()
83+
wait(CUDADevice(), ev, nothing, stream)
84+
CUDA.launch(stream) do
85+
notify(event)
8586
end
87+
wait(event)
8688
end
8789

8890
# Use this to synchronize between computation using the CuDefaultStream

0 commit comments

Comments
 (0)