Don't wait for work that an event already ordered

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
@@ -77,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

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()

CUDACore/lib/cudadrv/stream.jl

@@ -3,8 +3,31 @@
 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
+
+    # 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.)
@@ -19,7 +42,63 @@ mutable struct StreamOrder
 
     const lock::ReentrantLock
 
-    StreamOrder() = new(nothing, nothing, ReentrantLock())
+    StreamOrder() = new(1, 0, nothing, false, WeakKeyDict{StreamOrder,UInt64}(), 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
+
+# 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
 
 """
@@ -111,6 +190,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
 

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

CUDACore/src/compiler/execution.jl

@@ -445,13 +445,26 @@ 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
         end
     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,