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)