Hand off arrays between streams on the GPU instead of the host

CUDACore/lib/cudadrv/events.jl

@@ -79,6 +79,30 @@ Make a stream wait on a event. This only makes the stream wait, and not the host
 wait(e::CuEvent, stream::CuStream=stream()) =
     capture_submission(() -> cuStreamWaitEvent(stream, e, 0), stream)
 
+# 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
+# captured without an event that covers the work submitted to it before, as recording an
+# event on it would then only add a node to the graph.
+function stream_wait(stream::CuStream, source::CuStream)
+    order = source.order::StreamOrder
+    Base.@lock order.lock begin
+        event = order.capture_event
+        if event === nothing
+            in_capture(source) && return false
+            event = order.event
+            if event === nothing
+                event = context!(source.ctx) do
+                    CuEvent(EVENT_DISABLE_TIMING)
+                end
+                order.event = event
+            end
+            cuEventRecord(event::CuEvent, source)
+        end
+        cuStreamWaitEvent(stream, event::CuEvent, 0)
+    end
+    return true
+end
+
 """
     elapsed(start::CuEvent, stop::CuEvent)
 

CUDACore/lib/cudadrv/state.jl

@@ -47,14 +47,19 @@ mutable struct TaskLocalState
     math_mode::MathMode
     math_precision::Symbol
 
+    # the stream of the operation CUDA.jl is submitting, if any, and whether it's being
+    # captured (see `with_managed`)
+    operation_stream::Union{Nothing,CuStream}
+    operation_capturing::Bool
+
     function TaskLocalState(dev::CuDevice=something(default_device[], CuDevice(0)),
                             ctx::CuContext = context(dev))
         math_mode = something(default_math_mode[],
                               Base.JLOptions().fast_math==1 ? FAST_MATH : DEFAULT_MATH)
         math_precision = something(default_math_precision[], :TensorFloat32)
         new(dev, ctx, Union{Nothing,CuStream}[nothing for _ in 1:ndevices()],
             Cint[0 for _ in 1:ndevices()],
-            math_mode, math_precision)
+            math_mode, math_precision, nothing, false)
     end
 end
 
@@ -103,9 +108,8 @@ end
 
 # convenience function to get all relevant state
 # without querying task local storage multiple times
-@inline function active_state()
+@inline function active_state(state::TaskLocalState=task_local_state!())
     # inline to remove unused state properties
-    state = task_local_state!()
     return (device=state.device, context=state.context, stream=stream(state),
             math_mode=state.math_mode, math_precision=state.math_precision)
 end

CUDACore/lib/cudadrv/stream.jl

@@ -3,6 +3,25 @@
 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`).
+mutable struct StreamOrder
+    # 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.)
+    event::Any
+
+    # while `capture` is capturing on the stream, recording an event on it would only add a
+    # node to the graph. instead, other streams wait for an event that was recorded right
+    # before the capture began, which covers the work that was submitted to the stream
+    # before. it's set while holding the lock, so that no event is recorded on the stream
+    # after the capture began (see `capture(; stream)`). also a `CuEvent`, or `nothing`.
+    capture_event::Any
+
+    const lock::ReentrantLock
+
+    StreamOrder() = new(nothing, nothing, ReentrantLock())
+end
+
 """
     CuStream(; flags=STREAM_DEFAULT, priority=nothing)
 
@@ -22,6 +41,9 @@ mutable struct CuStream
     # happens when it is idle, so work submitted during earlier generations has finished.
     Base.@atomic generation::Int
 
+    # the special streams don't have one, as they implicitly synchronize with other streams
+    const order::Union{Nothing,StreamOrder}
+
     function CuStream(; flags::CUstream_flags=STREAM_DEFAULT,
                         priority::Union{Nothing,Integer}=nothing)
         handle_ref = Ref{CUstream}()
@@ -35,18 +57,18 @@ mutable struct CuStream
         ctx = current_context()
         # make sure memory last used on this stream can be released after destroying it
         disposal_stream(ctx)
-        obj = new(handle_ref[], true, ctx, nothing, 0)
+        obj = new(handle_ref[], true, ctx, nothing, 0, StreamOrder())
         # destroying a stream from a finalizer is deferred, as with all resources
         # (see `retire!`), also so that memory that was last used on it can be freed first
         resource_finalizer(obj)
         return obj
     end
 
