From 707c1922a083df96d338281fd1d481ef28f4e351 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 29 Sep 2026 17:29:52 +0200 Subject: [PATCH 1/5] Keep wrapped CPU arrays alive in unsafe_wrap(CuArray, ...) `unsafe_wrap(CuArray, a::Array)` did not retain `a`, so the CuArray could outlive the memory it wraps once the Array was garbage collected, both for HMM-backed unified memory and for registered host memory (whose finalizer would then unregister freed memory). Capture the Array in the DataRef finalizer so that it lives as long as the wrapper and any views derived from it. Wrapping a raw pointer keeps its existing semantics. Releasing the wrapper now also waits for outstanding work on it before letting go of the Array or unregistering the memory, so dropping the last reference while a kernel is still using it is safe. If that wait fails, the memory is leaked rather than released while it may still be in use. Also document that the reverse, `unsafe_wrap(Array, ::CuArray)`, does not keep the CuArray alive, and that captured graphs don't keep wrapped memory alive. --- CUDACore/src/array.jl | 118 ++++++++++++++++++++++++++++++++++++------ test/core/array.jl | 47 +++++++++++++++++ 2 files changed, 149 insertions(+), 16 deletions(-) diff --git a/CUDACore/src/array.jl b/CUDACore/src/array.jl index 2bf2e21fca..b0aaeb976e 100644 --- a/CUDACore/src/array.jl +++ b/CUDACore/src/array.jl @@ -220,6 +220,20 @@ Wrap a `CuArray` object around the data at the address given by the CUDA-managed (for a 1d array) or a tuple of the array dimensions. `own` optionally specified whether Julia should take ownership of the memory, calling `cudaFree` when the array is no longer referenced. The `ctx` argument determines the CUDA context where the data is allocated in. + +When wrapping a CPU `Array`, the resulting `CuArray` keeps it alive. When wrapping a CPU +pointer, the caller has to make sure the memory stays valid for as long as the `CuArray` +is used. Either way, the memory must not be freed or reallocated (e.g., by `resize!`) +while it is wrapped, and operations on the `CuArray` execute asynchronously, so +synchronize before accessing the memory from the CPU. Like other CuArrays, a wrapped array +is not kept alive by CUDA graphs that captured operations on it: replaying such a graph +after the array has been freed accesses invalid memory. + +!!! warning + + `unsafe_wrap(Array, a::CuArray)` does not keep `a` alive: the returned `Array` is only + valid for as long as the caller keeps a reference to `a` (or to the memory it wraps). + Using it after `a` has been freed results in undefined behavior. """ unsafe_wrap @@ -301,8 +315,69 @@ end # unmanaged pointer to CuArray supports_hmm(dev) = driver_version() >= v"12.2" && attribute(dev, DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS) == 1 -function Base.unsafe_wrap(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,Int}; - ctx::CuContext=context()) where {T,N,M<:AbstractMemory} + +# wrapped memory whose outstanding work could not be waited for. rather than releasing it +# while it may still be in use, we leak it by keeping its owner alive. +const leaked_wrapped_memory = Any[] +const leaked_wrapped_memory_lock = ReentrantLock() + +# wait for outstanding work on wrapped system memory, before it can be released. returns +# whether that succeeded. this runs from a finalizer, so it must not switch tasks (hence the +# driver's blocking synchronization instead of `synchronize`) and has to perform its own +# error handling. +function synchronize_wrapped(managed::Managed, wrap_ctx::CuContext) + Base.@lock managed.lock begin + (managed.dirty || managed.captured) || return true + try + stream = managed.stream + if stream.ctx !== nothing && isvalid(stream) + context!(stream.ctx) do + cuStreamSynchronize(stream) + end + else + # the stream has been destroyed, but its work may still be executing, or it is one + # of the special streams that aren't tied to a context. either way, we can't tell + # where the work was submitted, so wait for every device's primary context, as well + # as the context the memory was wrapped in (which may be a non-primary one). + synchronized = CuContext[] + for dev in devices() + pctx = CuPrimaryContext(dev) + isactive(pctx) || continue + CuContext(pctx) do ctx + context!(ctx) do + cuCtxSynchronize() + end + push!(synchronized, ctx) + end + end + # a destroyed context has no outstanding work + if !(wrap_ctx in synchronized) && isvalid(wrap_ctx) + context!(wrap_ctx) do + cuCtxSynchronize() + end + end + end + managed.dirty = false + return true + catch ex + Base.showerror_nostdio(ex, + "WARNING: Error while waiting for work on wrapped host memory; leaking that memory") + Base.show_backtrace(Core.stdout, catch_backtrace()) + Core.println() + return false + end + end +end + +function leak_wrapped_memory(owner) + owner === nothing && return + Base.@lock leaked_wrapped_memory_lock push!(leaked_wrapped_memory, owner) + return +end + +# `owner` is kept alive for as long as the wrapper +function wrap_system_memory(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,Int}, + owner=nothing; ctx::CuContext=context()) where {T,N,M<:AbstractMemory} isbitstype(T) || throw(ArgumentError("Can only unsafe_wrap a pointer to a bits type")) sz = prod(dims) * aligned_sizeof(T) @@ -311,15 +386,23 @@ function Base.unsafe_wrap(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,Int} supports_hmm(device(ctx)) || throw(ArgumentError("Cannot wrap system memory as unified memory on your system")) mem = UnifiedMemory(ctx, reinterpret(CuPtr{Nothing}, p), sz) - DataRef(Returns(nothing), Managed(mem)) + DataRef(Managed(mem)) do managed, args... + synchronize_wrapped(managed, ctx) || leak_wrapped_memory(owner) + GC.@preserve owner nothing + end elseif M == HostMemory # register as device-accessible host memory mem = context!(ctx) do register(HostMemory, p, sz, MEMHOSTREGISTER_DEVICEMAP) end - DataRef(Managed(mem)) do args... - context!(ctx) do - unregister(mem) + DataRef(Managed(mem)) do managed, args... + if synchronize_wrapped(managed, ctx) + GC.@preserve owner context!(ctx) do + unregister(mem) + end + else + # keep the memory registered (and its owner alive), as it may still be in use + leak_wrapped_memory(owner) end end else @@ -328,26 +411,29 @@ function Base.unsafe_wrap(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,Int} CuArray{T,N}(data, dims) end -function Base.unsafe_wrap(::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,N}}}, - p::Ptr{T}, dims::NTuple{N,Int}; ctx::CuContext=context()) where {T,N} - if supports_hmm(device(ctx)) - Base.unsafe_wrap(CuArray{T,N,UnifiedMemory}, p, dims; ctx) - else - Base.unsafe_wrap(CuArray{T,N,HostMemory}, p, dims; ctx) - end +function wrap_system_memory(::Type{CuArray{T,N}}, p::Ptr{T}, dims::NTuple{N,Int}, + owner=nothing; ctx::CuContext=context()) where {T,N} + M = supports_hmm(device(ctx)) ? UnifiedMemory : HostMemory + wrap_system_memory(CuArray{T,N,M}, p, dims, owner; ctx) end +Base.unsafe_wrap(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,Int}; + ctx::CuContext=context()) where {T,N,M<:AbstractMemory} = + wrap_system_memory(CuArray{T,N,M}, p, dims; ctx) +Base.unsafe_wrap(::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,N}}}, + p::Ptr{T}, dims::NTuple{N,Int}; ctx::CuContext=context()) where {T,N} = + wrap_system_memory(CuArray{T,N}, p, dims; ctx) # integer size input Base.unsafe_wrap(::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,1}}}, p::Ptr{T}, dim::Int) where {T} = unsafe_wrap(CuArray{T,1}, p, (dim,)) Base.unsafe_wrap(::Type{CuArray{T,1,M}}, p::Ptr{T}, dim::Int) where {T,M} = unsafe_wrap(CuArray{T,1,M}, p, (dim,)) -# array input +# array input: keep the array alive for as long as the wrapper Base.unsafe_wrap(::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,N}}}, a::Array{T,N}) where {T,N} = - unsafe_wrap(CuArray{T,N}, pointer(a), size(a)) + wrap_system_memory(CuArray{T,N}, pointer(a), size(a), a) Base.unsafe_wrap(::Type{CuArray{T,N,M}}, a::Array{T,N}) where {T,N,M} = - unsafe_wrap(CuArray{T,N,M}, pointer(a), size(a)) + wrap_system_memory(CuArray{T,N,M}, pointer(a), size(a), a) ## array interface diff --git a/test/core/array.jl b/test/core/array.jl index 29bdfbbb5e..efd5be1a10 100644 --- a/test/core/array.jl +++ b/test/core/array.jl @@ -80,6 +80,17 @@ end CUDA.enable_synchronization!(a) end +function slow_sum(out, a, cycles) + t0 = clock(UInt64) + while clock(UInt64) - t0 < cycles end + acc = 0 + for x in a + acc += x + end + out[] = acc + return +end + @testset "unsafe_wrap" begin # managed memory -> CuArray for a in [cu([1]; device=true), cu([1]; unified=true)] @@ -158,6 +169,42 @@ end @test Array(b) == a end end + + # wrapping an Array keeps it alive + memtypes = [] + CUDA.supports_hmm(device()) && push!(memtypes, CUDA.UnifiedMemory) + can_register && push!(memtypes, CUDA.HostMemory) + for M in memtypes + b = unsafe_wrap(CuArray{Int,1,M}, collect(1:1024)) + GC.gc(true) + @test Array(b) == 1:1024 + b .+= 1 + @test Array(b) == 2:1025 + end + + # releasing a wrapper waits for outstanding work before the Array can be freed + for M in memtypes + out = CuArray([0]) + let b = unsafe_wrap(CuArray{Int,1,M}, collect(1:1024)) + @cuda slow_sum(out, b, 500_000_000) + end + GC.gc(true) + @test CUDA.isdone(stream()) + @test Array(out)[] == sum(1:1024) + + # also when the stream the work was submitted to has been destroyed since. + # the result is read from the CPU directly, as synchronizing `out` would try to + # use that stream. + outh = [0] + out = unsafe_wrap(CuArray{Int,1,M}, outh) + s = CuStream() + let b = unsafe_wrap(CuArray{Int,1,M}, collect(1:1024)) + @cuda stream=s slow_sum(out, b, 500_000_000) + end + finalize(s) + GC.gc(true) + @test outh[] == sum(1:1024) + end end # errors From 3be5c986a91ff17dda8dd296748da80c34eeba4e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 29 Sep 2026 22:07:06 +0200 Subject: [PATCH 2/5] Release wrapped system memory without a global leak list Unregistering host memory already waits for outstanding work using it, so do that directly, with the thread's capture mode relaxed so that it neither fails nor invalidates a graph capture that is in progress. Unified memory has no such operation, so have the stream signal when the work is done, by launching a host function. The task waiting for that signal keeps the Array alive, so failing to signal it leaks the memory without needing a global list. When the stream is being captured, retry after the capture has finished. --- CUDACore/src/array.jl | 119 ++++++++++++++++++++++++------------------ test/core/array.jl | 45 ++++++++++++++-- 2 files changed, 109 insertions(+), 55 deletions(-) diff --git a/CUDACore/src/array.jl b/CUDACore/src/array.jl index b0aaeb976e..0af2c9daca 100644 --- a/CUDACore/src/array.jl +++ b/CUDACore/src/array.jl @@ -316,62 +316,83 @@ end supports_hmm(dev) = driver_version() >= v"12.2" && attribute(dev, DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS) == 1 -# wrapped memory whose outstanding work could not be waited for. rather than releasing it -# while it may still be in use, we leak it by keeping its owner alive. -const leaked_wrapped_memory = Any[] -const leaked_wrapped_memory_lock = ReentrantLock() - -# wait for outstanding work on wrapped system memory, before it can be released. returns -# whether that succeeded. this runs from a finalizer, so it must not switch tasks (hence the -# driver's blocking synchronization instead of `synchronize`) and has to perform its own -# error handling. -function synchronize_wrapped(managed::Managed, wrap_ctx::CuContext) +# returns an async condition that releases `owner` once signalled. until then, the task +# waiting for the condition keeps it alive, so never signalling it leaks the memory. that +# task is not affected by cancellation of the scope that wrapped the memory. +function release_on_signal(owner) + Base.AsyncCondition() do cond + GC.@preserve owner close(cond) + end +end + +# signal `cond` once outstanding work on wrapped memory has finished, without waiting for +# it. this runs from a finalizer, so it must not switch tasks and has to perform its own +# error handling. if the work can't be waited for, `cond` is never signalled. +function signal_when_done(managed::Managed, cond::Base.AsyncCondition, wrap_ctx::CuContext) Base.@lock managed.lock begin - (managed.dirty || managed.captured) || return true try stream = managed.stream - if stream.ctx !== nothing && isvalid(stream) - context!(stream.ctx) do - cuStreamSynchronize(stream) + if !(managed.dirty || managed.captured) + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) + elseif isvalid(stream) + # special streams aren't tied to a context, and resolve against the current one + ctx = something(stream.ctx, wrap_ctx) + capturing = context!(ctx) do + # a host function launched on a capturing stream would become part of the graph, + # and waiting for such a stream is not allowed, so retry after the capture. + is_capturing(stream) && return true + cuLaunchHostFunc(stream, cglobal(:uv_async_send), cond) + return false end - else - # the stream has been destroyed, but its work may still be executing, or it is one - # of the special streams that aren't tied to a context. either way, we can't tell - # where the work was submitted, so wait for every device's primary context, as well - # as the context the memory was wrapped in (which may be a non-primary one). - synchronized = CuContext[] - for dev in devices() - pctx = CuPrimaryContext(dev) - isactive(pctx) || continue - CuContext(pctx) do ctx - context!(ctx) do - cuCtxSynchronize() + if capturing + # finalizers can't switch tasks, but they can schedule them + @async begin + while context!(() -> is_capturing(stream), ctx) + sleep(0.01) end - push!(synchronized, ctx) + signal_when_done(managed, cond, wrap_ctx) end end - # a destroyed context has no outstanding work - if !(wrap_ctx in synchronized) && isvalid(wrap_ctx) - context!(wrap_ctx) do - cuCtxSynchronize() - end + elseif isvalid(stream.ctx) + # the stream has been destroyed, but its work may still be executing + context!(stream.ctx) do + cuCtxSynchronize() end + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) + else + # a destroyed context has no outstanding work + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) end - managed.dirty = false - return true catch ex Base.showerror_nostdio(ex, - "WARNING: Error while waiting for work on wrapped host memory; leaking that memory") + "WARNING: Error while releasing wrapped system memory; leaking that memory") Base.show_backtrace(Core.stdout, catch_backtrace()) Core.println() - return false end end + return end -function leak_wrapped_memory(owner) - owner === nothing && return - Base.@lock leaked_wrapped_memory_lock push!(leaked_wrapped_memory, owner) +# unregister wrapped host memory. that waits for outstanding work using it, so its owner +# can be released afterwards. this runs from a finalizer, possibly while a graph is being +# captured, in which case the default capture mode would reject unregistering memory (and +# invalidate the capture), so temporarily relax this thread's capture mode. +function unregister_wrapped(mem::HostMemory, ctx::CuContext) + mode = Ref(STREAM_CAPTURE_MODE_RELAXED) + try + cuThreadExchangeStreamCaptureMode(mode) + try + context!(ctx) do + unregister(mem) + end + finally + cuThreadExchangeStreamCaptureMode(mode) + end + catch ex + Base.showerror_nostdio(ex, "WARNING: Error while unregistering wrapped host memory") + Base.show_backtrace(Core.stdout, catch_backtrace()) + Core.println() + end return end @@ -386,9 +407,14 @@ function wrap_system_memory(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,In supports_hmm(device(ctx)) || throw(ArgumentError("Cannot wrap system memory as unified memory on your system")) mem = UnifiedMemory(ctx, reinterpret(CuPtr{Nothing}, p), sz) - DataRef(Managed(mem)) do managed, args... - synchronize_wrapped(managed, ctx) || leak_wrapped_memory(owner) - GC.@preserve owner nothing + if owner === nothing + DataRef(Returns(nothing), Managed(mem)) + else + # nothing tells us when the GPU is done with this memory, so have the stream do so + cond = release_on_signal(owner) + DataRef(Managed(mem)) do managed, args... + signal_when_done(managed, cond, ctx) + end end elseif M == HostMemory # register as device-accessible host memory @@ -396,14 +422,7 @@ function wrap_system_memory(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,In register(HostMemory, p, sz, MEMHOSTREGISTER_DEVICEMAP) end DataRef(Managed(mem)) do managed, args... - if synchronize_wrapped(managed, ctx) - GC.@preserve owner context!(ctx) do - unregister(mem) - end - else - # keep the memory registered (and its owner alive), as it may still be in use - leak_wrapped_memory(owner) - end + GC.@preserve owner unregister_wrapped(mem, ctx) end else throw(ArgumentError("Cannot wrap system memory as $M")) diff --git a/test/core/array.jl b/test/core/array.jl index efd5be1a10..d71d7ee6c6 100644 --- a/test/core/array.jl +++ b/test/core/array.jl @@ -182,15 +182,50 @@ end @test Array(b) == 2:1025 end - # releasing a wrapper waits for outstanding work before the Array can be freed + # the Array is only released once outstanding work on the wrapper has finished + function launch_tracked(out, M, collected, cycles) + a = collect(1:1024) + finalizer(_ -> collected[] = true, a) + b = unsafe_wrap(CuArray{Int,1,M}, a) + @cuda slow_sum(out, b, cycles) + return + end + function wait_collected(collected) + t = time() + while !collected[] && time() - t < 10 + GC.gc(true) + sleep(0.01) + end + return collected[] + end for M in memtypes out = CuArray([0]) - let b = unsafe_wrap(CuArray{Int,1,M}, collect(1:1024)) - @cuda slow_sum(out, b, 500_000_000) - end + collected = Threads.Atomic{Bool}(false) + launch_tracked(out, M, collected, 1_000_000_000) GC.gc(true) - @test CUDA.isdone(stream()) + if M == CUDA.HostMemory + # unregistering the memory waits for the outstanding work + @test CUDA.isdone(stream()) + else + # the stream signals when the Array can be released + @test !collected[] + end + @test Array(out)[] == sum(1:1024) + @test wait_collected(collected) + + # also when collected while capturing a graph on the stream that used it, + # without breaking that capture + x = CuArray([0]) + x .+= 1 + collected[] = false + launch_tracked(out, M, collected, 500_000_000) + graph = CUDA.capture() do + x .+= 1 + GC.gc(true) + end + @test graph isa CuGraph @test Array(out)[] == sum(1:1024) + @test wait_collected(collected) # also when the stream the work was submitted to has been destroyed since. # the result is read from the CPU directly, as synchronizing `out` would try to From 213ebd9f7804930a17b74d7b0712af3cba1620ab Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 29 Sep 2026 22:13:01 +0200 Subject: [PATCH 3/5] Actually collect the wrapper during the capture in the test CUDA.capture disables the GC, so re-enable it while collecting. --- test/core/array.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/test/core/array.jl b/test/core/array.jl index d71d7ee6c6..e5930b3756 100644 --- a/test/core/array.jl +++ b/test/core/array.jl @@ -214,14 +214,18 @@ end @test wait_collected(collected) # also when collected while capturing a graph on the stream that used it, - # without breaking that capture + # without breaking that capture. `capture` disables the GC, but captures made + # by other means don't. x = CuArray([0]) x .+= 1 collected[] = false + GC.gc(true) # other garbage may not be freeable during the capture launch_tracked(out, M, collected, 500_000_000) graph = CUDA.capture() do x .+= 1 + gc_state = GC.enable(true) GC.gc(true) + GC.enable(gc_state) end @test graph isa CuGraph @test Array(out)[] == sum(1:1024) From 15afbfd647aaf322b0e84fc0ba8812cda5f592be Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 07:17:59 +0200 Subject: [PATCH 4/5] Launch the destroyed-stream test's kernel from a function A `let` at the top level of a test loop doesn't make the wrapper unreachable, so it was never collected and the test only passed if the kernel happened to finish before the GC did. --- test/core/array.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/core/array.jl b/test/core/array.jl index e5930b3756..b260ca6f89 100644 --- a/test/core/array.jl +++ b/test/core/array.jl @@ -183,11 +183,11 @@ end end # the Array is only released once outstanding work on the wrapper has finished - function launch_tracked(out, M, collected, cycles) + function launch_tracked(out, M, collected, cycles; stream=stream()) a = collect(1:1024) finalizer(_ -> collected[] = true, a) b = unsafe_wrap(CuArray{Int,1,M}, a) - @cuda slow_sum(out, b, cycles) + @cuda stream=stream slow_sum(out, b, cycles) return end function wait_collected(collected) @@ -237,12 +237,12 @@ end outh = [0] out = unsafe_wrap(CuArray{Int,1,M}, outh) s = CuStream() - let b = unsafe_wrap(CuArray{Int,1,M}, collect(1:1024)) - @cuda stream=s slow_sum(out, b, 500_000_000) - end + collected[] = false + launch_tracked(out, M, collected, 500_000_000; stream=s) finalize(s) GC.gc(true) @test outh[] == sum(1:1024) + @test wait_collected(collected) end end From 27edd8495a22a0010d401e4e6b399fa194514286 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 12:23:04 +0200 Subject: [PATCH 5/5] Release wrapped system memory from a task, after waiting for its work Both memory kinds now share one mechanism: the wrapper's finalizer wakes a task, created along with the wrapper (so it isn't affected by cancellation), which waits for outstanding work on the stream that last used the memory, unregisters page-locked memory, and releases the Array. An explicit unsafe_free! waits for that task, so the Array can be wrapped again right away. This no longer relies on cuMemHostUnregister implicitly waiting for the GPU, launches nothing on the stream, doesn't block in finalizers, and removes the capture-mode swap and the polling task. GC during a capture is left unsupported, like elsewhere in CUDA.jl, so the test that re-enabled it (and freed unrelated memory on the capturing stream) is removed. The ordering tests now use explicit frees, reading the kernel's result from wrapped host memory, which doesn't depend on GC timing. --- CUDACore/src/array.jl | 154 +++++++++++++++++++----------------------- test/core/array.jl | 94 ++++++++++++-------------- 2 files changed, 113 insertions(+), 135 deletions(-) diff --git a/CUDACore/src/array.jl b/CUDACore/src/array.jl index 0af2c9daca..4969563be7 100644 --- a/CUDACore/src/array.jl +++ b/CUDACore/src/array.jl @@ -221,9 +221,9 @@ Wrap a `CuArray` object around the data at the address given by the CUDA-managed Julia should take ownership of the memory, calling `cudaFree` when the array is no longer referenced. The `ctx` argument determines the CUDA context where the data is allocated in. -When wrapping a CPU `Array`, the resulting `CuArray` keeps it alive. When wrapping a CPU -pointer, the caller has to make sure the memory stays valid for as long as the `CuArray` -is used. Either way, the memory must not be freed or reallocated (e.g., by `resize!`) +When wrapping a CPU `Array`, the resulting `CuArray` keeps it alive, until the GPU is done +using it. When wrapping a CPU pointer, the caller has to make sure the memory stays valid +for as long as the `CuArray` is used. Either way, the memory must not be freed or reallocated (e.g., by `resize!`) while it is wrapped, and operations on the `CuArray` execute asynchronously, so synchronize before accessing the memory from the CPU. Like other CuArrays, a wrapped array is not kept alive by CUDA graphs that captured operations on it: replaying such a graph @@ -316,117 +316,105 @@ end supports_hmm(dev) = driver_version() >= v"12.2" && attribute(dev, DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS) == 1 -# returns an async condition that releases `owner` once signalled. until then, the task -# waiting for the condition keeps it alive, so never signalling it leaks the memory. that -# task is not affected by cancellation of the scope that wrapped the memory. -function release_on_signal(owner) - Base.AsyncCondition() do cond - GC.@preserve owner close(cond) - end +# whether the stream that last used wrapped memory is being captured +function stream_capturing(managed::Managed) + stream = managed.stream + isvalid(stream) && stream.ctx !== nothing && + context!(() -> is_capturing(stream), stream.ctx) end -# signal `cond` once outstanding work on wrapped memory has finished, without waiting for -# it. this runs from a finalizer, so it must not switch tasks and has to perform its own -# error handling. if the work can't be waited for, `cond` is never signalled. -function signal_when_done(managed::Managed, cond::Base.AsyncCondition, wrap_ctx::CuContext) +# wait for outstanding work on wrapped memory, without blocking the thread. this doesn't +# use `synchronize`, which would also report (and reset) unrelated kernel exceptions. +function wait_for_work(managed::Managed, wrap_ctx::CuContext) Base.@lock managed.lock begin - try - stream = managed.stream - if !(managed.dirty || managed.captured) - ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) - elseif isvalid(stream) - # special streams aren't tied to a context, and resolve against the current one - ctx = something(stream.ctx, wrap_ctx) - capturing = context!(ctx) do - # a host function launched on a capturing stream would become part of the graph, - # and waiting for such a stream is not allowed, so retry after the capture. - is_capturing(stream) && return true - cuLaunchHostFunc(stream, cglobal(:uv_async_send), cond) - return false - end - if capturing - # finalizers can't switch tasks, but they can schedule them - @async begin - while context!(() -> is_capturing(stream), ctx) - sleep(0.01) - end - signal_when_done(managed, cond, wrap_ctx) - end - end - elseif isvalid(stream.ctx) - # the stream has been destroyed, but its work may still be executing - context!(stream.ctx) do - cuCtxSynchronize() - end - ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) - else - # a destroyed context has no outstanding work - ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) + (managed.dirty || managed.captured) || return + stream = managed.stream + if stream.ctx === nothing + # special streams aren't tied to a context or a thread, so wait for the whole context + isvalid(wrap_ctx) && nonblocking_synchronize(wrap_ctx) + elseif isvalid(stream) + # waiting for a stream that is being captured is not allowed + while stream_capturing(managed) + sleep(0.01) end - catch ex - Base.showerror_nostdio(ex, - "WARNING: Error while releasing wrapped system memory; leaking that memory") - Base.show_backtrace(Core.stdout, catch_backtrace()) - Core.println() + nonblocking_synchronize(stream) + elseif isvalid(stream.ctx) + # the stream has been destroyed, but its work may still be executing + nonblocking_synchronize(stream.ctx) end + # a destroyed context has no outstanding work + managed.dirty = false end - return end -# unregister wrapped host memory. that waits for outstanding work using it, so its owner -# can be released afterwards. this runs from a finalizer, possibly while a graph is being -# captured, in which case the default capture mode would reject unregistering memory (and -# invalidate the capture), so temporarily relax this thread's capture mode. -function unregister_wrapped(mem::HostMemory, ctx::CuContext) - mode = Ref(STREAM_CAPTURE_MODE_RELAXED) - try - cuThreadExchangeStreamCaptureMode(mode) +# release wrapped system memory, and its owner, once outstanding work using it has finished +function release_after_work(managed::Managed, owner, ctx::CuContext) + GC.@preserve owner begin try + wait_for_work(managed, ctx) + catch ex + # waiting only fails if the work can't complete anymore, e.g., after a fatal error + @error "Error while waiting for GPU work on wrapped system memory" exception=(ex, catch_backtrace()) + end + # a destroyed context has no registrations left + if managed.mem isa HostMemory && isvalid(ctx) context!(ctx) do - unregister(mem) + unregister(managed.mem) end - finally - cuThreadExchangeStreamCaptureMode(mode) end - catch ex - Base.showerror_nostdio(ex, "WARNING: Error while unregistering wrapped host memory") - Base.show_backtrace(Core.stdout, catch_backtrace()) - Core.println() end return end -# `owner` is kept alive for as long as the wrapper +# `owner` is kept alive for as long as the wrapper, and until the GPU is done with it function wrap_system_memory(::Type{CuArray{T,N,M}}, p::Ptr{T}, dims::NTuple{N,Int}, owner=nothing; ctx::CuContext=context()) where {T,N,M<:AbstractMemory} isbitstype(T) || throw(ArgumentError("Can only unsafe_wrap a pointer to a bits type")) sz = prod(dims) * aligned_sizeof(T) - data = if M == UnifiedMemory + mem = if M == UnifiedMemory # HMM extends unified memory to include system memory supports_hmm(device(ctx)) || throw(ArgumentError("Cannot wrap system memory as unified memory on your system")) - mem = UnifiedMemory(ctx, reinterpret(CuPtr{Nothing}, p), sz) - if owner === nothing - DataRef(Returns(nothing), Managed(mem)) - else - # nothing tells us when the GPU is done with this memory, so have the stream do so - cond = release_on_signal(owner) - DataRef(Managed(mem)) do managed, args... - signal_when_done(managed, cond, ctx) - end - end + UnifiedMemory(ctx, reinterpret(CuPtr{Nothing}, p), sz) elseif M == HostMemory # register as device-accessible host memory - mem = context!(ctx) do + context!(ctx) do register(HostMemory, p, sz, MEMHOSTREGISTER_DEVICEMAP) end - DataRef(Managed(mem)) do managed, args... - GC.@preserve owner unregister_wrapped(mem, ctx) - end else throw(ArgumentError("Cannot wrap system memory as $M")) end + managed = Managed(mem) + + # wrapped unified memory without an owner doesn't need to be released + if M == UnifiedMemory && owner === nothing + return CuArray{T,N}(DataRef(Returns(nothing), managed), dims) + end + + # waiting for the GPU requires switching tasks, which isn't possible from a finalizer, so + # the memory is released by a task that's woken up when the wrapper is freed. it is + # created here, so it isn't affected by cancellation of the scope the wrapper is freed in. + released = Base.Event() + cond = Base.AsyncCondition() do cond + try + release_after_work(managed, owner, ctx) + catch ex + @error "Error while releasing wrapped system memory" exception=(ex, catch_backtrace()) + finally + close(cond) + notify(released) + end + end + data = DataRef(managed) do managed, args... + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) + # when possible, wait for that, so that the memory can be wrapped again right away. + # that's not possible from a finalizer, or while a capture of the stream that last + # used the memory, which this task may have to end, is holding up the release. + if !GC.in_finalizer() && !stream_capturing(managed) + Base.wait(released) + end + end CuArray{T,N}(data, dims) end diff --git a/test/core/array.jl b/test/core/array.jl index b260ca6f89..3a585bfb1f 100644 --- a/test/core/array.jl +++ b/test/core/array.jl @@ -180,69 +180,59 @@ end @test Array(b) == 1:1024 b .+= 1 @test Array(b) == 2:1025 + + # also by arrays derived from the wrapper + b = reshape(unsafe_wrap(CuArray{Int,1,M}, collect(1:1024)), 32, 32) + GC.gc(true) + @test vec(Array(b)) == 1:1024 end # the Array is only released once outstanding work on the wrapper has finished - function launch_tracked(out, M, collected, cycles; stream=stream()) + for M in memtypes + # the kernel's output is written to wrapped host memory, which can be checked + # from the CPU without synchronizing + outh = [0] + out = unsafe_wrap(CuArray{Int,1,M}, outh) a = collect(1:1024) - finalizer(_ -> collected[] = true, a) b = unsafe_wrap(CuArray{Int,1,M}, a) - @cuda stream=stream slow_sum(out, b, cycles) - return + @cuda slow_sum(out, b, 500_000_000) + CUDA.unsafe_free!(b) + @test outh[] == sum(1:1024) + + # the memory is released right away, so it can be wrapped again + b = unsafe_wrap(CuArray{Int,1,M}, a) + @test Array(b) == 1:1024 + + # also when the stream the work was submitted to has been destroyed since + outh[] = 0 + s = CuStream() + @cuda stream=s slow_sum(out, b, 500_000_000) + finalize(s) + CUDA.unsafe_free!(b) + @test outh[] == sum(1:1024) end - function wait_collected(collected) - t = time() - while !collected[] && time() - t < 10 - GC.gc(true) - sleep(0.01) + + # when the GC frees the wrapper, a task releases the Array + function launch_tracked(outh, M, released) + a = collect(1:1024) + finalizer(a) do _ + released[] = outh[] == sum(1:1024) ? 1 : -1 end - return collected[] + out = unsafe_wrap(CuArray{Int,1,M}, outh) + b = unsafe_wrap(CuArray{Int,1,M}, a) + @cuda slow_sum(out, b, 500_000_000) + return end for M in memtypes - out = CuArray([0]) - collected = Threads.Atomic{Bool}(false) - launch_tracked(out, M, collected, 1_000_000_000) - GC.gc(true) - if M == CUDA.HostMemory - # unregistering the memory waits for the outstanding work - @test CUDA.isdone(stream()) - else - # the stream signals when the Array can be released - @test !collected[] - end - @test Array(out)[] == sum(1:1024) - @test wait_collected(collected) - - # also when collected while capturing a graph on the stream that used it, - # without breaking that capture. `capture` disables the GC, but captures made - # by other means don't. - x = CuArray([0]) - x .+= 1 - collected[] = false - GC.gc(true) # other garbage may not be freeable during the capture - launch_tracked(out, M, collected, 500_000_000) - graph = CUDA.capture() do - x .+= 1 - gc_state = GC.enable(true) + outh = [0] + released = Threads.Atomic{Int}(0) + launch_tracked(outh, M, released) + t = time() + while released[] == 0 && time() - t < 10 GC.gc(true) - GC.enable(gc_state) + sleep(0.01) end - @test graph isa CuGraph - @test Array(out)[] == sum(1:1024) - @test wait_collected(collected) - - # also when the stream the work was submitted to has been destroyed since. - # the result is read from the CPU directly, as synchronizing `out` would try to - # use that stream. - outh = [0] - out = unsafe_wrap(CuArray{Int,1,M}, outh) - s = CuStream() - collected[] = false - launch_tracked(out, M, collected, 500_000_000; stream=s) - finalize(s) - GC.gc(true) - @test outh[] == sum(1:1024) - @test wait_collected(collected) + @test released[] == 1 end end