diff --git a/CUDACore/lib/cudadrv/events.jl b/CUDACore/lib/cudadrv/events.jl index 634c1d6373..bfb6f0e075 100644 --- a/CUDACore/lib/cudadrv/events.jl +++ b/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) 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..04fbcfc930 100644 --- a/CUDACore/lib/cudadrv/stream.jl +++ b/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 """ 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..df05f586de 100644 --- a/CUDACore/src/compiler/execution.jl +++ b/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 diff --git a/CUDACore/src/graph.jl b/CUDACore/src/graph.jl index b1c67d9dda..24ca3eb6d0 100644 --- a/CUDACore/src/graph.jl +++ b/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 diff --git a/CUDACore/src/resources/managed.jl b/CUDACore/src/resources/managed.jl index 0bf32ed58b..25c17bc153 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 @@ -100,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) @@ -150,24 +173,42 @@ 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) managed.attachment = HOST_ATTACHED managed.dirty = false + managed.waiting = false check_exceptions() else maybe_synchronize(managed) @@ -175,12 +216,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 +268,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 +304,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 +338,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 ############################################################################################ 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