Keep wrapped CPU arrays alive in unsafe_wrap(CuArray, ...)

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

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