diff --git a/CUDACore/src/array.jl b/CUDACore/src/array.jl index 2bf2e21fca..4969563be7 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, 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 +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,53 +315,132 @@ 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} + +# 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 + +# 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 + (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 + 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 +end + +# 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(managed.mem) + end + end + end + return +end + +# `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) - DataRef(Returns(nothing), Managed(mem)) + 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 args... - context!(ctx) do - unregister(mem) - end - 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 -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..3a585bfb1f 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,71 @@ 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 + + # 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 + 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) + b = unsafe_wrap(CuArray{Int,1,M}, a) + @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 + + # 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 + 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 + outh = [0] + released = Threads.Atomic{Int}(0) + launch_tracked(outh, M, released) + t = time() + while released[] == 0 && time() - t < 10 + GC.gc(true) + sleep(0.01) + end + @test released[] == 1 + end end # errors