From 0de7b0d6b627e4832ae19bed12545ce468e5f8f5 Mon Sep 17 00:00:00 2001 From: SimonDanisch Date: Tue, 3 Nov 2026 16:42:02 +0100 Subject: [PATCH] Don't allocate when `cooperative_wait` waits on a worker thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- Project.toml | 2 +- src/synchronization.jl | 111 ++++++++++++++++---- test/synchronization.jl | 221 ++++++++++++++++++++++++++++++++++------ 3 files changed, 282 insertions(+), 52 deletions(-) diff --git a/Project.toml b/Project.toml index 76f41ab..1a6d350 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "GPUToolbox" uuid = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc" -version = "3.5.0" +version = "3.5.1" [compat] julia = "1.10" diff --git a/src/synchronization.jl b/src/synchronization.jl index cb89b62..02f6a1d 100644 --- a/src/synchronization.jl +++ b/src/synchronization.jl @@ -65,15 +65,44 @@ end # deadlock if what the finalizer waits for depends on that. finalizers that become pending # on a worker thread (e.g., when it triggers a collection) run on another thread instead. -mutable struct WaitRequest - const wait::Any - const obj::Any +# requests are reused, so that waiting on a worker does not allocate. they are typed by the +# function, the object and the inferred result type, so that those are not boxed either. to +# make that type exact, the worker calls `wait` in the caller's world, like a direct call. +# the arguments and the result are kept in vectors, which store immutable values inline +# (unlike a `Union{Nothing,T}` field when `T` contains references), and are emptied when the +# request is released, so that idle requests do not keep the driver's objects alive. +# +# only the caller releases a request, once the worker has signalled it (its last access to +# the request), or when polling found the object to have completed before submitting it. +# requests of interrupted or cancelled waits are not reused. +mutable struct WaitRequest{W,O,R} + const call::Vector{Tuple{W,O}} # `(wait, obj)`, while in use + const result::Vector{R} # empty if `wait` threw + exception::Any + world::Any # boxed, as passed to `Core._call_in_world` const done::Base.Event - result::Any - failed::Bool - submitted::Bool # whether a worker has been handed the request + submitted::Bool # whether a worker has been handed the request + + WaitRequest{W,O,R}() where {W,O,R} = + new(sizehint!(Tuple{W,O}[], 1), sizehint!(R[], 1), nothing, nothing, Base.Event(), + false) +end + +function perform!(request::WaitRequest{W,O,R}) where {W,O,R} + wait, obj = request.call[1] + try + # inferred in the same world, so this only fails if inference is wrong + push!(request.result, wait(obj)::R) + catch err + request.exception = err + end + return +end - WaitRequest(wait, obj) = new(wait, obj, Base.Event(), nothing, false, false) +@static if isdefined(Base, :tls_world_age) + tls_world_age() = Base.tls_world_age() +else + tls_world_age() = ccall(:jl_get_tls_world_age, UInt, ()) end mutable struct WaitWorker @@ -93,9 +122,32 @@ const idle_wait_workers = WaitWorker[] const idle_overflow_wait_workers = WaitWorker[] const wait_worker_lock = ReentrantLock() +# idle requests, by type. the values are untyped, as checking them against `Vector` (a +# `UnionAll`) is much slower than against the expected `Vector{T}`. +const request_pools = IdDict{DataType,Any}() +const request_pool_lock = Threads.SpinLock() + # how long to poll before checking again whether a worker has become available const POLL_RETRY_NS = 100_000 +function acquire_request(::Type{T}) where {T<:WaitRequest} + @lock request_pool_lock begin + pool = get(request_pools, T, nothing) + pool isa Vector{T} && !isempty(pool) && return pop!(pool) + end + return T() +end + +function release_request(request::T) where {T<:WaitRequest} + empty!(request.call) + empty!(request.result) + request.exception = nothing + request.submitted = false + reset(request.done) + @lock request_pool_lock push!(get!(Vector{T}, request_pools, T)::Vector{T}, request) + return +end + function wait_worker_loop(data::Ptr{Cvoid}) worker = unsafe_pointer_to_objref(data)::WaitWorker # never run finalizers on this thread (see above). this is never undone, so that also @@ -106,13 +158,7 @@ function wait_worker_loop(data::Ptr{Cvoid}) request = worker.request::WaitRequest worker.request = nothing - try - # the worker was started in an older world - request.result = Base.invokelatest(request.wait, request.obj) - catch err - request.result = err - request.failed = true - end + Core._call_in_world(request.world, perform!, request) # become available before waking the caller, so that it finds this worker when it # immediately waits again @@ -370,9 +416,9 @@ Wait for `obj` to complete without blocking the calling thread, so that other ta on it in the meantime. `wait(obj)` should perform a blocking wait for `obj`. By default, it is executed on a -separate thread, so it should block in a GC-safe manner (e.g., using -[`@gcsafe_ccall`](@ref)) and set up any thread-local state it relies on (e.g., the active -context). +separate thread (in the caller's world, i.e., seeing the same method definitions as a +direct call), so it should block in a GC-safe manner (e.g., using [`@gcsafe_ccall`](@ref)) +and set up any thread-local state it relies on (e.g., the active context). Alternatively, if the driver can notify completion by calling back, pass a function `subscribe(obj, payload::Ptr{Cvoid})` that registers such a callback, which in turn calls @@ -436,9 +482,9 @@ immediately, while the operation may still be executing. In finalizers, where it is not possible to switch tasks, and while generating output (e.g., during precompilation), `wait(obj)` is called on the calling thread instead. """ -function cooperative_wait(wait::W, obj; subscribe::S=nothing, isdone::D=nothing, +function cooperative_wait(wait::W, obj::O; subscribe::S=nothing, isdone::D=nothing, spin::Union{Bool,Real}=true, - cancellable::Bool=false) where {W,S,D} + cancellable::Bool=false) where {W,O,S,D} spin isa Bool || spin >= 0 || throw(ArgumentError("spin duration must be non-negative")) if GC.in_finalizer() || generating_output() return Some(wait(obj)) @@ -470,12 +516,33 @@ function cooperative_wait(wait::W, obj; subscribe::S=nothing, isdone::D=nothing, # slow path: wait for a completion notification, or hand the wait to a worker thread if subscribe === nothing - request = WaitRequest(wait, obj) + R = Base.promote_op(wait, O) + # an interrupt while spinning on the pool's lock is deferred like above + acquired = nothing + while acquired === nothing + try + acquired = acquire_request(WaitRequest{W,O,R}) + catch err + (cancellable || !(err isa InterruptException)) && rethrow() + interrupt === nothing && (interrupt = err) + end + end + request = acquired::WaitRequest{W,O,R} + push!(request.call, (wait, obj)) + # only box the world when it changed, which is rare + world = tls_world_age() + old_world = request.world + if !(old_world isa UInt && old_world == world) + request.world = world + end slow_wait = () -> wait_cooperatively(request, obj, isdone) polled = cancellable ? slow_wait() : wait_uninterrupted(slow_wait, interrupt) + exception = request.exception + result = isempty(request.result) ? nothing : Some{R}(request.result[1]) + release_request(request) polled && return nothing - request.failed && throw(request.result) - return Some(request.result) + result === nothing && throw(exception) + return result else c = acquire_completion() try diff --git a/test/synchronization.jl b/test/synchronization.jl index 0413b1e..2a34461 100644 --- a/test/synchronization.jl +++ b/test/synchronization.jl @@ -16,6 +16,62 @@ function blocking_wait(op::Operation) return :waited end +# bits-type argument and result, both of which would be boxed by an untyped request +struct Deadline + time::UInt64 +end +Deadline(secs::Real) = Deadline(time_ns() + round(UInt64, secs * 1e9)) +passed(d::Deadline) = time_ns() >= d.time +function sleep_until(d::Deadline) + while !passed(d) + @gcsafe_ccall uv_sleep(1::Cuint)::Cvoid + end + return d.time +end + +# the same wait, like CUDA.jl's: on an immutable object containing references, returning a +# status code (whose boxes, unlike those of small integers, are not cached) +primitive type Status 32 end +const WAITED = reinterpret(Status, Int32(0)) +const wait_context = Ref(0) +status_wait((op, _)) = (blocking_wait(op); WAITED) +status_isdone((op, _)) = isdone(op) + +worker_wait(d::Deadline) = cooperative_wait(sleep_until, d; isdone=passed, spin=false) +# not inlined, so that the result is returned (as with CUDA.jl's `synchronize`) +@noinline tuple_wait(op::Operation) = + cooperative_wait(status_wait, (op, wait_context); isdone=status_isdone, spin=false) + +# create deadlines after compilation, so measurements include actual blocking waits +Base.@constprop :none function worker_wait_allocations(n) + bytes = 0 + for _ in 1:n + d = Deadline(1e-3) + bytes += @allocated worker_wait(d) + end + return bytes +end + +Base.@constprop :none function tuple_wait_allocations(op, n) + bytes = 0 + for _ in 1:n + bytes += @allocated something(tuple_wait(op), WAITED)::Status + end + return bytes +end + +# a wait whose function, object and result are immutable but contain references +@noinline function pooled_wait_weakrefs() + owner, obj = Ref(0), Ref(0) + refs = (WeakRef(owner), WeakRef(obj)) + wait = ((o, _),) -> (owner[] += o[]; (owner, 1)) + @test cooperative_wait(wait, (obj, 1); spin=false) == Some((owner, 1)) + return refs +end + +# redefined during a wait, to return something else +changing_wait(_) = :before + function short_waits() for _ in 1:10 cooperative_wait(blocking_wait, complete_after!(Operation(), 0.001); isdone) @@ -50,6 +106,35 @@ function signal_later(payload, ms) return end +# run `f` on a thread that is not managed by Julia +mutable struct ForeignCall + const f::Any + @atomic done::Bool + @atomic failed::Bool +end +const foreign_calls = ForeignCall[] # keep them alive +function foreign_call(arg::Ptr{Cvoid}) + call = unsafe_pointer_to_objref(arg)::ForeignCall + try + call.f() + catch + @atomic call.failed = true + end + @atomic call.done = true + return +end +function on_foreign_thread(f) + call = ForeignCall(f, false, false) + push!(foreign_calls, call) + tid = Ref{NTuple{32, UInt8}}(ntuple(_ -> 0x0, 32)) + cb = @cfunction(foreign_call, Cvoid, (Ptr{Cvoid},)) + err = ccall(:uv_thread_create, Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), + tid, cb, pointer_from_objref(call)) + @test err == 0 + ccall(:uv_thread_detach, Cint, (Ptr{Cvoid},), tid) + return call +end + @testset "cooperative_wait" begin # completed operations are detected by polling @test cooperative_wait(blocking_wait, complete!(Operation()); isdone) === nothing @@ -68,6 +153,20 @@ end @test @allocated(fast_wait(op)) == 0 end + # requests are reused, so that waiting on a worker does not allocate either, also + # not to box the object or the result + if !isdefined(Base, :CANCEL_TOKEN) + worker_wait_allocations(4) + op = complete!(Operation()) + tuple_wait_allocations(op, 4) + # the scheduler can allocate when the thread goes idle, so take the best batch + @test minimum(worker_wait_allocations(32) for _ in 1:5) == 0 + @test minimum(tuple_wait_allocations(op, 32) for _ in 1:5) == 0 + d = Deadline(1e-3) + @test worker_wait(d) == Some(d.time) + @test tuple_wait(op) === Some(WAITED) + end + # other tasks on this thread keep running while waiting on a worker for polled in (false, true) op = complete_after!(Operation(), 30) # in case the thread is blocked @@ -91,10 +190,69 @@ end @test_throws ErrorException("oops") cooperative_wait(blocking_wait, Operation(); isdone=_ -> error("oops")) + # an interrupt after the worker signalled the request resumes waiting, which returns + # immediately + let T = GPUToolbox.WaitRequest{typeof(identity),Nothing,Nothing}, + request = GPUToolbox.acquire_request(T), + attempts = 0 + request.submitted = true + notify(request.done) + interrupted_wait = function () + GPUToolbox.wait_cooperatively(request, nothing, nothing) + attempts += 1 + attempts == 1 && throw(InterruptException()) + end + t = @async GPUToolbox.wait_uninterrupted(interrupted_wait, nothing) + finished = timedwait(() -> istaskdone(t), 10) === :ok + @test finished + finished || notify(request.done) + @test_throws TaskFailedException wait(t) + @test t.exception isa InterruptException + @test attempts == 2 + GPUToolbox.release_request(request) + end + + # reused requests do not carry over a previous result or exception + let value = Ref{Any}(:success) + maybe_fail = _ -> value[] === :success ? :waited : throw(value[]) + value[] = nothing + @test_throws Nothing cooperative_wait(maybe_fail, nothing; spin=false) + value[] = :success + @test cooperative_wait(maybe_fail, nothing; spin=false) === Some(:waited) + value[] = false + @test_throws Bool cooperative_wait(maybe_fail, nothing; spin=false) + value[] = :success + @test cooperative_wait(maybe_fail, nothing; spin=false) === Some(:waited) + end + + for value in (nothing, UInt64(0xf123456789abcdef), "result", Ref(1), (Ref(1), 2)) + @test something(cooperative_wait(identity, value; spin=false)) === value + end + + # idle requests do not keep the function, the object or the result alive + refs = pooled_wait_weakrefs() + GC.gc() + @test all(r -> r.value === nothing, refs) + # functions defined after the workers were started can be used f = @eval _ -> :new_function @test cooperative_wait(f, nothing) == Some(:new_function) + # the function is called in the caller's world, like a direct call would be + let redefined = Ref(false), world = Base.get_world_counter() + function redefining_isdone(_) + redefined[] || (redefined[] = true; @eval changing_wait(_) = :after) + return false + end + @test cooperative_wait(changing_wait, nothing; isdone=redefining_isdone) == + Some(:before) + @test redefined[] + # also when requests are reused across worlds + @test Base.invokelatest(cooperative_wait, changing_wait, nothing) == Some(:after) + @test Base.invoke_in_world(world, cooperative_wait, changing_wait, nothing) == + Some(:before) + end + # a long wait does not delay other ones, even with more waits than worker threads for polled in (false, true) long = Operation() @@ -203,6 +361,40 @@ end @test t.exception isa InterruptException end + # and for interrupts while spinning on the lock of the request pool (see below for the + # conditions under which SIGINT can be tested) + if Sys.islinux() && Threads.threadid() == 1 && !isdefined(Base, :CANCEL_TOKEN) && + !(v"1.12-" <= VERSION < v"1.13-") + Base.exit_on_sigint(false) + acquiring = Threads.Atomic{Bool}(false) + holder = nothing + try + holder = on_foreign_thread() do + @lock GPUToolbox.request_pool_lock begin + while !acquiring[] + @gcsafe_ccall uv_sleep(1::Cuint)::Cvoid + end + # give the main thread time to start spinning on the lock + @gcsafe_ccall uv_sleep(100::Cuint)::Cvoid + ccall(:kill, Cint, (Cint, Cint), getpid(), Base.SIGINT) + @gcsafe_ccall uv_sleep(100::Cuint)::Cvoid + end + end + @test timedwait(() -> islocked(GPUToolbox.request_pool_lock), 30) === :ok + op = complete_after!(Operation(), 0.5) + @test_throws InterruptException begin + acquiring[] = true + cooperative_wait(blocking_wait, op; spin=false) + end + @test isdone(op) + finally + acquiring[] = true + @test timedwait(() -> @atomic(holder.done), 30) === :ok + @test !@atomic(holder.failed) + Base.exit_on_sigint(true) + end + end + # the same goes for task cancellation if isdefined(Base, :CancellationTokenSource) src = Base.CancellationTokenSource() @@ -367,35 +559,6 @@ end @test isassigned(ret) && ret[] == Some(:waited) end -# run `f` on a thread that is not managed by Julia -mutable struct ForeignCall - const f::Any - @atomic done::Bool - @atomic failed::Bool -end -const foreign_calls = ForeignCall[] # keep them alive -function foreign_call(arg::Ptr{Cvoid}) - call = unsafe_pointer_to_objref(arg)::ForeignCall - try - call.f() - catch - @atomic call.failed = true - end - @atomic call.done = true - return -end -function on_foreign_thread(f) - call = ForeignCall(f, false, false) - push!(foreign_calls, call) - tid = Ref{NTuple{32, UInt8}}(ntuple(_ -> 0x0, 32)) - cb = @cfunction(foreign_call, Cvoid, (Ptr{Cvoid},)) - err = ccall(:uv_thread_create, Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), - tid, cb, pointer_from_objref(call)) - @test err == 0 - ccall(:uv_thread_detach, Cint, (Ptr{Cvoid},), tid) - return call -end - @testset "completion lock" begin l = GPUToolbox.YieldingSpinLock() @test_throws ConcurrencyViolationError Base.assert_havelock(l)