-    global default_stream() = new(convert(CUstream, C_NULL), true, nothing, nothing, 0)
+    global default_stream() = new(convert(CUstream, C_NULL), true, nothing, nothing, 0, nothing)
 
-    global legacy_stream() = new(convert(CUstream, 1), true, nothing, nothing, 0)
+    global legacy_stream() = new(convert(CUstream, 1), true, nothing, nothing, 0, nothing)
 
-    global per_thread_stream() = new(convert(CUstream, 2), true, nothing, nothing, 0)
+    global per_thread_stream() = new(convert(CUstream, 2), true, nothing, nothing, 0, nothing)
 end
 
 """

CUDACore/src/CUDAKernels.jl

@@ -49,7 +49,7 @@ Adapt.adapt_storage(::KA.CPU, a::Union{CuArray,GPUArrays.AbstractGPUSparseArray}
 ## memory operations
 
 function KA.copyto!(::CUDABackend, A, B)
-    GC.@preserve A B begin
+    GC.@preserve A B CUDACore.with_managed_arrays(A, B) do
         destptr = pointer(A)
         srcptr  = pointer(B)
         N       = length(A)

CUDACore/src/array.jl

@@ -524,6 +524,11 @@ Base.unsafe_convert(typ::Type{CuPtr{T}}, x::CuArray{T}) where {T} =
   return data
 end
 
+# the address of an array's memory, for use by an operation that already took ownership of
+# it (see `with_managed_arrays`)
+unmanaged_pointer(x::CuArray{T}, i::Integer=1) where {T} =
+  convert(CuPtr{T}, x.data[].mem) + x.offset + Base._memory_offset(x, i)
+
 
 ## indexing
 
@@ -575,6 +580,31 @@ end
 
 ## memory copying
 
+# submit an operation on arrays, or views of them, to the current stream (see
+# `with_managed`). host arrays are ignored.
+function with_managed_arrays(f, a::AbstractArray)
+  ma = managed_memory(a)
+  ma === nothing ? f() : with_ordered_managed(f, (ma,))
+end
+function with_managed_arrays(f, a::AbstractArray, b::AbstractArray)
+  ma, mb = managed_memory(a), managed_memory(b)
+  if mb === nothing || mb === ma
+    ma === nothing ? f() : with_ordered_managed(f, (ma,))
+  elseif ma === nothing
+    with_ordered_managed(f, (mb,))
+  elseif objectid(ma.lock) < objectid(mb.lock)  # see `locking_order`
+    with_ordered_managed(f, (ma, mb))
+  else
+    with_ordered_managed(f, (mb, ma))
+  end
+end
+
+managed_memory(a::CuArray) = check_capture(a.data)[]
+function managed_memory(a::AbstractArray)
+  p = parent(a)
+  p === a ? nothing : managed_memory(p)
+end
+
 if VERSION >= v"1.11.0-DEV.753"
 function typetagdata(a::Array, i=1)
   ptr_or_offset = Int(a.ref.ptr_or_offset)
@@ -653,12 +683,12 @@ function Base.unsafe_copyto!(dest::DenseCuArray{T}, doffs,
       # synchronization here, but the exact cases are hard to know and detect (e.g.,
       # unpinned memory normally blocks, but not for all sizes, and not on all memory
       # architectures).
-      GC.@preserve src dest begin
+      GC.@preserve src dest with_managed_arrays(dest) do
         # semantically, it is not safe for this operation to execute asynchronously, because
         # the Array may be collected before the copy starts executing. However, when using
         # unpinned memory, CUDA first stages a copy to a pinned buffer that will outlive
         # the source array, making this operation safe.
-        unsafe_copyto!(pointer(dest, doffs), pointer(src, soffs), n; async=true)
+        unsafe_copyto!(unmanaged_pointer(dest, doffs), pointer(src, soffs), n; async=true)
         if Base.isbitsunion(T)
           unsafe_copyto!(typetagdata(dest, doffs), typetagdata(src, soffs), n; async=true)
         end
@@ -694,8 +724,9 @@ function Base.unsafe_copyto!(dest::DenseCuArray{T}, doffs,
     # use direct device-to-device copy
     context!(context(src)) do
       capture_submission() do
-        GC.@preserve src dest begin
-          unsafe_copyto!(pointer(dest, doffs), pointer(src, soffs), n; async=true)
+        GC.@preserve src dest with_managed_arrays(dest, src) do
+          unsafe_copyto!(unmanaged_pointer(dest, doffs), unmanaged_pointer(src, soffs), n;
+                         async=true)
           if Base.isbitsunion(T)
             unsafe_copyto!(typetagdata(dest, doffs), typetagdata(src, soffs), n; async=true)
           end
@@ -756,8 +787,9 @@ end
 function Base.unsafe_copyto!(dest::DenseCuArray{T,<:Any,<:Union{UnifiedMemory,HostMemory}}, doffs,
                              src::DenseCuArray{T}, soffs, n) where T
   context!(context(src)) do
-    GC.@preserve src dest begin
-      unsafe_copyto!(pointer(dest, doffs), pointer(src, soffs), n; async=true)
+    GC.@preserve src dest with_managed_arrays(dest, src) do
+      unsafe_copyto!(unmanaged_pointer(dest, doffs), unmanaged_pointer(src, soffs), n;
+                     async=true)
       if Base.isbitsunion(T)
         unsafe_copyto!(typetagdata(dest, doffs), typetagdata(src, soffs), n; async=true)
       end
@@ -769,8 +801,9 @@ end
 function Base.unsafe_copyto!(dest::DenseCuArray{T}, doffs,
                              src::DenseCuArray{T,<:Any,<:Union{UnifiedMemory,HostMemory}}, soffs, n) where T
   context!(context(dest)) do
-    GC.@preserve src dest begin
-      unsafe_copyto!(pointer(dest, doffs), pointer(src, soffs), n; async=true)
+    GC.@preserve src dest with_managed_arrays(dest, src) do
+      unsafe_copyto!(unmanaged_pointer(dest, doffs), unmanaged_pointer(src, soffs), n;
+                     async=true)
       if Base.isbitsunion(T)
         unsafe_copyto!(typetagdata(dest, doffs), typetagdata(src, soffs), n; async=true)
       end
@@ -781,8 +814,9 @@ end
 
 function Base.unsafe_copyto!(dest::DenseCuArray{T,<:Any,<:Union{UnifiedMemory,HostMemory}}, doffs,
                              src::DenseCuArray{T,<:Any,<:Union{UnifiedMemory,HostMemory}}, soffs, n) where T
-  GC.@preserve src dest begin
-    unsafe_copyto!(pointer(dest, doffs), pointer(src, soffs), n; async=true)
+  GC.@preserve src dest with_managed_arrays(dest, src) do
+    unsafe_copyto!(unmanaged_pointer(dest, doffs), unmanaged_pointer(src, soffs), n;
+                   async=true)
     if Base.isbitsunion(T)
       unsafe_copyto!(typetagdata(dest, doffs), typetagdata(src, soffs), n; async=true)
     end
@@ -909,7 +943,9 @@ function Base.fill!(A::DenseCuArray{T}, x) where T <: MemsetCompatTypes
   y = reinterpret(U, convert(T, x))
   context!(context(A)) do
     capture_submission() do
-      GC.@preserve A memset(convert(CuPtr{U}, pointer(A)), y, length(A))
+      GC.@preserve A with_managed_arrays(A) do
+        memset(convert(CuPtr{U}, unmanaged_pointer(A)), y, length(A))
+      end
     end
   end
   A

CUDACore/src/compiler/execution.jl

@@ -411,64 +411,57 @@ function locking_order(managed::AbstractVector{<:Managed})
     return ordered
 end
 
-function lock_managed(managed::AbstractVector{<:Managed})
-    locked = locking_order(managed)
-    for memory in locked
-        lock(memory.lock)
-    end
-    return locked
-end
-
-function unlock_managed(locked::AbstractVector{<:Managed})
-    for memory in Iterators.reverse(locked)
-        unlock(memory.lock)
-    end
-    return
-end
+"""
+    with_managed(f, managed::AbstractVector{<:Managed}; stream=stream())
 
