From 23aab4ffe2662b49b723f866b49caa9eef9f70c8 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 6 Oct 2026 11:49:24 +0200 Subject: [PATCH 1/3] Lock arrays until copies and fill! have been submitted Kernel launches keep the memory they use locked until the launch has been submitted, so that another task can't see the new owner of the memory and wait for its stream before the work it needs to wait for is on there. Copies and fill! took ownership when converting their arrays to pointers, but released the lock again before submitting the operation. Keep their arrays locked until the operation has been submitted too. --- CUDACore/src/CUDAKernels.jl | 2 +- CUDACore/src/array.jl | 58 +++++++++++++++++++++------ CUDACore/src/compiler/execution.jl | 63 ++++++++++-------------------- 3 files changed, 68 insertions(+), 55 deletions(-) diff --git a/CUDACore/src/CUDAKernels.jl b/CUDACore/src/CUDAKernels.jl index 4aec563fa1..4527f31d76 100644 --- a/CUDACore/src/CUDAKernels.jl +++ b/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) diff --git a/CUDACore/src/array.jl b/CUDACore/src/array.jl index 7714a80e4a..72e9da8912 100644 --- a/CUDACore/src/array.jl +++ b/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 diff --git a/CUDACore/src/compiler/execution.jl b/CUDACore/src/compiler/execution.jl index c88cb54a74..f3b7bcef6a 100644 --- a/CUDACore/src/compiler/execution.jl +++ b/CUDACore/src/compiler/execution.jl @@ -411,64 +411,41 @@ 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. +""" 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() 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() finally - unlock_managed(locked) - 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 - try - for memory in ordered - take_ownership!(memory; state, stream, capturing) + for memory in Iterators.reverse(ordered) + unlock(memory.lock) + end end - return f() - finally - unlock_managed(ordered) end end From 1b95c001790e5af263fd24eac5b5e30c8c17c67f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 6 Oct 2026 11:50:59 +0200 Subject: [PATCH 2/3] Hand off memory between streams on the device When an array is used on another stream than the one that last used it, e.g., because it was passed to another task, CUDA.jl synchronizes the previous stream from the host. That makes it safe to share arrays between tasks, but the task that takes over the array blocks until all work on the other stream has finished, including work that has nothing to do with the array. For operations that CUDA.jl submits itself (kernel launches, graph launches, copies and fills), make the new stream wait for the previous one on the device instead, using an event that's cached per stream. The operations still execute in order, but the host isn't blocked. Pointers to memory that are taken outside of such an operation, e.g., by a library or MPI, may be used from the host or other streams, so these conversions still synchronize on the host. That also applies when the memory was already handed off to the current stream: the memory remembers that its stream waits for others, and a pointer conversion then synchronizes that stream. Only device memory moving between ordinary streams in the same context is handed off on the device. Unified and host memory, memory used by captures that CUDA.jl doesn't know about, and the special streams keep using host synchronization. So does a stream that is being captured on explicitly, as recording an event on it would only add a node to the graph; `capture` marks the stream under the lock that hand-offs hold while recording, before it begins capturing, so that no hand-off can record an event on it after that. Captures don't take ownership of the memory they use, launching the graph does, so a graph launch is ordered after earlier uses of its memory like any other operation. --- CUDACore/lib/cudadrv/events.jl | 20 ++++++++ CUDACore/lib/cudadrv/state.jl | 10 ++-- CUDACore/lib/cudadrv/stream.jl | 27 ++++++++-- CUDACore/src/compiler/execution.jl | 22 ++++++-- CUDACore/src/graph.jl | 5 ++ CUDACore/src/resources/managed.jl | 81 ++++++++++++++++++++++++++---- NEWS.md | 13 +++++ docs/src/usage/multitasking.md | 29 +++++++---- test/core/cudadrv.jl | 78 ++++++++++++++++++++++++++++ 9 files changed, 256 insertions(+), 29 deletions(-) diff --git a/CUDACore/lib/cudadrv/events.jl b/CUDACore/lib/cudadrv/events.jl index 634c1d6373..17da44913f 100644 --- a/CUDACore/lib/cudadrv/events.jl +++ b/CUDACore/lib/cudadrv/events.jl @@ -79,6 +79,26 @@ 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, 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 + (order.capturing || 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) + cuStreamWaitEvent(stream, event::CuEvent, 0) + end + return true +end + """ elapsed(start::CuEvent, stop::CuEvent) diff --git a/CUDACore/lib/cudadrv/state.jl b/CUDACore/lib/cudadrv/state.jl index 542adc7a80..8dbdf60535 100644 --- a/CUDACore/lib/cudadrv/state.jl +++ b/CUDACore/lib/cudadrv/state.jl @@ -47,6 +47,11 @@ 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[], @@ -54,7 +59,7 @@ mutable struct TaskLocalState 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 diff --git a/CUDACore/lib/cudadrv/stream.jl b/CUDACore/lib/cudadrv/stream.jl index 0c9732d61d..3e054e78d4 100644 --- a/CUDACore/lib/cudadrv/stream.jl +++ b/CUDACore/lib/cudadrv/stream.jl @@ -3,6 +3,22 @@ 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 + + # whether `capture` is capturing on the stream. it's set while holding the lock before + # the capture begins, so that no event is recorded on the stream after that. + capturing::Bool + + const lock::ReentrantLock + + StreamOrder() = new(nothing, false, ReentrantLock()) +end + """ CuStream(; flags=STREAM_DEFAULT, priority=nothing) @@ -22,6 +38,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 +54,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 """ diff --git a/CUDACore/src/compiler/execution.jl b/CUDACore/src/compiler/execution.jl index f3b7bcef6a..df05f586de 100644 --- a/CUDACore/src/compiler/execution.jl +++ b/CUDACore/src/compiler/execution.jl @@ -416,7 +416,9 @@ end 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. +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} @@ -431,7 +433,8 @@ function with_ordered_managed(f::F, ordered::Union{AbstractVector{<: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) for memory in ordered lock(memory.lock) @@ -440,7 +443,7 @@ function with_ordered_managed(f::F, ordered::Union{AbstractVector{<:Managed}, for memory in ordered take_ownership!(memory; state, stream, capturing) end - return f() + return with_operation_stream(f, tls, stream, capturing) finally for memory in Iterators.reverse(ordered) unlock(memory.lock) @@ -449,6 +452,19 @@ function with_ordered_managed(f::F, ordered::Union{AbstractVector{<:Managed}, end 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 + return f() + finally + tls.operation_stream, tls.operation_capturing = old_stream, old_capturing + end +end + function managed_kernel_launch(backend, kernel, arguments, managed; kwargs...) # tasks that are part of a capture submit while holding its lock (see `CaptureScope`) current_capture_scope() === nothing || diff --git a/CUDACore/src/graph.jl b/CUDACore/src/graph.jl index b1c67d9dda..52645423ca 100644 --- a/CUDACore/src/graph.jl +++ b/CUDACore/src/graph.jl @@ -398,9 +398,14 @@ function capture_stream(f, graph::Union{Nothing,CuGraph}, deps::Vector{CUgraphNo end else reserve_capture_stream(stream, ctx) + order = stream.order::StreamOrder try + # other tasks wait for the stream by recording an event on it (see + # `stream_wait`), which they can't do anymore once the capture has begun + @lock order.lock order.capturing = true capture_on(f, stream, ctx, graph, deps; mode, throw_error) finally + @lock order.lock order.capturing = false @lock capture_streams_lock delete!(busy_capture_streams, stream) end end diff --git a/CUDACore/src/resources/managed.jl b/CUDACore/src/resources/managed.jl index 0bf32ed58b..b5e8ed036e 100644 --- a/CUDACore/src/resources/managed.jl +++ b/CUDACore/src/resources/managed.jl @@ -47,6 +47,10 @@ mutable struct Managed{M} # (only for captures that CUDA.jl does not know about) captured::Bool + # whether `stream` waits on the device for operations on other streams that used the + # memory, which therefore haven't been synchronized either (implies `dirty`) + waiting::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 @@ -66,7 +70,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, owned_allocation, attachment, 0, nothing) + synchronizing, dirty, captured, false, owned_allocation, attachment, + 0, nothing) end end @@ -91,6 +96,7 @@ function synchronize(managed::Managed) end end managed.dirty = false + managed.waiting = false end end @@ -168,6 +174,7 @@ function prepare_host_access!(managed::Managed) synchronize(event) managed.attachment = HOST_ATTACHED managed.dirty = false + managed.waiting = false check_exceptions() else maybe_synchronize(managed) @@ -175,12 +182,50 @@ function prepare_host_access!(managed::Managed) return end -# Transfer stream ownership of an allocation and mark it dirty in anticipation of a -# device-side operation. The caller must hold `managed.lock` until that operation has been -# submitted to `stream`, so the recorded owner cannot become visible before its submission. +# Memory that moves to another stream needs to be ordered after the operations that used it +# before. For an operation that CUDA.jl submits to that stream itself, it suffices to make +# the stream wait for the previous one on the device. Other consumers, like a library that +# is passed a pointer, may access the memory from the host or from streams we don't know +# of, so for them the previous stream is synchronized from the host instead. Either way, +# this relies on the previous operations having been submitted already. + +# whether memory can move from its stream to `stream` by waiting on the device. that needs +# two ordinary streams in the same context. memory used during unknown captures may also be +# used by graph launches we don't know of, and other kinds of memory can be accessed from the +# CPU in ways that we don't see. (known captures don't take ownership, see `take_ownership!`) +function can_handoff(managed::Managed{M}, stream::CuStream) where {M} + source = managed.stream + M == DeviceMemory && managed.synchronizing && !managed.captured && + source.order !== nothing && stream.order !== nothing && source.ctx == stream.ctx +end + +# make `stream` wait on the device for the operations on the memory's stream, returning +# whether that was possible (see `stream_wait`) +function handoff!(managed::Managed, stream::CuStream) + managed.dirty || return true + source = managed.stream + # (holding the lock that recycling streams takes, see `pending_work`) + @lock stream_disposal_lock begin + if recycled(managed) + # the stream was handed to another task, which only happens when it was idle + managed.dirty = false + managed.waiting = false + return true + end + isvalid(source) && stream_wait(stream, source) || return false + end + managed.waiting = true + return true +end + +# Transfer stream ownership of an allocation and mark it dirty in anticipation of an +# operation on it. The caller must hold `managed.lock` until that operation has been +# submitted, so the recorded owner cannot become visible before its submission. Pass +# `external=true` if the operation isn't submitted to `stream` by CUDA.jl (see above). function take_ownership!(managed::Managed{M}; state=active_state(), stream::CuStream=state.stream, - capturing::Bool=is_capturing(stream)) where {M} + capturing::Bool=is_capturing(stream), + external::Bool=false) where {M} sizeof(managed) == 0 && return managed M == UnifiedMemory && attach_globally!(managed, stream, capturing) @@ -189,7 +234,8 @@ function take_ownership!(managed::Managed{M}; state=active_state(), capture = current_capture(stream) if capture !== nothing # captured operations don't execute until the graph is launched, so only record the - # use of the memory. the graph will take ownership of it when it is launched. + # use of the memory. the graph will take ownership of it when it is launched, which + # orders the launch after earlier uses of the memory. record!(capture, managed) return managed end @@ -224,12 +270,21 @@ 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 - maybe_synchronize(managed) + if !(!external && can_handoff(managed, stream) && handoff!(managed, stream)) + maybe_synchronize(managed) + end managed.stream = stream managed.stream_ctx = state.context end managed.generation = generation(managed.stream) + # an external consumer needs operations on other streams to have finished, also when the + # stream already waits for them. operations on the stream itself are, as before, assumed + # to be ordered by the consumer. (during an unknown capture, nothing is waited for) + if external && managed.waiting && !capturing + synchronize(managed) + end + # prefetch unified memory as we're likely to use it on the GPU if M == UnifiedMemory can_prefetch = !capturing @@ -249,8 +304,16 @@ function Base.convert(::Type{CuPtr{T}}, managed::Managed{M}) where {T,M} ptr = convert(CuPtr{T}, managed.mem) ptr == CU_NULL && return ptr - state = active_state() - take_ownership!(managed; state, stream=state.stream) + # 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. + tls = task_local_state!() + state = active_state(tls) + stream = tls.operation_stream + if stream === nothing + take_ownership!(managed; state, stream=state.stream, external=true) + else + take_ownership!(managed; state, stream, capturing=tls.operation_capturing) + end return ptr end end diff --git a/NEWS.md b/NEWS.md index ccc2baf5a5..7e2bcac986 100644 --- a/NEWS.md +++ b/NEWS.md @@ -34,6 +34,12 @@ are listed as subsections of the minor release they belong to. each hold on to device memory and slow down memory allocation. Code that uses a task's stream after the task has finished should create its own stream with `CuStream()` instead. +- Pointers that are taken from GPU memory inside the function passed to + `CUDA.with_managed` are assumed to be used by an operation on the stream + passed to `with_managed`, and no longer wait on the CPU for other streams + that used the memory before. Code that passes such pointers to libraries, + other streams or the CPU should take them outside of `with_managed` + ([#3329](https://github.com/JuliaGPU/CUDA.jl/pull/3329)). *New features*: @@ -56,6 +62,13 @@ are listed as subsections of the minor release they belong to. perform them, instead of being executed right away. These tasks need to finish before the capture ends. Changing the stream of a task that takes part in a capture, using `stream!`, is not supported and throws an error. +- Handing an array to another task no longer blocks the CPU until the stream + that last used the array has finished all of its work. For the operations + that CUDA.jl submits itself (kernel launches, including broadcasts and + KernelAbstractions kernels, copies, `fill!` and graph launches), the stream + of the new task now waits for the previous one on the GPU. Library calls and + pointers that are passed to other code still wait on the CPU + ([#3329](https://github.com/JuliaGPU/CUDA.jl/pull/3329)). *Bug fixes*: diff --git a/docs/src/usage/multitasking.md b/docs/src/usage/multitasking.md index d70ad0fdb5..d2bcfa0f30 100644 --- a/docs/src/usage/multitasking.md +++ b/docs/src/usage/multitasking.md @@ -51,18 +51,26 @@ GPU arrays filled with random numbers: function main(N=1024) a = CUDA.rand(N,N) b = CUDA.rand(N,N) - - # make sure this data can be used by other tasks! - synchronize() - run(a, b) end ``` -The `main` function illustrates how we need to take care when sharing data between tasks: -GPU operations typically execute asynchronously, queued on an execution stream, so if we -switch tasks and thus switch execution streams we need to `synchronize()` to ensure the data -is actually available. +The arrays are initialized on the stream of the main task, and used from the streams of the +two tasks that `run` spawns. GPU operations execute asynchronously, so that initialization +may not have finished when the tasks start. CUDA.jl takes care of that: it remembers which +stream last used each array, and orders work on the array from another stream after all the +work that was queued on that stream before: + +- for operations that CUDA.jl submits itself, i.e., kernel launches (including broadcasts + and KernelAbstractions kernels), copies, `fill!` and graph launches, the new stream waits + for the previous one on the GPU, without blocking the task; +- for other uses of the array's memory, like library calls (the matrix multiplication in + `compute` uses CUBLAS) or passing a pointer to C code or MPI, the task first waits until + the previous stream has finished its work, because CUDA.jl doesn't know which stream, or + whether the CPU, will access the memory. + +Unified and host memory, and memory that is used on another device, are always waited for +by the task. Using Nsight Systems, we can visualize the execution of this example: @@ -125,8 +133,9 @@ high-priority stream; `CUDA.priority!(:normal)` selects the normal priority agai The scoped form, `CUDA.priority!(:high) do ... end`, restores the previous stream when the block ends. CUDA treats stream priority as a scheduling hint for pending kernels, and existing work is not interrupted. Repeated priority changes reuse the task's streams. -Using an array from before the switch may wait on the CPU for pending work on its -previous stream; switch before queuing long work when that wait matters. +Using an array from before the switch then waits for pending work on its previous +stream, as described above; for library calls, that blocks the task, so switch before +queuing long work when that wait matters. 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. diff --git a/test/core/cudadrv.jl b/test/core/cudadrv.jl index 2dc2b5b489..d950665d2a 100644 --- a/test/core/cudadrv.jl +++ b/test/core/cudadrv.jl @@ -1188,6 +1188,84 @@ let long = CuStream() end end +# memory used by CUDA.jl operations on another stream is handed off on the device: the new +# stream waits for the old one, without blocking the host +let a = CUDA.zeros(Int, 1), b = CUDA.zeros(Int, 1), u = cu([0]; unified=true), + s = stream(), other = CuStream(), third = CuStream() + # warm up, as compiling kernels waits for the GPU to become idle + for x in (a, u) + x .+= 1 + fetch(Threads.@spawn (stream!(other); x .*= 2; copyto!(b, a))) + fetch(Threads.@spawn (stream!(third); x .-= 1)) + synchronize(third) + x .= 1 + end + + # kernels and copies don't block, and still execute in order + @test gated(s) do + a .+= 1 + fetch(Threads.@spawn begin + stream!(other) + a .*= 2 + fetch(Threads.@spawn (stream!(third); a .-= 1; copyto!(b, a))) + !CUDA.isdone(s) + end) + end + synchronize(third) + @test Array(a) == [3] + @test Array(b) == [3] + + # synchronizing the memory also waits for the stream it was handed off from + @test gated(s) do + a .+= 1 + fetch(Threads.@spawn (stream!(other); a .*= 2)) + open_gate_during(() -> synchronize(a)) + end + @test Array(a) == [8] + + # other consumers of a pointer, like a library, get the memory ready on the host, also + # when the stream that hands it to them already waits for the previous use + for hops in (1, 2) + @test gated(s) do + a .+= 1 + fetch(Threads.@spawn begin + stream!(other) + a .*= 2 + if hops == 2 + fetch(Threads.@spawn (stream!(third); a .-= 1)) + end + open_gate_during(() -> Base.unsafe_convert(CuPtr{Int}, a)) + end) + end + end + @test Array(a) == [37] + + # capturing doesn't wait for memory that was handed off, but launching the graph takes + # ownership of it, so a launch on yet another stream waits for it on the device + @test gated(s) do + a .+= 1 + fetch(Threads.@spawn begin + stream!(other) + a .*= 2 + graph = capture(() -> (a .-= 1)) + launch(instantiate(graph), third) + !CUDA.isdone(s) && !CUDA.isdone(third) + end) + end + synchronize(third) + @test Array(a) == [75] + + # memory that isn't device memory is still synchronized on the host + @test gated(s) do + u .+= 1 + open_gate_during() do + fetch(Threads.@spawn (stream!(other); u .*= 2)) + end + end + synchronize(other) + @test Array(u) == [4] +end + end ############################################################################################ From 00a2218a406e75ff7ffa48cbe05234a55629bd64 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 6 Oct 2026 14:46:09 +0200 Subject: [PATCH 3/3] Let other tasks use memory last used on a stream that is being captured When capturing on a specific stream, memory may have been used on that stream before the capture. Another task that wants to use that memory needs to wait for the work that was submitted to the stream, but can't do so by recording an event on it: during the capture, that would only add a node to the graph. So such uses failed with a CaptureError until the capture had ended. Record an event on the stream right before beginning the capture, and have other tasks wait for that event instead, when handing memory off on the device, when synchronizing it from the host, and when attaching unified memory to the host before CPU access (which then happens on another stream). Captured operations don't take ownership of memory, so the event covers the last use of all memory owned by the stream. It is published under the lock that is held while recording events on the stream, so that no event is recorded on it after the capture began. Each capture records a new event, as tasks may still be waiting for the one of an earlier capture. The capture itself keeps failing to synchronize such memory, as it may have captured operations on it that the event doesn't cover. --- CUDACore/lib/cudadrv/events.jl | 18 ++++---- CUDACore/lib/cudadrv/stream.jl | 11 +++-- CUDACore/src/graph.jl | 18 +++++--- CUDACore/src/resources/managed.jl | 68 +++++++++++++++++++++++-------- test/core/graph.jl | 21 ++++++++++ 5 files changed, 102 insertions(+), 34 deletions(-) diff --git a/CUDACore/lib/cudadrv/events.jl b/CUDACore/lib/cudadrv/events.jl index 17da44913f..bfb6f0e075 100644 --- a/CUDACore/lib/cudadrv/events.jl +++ b/CUDACore/lib/cudadrv/events.jl @@ -81,19 +81,23 @@ wait(e::CuEvent, stream::CuStream=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, as recording an event on it would then only add a node to the graph. +# 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 - (order.capturing || in_capture(source)) && return false - event = order.event + event = order.capture_event if event === nothing - event = context!(source.ctx) do - CuEvent(EVENT_DISABLE_TIMING) + 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 - order.event = event + cuEventRecord(event::CuEvent, source) end - cuEventRecord(event::CuEvent, source) cuStreamWaitEvent(stream, event::CuEvent, 0) end return true diff --git a/CUDACore/lib/cudadrv/stream.jl b/CUDACore/lib/cudadrv/stream.jl index 3e054e78d4..04fbcfc930 100644 --- a/CUDACore/lib/cudadrv/stream.jl +++ b/CUDACore/lib/cudadrv/stream.jl @@ -10,13 +10,16 @@ mutable struct StreamOrder # which is defined after streams.) event::Any - # whether `capture` is capturing on the stream. it's set while holding the lock before - # the capture begins, so that no event is recorded on the stream after that. - capturing::Bool + # 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, false, ReentrantLock()) + StreamOrder() = new(nothing, nothing, ReentrantLock()) end """ diff --git a/CUDACore/src/graph.jl b/CUDACore/src/graph.jl index 52645423ca..24ca3eb6d0 100644 --- a/CUDACore/src/graph.jl +++ b/CUDACore/src/graph.jl @@ -400,12 +400,18 @@ function capture_stream(f, graph::Union{Nothing,CuGraph}, deps::Vector{CUgraphNo reserve_capture_stream(stream, ctx) order = stream.order::StreamOrder try - # other tasks wait for the stream by recording an event on it (see - # `stream_wait`), which they can't do anymore once the capture has begun - @lock order.lock order.capturing = true + # 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.capturing = false + @lock order.lock order.capture_event = nothing @lock capture_streams_lock delete!(busy_capture_streams, stream) end end @@ -604,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 diff --git a/CUDACore/src/resources/managed.jl b/CUDACore/src/resources/managed.jl index b5e8ed036e..25c17bc153 100644 --- a/CUDACore/src/resources/managed.jl +++ b/CUDACore/src/resources/managed.jl @@ -106,13 +106,30 @@ end function pending_work(managed::Managed) @lock stream_disposal_lock begin recycled(managed) && return nothing - check_capture(managed.stream) - relaxed_capture_mode(() -> isdone(managed.stream)) && return nothing - event = CuEvent(EVENT_DISABLE_TIMING) - cuEventRecord(event, managed.stream) - return event + stream = managed.stream + order = stream.order + order === nothing && return record_pending_work(stream) + # a stream that `capture` is capturing on has an event that covers the work submitted + # to it before, including the last use of the memory (captured operations don't take + # ownership of memory). that's not enough for the capture itself, which may have + # captured operations on the memory. the lock keeps a capture from beginning meanwhile. + @lock order.lock begin + event = order.capture_event + if event === nothing || CUDACore.stream() == stream + record_pending_work(stream) + else + event::CuEvent + end + end end end +function record_pending_work(stream::CuStream) + check_capture(stream) + relaxed_capture_mode(() -> isdone(stream)) && return nothing + event = CuEvent(EVENT_DISABLE_TIMING) + cuEventRecord(event, stream) + return event +end function maybe_synchronize(managed::Managed) Base.@lock managed.lock begin if managed.synchronizing && (managed.dirty || managed.captured) @@ -156,19 +173,36 @@ function prepare_host_access!(managed::Managed) if managed.attachment == GLOBALLY_ATTACHED && managed.synchronizing && !managed.captured && managed.stream.ctx !== nothing event = @lock stream_disposal_lock begin - if !recycled(managed) && isvalid(managed.stream) - check_capture(managed.stream) - end - # this also handles owners that have been recycled or destroyed - stream, ctx = release_stream(managed.stream, managed.stream_ctx, managed.generation) - context!(ctx) do - # (attaching is prohibited while another thread captures in global mode) - relaxed_capture_mode() do - cuStreamAttachMemAsync(stream, managed.mem, 0, MEM_ATTACH_HOST) + # (holding the lock under which a capture on the stream begins, see `pending_work`) + order = managed.stream.order + order === nothing || lock(order.lock) + try + fence = order === nothing || recycled(managed) ? nothing : order.capture_event + if fence !== nothing && CUDACore.stream() != managed.stream + # the stream is being captured by another task, so attach on another stream, + # after the work that was submitted before the capture began + ctx = something(managed.stream.ctx) + stream = disposal_stream(ctx) + else + fence = nothing + if !recycled(managed) && isvalid(managed.stream) + check_capture(managed.stream) + end + # this also handles owners that have been recycled or destroyed + stream, ctx = release_stream(managed.stream, managed.stream_ctx, managed.generation) + end + context!(ctx) do + fence === nothing || cuStreamWaitEvent(stream, fence::CuEvent, 0) + # (attaching is prohibited while another thread captures in global mode) + relaxed_capture_mode() do + cuStreamAttachMemAsync(stream, managed.mem, 0, MEM_ATTACH_HOST) + end + event = CuEvent(EVENT_DISABLE_TIMING) + cuEventRecord(event, stream) + event end - event = CuEvent(EVENT_DISABLE_TIMING) - cuEventRecord(event, stream) - event + finally + order === nothing || unlock(order.lock) end end synchronize(event) diff --git a/test/core/graph.jl b/test/core/graph.jl index 4dc469415c..408cc3a012 100644 --- a/test/core/graph.jl +++ b/test/core/graph.jl @@ -370,6 +370,27 @@ end graph = capture(() -> a .+= 1; stream=s) instantiate(graph)() @test Array(a) == fill(4f0, 4) + + # other tasks can use memory that was last used on the stream before capturing on it + for (use, expected) in ((x -> Array(x), 1f0), (x -> Array(x .+ 1), 2f0)) + x = CUDA.zeros(Float32, 1 << 16) + stream!(s) do + x .+= 1 + end + go = Base.Event() + other = @async (wait(go); use(x)) + capture(; stream=s) do + notify(go) + wait(other) + end + @test fetch(other) == fill(expected, 1 << 16) + + # but the capture itself can't wait for it, as it may have captured uses of it + @test_throws CaptureError capture(; stream=s) do + x .+= 1 + Array(x) + end + end end @testset "using memory from a capturing task" begin