Repository navigation
Don't allocate when cooperative_wait waits on a worker thread - #39
Merged
Merged
Conversation
Waits handed to a worker thread (CUDA.jl, AMDGPU.jl, OpenCL on GPUs) allocated a request with its event, and boxed the object and the result: 224 to 288 bytes per wait. Whether a wait takes that path depends on the operation's duration and on host load, so allocation counts weren't deterministic, as with completion notifications before #32. Requests are now reused, from pools per type of function, object and inferred result, so that those are stored unboxed. The arguments and the result are kept in vectors that are emptied on release, which also store immutable objects with references (like CUDA.jl's `(stream, context)`) inline, and keep idle requests from retaining them. Requests of interrupted or cancelled waits are not reused. The worker now calls the function in the caller's world instead of the latest one, as a direct call would, so that its result always has the inferred type. Otherwise the return type would need to include another, boxed, `Some`, which makes returning the result allocate when `cooperative_wait` isn't inlined. The pools use a spin lock of their own, which defers interrupts like the rest of the wait. A worker round trip takes about as long as before (1.1-1.3 µs on Linux, Julia 1.12). Bumps the version to 3.5.1.
maleadt
force-pushed
the
sd/pool-worker-wait-requests
branch
from
November 5, 2026 08:48
f8a5424 to
0de7b0d
Compare
cooperative_wait waits on a worker thread
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.
Waits that
cooperative_waithands to a worker thread (CUDA.jl, AMDGPU.jl, OpenCL on GPUs) allocate a request with its event, and box the object and the result: 224 to 288 bytes per wait. Whether a wait takes that path depends on how long the operation runs and how loaded the host is, so allocation counts of e.g. a kernel launch followed bysynchronize()aren't deterministic. #32 fixed the same problem for waits on completion notifications (PoCL).Requests are now reused from pools, one per type of wait function, object and inferred result type (
promote_op), so the object and the result are stored without boxing. They are kept in 1-element vectors that are emptied on release. That stores immutable objects containing references inline, like CUDA.jl's(stream, context)tuple, which aUnion{Nothing,T}field would box. It also keeps idle requests from retaining the driver's objects. Requests of interrupted or cancelled waits are dropped rather than reused, since the worker may still be using them.The worker now calls
waitin the caller's world rather than the latest one, so it sees the same methods as a direct call (and as the in-finalizer path, which already callswaitdirectly). That keeps the result of the inferred type, and the return type atUnion{Nothing,Some{R}}. Supporting a latest-world result of another type would need aSome{Any}in the return type, which boxes every returned result whencooperative_waitisn't inlined. The only behavior change: a method defined after the caller's world, e.g. a redefinition while the wait is in progress, is no longer picked up.The pools have their own spin lock. Sharing the worker lock cost about 0.6 µs per round trip, since the worker takes that lock just before waking the caller. An interrupt while spinning on the pool lock is deferred, like any other interrupt during the wait.
Per wait on a worker thread (Linux, Julia 1.12):
nothingInt32-like(mutable, mutable)tuple,Int32-likeWith CUDA.jl on an RTX 5080,
synchronize()after a long kernel goes from 672 to 400 bytes. The rest is allocated by CUDA.jl itself, mostly when its worker function sets up the context. A worker round trip takes about as long as before (1.0–1.2 µs, against 0.9–1.1). Tested on Julia 1.10–1.14 on Linux, and 1.10 and 1.13 on macOS.Bumps the version to 3.5.1.