Transfer buffer ownership at conversion time instead of a per-launch locked phase

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)

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]

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))

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()

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