+Submit an operation on `managed` memory to `stream` by calling `f`. The memory is made
+ready for use on `stream`, and locked so that it can't move to another stream before the
+operation has been submitted. Pointers to memory that are taken within `f` (e.g., by
+converting an array) are assumed to be used by an operation on `stream` too, so they don't
+wait for other streams that used the memory before.
+"""
 function with_managed(f::F, managed::AbstractVector{<:Managed};
                       stream::CuStream=stream()) where {F}
+    # (a single memory is trivially in locking order)
+    ordered = length(managed) == 1 ? managed : locking_order(managed)
+    with_ordered_managed(f, ordered; stream)
+end
+
+# like `with_managed`, but for memory that's already in locking order
+function with_ordered_managed(f::F, ordered::Union{AbstractVector{<:Managed},
+                                                   Tuple{Vararg{Managed}}};
+                              stream::CuStream=stream()) where {F}
     # (taking ownership of the memory and using it is a single submission)
     capture_submission(stream) do
-        state = active_state()
+        tls = task_local_state!()
+        state = active_state(tls)
         capturing = is_capturing(stream)
-        if length(managed) == 1
-            memory = @inbounds managed[1]
+        for memory in ordered
             lock(memory.lock)
-            try
-                take_ownership!(memory; state, stream, capturing)
-                return f()
-            finally
-                unlock(memory.lock)
-            end
         end
-        locked = lock_managed(managed)
         try
-            for memory in locked
+            for memory in ordered
                 take_ownership!(memory; state, stream, capturing)
             end
-            return f()
+            return with_operation_stream(f, tls, stream, capturing)
         finally
-            unlock_managed(locked)
+            for memory in Iterators.reverse(ordered)
+                unlock(memory.lock)
+            end
         end
     end
 end
 
-# like `with_managed`, but for memory that's already in locking order
-function with_ordered_managed(f::F, ordered::AbstractVector{<:Managed};
-                              stream::CuStream=stream()) where {F}
-    state = active_state()
-    capturing = is_capturing(stream)
-    for memory in ordered
-        lock(memory.lock)
-    end
+# 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,
+                                       capturing::Bool) where {F}
+    old_stream, old_capturing = tls.operation_stream, tls.operation_capturing
+    tls.operation_stream, tls.operation_capturing = stream, capturing
     try
-        for memory in ordered
-            take_ownership!(memory; state, stream, capturing)
-        end
         return f()
     finally
-        unlock_managed(ordered)
+        tls.operation_stream, tls.operation_capturing = old_stream, old_capturing
     end
 end
 

CUDACore/src/graph.jl

@@ -398,9 +398,20 @@ function capture_stream(f, graph::Union{Nothing,CuGraph}, deps::Vector{CUgraphNo
         end
     else
         reserve_capture_stream(stream, ctx)
+        order = stream.order::StreamOrder
         try
+            # memory may have been used on the stream before. other tasks wait for that by
+            # recording an event on the stream (see `stream_wait`), which they can't do once
+            # the capture has begun, so record one for them. (a new one, as tasks may still
+            # be waiting for the event of an earlier capture.)
+            event = CuEvent(EVENT_DISABLE_TIMING)
+            @lock order.lock begin
+                cuEventRecord(event, stream)
+                order.capture_event = event
+            end
             capture_on(f, stream, ctx, graph, deps; mode, throw_error)
         finally
+            @lock order.lock order.capture_event = nothing
             @lock capture_streams_lock delete!(busy_capture_streams, stream)
         end
     end
@@ -599,8 +610,8 @@ affected by the capture.
 
 To capture on a specific stream instead, pass it as the `stream` keyword argument. That
 needs to be a stream that was created with `flags=STREAM_NON_BLOCKING` in the current
-context, and that isn't used by anything else while capturing. Other tasks cannot wait for
-work that was submitted to that stream before the capture until the capture has ended.
+context, and that isn't used by anything else while capturing. Other tasks can still use
+memory that was last used on that stream before the capture.
 
 CUDA.jl checks the operations it performs itself, like waiting for the GPU, but by default
 doesn't ask the driver to prohibit other operations that are potentially unsafe during