Add cooperative_wait, for waiting on GPU operations without blocking the thread - #23
Merged
Merged
Conversation
…the thread. Generalizes the nonblocking synchronization from CUDA.jl, of which copies exist in oneAPI.jl and in the works for OpenCL.jl: short operations are detected by polling, longer ones are waited for on a pool of foreign worker threads while the calling task yields. Workers serve a single request at a time and are reused last-in-first-out; when all are busy, pollable objects are polled until a worker becomes available. Interrupts and cancellation are deferred until the operation has completed, unless the wait is cancellable.
This was referenced Sep 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Waiting for the GPU, e.g. in
CUDA.synchronize(), shouldn't block the Julia thread it's called from, because other tasks may need that thread in the meantime (to launch work on another stream, do I/O, etc.). CUDA.jl has done this for a while with its "nonblocking synchronization", oneAPI.jl carries a verbatim copy of it, and OpenCL.jl and KernelAbstractions' POCL back-end were about to grow their own. This PR adds a reworked version to GPUToolbox ascooperative_wait, so that all back-ends can share one implementation.A back-end provides a blocking wait, which will run on another thread, and optionally a non-blocking query. For example, CUDA.jl would do:
How it works
isdone: a few spins without yielding, then yielding to other tasks in between polls. This is the same heuristic CUDA.jl uses today.Event. A worker serves one request at a time, and idle workers are reused last-in-first-out, so the one you get is usually still awake.In finalizers and during precompilation, the blocking wait is simply called on the current thread. An
InterruptException, or task cancellation on Julia 1.14, doesn't make the wait return early: the GPU may still be using memory the caller would free when unwinding, so it keeps waiting and throws once the operation has completed.cancellable=trueopts out of that.What changes compared to CUDA.jl
CUDA.jl assigns each task one of its 4 workers (round-robin, and then sticky) and talks to it through a custom channel. As a result, a long wait blocks every other task that happens to share its worker. With 9 tasks each synchronizing their own stream, one of them running 20 ms kernels and the others 1 ms kernels, the tasks sharing a worker with the slow one saw a p90 overhead of 18.8 ms, against ~25 µs for the others. With a prototype of this design, all of them stayed below 30 µs. The tests here check the same thing deterministically: short waits have to complete while a long one is pending, even when there are more waits than workers.
It also defers interrupts until the operation has completed, as described above (previously that depended on where the interrupt hit), and calls the back-end's wait function with
invokelatest, since the worker threads live in the world age they were started in.Latency is unchanged. Comparing both implementations in the same process, alternating between them to cancel out noise (RTX 5080, Ryzen 9 9950X, Julia 1.12), the median overhead of a launch plus synchronization, minus the kernel duration:
(This was measured on a heavily loaded machine, hence the ~2 µs higher baseline than on an idle one.) Past the polling phase, both are dominated by the time it takes to wake up the waiting thread, and were within noise of each other. The blocking column is lower because the driver spins on the calling thread, which is exactly what we want to avoid.
I also tried a time-based polling budget, since the iteration count seemed arbitrary: how long it lasts depends on the CPU and on how many tasks are runnable. That turned out to be worse on a loaded machine. The iteration budget stretches when yielding gets slow, whereas a fixed time budget hands the wait to a worker that has gone to sleep in the meantime, and waking it up then costs milliseconds.
Alternatives considered
cuLaunchHostFunc,clSetEventCallback) waking anAsyncCondition, as AMDGPU.jl does. This added 20–75 µs over a blocking wait, compared to 2–4 µs for handing off to a worker. CUDA host functions also stall the stream and aren't called when the context has an error, and NVIDIA's OpenCL driver delivered ~3.5% of callbacks about 20 ms late. I also tried callbacks only as a last resort for when all workers are busy, but they didn't earn their complexity there either.@threadcallruns on libuv's pool, which is 4 threads shared with filesystem work, so it has the same head-of-line blocking, and it wakes the caller through the event loop.Julia itself doesn't offer anything better to build on: there's no way to hook polling into the scheduler, and
@threadcallhas the limitations above.CUDA.jl, OpenCL.jl and KernelAbstractions' POCL back-end will switch to this once it's released, so this bumps the version to 3.2.0.