From 11d1bdeb06b3336cd11126a082e07a0e4376d03f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 3 Oct 2026 21:38:28 +0200 Subject: [PATCH 1/2] Don't wait for memory whose last use is known to have completed When memory moves to another stream, CUDA.jl orders the new stream after the one that last used the memory, by waiting on the device or synchronizing on the host. That's unnecessary when the host already knows that the last use has completed, e.g., because it synchronized an event recorded after it. Divide each stream's work into epochs, stamp memory with the epoch of its last use, and close an epoch when recording an event or synchronizing the stream. Synchronizing a stream, or observing that an event has completed, then tells us that all accesses up to the closed epoch have finished, so memory last used in them doesn't need to be waited for anymore. Waiting for an event doesn't hold its lock, so other tasks can record or query it in the meantime; what the wait covers is only learned if the event wasn't recorded again. Operations that CUDA.jl submits itself (kernel launches, copies, fill!, graph launches) stamp memory when taking its pointer, and again once the operation has been submitted, also when that fails. Between the two, an epoch closed by another task could appear to cover the operation, so only the task submitting work to a stream closes epochs that count. Epochs closed by another task, or on a stream that several live tasks submit work to, therefore don't cover anything. A stream whose task has finished can be taken over by another one, as happens when streams are recycled. A pointer taken outside of such an operation, e.g. with `pointer(a)` or when passing an array to a library call, is only stamped once and can be used by work that is submitted after its epoch was closed. Such memory isn't considered complete until it moves to another stream, which waits for all work submitted to the old one so far. Recording an event on a stream that is being captured only adds a node to the graph, so it doesn't close an epoch, and captured operations don't stamp the memory they use (launching the graph does). Memory used by captures that CUDA.jl doesn't know about, and memory used on the special streams, are never considered complete. --- CUDACore/lib/cudadrv/events.jl | 49 ++++++++- CUDACore/lib/cudadrv/state.jl | 3 + CUDACore/lib/cudadrv/stream.jl | 64 +++++++++++- CUDACore/lib/cudadrv/synchronization.jl | 30 +++++- CUDACore/src/compiler/execution.jl | 13 +++ CUDACore/src/resources/managed.jl | 49 ++++++++- NEWS.md | 8 ++ docs/src/usage/multitasking.md | 7 ++ test/core/cudadrv.jl | 128 ++++++++++++++++++++++++ 9 files changed, 339 insertions(+), 12 deletions(-) diff --git a/CUDACore/lib/cudadrv/events.jl b/CUDACore/lib/cudadrv/events.jl index bfb6f0e075..af1f6210e2 100644 --- a/CUDACore/lib/cudadrv/events.jl +++ b/CUDACore/lib/cudadrv/events.jl @@ -15,12 +15,20 @@ mutable struct CuEvent handle::CUevent ctx::CuContext + # the stream and epoch covered by the most recent recording of this event, if tracked. + # each tracked recording closes a different epoch, so this also identifies the + # recording. protected by `lock`, which is held while recording the event, so that + # what we learn from querying or waiting for it applies to the recording it was told + # about. + source::Union{Nothing,Tuple{StreamOrder,UInt64}} + const lock::ReentrantLock + function CuEvent(flags=EVENT_DEFAULT) handle_ref = Ref{CUevent}() cuEventCreate(handle_ref, flags) ctx = current_context() - obj = new(handle_ref[], ctx) + obj = new(handle_ref[], ctx, nothing, ReentrantLock()) # destroying an event from a finalizer is deferred, as with all resources resource_finalizer(obj) return obj @@ -44,7 +52,27 @@ Base.hash(e::CuEvent, h::UInt) = hash(e.handle, h) Record an event on a stream. """ record(e::CuEvent, stream::CuStream=stream()) = - capture_submission(() -> cuEventRecord(e, stream), stream) + capture_submission(() -> record_tracked(e, stream), stream) +function record_tracked(e::CuEvent, stream::CuStream) + Base.@lock e.lock begin + # recording during capture only adds a node to the graph, which neither completes + # nor orders anything before the graph is launched + order = stream.order + epoch = order === nothing || is_capturing(stream) ? nothing : close_epoch!(order) + e.source = nothing + cuEventRecord(e, stream) + epoch === nothing || (e.source = (order, epoch)) + end + return +end + +# the recording of the event that covers `source` has completed, so the work it covers has +# too. that's only known if the event hasn't been recorded again since. (caller holds +# `e.lock`) +function mark_completed!(e::CuEvent, source) + source === nothing || e.source !== source || mark_completed!(source...) + return +end """ synchronize(e::CuEvent) @@ -60,6 +88,23 @@ Return `false` if there is outstanding work preceding the most recent call to `record(e)` and `true` if all captured work has been completed. """ function isdone(e::CuEvent) + Base.@lock e.lock begin + done = unsafe_isdone(e) + source = e.source + if done && source !== nothing && !is_completed(source...) + # work that is known to be done isn't waited for anymore, so this needs the + # sanitizer workaround in `synchronize_completed`. a tracked recording can't + # change while we hold the lock, so that doesn't block. + res = synchronize_completed(e) + res == SUCCESS || throw_api_error(res) + mark_completed!(e, source) + end + return done + end +end + +# query the event without taking its lock, e.g., from a worker thread +function unsafe_isdone(e::CuEvent) res = unchecked_cuEventQuery(e) if res == ERROR_NOT_READY return false diff --git a/CUDACore/lib/cudadrv/state.jl b/CUDACore/lib/cudadrv/state.jl index 8dbdf60535..1a990c1f37 100644 --- a/CUDACore/lib/cudadrv/state.jl +++ b/CUDACore/lib/cudadrv/state.jl @@ -344,6 +344,7 @@ const stream_pool_lock = ReentrantLock() push!(stream_pools[key], PooledStream(stream, WeakRef(task))) end end + bind!(stream) # register the name of this task # XXX: do this when the user has imported NVTX.jl (using weak dependencies?) @@ -525,6 +526,7 @@ function stream!(stream::CuStream) check_stream_capture() state = task_local_state!() devidx = deviceid(state.device)+1 + bind!(stream) state.streams[devidx] = stream return end @@ -535,6 +537,7 @@ function stream!(f::Function, stream::CuStream) devidx = deviceid(state.device)+1 old_stream = state.streams[devidx] old_priority = state.priorities[devidx] + bind!(stream) state.streams[devidx] = stream try f() diff --git a/CUDACore/lib/cudadrv/stream.jl b/CUDACore/lib/cudadrv/stream.jl index 04fbcfc930..660726e98a 100644 --- a/CUDACore/lib/cudadrv/stream.jl +++ b/CUDACore/lib/cudadrv/stream.jl @@ -3,8 +3,26 @@ export CuStream, default_stream, legacy_stream, per_thread_stream, unique_id, priority, priority_range, synchronize, device_synchronize -# State used to order work on other streams after the work on a stream (see `stream_wait`). +# What the host knows about the work submitted to a stream, so that memory moving between +# streams only needs to be ordered when that's actually necessary (see `take_ownership!`). +# The stream's work is divided in epochs: memory accesses are stamped with the current +# epoch, and recording an event or synchronizing the stream closes it. An event or +# synchronization thus covers all accesses stamped with an epoch up to and including the +# one it closed. Recording an event on a stream that is being captured only adds a node to +# the graph, so it doesn't close anything. mutable struct StreamOrder + # the epoch accesses are currently stamped with + Base.@atomic epoch::UInt64 + + # all accesses stamped with this epoch, or an earlier one, have completed + Base.@atomic completed::UInt64 + + # the task submitting work to this stream, and whether other tasks do so too. a stream + # stays shared: tasks that activated it with `stream!` may still use it after the task + # it was created for has finished and it has been recycled. + Base.@atomic owner::Union{Nothing,WeakRef} + Base.@atomic shared::Bool + # an event to make other streams wait for this one, created when first needed. the lock # is held from recording the event until the other stream waited for it. (a `CuEvent`, # which is defined after streams.) @@ -19,9 +37,49 @@ mutable struct StreamOrder const lock::ReentrantLock - StreamOrder() = new(nothing, nothing, ReentrantLock()) + StreamOrder() = new(1, 0, nothing, false, nothing, nothing, ReentrantLock()) end +current_epoch(order::StreamOrder) = Base.@atomic order.epoch + +# close the current epoch, returning it when the operation about to be submitted (an event +# or synchronization) will cover all accesses stamped with it. the epoch needs to be closed +# *before* submitting that operation, so that accesses stamped after submitting their own +# work aren't covered. however, an access may also be stamped when taking a pointer, before +# submitting the work that uses it. only the task submitting work to the stream knows that +# has happened, so an epoch closed by another task, or when several tasks submit work to +# the stream, doesn't cover anything. +function close_epoch!(order::StreamOrder) + epoch = (Base.@atomic order.epoch += 1) - 1 + owner = Base.@atomic order.owner + if Base.@atomic(order.shared) || (owner !== nothing && owner.value !== current_task()) + return nothing + end + return epoch +end + +# note that the current task submits work to the stream. a task that has finished can't be +# about to submit work anymore, so its stream can be handed to another task, e.g., when +# streams are recycled, without considering it shared. +function bind!(order::StreamOrder) + task = current_task() + while true + owner = Base.@atomic order.owner + previous = owner === nothing ? nothing : owner.value + previous === task && return + if previous === nothing || istaskdone(previous) + _, bound = Base.@atomicreplace order.owner owner => WeakRef(task) + bound && return + else + Base.@atomic order.shared = true + return + end + end +end + +mark_completed!(order::StreamOrder, epoch::UInt64) = (Base.@atomic order.completed max epoch; return) +is_completed(order::StreamOrder, epoch::UInt64) = epoch <= Base.@atomic order.completed + """ CuStream(; flags=STREAM_DEFAULT, priority=nothing) @@ -111,6 +169,8 @@ per_thread_stream() Base.unsafe_convert(::Type{CUstream}, s::CuStream) = s.handle +bind!(s::CuStream) = (s.order === nothing || bind!(s.order); return) + # only bumped while holding `stream_pool_lock` and `stream_disposal_lock`, but read without them generation(s::CuStream) = Base.@atomic :acquire s.generation diff --git a/CUDACore/lib/cudadrv/synchronization.jl b/CUDACore/lib/cudadrv/synchronization.jl index ff05a4e4eb..e9162e7249 100644 --- a/CUDACore/lib/cudadrv/synchronization.jl +++ b/CUDACore/lib/cudadrv/synchronization.jl @@ -20,7 +20,12 @@ function worker_synchronize((obj, ctx)) unchecked_synchronize(obj) end end -worker_isdone((obj, ctx)) = isdone(obj) +worker_isdone((obj, ctx)) = poll(obj) + +# whether the object is done, without taking locks, so that this can be called from any +# thread. (`isdone` on an event also learns what the work it covers is.) +poll(obj::CuStream) = isdone(obj) +poll(obj::CuEvent) = unsafe_isdone(obj) function synchronize_object(obj::SyncObject; blocking::Bool, spin::Bool) obj isa CuEvent || check_capture(obj) @@ -30,7 +35,7 @@ function synchronize_object(obj::SyncObject; blocking::Bool, spin::Bool) pollable = !(obj isa CuContext) # if we're about to wait, now may be a good time for a GC pause - if !pollable || !isdone(obj) + if !pollable || !poll(obj) maybe_collect(true) end @@ -65,7 +70,9 @@ end # XXX: compute-sanitizer doesn't treat a successful query of an event as synchronization, and # reports races with the work that the event ordered (#3346). so when synchronizing an # event finds it done by polling, synchronize it too. that only blocks if the event was -# recorded again since. other queries of events, e.g. with `isdone`, aren't covered. +# recorded again since. `isdone` also does so when it learns that tracked work has +# completed. other queries of events, e.g. the ones that gate releasing held resources +# and reusing cached memory, aren't covered. # synchronizing is prohibited while another thread captures in global mode, even though # it doesn't affect that capture, so relax the capture mode. synchronize_completed(event::CuEvent) = @@ -77,6 +84,14 @@ function device_synchronize(; blocking::Bool=false, spin::Bool=true) end function synchronize(stream::CuStream=stream(); blocking::Bool=false, spin::Bool=true) + order = stream.order + epoch = if order === nothing + nothing + else + # (a stream that is being captured can't be synchronized, so don't close its epoch) + check_capture(stream) + close_epoch!(order) + end if stream.handle == CU_STREAM_PER_THREAD && !blocking && use_nonblocking_synchronization # the per-thread stream is specific to the calling thread, so it can't be # synchronized from a worker thread. wait for an event recorded on it instead. @@ -90,8 +105,15 @@ function synchronize(stream::CuStream=stream(); blocking::Bool=false, spin::Bool else synchronize_object(stream; blocking, spin) end + epoch === nothing || mark_completed!(order, epoch) check_exceptions() end -synchronize(event::CuEvent; blocking::Bool=false, spin::Bool=true) = +# the event's lock isn't held while waiting, so that other tasks can record or query it in +# the meantime. what the wait covers is only learned if the event hasn't been recorded again. +function synchronize(event::CuEvent; blocking::Bool=false, spin::Bool=true) + source = Base.@lock event.lock event.source synchronize_object(event; blocking, spin) + source === nothing || Base.@lock event.lock mark_completed!(event, source) + return +end diff --git a/CUDACore/src/compiler/execution.jl b/CUDACore/src/compiler/execution.jl index df05f586de..83808b97c0 100644 --- a/CUDACore/src/compiler/execution.jl +++ b/CUDACore/src/compiler/execution.jl @@ -445,6 +445,12 @@ function with_ordered_managed(f::F, ordered::Union{AbstractVector{<:Managed}, end return with_operation_stream(f, tls, stream, capturing) finally + if !capturing + epoch = stream_epoch(stream) + for memory in ordered + restamp!(memory, stream, epoch) + end + end for memory in Iterators.reverse(ordered) unlock(memory.lock) end @@ -452,6 +458,13 @@ function with_ordered_managed(f::F, ordered::Union{AbstractVector{<:Managed}, end end +# memory is stamped with the epoch of its stream once the operation using it has been +# submitted, also when `f` throws, as it may have submitted work first (see `StreamOrder`). +# only memory that `take_ownership!` actually moved to `stream` is stamped, which captured +# operations don't do (they execute when the graph is launched, which stamps the memory). +restamp!(memory::Managed, stream::CuStream, epoch::UInt64) = + (memory.stream == stream && (memory.epoch = epoch); return) + # pointers taken during `f` are used by an operation that CUDA.jl submits to `stream` (see # `convert(::Type{CuPtr}, ::Managed)`) @inline function with_operation_stream(f::F, tls::TaskLocalState, stream::CuStream, diff --git a/CUDACore/src/resources/managed.jl b/CUDACore/src/resources/managed.jl index 25c17bc153..1effd65a22 100644 --- a/CUDACore/src/resources/managed.jl +++ b/CUDACore/src/resources/managed.jl @@ -37,6 +37,9 @@ mutable struct Managed{M} stream_ctx::CuContext generation::Int + # the epoch of `stream` during which the memory was last used (see `StreamOrder`) + epoch::UInt64 + # whether accessing this memory can cause implicit synchronization synchronizing::Bool @@ -51,6 +54,10 @@ mutable struct Managed{M} # memory, which therefore haven't been synchronized either (implies `dirty`) waiting::Bool + # whether a pointer was taken outside of an operation that CUDA.jl submits, and may thus + # be used by work submitted to `stream` after it was stamped with `epoch` + escaped::Bool + # whether the memory was allocated by CUDA.jl, as opposed to imported using `unsafe_wrap`. # only such memory counts towards our memory usage, and can be reused once freed. const owned_allocation::Bool @@ -70,8 +77,8 @@ mutable struct Managed{M} # NOTE: memory starts as dirty, because stream-ordered allocations are only # guaranteed to be physically allocated at a synchronization event. new{typeof(mem)}(mem, ReentrantLock(), stream, mem.ctx, generation(stream), - synchronizing, dirty, captured, false, owned_allocation, attachment, - 0, nothing) + stream_epoch(stream), synchronizing, dirty, captured, false, false, + owned_allocation, attachment, 0, nothing) end end @@ -80,6 +87,21 @@ recycled(managed::Managed) = managed.generation != generation(managed.stream) Base.sizeof(managed::Managed) = sizeof(managed.mem) +# the epoch to stamp an access on `stream` with. accesses on untracked streams are never +# covered by an event or synchronization. +function stream_epoch(stream::CuStream) + order = stream.order + order === nothing ? typemax(UInt64) : current_epoch(order) +end + +# whether the last use of memory is known to have completed. memory whose pointer escaped +# may be used by work that was submitted after the epoch it was stamped with was closed. +function is_completed(managed::Managed) + order = managed.stream.order + !managed.captured && !managed.escaped && order !== nothing && + is_completed(order, managed.epoch) +end + # wait for the current owner of memory to finish processing function synchronize(managed::Managed) Base.@lock managed.lock begin @@ -133,7 +155,12 @@ end function maybe_synchronize(managed::Managed) Base.@lock managed.lock begin if managed.synchronizing && (managed.dirty || managed.captured) - synchronize(managed) + if is_completed(managed) + managed.dirty = false + managed.waiting = false + else + synchronize(managed) + end end end end @@ -237,6 +264,11 @@ end # whether that was possible (see `stream_wait`) function handoff!(managed::Managed, stream::CuStream) managed.dirty || return true + if is_completed(managed) + managed.dirty = false + managed.waiting = false + return true + end source = managed.stream # (holding the lock that recycling streams takes, see `pending_work`) @lock stream_disposal_lock begin @@ -309,6 +341,9 @@ function take_ownership!(managed::Managed{M}; state=active_state(), end managed.stream = stream managed.stream_ctx = state.context + # the work submitted to the previous stream has been waited for, which is all that's + # known about uses of pointers that escaped (see `convert(::Type{CuPtr}, ::Managed)`) + managed.escaped = false end managed.generation = generation(managed.stream) @@ -328,6 +363,7 @@ function take_ownership!(managed::Managed{M}; state=active_state(), can_prefetch && prefetch(managed.mem; device=state.device, stream) end + managed.epoch = stream_epoch(stream) managed.dirty = true return managed end @@ -339,12 +375,17 @@ function Base.convert(::Type{CuPtr{T}}, managed::Managed{M}) where {T,M} ptr == CU_NULL && return ptr # within an operation that CUDA.jl submits (see `with_managed`), the pointer is used on - # that operation's stream. otherwise, we don't know who is going to use it. + # that operation's stream. otherwise, we don't know who is going to use it, or when. + # it's assumed to be used on the task's stream, until the memory is used on another + # stream: work submitted through the pointer after that isn't waited for. that work + # may also be submitted after the stream was synchronized or an event was recorded, + # so neither tells anything about the memory anymore (see `is_completed`). tls = task_local_state!() state = active_state(tls) stream = tls.operation_stream if stream === nothing take_ownership!(managed; state, stream=state.stream, external=true) + managed.escaped = true else take_ownership!(managed; state, stream, capturing=tls.operation_capturing) end diff --git a/NEWS.md b/NEWS.md index 7e2bcac986..9a9f25d2fd 100644 --- a/NEWS.md +++ b/NEWS.md @@ -91,6 +91,14 @@ are listed as subsections of the minor release they belong to. (Windows, and Jetson boards up to Orin). Previously, that crashed with a memory access error. +*Minor changes*: + +- An array that another task uses doesn't wait for work that the host already + waited for, by synchronizing the stream or an event recorded after that work + on the task that submitted it. Work submitted through a pointer taken outside + of a CUDA.jl operation (e.g. with `pointer(a)`) is only waited for if it was + submitted from that task's stream before the array is used on another stream. + ## v6.4 (September 2026) diff --git a/docs/src/usage/multitasking.md b/docs/src/usage/multitasking.md index d2bcfa0f30..a325291ec0 100644 --- a/docs/src/usage/multitasking.md +++ b/docs/src/usage/multitasking.md @@ -140,6 +140,13 @@ New tasks start at normal priority, even when spawned from a high-priority task. If you selected an explicit stream with `stream!`, use the scoped form of `priority!` to restore that stream after the priority change. +When another task uses an array, CUDA.jl makes it wait for the work that was submitted for +that array before. That includes work submitted through a pointer to the array's memory, +e.g., one obtained with `pointer(a)` and passed to a library or a kernel, but only if it was +submitted from the task that took the pointer, and before the array was used by another +task. Work submitted through such a pointer later is not ordered with the other task's +work, so take the pointer again instead. + ## Multithreading diff --git a/test/core/cudadrv.jl b/test/core/cudadrv.jl index d950665d2a..39e5fedd53 100644 --- a/test/core/cudadrv.jl +++ b/test/core/cudadrv.jl @@ -1266,6 +1266,134 @@ let a = CUDA.zeros(Int, 1), b = CUDA.zeros(Int, 1), u = cu([0]; unified=true), @test Array(u) == [4] end +# memory whose last use is known to have completed doesn't need to be waited for. as the +# kernels above don't block anyway, check that with the pointer conversion of a library. +unsafe_use(a) = (Base.unsafe_convert(CuPtr{Int}, a); nothing) +let a = CUDA.zeros(Int, 1), s = stream(), other = CuStream(), third = CuStream() + a .+= 1 + fetch(Threads.@spawn (stream!(other); unsafe_use(a))) + synchronize() + + # the host waited for an event recorded after the last use + a .+= 1 + e = CuEvent() + record(e) + @test gated(s) do + fetch(Threads.@spawn begin + stream!(other) + synchronize(e) + unsafe_use(a) + !CUDA.isdone(s) + end) + end + + # but not for a use after that event + e = CuEvent() + record(e) + @test gated(s) do + a .+= 1 + open_gate_during() do + fetch(Threads.@spawn (stream!(other); synchronize(e); unsafe_use(a))) + end + end + + # an event recorded by another task doesn't count, as it may have been recorded before + # the work that used the memory was submitted + a .+= 1 + e = fetch(Threads.@spawn (e = CuEvent(); record(e, s); e)) + @test gated(s) do + open_gate_during() do + fetch(Threads.@spawn (stream!(other); synchronize(e); unsafe_use(a))) + end + end + + # re-recording an event replaces the work it covers + a .+= 1 + e = CuEvent() + record(e) + record(e, third) + @test gated(s) do + open_gate_during() do + fetch(Threads.@spawn (stream!(other); synchronize(e); unsafe_use(a))) + end + end + + # a stream can be used by another task once the previous one has finished + s2 = CuStream() + wait(Threads.@spawn (stream!(s2); synchronize(s2))) + e = fetch(Threads.@spawn begin + stream!(s2) + a .+= 1 + e = CuEvent() + record(e) + e + end) + @test gated(s2) do + fetch(Threads.@spawn begin + stream!(other) + synchronize(e) + unsafe_use(a) + !CUDA.isdone(s2) + end) + end + synchronize(s2) + @test Array(a) == [6] +end + +# work submitted through a pointer that was taken before synchronizing the stream isn't +# covered by that synchronization, so other tasks using the memory still need to wait for it +store_kernel(x) = (@inbounds x[1] = 42; return) +raw_array(p::CuPtr{Int}) = + CuDeviceArray{Int,1,CUDA.AS.Global}(reinterpret(Core.LLVMPtr{Int,CUDA.AS.Global}, p), (1,)) +let a = CUDA.zeros(Int, 1), s = stream(), other = CuStream() + # warm up, as compiling kernels waits for the GPU to become idle + @cuda stream=s store_kernel(raw_array(pointer(a))) + fetch(Threads.@spawn (stream!(other); a .+= 1; unsafe_use(a))) + synchronize(other) + + # library pointer conversions on another task wait for the launch + p = pointer(a) + synchronize(s) + @test gated(s) do + @cuda stream=s store_kernel(raw_array(p)) + open_gate_during() do + fetch(Threads.@spawn (stream!(other); unsafe_use(a))) + end + end + + # and so do kernels on another task's stream + p = pointer(a) + synchronize(s) + @test gated(s) do + @cuda stream=s store_kernel(raw_array(p)) + fetch(Threads.@spawn begin + stream!(other) + a .+= 1 + !CUDA.isdone(other) + end) + end + synchronize(other) + @test Array(a) == [43] +end + +# an event can be recorded again while another task waits for it, which then doesn't tell +# anything about the work covered by the earlier recording +let a = CUDA.zeros(Int, 1), s = stream(), other = CuStream(), e = CuEvent() + @test gated(s) do + a .+= 1 + record(e) + waiter = Threads.@spawn synchronize(e) + for _ in 1:100 + yield() + end + record(e, other) + wait(waiter) + open_gate_during() do + fetch(Threads.@spawn (stream!(other); unsafe_use(a))) + end + end +end + end ############################################################################################ From 38d4a784772fc89636ccea1a029d44adfeb1200e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 3 Oct 2026 21:45:55 +0200 Subject: [PATCH 2/2] Don't wait for memory that a stream is already ordered after When a stream waits for an event, everything it executes afterwards is ordered after the work the event covers. Memory last used in that work can then move to that stream without waiting for the stream that used it. That's what KernelAbstractions' `@spawn` relies on: the spawned task waits for an event recorded by its parent, but its first use of the parent's arrays currently also waits for everything the parent queued after the event. Have each stream remember, per stream it waited for, up to which epoch its work is ordered, and don't wait again for memory whose last use is covered by that. This applies to the memory that could be handed off on the device anyway, and not to memory whose pointer escaped. The memory remains marked as waiting on another stream, so that a pointer conversion for a consumer outside of CUDA.jl, like MPI, synchronizes the new stream. That only waits for the work the event covered, and not for everything the other stream queued since, as synchronizing that stream would. --- CUDACore/lib/cudadrv/events.jl | 17 +++++++- CUDACore/lib/cudadrv/stream.jl | 23 ++++++++++- CUDACore/src/resources/managed.jl | 15 ++++++- NEWS.md | 4 ++ test/core/cudadrv.jl | 65 +++++++++++++++++++++++++++++++ 5 files changed, 121 insertions(+), 3 deletions(-) diff --git a/CUDACore/lib/cudadrv/events.jl b/CUDACore/lib/cudadrv/events.jl index af1f6210e2..daa1e13a23 100644 --- a/CUDACore/lib/cudadrv/events.jl +++ b/CUDACore/lib/cudadrv/events.jl @@ -122,7 +122,22 @@ Make a stream wait on a event. This only makes the stream wait, and not the host [`synchronize(::CuEvent)`](@ref) for that. """ wait(e::CuEvent, stream::CuStream=stream()) = - capture_submission(() -> cuStreamWaitEvent(stream, e, 0), stream) + capture_submission(() -> wait_tracked(e, stream), stream) +function wait_tracked(e::CuEvent, stream::CuStream) + Base.@lock e.lock begin + cuStreamWaitEvent(stream, e, 0) + + # only learn about the order once the wait has been submitted, so that work isn't + # believed to be ordered after the event before it actually is. a wait that is + # captured only orders the graph's operations, once it is launched. + source = e.source + order = stream.order + if source !== nothing && order !== nothing && !is_capturing(stream) + mark_ordered!(order, source...) + end + end + return +end # make the work submitted to `stream` from now on wait for the work submitted to `source` so # far, on the device. returns false if that isn't possible because `source` is being diff --git a/CUDACore/lib/cudadrv/stream.jl b/CUDACore/lib/cudadrv/stream.jl index 660726e98a..0af123511b 100644 --- a/CUDACore/lib/cudadrv/stream.jl +++ b/CUDACore/lib/cudadrv/stream.jl @@ -23,6 +23,11 @@ mutable struct StreamOrder Base.@atomic owner::Union{Nothing,WeakRef} Base.@atomic shared::Bool + # work submitted to this stream from now on is ordered after the accesses to other + # streams up to these epochs, because the stream waited for an event covering them. + # weakly keyed, as there's nothing to learn about streams that are gone. + const waited::WeakKeyDict{StreamOrder,UInt64} + # an event to make other streams wait for this one, created when first needed. the lock # is held from recording the event until the other stream waited for it. (a `CuEvent`, # which is defined after streams.) @@ -37,7 +42,8 @@ mutable struct StreamOrder const lock::ReentrantLock - StreamOrder() = new(1, 0, nothing, false, nothing, nothing, ReentrantLock()) + StreamOrder() = new(1, 0, nothing, false, WeakKeyDict{StreamOrder,UInt64}(), nothing, + nothing, ReentrantLock()) end current_epoch(order::StreamOrder) = Base.@atomic order.epoch @@ -80,6 +86,21 @@ end mark_completed!(order::StreamOrder, epoch::UInt64) = (Base.@atomic order.completed max epoch; return) is_completed(order::StreamOrder, epoch::UInt64) = epoch <= Base.@atomic order.completed +# record that work submitted to `order` from now on runs after `source`'s accesses up to `epoch` +function mark_ordered!(order::StreamOrder, source::StreamOrder, epoch::UInt64) + Base.@lock order.waited begin + order.waited[source] = max(get(order.waited, source, UInt64(0)), epoch) + end + return +end + +# whether work submitted to `order` now runs after `source`'s accesses up to `epoch` +function is_ordered(order::StreamOrder, source::StreamOrder, epoch::UInt64) + Base.@lock order.waited begin + epoch <= get(order.waited, source, UInt64(0)) + end +end + """ CuStream(; flags=STREAM_DEFAULT, priority=nothing) diff --git a/CUDACore/src/resources/managed.jl b/CUDACore/src/resources/managed.jl index 1effd65a22..a982d2bc59 100644 --- a/CUDACore/src/resources/managed.jl +++ b/CUDACore/src/resources/managed.jl @@ -94,6 +94,14 @@ function stream_epoch(stream::CuStream) order === nothing ? typemax(UInt64) : current_epoch(order) end +# whether work submitted to `stream` is known to run after the last use of memory +function is_ordered(managed::Managed, stream::CuStream) + order = stream.order + source = managed.stream.order + !managed.captured && !managed.escaped && order !== nothing && source !== nothing && + is_ordered(order, source, managed.epoch) +end + # whether the last use of memory is known to have completed. memory whose pointer escaped # may be used by work that was submitted after the epoch it was stamped with was closed. function is_completed(managed::Managed) @@ -336,7 +344,12 @@ function take_ownership!(managed::Managed{M}; state=active_state(), # accessing memory on another stream: ensure the data is ready and take ownership. # (the default streams are specific to a context, so also check that.) if managed.stream != stream || managed.stream_ctx != state.context - if !(!external && can_handoff(managed, stream) && handoff!(managed, stream)) + on_device = can_handoff(managed, stream) + if on_device && is_ordered(managed, stream) + # `stream` already waits for the last use on the device. an external consumer + # synchronizes `stream` below, which only waits for what `stream` is ordered after. + managed.waiting = managed.dirty && !is_completed(managed) + elseif !(on_device && !external && handoff!(managed, stream)) maybe_synchronize(managed) end managed.stream = stream diff --git a/NEWS.md b/NEWS.md index 9a9f25d2fd..fb22f8a540 100644 --- a/NEWS.md +++ b/NEWS.md @@ -98,6 +98,10 @@ are listed as subsections of the minor release they belong to. on the task that submitted it. Work submitted through a pointer taken outside of a CUDA.jl operation (e.g. with `pointer(a)`) is only waited for if it was submitted from that task's stream before the array is used on another stream. +- A task whose stream waited for an event with `CUDA.wait(event)` uses arrays + last used before that event without waiting for other work that was queued + after it. This lets tasks spawned with `KernelAbstractions.@spawn` overlap + with their parent. ## v6.4 (September 2026) diff --git a/test/core/cudadrv.jl b/test/core/cudadrv.jl index 39e5fedd53..4bad735429 100644 --- a/test/core/cudadrv.jl +++ b/test/core/cudadrv.jl @@ -1394,6 +1394,71 @@ let a = CUDA.zeros(Int, 1), s = stream(), other = CuStream(), e = CuEvent() end end +# memory whose last use is known to come before the work on a stream, because that stream +# waited for an event recorded after it, doesn't need to be handed off. this is the pattern +# `KernelAbstractions.@spawn` uses. +let a = CUDA.zeros(Int, 1), s = stream(), other = CuStream() + a .+= 1 + fetch(Threads.@spawn (stream!(other); a .+= 1; unsafe_use(a))) + synchronize() + + # kernels on the waiting stream don't wait for work queued after the event + a .+= 1 + e = CuEvent() + record(e) + @test gated(s) do + fetch(Threads.@spawn begin + stream!(other) + CUDA.wait(e) + a .+= 1 + synchronize(other) + !CUDA.isdone(s) + end) + end + + # library pointer conversions only wait for the waiting stream + a .+= 1 + e = CuEvent() + record(e) + @test gated(s) do + fetch(Threads.@spawn begin + stream!(other) + CUDA.wait(e) + unsafe_use(a) + !CUDA.isdone(s) + end) + end + + # memory used after the event still needs to be waited for + e = CuEvent() + record(e) + @test gated(s) do + a .+= 1 + fetch(Threads.@spawn begin + stream!(other) + CUDA.wait(e) + a .+= 1 + open_gate_during(() -> synchronize(other)) + end) + end + + # and so is work submitted through a pointer that was taken before the event + p = pointer(a) + e = CuEvent() + record(e) + @test gated(s) do + @cuda stream=s store_kernel(raw_array(p)) + fetch(Threads.@spawn begin + stream!(other) + CUDA.wait(e) + a .+= 1 + open_gate_during(() -> synchronize(other)) + end) + end + synchronize(other) + @test Array(a) == [43] +end + end ############################################################################################