diff --git a/src/array.jl b/src/array.jl index fe680f07d..7eac39eda 100644 --- a/src/array.jl +++ b/src/array.jl @@ -410,7 +410,7 @@ end function Adapt.adapt_storage(to::Runtime.Adaptor, x::ROCArray{T,N}) where {T,N} managed = x.buf[] - push!(to.managed, managed) + to.stream === nothing || take_ownership_fast!(managed, to.stream) buf = managed.mem ptr = convert(Ptr{T}, typeof(buf) <: Mem.HIPBuffer ? buf : buf.dev_ptr) llvm_ptr = AMDGPU.LLVMPtr{T,AS.Global}(ptr + x.offset) diff --git a/src/highlevel.jl b/src/highlevel.jl index c045bc06b..5fc084b21 100644 --- a/src/highlevel.jl +++ b/src/highlevel.jl @@ -90,8 +90,8 @@ input object `x` as-is. Do not add methods to this function, but instead extend the underlying Adapt.jl package and register methods for the the `AMDGPU.Adaptor` type. """ -rocconvert(arg, managed::Vector{Managed}=Managed[]) = - adapt(Runtime.Adaptor(managed), arg) +rocconvert(arg) = adapt(Runtime.Adaptor(nothing), arg) +rocconvert(arg, stream::HIPStream) = adapt(Runtime.Adaptor(stream), arg) const MACRO_KWARGS = [:launch] const COMPILER_KWARGS = [:name, :unsafe_fp_atomics, :wavefrontsize64] diff --git a/src/memory.jl b/src/memory.jl index dd581b4f4..a1a7834b1 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -445,6 +445,13 @@ function take_ownership!(managed::Managed; stream::HIPStream=AMDGPU.stream()) return managed end +# Fast-path ownership transfer for the kernel-launch path +@inline function take_ownership_fast!(managed::Managed, stream::HIPStream) + (managed.stream === stream && managed.dirty) && return + Base.@lock managed.lock take_ownership!(managed; stream) + return +end + function lock_managed(managed::AbstractVector{<:Managed}) locked = unique(managed) sort!(locked; by=m -> objectid(m.lock)) diff --git a/src/runtime/Runtime.jl b/src/runtime/Runtime.jl index 6004aa23f..b9b981185 100644 --- a/src/runtime/Runtime.jl +++ b/src/runtime/Runtime.jl @@ -13,8 +13,11 @@ import ..AMDGPU import ..AMDGPU: LockedObject import .HIP: HIPDevice -struct Adaptor{M} - managed::M +# Carries the launching stream so that argument conversion can transfer +# buffer ownership inline (`take_ownership_fast!`); `nothing` skips ownership +# handling (bare `rocconvert` for reflection/type computation). +struct Adaptor{S} + stream::S end const RT_LOCK = Threads.ReentrantLock() diff --git a/src/runtime/hip-execution.jl b/src/runtime/hip-execution.jl index 86571fe4e..d011c8c0e 100644 --- a/src/runtime/hip-execution.jl +++ b/src/runtime/hip-execution.jl @@ -56,12 +56,9 @@ function (ker::HIPKernel{F, TT})( ) where {F, TT, N} # Check if previous kernels threw an exception. AMDGPU.throw_if_exception(stream.device) - managed = AMDGPU.Managed[] GC.@preserve args begin - converted = map(arg -> AMDGPU.rocconvert(arg, managed), args) - AMDGPU.with_managed(managed; stream) do - call(ker, converted...; stream, call_kwargs...) - end + converted = map(arg -> AMDGPU.rocconvert(arg, stream), args) + call(ker, converted...; stream, call_kwargs...) end end