Skip to content

Add cooperative_wait, for waiting on GPU operations without blocking the thread - #23

Merged
maleadt merged 1 commit into
mainfrom
tb/cooperative-wait
Sep 30, 2026
Merged

maleadt merged 1 commit into
mainfrom
tb/cooperative-wait

Conversation

@maleadt

@maleadt maleadt commented Sep 30, 2026

Copy link
Copy Markdown
Member

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 as cooperative_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:

function synchronize(stream::CuStream)
    res = cooperative_wait(stream; isdone) do stream
        context!(() -> unchecked_cuStreamSynchronize(stream), stream.ctx)
    end
    # `nothing` means polling saw the stream finish; synchronize again to get the status
    res = @something res unchecked_cuStreamSynchronize(stream)
    res == SUCCESS || throw_api_error(res)
end

How it works

  1. Short operations are caught by polling isdone: a few spins without yielding, then yielding to other tasks in between polls. This is the same heuristic CUDA.jl uses today.
  2. Longer waits are handed to a pool of at most 4 foreign threads (created with libuv and adopted by Julia), while the calling task waits on an 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.
  3. When all workers are busy, pollable objects keep being polled while yielding, and grab a worker as soon as one frees up. Objects that can't be polled queue for a worker.

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=true opts 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:

kernel blocking CUDA.jl this PR
0 µs 6.6 µs 6.8 µs 6.8 µs
50 µs 6.7 µs 7.0 µs 7.0 µs
100 µs 7.3 µs 7.7 µs 7.8 µs

(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

  • Driver completion callbacks (cuLaunchHostFunc, clSetEventCallback) waking an AsyncCondition, 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.
  • @threadcall runs 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.
  • A single thread polling all pending operations avoids head-of-line blocking, but any hiccup of that thread delays every waiter, and it can't wait for a whole CUDA context.
  • Only polling, as Metal.jl does, has the best latency, but keeps the waiting task, and thus its thread, busy.

Julia itself doesn't offer anything better to build on: there's no way to hook polling into the scheduler, and @threadcall has 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.

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant