Keep track of stream usage for arrays (#633)

- Add Managed struct that keeps track of TLS streams and correctly
    switches between them. This in turn allows us to **not** use global
    stream in finalizers.
- Remove REPL hook.

Author: pxl-th

src/AMDGPU.jl

@@ -72,7 +72,7 @@ include("runtime/Runtime.jl")
 import .Runtime
 import .Runtime: Mem, ROCDim, ROCDim3
 
-include("stats.jl")
+include("memory.jl")
 
 const ci_cache = GPUCompiler.CodeCache()
 Base.Experimental.@MethodTable(method_table)
@@ -217,12 +217,6 @@ function __init__()
             @warn "$name is unavailable, functionality will be disabled."
         end
     end
-
-    # Ensure that operations executed by the REPL backend finish before returning.
-    # Displaying values happens on a different task.
-    if isdefined(Base, :active_repl_backend)
-        push!(Base.active_repl_backend.ast_transforms, synchronize_rocm_tasks)
-    end
 end
 
 end

src/array.jl

@@ -1,45 +1,36 @@
 mutable struct ROCArray{T, N, B} <: AbstractGPUArray{T, N}
-    buf::DataRef{B}
+    buf::DataRef{Managed{B}}
     dims::Dims{N}
     offset::Int # Offset is in number of elements (not bytes).
 
     function ROCArray{T, N, B}(
         ::UndefInitializer, dims::Dims{N},
     ) where {T, N, B <: Union{Mem.HIPBuffer, Mem.HostBuffer}}
         @assert isbitstype(T) "ROCArray only supports bits types"
-        buf = B(prod(dims) * sizeof(T); stream=stream())
-        xs = new{T, N, B}(DataRef(_free_buf, buf), dims, 0)
-        finalizer(unsafe_finalize!, xs)
+        data = DataRef(pool_free, pool_alloc(B, prod(dims) * sizeof(T)))
+        xs = new{T, N, B}(data, dims, 0)
+        finalizer(unsafe_free!, xs)
         return xs
     end
 
     function ROCArray{T, N}(
-        buf::DataRef{B}, dims::Dims{N}; offset::Integer = 0,
+        buf::DataRef{Managed{B}}, dims::Dims{N}; offset::Integer = 0,
     ) where {T, N, B <: Union{Mem.HIPBuffer, Mem.HostBuffer}}
         @assert isbitstype(T) "ROCArray only supports bits types"
         xs = new{T, N, B}(buf, dims, offset)
-        finalizer(unsafe_finalize!, xs)
+        finalizer(unsafe_free!, xs)
         return xs
     end
 end
 
-# Passed to `DataRef` to handle freeing.
-function _free_buf(buf, stream_ordered::Bool)
-    context!(buf.ctx) do
-        s = stream_ordered ? AMDGPU.stream() : AMDGPU.default_stream()
-        Mem.free(buf; stream=s)
-    end
-end
-
-unsafe_free!(x::ROCArray) = GPUArrays.unsafe_free!(x.buf, true)
-unsafe_finalize!(x::ROCArray) = GPUArrays.unsafe_free!(x.buf, false) # Use global stream.
+unsafe_free!(x::ROCArray) = GPUArrays.unsafe_free!(x.buf)
 
 """
     device(A::ROCArray) -> HIPDevice
 
 Return the device associated with the array `A`.
 """
-device(A::ROCArray) = A.buf[].device
+device(A::ROCArray) = A.buf[].mem.device
 
 buftype(x::ROCArray) = buftype(typeof(x))
 buftype(::Type{<:ROCArray{<:Any, <:Any, B}}) where B = B # TODO check `@isdefined`?
@@ -73,8 +64,7 @@ AnyROCVecOrMat{T} = Union{AnyROCVector{T}, AnyROCMatrix{T}}
 
 # type and dimensionality specified, accepting dims as tuples of Ints
 function ROCArray{T,N}(::UndefInitializer, dims::Dims{N}) where {T,N}
-    buf = Mem.HIPBuffer(prod(dims) * sizeof(T); stream=stream())
-    ROCArray{T, N}(DataRef(_free_buf, buf), dims)
+    ROCArray{T, N, Mem.HIPBuffer}(undef, dims)
 end
 
 # buffer, type and dimensionality specified
@@ -98,7 +88,8 @@ ROCArray{T}(::UndefInitializer, dims::Vararg{Integer, N}) where {T, N} =
 # from Base arrays
 function ROCArray{T,N}(x::Array{T,N}, dims::Dims{N}) where {T,N}
     r = ROCArray{T,N}(undef, dims)
-    Mem.upload!(r.buf[], pointer(x), sizeof(x); stream=stream())
+    Mem.upload!(convert(Mem.AbstractAMDBuffer, r.buf[]),
+        pointer(x), sizeof(x); stream=stream())
     return r
 end
 
@@ -137,6 +128,8 @@ Base.convert(::Type{T}, x::T) where T <: ROCArray = x
 
 ## memory operations
 
+# TODO rework, to pass pointers, instead of accessing .mem
+
 function Base.copyto!(
     dest::Array{T}, d_offset::Integer,
     source::ROCArray{T}, s_offset::Integer, amount::Integer;
@@ -145,11 +138,11 @@ function Base.copyto!(
     amount == 0 && return dest
     @boundscheck checkbounds(dest, d_offset + amount - 1)
     @boundscheck checkbounds(source, s_offset + amount - 1)
-    strm = stream()
     Mem.download!(
         pointer(dest, d_offset),
-        Mem.view(source.buf[], (source.offset + s_offset - 1) * sizeof(T)),
-        amount * sizeof(T); stream=strm, async)
+        Mem.view(convert(Mem.AbstractAMDBuffer, source.buf[]),
+            (source.offset + s_offset - 1) * sizeof(T)),
+        amount * sizeof(T); stream=stream(), async)
     dest
 end
 
@@ -161,7 +154,8 @@ function Base.copyto!(
     @boundscheck checkbounds(dest, d_offset + amount - 1)
     @boundscheck checkbounds(source, s_offset + amount - 1)
     Mem.upload!(
-        Mem.view(dest.buf[], (dest.offset + d_offset - 1) * sizeof(T)),
+        Mem.view(convert(Mem.AbstractAMDBuffer, dest.buf[]),
+            (dest.offset + d_offset - 1) * sizeof(T)),
         pointer(source, s_offset), amount * sizeof(T); stream=stream())
     dest
 end
@@ -174,8 +168,10 @@ function Base.copyto!(
     @boundscheck checkbounds(dest, d_offset + amount - 1)
     @boundscheck checkbounds(source, s_offset + amount - 1)
     Mem.transfer!(
-        Mem.view(dest.buf[], (dest.offset + d_offset - 1) * sizeof(T)),
-        Mem.view(source.buf[], (source.offset + s_offset - 1) * sizeof(T)),
+        Mem.view(convert(Mem.AbstractAMDBuffer, dest.buf[]),
+            (dest.offset + d_offset - 1) * sizeof(T)),
+        Mem.view(convert(Mem.AbstractAMDBuffer, source.buf[]),
+            (source.offset + s_offset - 1) * sizeof(T)),
         amount * sizeof(T); stream=stream())
     dest
 end
@@ -197,7 +193,7 @@ function Base.unsafe_wrap(
     buf = lock ?
         Mem.HostBuffer(Ptr{Cvoid}(ptr), sz) :
         Mem.HIPBuffer(Ptr{Cvoid}(ptr), sz)
-    ROCArray{T, N}(DataRef(_free_buf, buf), dims)
+    ROCArray{T, N}(DataRef(pool_free, Managed(buf)), dims)
 end
 
 Base.unsafe_wrap(::Type{ROCArray{T}}, ptr::Ptr, dims; kwargs...) where T =
@@ -237,12 +233,8 @@ Adapt.adapt_storage(::Float32Adaptor, xs::AbstractArray{Float16}) =
 
 roc(xs) = adapt(Float32Adaptor(), xs)
 
-function Base.unsafe_convert(::Type{Ptr{T}}, x::ROCArray{T}) where T
-    # TODO have specialized convert function for buffers:
-    # convert(hipPtr, buf) -> dev ptr
-    tmp = typeof(x.buf[]) <: Mem.HIPBuffer ? x.buf[] : x.buf[].dev_ptr
-    Base.unsafe_convert(Ptr{T}, tmp) + x.offset * sizeof(T)
-end
+Base.unsafe_convert(typ::Type{Ptr{T}}, x::ROCArray{T}) where T =
+    convert(typ, x.buf[]) + x.offset * sizeof(T)
 
 # some nice utilities
 
@@ -280,13 +272,14 @@ function Base.resize!(A::ROCVector{T}, n::Integer) where T
 
     copy_size = min(length(A), n) * sizeof(T)
     if copy_size > 0
-        Mem.transfer!(new_buf, A.buf[], copy_size; stream=stream())
+        Mem.transfer!(new_buf, convert(Mem.AbstractAMDBuffer, A.buf[]),
+            copy_size; stream=stream())
     end
 
     # Free old buffer.
     unsafe_free!(A)
 
-    A.buf = DataRef(_free_buf, new_buf)
+    A.buf = DataRef(pool_free, Managed(new_buf))
     A.dims = (n,)
     A.offset = 0
     return A

src/blas/wrappers.jl

@@ -580,18 +580,19 @@ end
 # helper function to get a device array of device pointers
 function device_batch(batch::Array{T}) where T <: ROCArray
     E = eltype(T)
-    ROCArray([Base.unsafe_convert(Ptr{E}, arr.buf[]) for arr in batch])
+    ROCArray([convert(Ptr{E}, arr.buf[]) for arr in batch])
 end
 
 function device_batch(x::AnyROCArray{T, 3}) where T
     shift = size(x, 1) * size(x, 2) * sizeof(T)
+    buf = convert(AMDGPU.Mem.AbstractAMDBuffer, x.buf[])
     ROCArray([
-        Base.unsafe_convert(Ptr{T}, AMDGPU.Mem.view(x.buf[], shift * (i - 1)))
+        convert(Ptr{T}, AMDGPU.Mem.view(buf, shift * (i - 1)))
         for i in 1:size(x, 3)])
 end
 
 function device_batch(x::AnyROCArray{T, 3}, batch_count::Int) where T
-    ptr = Base.unsafe_convert(Ptr{T}, x.buf[])
+    ptr = convert(Ptr{T}, x.buf[])
     ROCArray([ptr for i in 1:batch_count])
 end
 

src/compiler/dynamic_memory.jl

@@ -13,8 +13,7 @@ function create_malloc_hostcall!()
         # Create host pinned memory and store HostCall in it.
         # It will be then accessed by kernels from kernel state.
         buf = Mem.HostBuffer(sizeof(holder.hc), HIP.hipHostAllocDefault)
-        ptr = Base.unsafe_convert(
-            Ptr{Device.HostCall{Ptr{Cvoid}, Tuple{Csize_t}}}, buf)
+        ptr = convert(Ptr{Device.HostCall{Ptr{Cvoid}, Tuple{Csize_t}}}, buf)
         Base.unsafe_store!(ptr, holder.hc)
         return holder, buf
     end
@@ -34,8 +33,7 @@ function create_free_hostcall!()
         end
 
         buf = Mem.HostBuffer(sizeof(holder.hc), HIP.hipHostAllocDefault)
-        ptr = Base.unsafe_convert(
-            Ptr{Device.HostCall{Nothing, Tuple{Ptr{Cvoid}}}}, buf)
+        ptr = convert(Ptr{Device.HostCall{Nothing, Tuple{Ptr{Cvoid}}}}, buf)
         Base.unsafe_store!(ptr, holder.hc)
         return holder, buf
     end

src/compiler/exceptions.jl

@@ -17,7 +17,7 @@ end
 # Check for exceptions on every synchronization.
 function check_exceptions()
     for (dev, buf) in _exception_flags
-        ptr = Base.unsafe_convert(Ptr{Int}, buf)
+        ptr = convert(Ptr{Int}, buf)
         flag = unsafe_load(ptr)
         if flag != 0
             unsafe_store!(ptr, 0)

src/compiler/output_context.jl

@@ -15,7 +15,7 @@ function create_output_context!(#= TODO mod::HIP.HIPModule =#)
 
         # Pointer to HostCall to be read from device.
         buf = Mem.HostBuffer(sizeof(holder.hc), HIP.hipHostAllocDefault)
-        ptr = Base.unsafe_convert(Ptr{Device.OUTPUT_CONTEXT_TYPE}, buf)
+        ptr = convert(Ptr{Device.OUTPUT_CONTEXT_TYPE}, buf)
         Base.unsafe_store!(ptr, holder.hc)
         return holder, buf
     end
@@ -48,8 +48,7 @@ function create_printf_output_context!()
         end
         # Pointer to HostCall to be read from device.
         buf = Mem.HostBuffer(sizeof(holder.hc), HIP.hipHostAllocDefault)
-        ptr = Base.unsafe_convert(
-            Ptr{Device.PRINTF_OUTPUT_CONTEXT_TYPE}, buf)
+        ptr = convert(Ptr{Device.PRINTF_OUTPUT_CONTEXT_TYPE}, buf)
         Base.unsafe_store!(ptr, holder.hc)
         return holder, buf
     end

src/device/gcn/hostcall.jl

@@ -54,7 +54,7 @@ function HostCall(
     buf_len = max(sizeof(UInt64), buf_len) # make room for return buffer pointer
     buf = Mem.HostBuffer(buf_len, AMDGPU.HIP.hipHostAllocDefault)
 
-    buf_ptr = LLVMPtr{UInt8, AS.Global}(Base.unsafe_convert(Ptr{UInt8}, buf))
+    buf_ptr = LLVMPtr{UInt8, AS.Global}(convert(Ptr{UInt8}, buf))
     host_signal_store!(HSA.Signal(signal_handle), READY_SENTINEL)
     HostCall{RT, AT}(signal_handle, buf_ptr, buf_len)
 end
@@ -144,7 +144,7 @@ function HostCallHolder(
 
                     ret_ref = Ref{rettype}(ret)
                     GC.@preserve ret_ref begin
-                        ret_ptr = Base.unsafe_convert(Ptr{Cvoid}, ret_buf)
+                        ret_ptr = convert(Ptr{Cvoid}, ret_buf)
                         if sizeof(ret) > 0
                             src_ptr = reinterpret(Ptr{Cvoid},
                                 Base.unsafe_convert(Ptr{rettype}, ret_ref))

src/exception_handler.jl

@@ -71,13 +71,13 @@ end
 
 function has_exception(dev::HIPDevice)::Bool
     ex = exception_holder(dev)
-    ptr = Base.unsafe_convert(Ptr{Int}, ex.exception_flag)
+    ptr = convert(Ptr{Int}, ex.exception_flag)
     unsafe_load(ptr) != 0
 end
 
 function reset_exception_holder!(dev::HIPDevice)
     ex = exception_holder(dev)
-    ptr = Base.unsafe_convert(Ptr{Int}, ex.exception_flag)
+    ptr = convert(Ptr{Int}, ex.exception_flag)
     unsafe_store!(ptr, 0)
 
     fill!(ex.buffers_counter, 0)

src/gpuarrays.jl

@@ -58,8 +58,9 @@ function Base.convert(
     ::Type{ROCDeviceArray{T, N, AS.Global}}, a::ROCArray{T, N},
 ) where {T, N}
     # If HostBuffer, use device pointer.
-    ptr = Base.unsafe_convert(Ptr{T},
-        typeof(a.buf[]) <: Mem.HIPBuffer ? a.buf[] : a.buf[].dev_ptr)
+    buf = convert(Mem.AbstractAMDBuffer, a.buf[])
+    ptr = convert(Ptr{T}, typeof(buf) <: Mem.HIPBuffer ?
+        buf : buf.dev_ptr)
     llvm_ptr = AMDGPU.LLVMPtr{T,AS.Global}(ptr + a.offset * sizeof(T))
     ROCDeviceArray{T, N, AS.Global}(a.dims, llvm_ptr)
 end

src/hip/stream.jl

@@ -6,6 +6,8 @@ mutable struct HIPStream
     priority::Symbol
     device::HIPDevice
     ctx::HIPContext
+
+    Base.@atomic valid::Bool
 end
 
 """
@@ -24,16 +26,19 @@ function HIPStream(priority::Symbol = :normal)
     stream_ref = Ref{hipStream_t}()
     hipStreamCreateWithPriority(stream_ref, 0, priority_int) |> check
     d = device()
-    stream = HIPStream(stream_ref[], priority, d, HIPContext(d))
+    stream = HIPStream(stream_ref[], priority, d, HIPContext(d), true)
     finalizer(stream) do s
         AMDGPU.context!(s.ctx) do
             hipStreamDestroy(s.stream) |> check
         end
+        Base.@atomic s.valid = false
     end
     return stream
 end
 
-default_stream() = HIPStream(C_NULL, :normal, device(), HIPContext())
+isvalid(s::HIPStream) = s.valid
+
+default_stream() = HIPStream(C_NULL, :normal, device(), HIPContext(), true)
 
 """
     HIPStream(stream::hipStream_t)
@@ -43,7 +48,7 @@ Device is the default device that's currently in use.
 """
 function HIPStream(stream::hipStream_t)
     d = device()
-    HIPStream(stream, priority(stream), d, HIPContext(d))
+    HIPStream(stream, priority(stream), d, HIPContext(d), true)
 end
 
 function isdone(stream::HIPStream)