From cf47b89d3087cda9117c083f91a446b7a6ed3fa7 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 29 Sep 2026 17:33:28 +0200 Subject: [PATCH 1/3] Undo host registrations made by unsafe_wrap, and support wrapping Arrays Wrapping unregistered host memory with `unsafe_wrap(ROCArray, ptr)` page-locks it with a refcounted `hipHostRegister`, but with the default `own=false` the wrapper never undid that registration: its finalizer did nothing, and `free(::HostBuffer)` returns early for unowned buffers. The memory stayed pinned forever, and because the refcount is keyed by address, wrapping new memory that happens to reuse that address skipped registering it. The wrapper now drops its reference to the registration when freed, without freeing memory it doesn't own. Releasing wrapped host memory, owned or not, first waits for the device to stop using it, polling the stream so that the thread stays available to service hostcalls. Finalizers cannot yield, so when a wrapper is finalized the release is queued for a background task, started the first time host memory is wrapped. Explicit `unsafe_free!` still releases immediately. If waiting fails, the memory stays registered and rooted rather than risking a use after free. Also add `unsafe_wrap(ROCArray, ::Array)`, which keeps the array alive for as long as the wrapper exists, and `unsafe_wrap(Array, ::ROCArray)` for host-backed arrays, documenting that the latter does not keep the ROCArray alive. Copies now preserve their operands until the copy has been submitted, so that a wrapper can't be finalized (unregistering its memory) in between. --- src/array.jl | 156 +++++++++++++++++++++++++++++++++++-- test/core/rocarray_base.jl | 31 ++++++++ 2 files changed, 181 insertions(+), 6 deletions(-) diff --git a/src/array.jl b/src/array.jl index 7eac39eda..3feb4d307 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,11 +271,33 @@ 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. + +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 B = if memtype == HIP.hipMemoryTypeUnregistered @@ -289,17 +311,139 @@ 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)) + B == Mem.HostBuffer && start_host_release_task() + if B == Mem.HostBuffer && sz == 0 + # registering an empty range is invalid + buf = Mem.HostBuffer() + else + buf = B(Ptr{Cvoid}(ptr), sz; own) + end + 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 + managed -> release_host_memory(managed, own, owner) + elseif own + pool_free + else + managed -> GC.@preserve owner nothing + end + dref = DataRef(finalize_buffer, Managed(buf)) return ROCArray{T, N}(dref, dims) end +# Releasing wrapped host memory has to wait for the device to finish using it. That is not +# possible from a finalizer: finalizers cannot yield, and blocking could deadlock with a +# kernel waiting for the host to service a hostcall. So when finalized, the memory is +# queued for release by a background task instead. +const host_release_lock = Threads.SpinLock() +const host_release_queue = Tuple{Managed{Mem.HostBuffer}, Bool, Any}[] +const host_release_cond = Ref{Base.AsyncCondition}() +const host_release_start_lock = ReentrantLock() +# memory whose release failed, kept pinned and rooted rather than risking use after free +const host_release_leaked = Any[] + +function release_host_memory(managed::Managed{Mem.HostBuffer}, own::Bool, owner) + if GC.in_finalizer() + @lock host_release_lock begin + push!(host_release_queue, (managed, own, owner)) + end + # safe to call from finalizers and any thread + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), host_release_cond[]) + else + _release_host_memory(managed, own, owner) + end + return +end + +function _release_host_memory(managed::Managed{Mem.HostBuffer}, own::Bool, owner = nothing) + buf = managed.mem + # the owner has to stay alive until its memory has been released + GC.@preserve owner Base.@lock managed.lock begin + try + # wait without blocking the thread, regardless of the synchronization + # preference, as a kernel may be waiting for us to service a hostcall + if managed.dirty + spins = 0 + while !HIP.isdone(managed.stream) + spins < 100 ? yield() : sleep(0.001) + spins += 1 + end + managed.dirty = false + end + catch ex + @error "Error while waiting to release $(Base.format_bytes(buf.bytesize)) of wrapped host memory; leaking it" exception=(ex, catch_backtrace()) + @lock host_release_lock push!(host_release_leaked, (managed, owner)) + return + end + 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 + end + return +end + +# started when host memory is first wrapped, i.e., outside of a finalizer +function start_host_release_task() + isassigned(host_release_cond) && return + Base.@lock host_release_start_lock begin + isassigned(host_release_cond) && return + cond = Base.AsyncCondition() + errormonitor(Threads.@spawn begin + while true + wait(cond) + while true + entry = @lock host_release_lock begin + isempty(host_release_queue) ? nothing : popfirst!(host_release_queue) + end + entry === nothing && break + managed, own, owner = entry + _release_host_memory(managed, own, owner) + end + end + end) + host_release_cond[] = cond + 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/test/core/rocarray_base.jl b/test/core/rocarray_base.jl index 78644c669..7d9fee39d 100644 --- a/test/core/rocarray_base.jl +++ b/test/core/rocarray_base.jl @@ -162,6 +162,37 @@ 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 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 "Broadcasting different buffer types" begin x = rand(Float32, 4, 16, 16) xd = unsafe_wrap(ROCArray, pointer(x), size(x)) From d02902a858597d3facfc95f2317bd3ed9bbba3e9 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 29 Sep 2026 22:37:37 +0200 Subject: [PATCH 2/3] Release wrapped memory from a stream callback instead of a service task Replace the global release queue, its service task and the polling loop with a host function launched on the stream that last used the memory. It signals an async condition once the device is done with the memory, and the task waiting for that condition releases it. That task keeps the wrapped Array alive, so failing to signal it leaks the memory without needing a global list, and it isn't affected by task cancellation. Explicitly freeing a wrapper that the device is done with still releases it right away. --- src/array.jl | 136 ++++++++++++++++++------------------- test/core/rocarray_base.jl | 17 +++++ 2 files changed, 82 insertions(+), 71 deletions(-) diff --git a/src/array.jl b/src/array.jl index 3feb4d307..53d2802f9 100644 --- a/src/array.jl +++ b/src/array.jl @@ -282,7 +282,7 @@ memory or to host memory; host memory that is not yet page-locked is registered 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. +`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. @@ -311,104 +311,98 @@ function wrap_memory(ptr::Ptr{T}, dims::NTuple{N, <:Integer}, own::Bool, end sz = prod(dims) * aligned_sizeof(T) - B == Mem.HostBuffer && start_host_release_task() 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 - managed -> release_host_memory(managed, own, owner) + release_after_use(() -> release_host_memory(managed, own), owner) elseif own pool_free + elseif owner !== nothing + release_after_use(Returns(nothing), owner) else - managed -> GC.@preserve owner nothing + Returns(nothing) end - dref = DataRef(finalize_buffer, Managed(buf)) + dref = DataRef(finalize_buffer, managed) return ROCArray{T, N}(dref, dims) end -# Releasing wrapped host memory has to wait for the device to finish using it. That is not -# possible from a finalizer: finalizers cannot yield, and blocking could deadlock with a -# kernel waiting for the host to service a hostcall. So when finalized, the memory is -# queued for release by a background task instead. -const host_release_lock = Threads.SpinLock() -const host_release_queue = Tuple{Managed{Mem.HostBuffer}, Bool, Any}[] -const host_release_cond = Ref{Base.AsyncCondition}() -const host_release_start_lock = ReentrantLock() -# memory whose release failed, kept pinned and rooted rather than risking use after free -const host_release_leaked = Any[] - -function release_host_memory(managed::Managed{Mem.HostBuffer}, own::Bool, owner) - if GC.in_finalizer() - @lock host_release_lock begin - push!(host_release_queue, (managed, own, owner)) - end - # safe to call from finalizers and any thread - ccall(:uv_async_send, Cint, (Ptr{Cvoid},), host_release_cond[]) - else - _release_host_memory(managed, own, owner) +# 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() + cond = Base.AsyncCondition() do cond + close(cond) + GC.@preserve owner release_once() end - return + return managed -> release_when_done(managed, release_once, cond) end -function _release_host_memory(managed::Managed{Mem.HostBuffer}, own::Bool, owner = nothing) - buf = managed.mem - # the owner has to stay alive until its memory has been released - GC.@preserve owner Base.@lock managed.lock begin - try - # wait without blocking the thread, regardless of the synchronization - # preference, as a kernel may be waiting for us to service a hostcall - if managed.dirty - spins = 0 - while !HIP.isdone(managed.stream) - spins < 100 ? yield() : sleep(0.001) - spins += 1 +function release_when_done(managed::Managed, release, cond::Base.AsyncCondition) + if !managed.dirty + 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) end - managed.dirty = false - end - catch ex - @error "Error while waiting to release $(Base.format_bytes(buf.bytesize)) of wrapped host memory; leaking it" exception=(ex, catch_backtrace()) - @lock host_release_lock push!(host_release_leaked, (managed, owner)) - return - end - try - if own - pool_free(managed) + elseif !GC.in_finalizer() && HIP.isdone(stream) + # freed explicitly after the device is done with the memory: release it now + release() + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), cond) else - AMDGPU.context!(() -> Mem.unregister(buf.ptr), buf.ctx) + HIP.hipLaunchHostFunc(stream, cglobal(:uv_async_send), cond) end - catch ex - @error "Error while releasing $(Base.format_bytes(buf.bytesize)) of wrapped host memory" exception=(ex, catch_backtrace()) 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 -# started when host memory is first wrapped, i.e., outside of a finalizer -function start_host_release_task() - isassigned(host_release_cond) && return - Base.@lock host_release_start_lock begin - isassigned(host_release_cond) && return - cond = Base.AsyncCondition() - errormonitor(Threads.@spawn begin - while true - wait(cond) - while true - entry = @lock host_release_lock begin - isempty(host_release_queue) ? nothing : popfirst!(host_release_queue) - end - entry === nothing && break - managed, own, owner = entry - _release_host_memory(managed, own, owner) - end - end - end) - host_release_cond[] = cond +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 diff --git a/test/core/rocarray_base.jl b/test/core/rocarray_base.jl index 7d9fee39d..b5bb12fdf 100644 --- a/test/core/rocarray_base.jl +++ b/test/core/rocarray_base.jl @@ -186,6 +186,23 @@ end 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) + 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) From c5843a1ff17a59648ded72c8ef6d119382edad48 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 2 Oct 2026 14:05:27 +0200 Subject: [PATCH 3/3] Address review of wrapped memory release - Query and launch on the stream with a relaxed capture mode, so that freeing a wrapper while another stream is being captured doesn't invalidate the capture. - Release immediately if the stream was recycled to another task, since the work has finished and the stream may be captured by its new owner. - Unregister host memory while holding the pin lock, so a concurrent register doesn't see it as externally pinned. Querying the memory type also takes the lock, as HIP can crash when it races with hipHostUnregister. - Wait for the release when a wrapper is freed explicitly, so the memory can be wrapped again right away, and refuse to extend a registration still in use. - Fix the test that checks the wrapped array is released. --- src/array.jl | 36 ++++++++++++++++---- src/hip/stream.jl | 11 +++--- src/runtime/memory/hip.jl | 16 +++++---- test/core/rocarray_base.jl | 70 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 116 insertions(+), 17 deletions(-) diff --git a/src/array.jl b/src/array.jl index 53d2802f9..f57a52ccd 100644 --- a/src/array.jl +++ b/src/array.jl @@ -299,7 +299,7 @@ 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 @@ -347,15 +347,20 @@ end 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) + return managed -> release_when_done(managed, release_once, cond, done) end -function release_when_done(managed::Managed, release, cond::Base.AsyncCondition) - if !managed.dirty +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 @@ -375,14 +380,20 @@ function release_when_done(managed::Managed, release, cond::Base.AsyncCondition) while AMDGPU.context!(() -> HIP.is_capturing(stream), stream.ctx) sleep(0.01) end - release_when_done(managed, release, cond) + release_when_done(managed, release, cond, done) end - elseif !GC.in_finalizer() && HIP.isdone(stream) + 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 - HIP.hipLaunchHostFunc(stream, cglobal(:uv_async_send), cond) + # 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 @@ -393,6 +404,17 @@ function release_when_done(managed::Managed, release, cond::Base.AsyncCondition) 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 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 b5bb12fdf..18e2f0f37 100644 --- a/test/core/rocarray_base.jl +++ b/test/core/rocarray_base.jl @@ -188,6 +188,7 @@ end # ... 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) @@ -210,6 +211,75 @@ end @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))