diff --git a/src/array.jl b/src/array.jl index 7eac39eda..f57a52ccd 100644 --- a/src/array.jl +++ b/src/array.jl @@ -221,7 +221,7 @@ function Base.copyto!( @boundscheck checkbounds(dest, d_offset + amount - 1) @boundscheck checkbounds(source, s_offset + amount - 1) stm = stream() - Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stm) + GC.@preserve dest source Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stm) async || synchronize(stm) return dest end @@ -233,7 +233,7 @@ function Base.copyto!( amount == 0 && return dest @boundscheck checkbounds(dest, d_offset + amount - 1) @boundscheck checkbounds(source, s_offset + amount - 1) - Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stream()) + GC.@preserve dest source Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stream()) return dest end @@ -244,7 +244,7 @@ function Base.copyto!( amount == 0 && return dest @boundscheck checkbounds(dest, d_offset + amount - 1) @boundscheck checkbounds(source, s_offset + amount - 1) - Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stream()) + GC.@preserve dest source Mem.memcpy!(pointer(dest, d_offset), pointer(source, s_offset), amount * aligned_sizeof(T); stream=stream()) return dest end @@ -271,13 +271,35 @@ function Base.copy(X::ROCArray{T}) where T return Xnew end +""" + unsafe_wrap(ROCArray, ptr::Ptr{T}, dims; own=false) + unsafe_wrap(ROCArray, a::Array) + +Wrap a `ROCArray` around existing memory, without copying it. `ptr` can point to device +memory or to host memory; host memory that is not yet page-locked is registered with +`hipHostRegister` for as long as the wrapper exists, which can be slow. + +When wrapping an `Array`, the returned `ROCArray` keeps it alive. When wrapping a pointer, +the caller has to keep the memory valid for as long as the `ROCArray` is used. If `own` is +set, the memory is freed (or, for registered host memory, unregistered) when the +`ROCArray` is freed, as soon as the device is done using it. + +Device operations execute asynchronously, so synchronize (e.g., using +`AMDGPU.synchronize()`) before accessing wrapped host memory on the host. +""" function Base.unsafe_wrap( ::Type{<:ROCArray}, ptr::Ptr{T}, dims::NTuple{N, <:Integer}; own::Bool = false, ) where {T,N} - check_eltype("unsafe_wrap(CuArray, ...)", T) + return wrap_memory(ptr, dims, own) +end + +# `owner` is kept alive for as long as the wrapper +function wrap_memory(ptr::Ptr{T}, dims::NTuple{N, <:Integer}, own::Bool, + owner = nothing) where {T,N} + check_eltype("unsafe_wrap(ROCArray, ...)", T) - memtype = Mem.attributes(ptr).type + memtype = Mem.memory_type(Ptr{Cvoid}(ptr)) B = if memtype == HIP.hipMemoryTypeUnregistered Mem.HostBuffer elseif memtype == HIP.hipMemoryTypeHost @@ -289,17 +311,155 @@ function Base.unsafe_wrap( end sz = prod(dims) * aligned_sizeof(T) - buf = B(Ptr{Cvoid}(ptr), sz; own) - dref = DataRef(own ? pool_free : Returns(nothing), Managed(buf)) + if B == Mem.HostBuffer && sz == 0 + # registering an empty range is invalid + buf = Mem.HostBuffer() + else + buf = B(Ptr{Cvoid}(ptr), sz; own) + end + managed = Managed(buf) + finalize_buffer = if buf isa Mem.HostBuffer && buf.ptr != C_NULL && + (own || Mem.is_registered(buf.ptr)) + # constructing the buffer registered the memory (or took another reference to an + # existing registration), which needs to be undone even if we don't own it + release_after_use(() -> release_host_memory(managed, own), owner) + elseif own + pool_free + elseif owner !== nothing + release_after_use(Returns(nothing), owner) + else + Returns(nothing) + end + dref = DataRef(finalize_buffer, managed) return ROCArray{T, N}(dref, dims) end +# Returns a finalizer for the `Managed` memory of a wrapper, which calls `release` and lets +# go of `owner` once the device is done with that memory. +# +# Waiting for the device is not possible from a finalizer: finalizers cannot yield, and +# blocking could deadlock with a kernel waiting for the host to service a hostcall. Instead, +# a host function launched on the stream that last used the memory signals an async +# condition when the device reaches it, and the task waiting for that condition performs +# the release. That task keeps `owner` alive, so if the condition can't be signalled, the +# memory is leaked rather than released while it may still be in use. The task is not +# affected by cancellation of the scope that wrapped the memory. +function release_after_use(release, owner) + released = Threads.Atomic{Bool}(false) + release_once() = Threads.atomic_xchg!(released, true) || release() + done = Threads.Atomic{Bool}(false) + cond = Base.AsyncCondition() do cond + close(cond) + GC.@preserve owner release_once() + done[] = true + end + return managed -> release_when_done(managed, release_once, cond, done) +end + +function release_when_done(managed::Managed, release, cond::Base.AsyncCondition, + done::Threads.Atomic{Bool}) + # if the stream has been handed to another task since it last used the memory, that + # work has finished, and the stream may now be used (or captured) by the other task + if !managed.dirty || recycled(managed) + GC.in_finalizer() || release() + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) + return + end + try + stream = managed.stream + AMDGPU.context!(stream.ctx) do + # a destroyed stream may still have work in flight. our streams are blocking, + # so the default stream waits for that work. + HIP.isvalid(stream) || (stream = AMDGPU.default_stream()) + if HIP.is_capturing(stream) + # a host function launched on a capturing stream would become part of the + # graph, so retry after the capture. finalizers can't switch tasks, but + # they can schedule them. note that this can't be checked atomically, so a + # capture started concurrently by another thread may still record it. + @async begin + while AMDGPU.context!(() -> HIP.is_capturing(stream), stream.ctx) + sleep(0.01) + end + release_when_done(managed, release, cond, done) + end + elseif !GC.in_finalizer() && HIP.query(stream) == HIP.hipSuccess + # freed explicitly after the device is done with the memory: release it now + release() + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) + else + # another stream may be capturing, which launching on this one doesn't affect + HIP.relaxed_capture_mode() do + HIP.hipLaunchHostFunc(stream, cglobal(:uv_async_send), cond) + end + # when freed explicitly, wait for the release, so that the memory can be + # wrapped again right away + GC.in_finalizer() || wait_for_release(done, stream) + end + end + catch ex + Base.showerror_nostdio(ex, "WARNING: Error while releasing wrapped memory; leaking it") + Base.show_backtrace(Core.stdout, catch_backtrace()) + Core.println() + end + return +end + +function wait_for_release(done::Threads.Atomic{Bool}, stream::HIPStream) + while !done[] + # if the device fails, the host function never runs + HIP.isvalid(stream) || return + status = HIP.query(stream) + (status == HIP.hipSuccess || status == HIP.hipErrorNotReady) || return + sleep(0.001) + end + return +end + +function release_host_memory(managed::Managed{Mem.HostBuffer}, own::Bool) + buf = managed.mem + try + if own + pool_free(managed) + else + AMDGPU.context!(() -> Mem.unregister(buf.ptr), buf.ctx) + end + catch ex + @error "Error while releasing $(Base.format_bytes(buf.bytesize)) of wrapped host memory" exception=(ex, catch_backtrace()) + end + return +end + Base.unsafe_wrap(::Type{<:ROCArray}, ptr::Ptr, dim::Integer; own::Bool=false) = unsafe_wrap(ROCArray, ptr, (dim,); own) Base.unsafe_wrap(::Type{ROCArray{T}}, ptr::Ptr, dims::NTuple{N, <:Integer}; kwargs...) where {T, N} = unsafe_wrap(ROCArray, Base.unsafe_convert(Ptr{T}, ptr), dims; kwargs...) +# array input: keep the array alive for as long as the wrapper +Base.unsafe_wrap(::Union{Type{ROCArray}, Type{ROCArray{T}}, Type{ROCArray{T, N}}}, + a::Array{T, N}) where {T, N} = + wrap_memory(pointer(a), size(a), false, a) + +""" + unsafe_wrap(Array, a::ROCArray) + +Wrap an `Array` around the memory of a `ROCArray`, without copying it. This is only +possible for arrays backed by host memory, i.e., with buffer type `Mem.HostBuffer`. + +!!! warning + + The returned `Array` does **not** keep the `ROCArray` alive. The caller has to keep a + reference to the `ROCArray` for as long as the `Array`, or anything derived from it, + is used; otherwise the `Array` may end up referring to freed memory. Device operations + execute asynchronously, so synchronize before accessing the returned array. +""" +function Base.unsafe_wrap(::Type{Array}, a::ROCArray{T, N, Mem.HostBuffer}) where {T, N} + ptr = convert(Ptr{T}, a.buf[].mem.ptr) + a.offset + return unsafe_wrap(Array, ptr, size(a)) +end +Base.unsafe_wrap(::Type{Array}, a::ROCArray) = + throw(ArgumentError("Can only wrap an Array around a ROCArray backed by host memory")) + ## interop with CPU arrays # We don't convert isbits types in `adapt`, since they are already diff --git a/src/hip/stream.jl b/src/hip/stream.jl index 5ba8200a9..ed83d2b61 100644 --- a/src/hip/stream.jl +++ b/src/hip/stream.jl @@ -112,18 +112,21 @@ function claim_stream!(pool::Vector{PooledStream}, task::Task) return candidate.stream end -function query(s::HIPStream) - # querying a stream is prohibited while another one is being captured in global - # mode, even though it doesn't interfere with the capture, so temporarily relax that +# some API calls, like querying a stream, are prohibited while another stream is being +# captured in global mode, even when they don't interfere with the capture. `f` must not +# yield, because the capture mode is a property of the thread. +function relaxed_capture_mode(f) mode = Ref(hipStreamCaptureModeRelaxed) hipThreadExchangeStreamCaptureMode(mode) try - return unchecked_hipStreamQuery(s) + return f() finally hipThreadExchangeStreamCaptureMode(mode) end end +query(s::HIPStream) = relaxed_capture_mode(() -> unchecked_hipStreamQuery(s)) + # only bumped while holding `STREAM_POOL_LOCK`, but read without it generation(s::HIPStream) = Base.@atomic :acquire s.generation diff --git a/src/runtime/memory/hip.jl b/src/runtime/memory/hip.jl index e66273a4f..f897f042d 100644 --- a/src/runtime/memory/hip.jl +++ b/src/runtime/memory/hip.jl @@ -219,6 +219,10 @@ function register(ptr::Ptr{Cvoid}, sz::Integer) Base.@lock __pin_lock begin count = get(__pin_count, ptr, 0) if count > 0 + # the existing registration can't be extended + sz > __pinned_memory[ptr] && error(""" + Cannot register $(Base.format_bytes(sz)) at $ptr, which is still registered with a smaller size of $(Base.format_bytes(__pinned_memory[ptr])). + Make sure that previous wrappers of this memory have been freed, and that their release has finished (it is postponed while the stream that last used them is being captured).""") __pin_count[ptr] = count + 1 return end @@ -237,6 +241,9 @@ function register(ptr::Ptr{Cvoid}, sz::Integer) return end +# querying host memory that is concurrently being unregistered can crash HIP +memory_type(ptr::Ptr{Cvoid}) = Base.@lock __pin_lock attributes(ptr).type + """ unregister(ptr::Ptr{Cvoid}) @@ -246,23 +253,20 @@ Decrement the refcount for `ptr`. When it reaches zero the underlying function unregister(ptr::Ptr{Cvoid}) ptr == C_NULL && error("Cannot unregister `NULL` pointer.") - do_unregister = false Base.@lock __pin_lock begin count = get(__pin_count, ptr, 0) count == 0 && error("Cannot unregister untracked pointer $ptr.") if count == 1 + # unregister while holding the lock, or a concurrent `register` could see the + # memory as registered externally, and not track it delete!(__pinned_memory, ptr) delete!(__pin_count, ptr) - do_unregister = true + HIP.relaxed_capture_mode(() -> HIP.hipHostUnregister(ptr)) else __pin_count[ptr] = count - 1 end end - - if do_unregister - HIP.hipHostUnregister(ptr) - end return end diff --git a/test/core/rocarray_base.jl b/test/core/rocarray_base.jl index 78644c669..18e2f0f37 100644 --- a/test/core/rocarray_base.jl +++ b/test/core/rocarray_base.jl @@ -162,6 +162,124 @@ end @test AMDGPU.Mem.is_pinned(Ptr{Cvoid}(pointer(x))) == false end + @testset "Registration is undone when freeing" begin + x = zeros(Float32, 16) + xd = unsafe_wrap(ROCArray, pointer(x), size(x)) + @test AMDGPU.Mem.is_pinned(Ptr{Cvoid}(pointer(x))) + AMDGPU.unsafe_free!(xd) + @test !AMDGPU.Mem.is_pinned(Ptr{Cvoid}(pointer(x))) + @test !AMDGPU.Mem.is_registered(Ptr{Cvoid}(pointer(x))) + end + + @testset "Wrap Array" begin + for AT in [ROCArray, ROCArray{Float32}, ROCArray{Float32, 1}] + a = Float32[1, 2, 3] + b = unsafe_wrap(AT, a) + @test b isa ROCVector{Float32, AMDGPU.Mem.HostBuffer} + @test Array(b) == a + end + @test isempty(Array(unsafe_wrap(ROCArray, Float32[]))) + + # the wrapper keeps the array alive + xd = unsafe_wrap(ROCArray, fill(1f0, 1024)) + GC.gc(true) + AMDGPU.@sync xd .+= 1f0 + @test all(==(2f0), Array(xd)) + + # ... and lets go of it once the device is done using it + function wrap_tracked(collected) + local a, xd + a = fill(1f0, 1024) + finalizer(_ -> collected[] = true, a) + xd = unsafe_wrap(ROCArray, a) + xd .+= 1f0 + return + end + collected = Threads.Atomic{Bool}(false) + wrap_tracked(collected) + t = time() + while !collected[] && time() - t < 10 + GC.gc(true) + sleep(0.01) + end + @test collected[] + + # and the other way around + a = Float32[1, 2, 3] + b = unsafe_wrap(ROCArray, a) + @test pointer(unsafe_wrap(Array, b)) == pointer(a) + @test_throws ArgumentError unsafe_wrap(Array, AMDGPU.zeros(Float32, 3)) + end + + @testset "Re-wrapping after freeing" begin + # an explicit free waits for the release, so that the memory can be wrapped again + x = zeros(Float32, 1 << 20) + xd = unsafe_wrap(ROCArray, pointer(x), 16) + xd .+= 1f0 + AMDGPU.unsafe_free!(xd) + @test !AMDGPU.Mem.is_registered(Ptr{Cvoid}(pointer(x))) + xd = unsafe_wrap(ROCArray, pointer(x), length(x)) + xd .+= 1f0 + AMDGPU.synchronize() + @test sum(x) == 16 + length(x) + + # a registration can't be extended while it's still in use + @test_throws ErrorException unsafe_wrap(ROCArray, pointer(x), 2 * length(x)) + AMDGPU.unsafe_free!(xd) + end + + @testset "Freeing while capturing" begin + # last used on the stream of another task, which is still alive + a = zeros(Float32, 1024) + wrapped = Channel(1) + finish = Channel(1) + t = @async begin + xd = unsafe_wrap(ROCArray, a) + xd .+= 1f0 + put!(wrapped, xd) + take!(finish) + end + xd = take!(wrapped) + y = AMDGPU.zeros(Float32, 16) + graph = AMDGPU.capture() do + AMDGPU.unsafe_free!(xd) + y .+= 1f0 + end + @test graph !== nothing + @test !AMDGPU.Mem.is_registered(Ptr{Cvoid}(pointer(a))) + put!(finish, nothing) + wait(t) + @test all(==(1f0), a) + AMDGPU.HIP.launch(AMDGPU.HIP.instantiate(graph)) + @test all(==(1f0), Array(y)) + + # last used on a stream that has since been handed to another task, which is + # capturing it. that work has finished, so the memory is released right away. + # (simulated by bumping the stream's generation, as recycling isn't deterministic) + a = zeros(Float32, 1024) + s = HIPStream() + old_stream = AMDGPU.stream() + AMDGPU.stream!(s) + try + xd = unsafe_wrap(ROCArray, a) + xd .+= 1f0 + AMDGPU.synchronize() + Base.@atomic s.generation += 1 + z = AMDGPU.zeros(Float32, 16) + released = false + graph = AMDGPU.capture() do + AMDGPU.unsafe_free!(xd) + released = !AMDGPU.Mem.is_registered(Ptr{Cvoid}(pointer(a))) + z .+= 1f0 + end + @test graph !== nothing + @test released + finally + AMDGPU.stream!(old_stream) + end + @test all(==(1f0), a) + end + @testset "Broadcasting different buffer types" begin x = rand(Float32, 4, 16, 16) xd = unsafe_wrap(ROCArray, pointer(x), size(x))