Recycle the streams of finished tasks

docs/src/usage/multitasking.md

@@ -34,6 +34,8 @@ Because launches are asynchronous, synchronize before reading results back or ti
 AMDGPU.@sync @roc groupsize=256 gridsize=n kernel(args...)
 ```
 
+HIP streams are expensive: each one pins several MiB of host memory. To keep applications that spawn many short-lived tasks from accumulating streams, the stream of a task is recycled once the task has finished and all work on it has completed. Tasks that are running at the same time never share a stream, but a newly started task may get the stream of an earlier one. If you need a stream that outlives the task that uses it, e.g., to pass it on to other tasks, create one explicitly with `AMDGPU.HIPStream()` and activate it with `AMDGPU.stream!`.
+
 Streams also carry a priority (`:normal`, `:low`, `:high`) to bias scheduling. See [Streams](@ref) for stream priorities, synchronization details, and the blocking-vs-nonblocking preference.
 
 ## Using multiple GPUs

src/hip/stream.jl

@@ -8,6 +8,10 @@ mutable struct HIPStream
     ctx::HIPContext
 
     Base.@atomic valid::Bool
+
+    # bumped when the stream is handed to another task (see `task_stream`), which only
+    # happens when it is idle, so work submitted during earlier generations has finished.
+    Base.@atomic generation::Int
 end
 
 """
@@ -26,7 +30,7 @@ function HIPStream(priority::Symbol = :normal)
     stream_ref = Ref{hipStream_t}()
     hipStreamCreateWithPriority(stream_ref, 0, priority_int)
     d = device()
-    stream = HIPStream(stream_ref[], priority, d, HIPContext(d), true)
+    stream = HIPStream(stream_ref[], priority, d, HIPContext(d), true, 0)
     return finalizer(stream) do s
         Base.@atomic s.valid = false
         AMDGPU.context!(s.ctx) do
@@ -35,9 +39,97 @@ function HIPStream(priority::Symbol = :normal)
     end
 end
 
+# Every task gets its own default stream, but HIP streams are expensive: creating one
+# takes milliseconds and pins ~8 MiB of host memory. Since the GC is in no hurry to
+# collect finished tasks (and with them, their streams), code that spawns many
+# short-lived tasks would pile up thousands of streams. Instead, recycle the streams of
+# tasks that have finished, keeping up to `STREAM_POOL_IDLE` unused ones per device and
+# priority.
+const STREAM_POOL_IDLE = 32
+struct PooledStream
+    stream::HIPStream
+    owner::WeakRef
+end
+const STREAM_POOLS = Dict{Tuple{Int,Symbol}, Vector{PooledStream}}()
+const STREAM_POOL_LOCK = ReentrantLock()
+
+function task_stream(priority::Symbol = :normal)
+    # finalizers can't wait for the pool's lock, so give them a stream of their own
+    GC.in_finalizer() && return HIPStream(priority)
+
+    key = (device_id(device()), priority)
+    task = current_task()
+    stream = Base.@lock STREAM_POOL_LOCK begin
+        claim_stream!(get!(Vector{PooledStream}, STREAM_POOLS, key), task)
+    end
+    stream === nothing || return stream
+
+    # creating a stream can be slow, so don't make other tasks wait for it
+    stream = HIPStream(priority)
+    Base.@lock STREAM_POOL_LOCK begin
+        push!(STREAM_POOLS[key], PooledStream(stream, WeakRef(task)))
+    end
+    return stream
+end
+
+function claim_stream!(pool::Vector{PooledStream}, task::Task)
+    candidate = nothing
+    idle = 0
+    i = 1
+    while i <= length(pool)
+        entry = pool[i]
+        owner = entry.owner.value
+        keep = if owner === task
+            # a task that switches back and forth between priorities keeps its streams
+            isvalid(entry.stream) && return entry.stream
+            false
+        elseif owner !== nothing && !istaskdone(owner::Task)
+            true
+        elseif !isvalid(entry.stream)
+            false
+        else
+            status = query(entry.stream)
+            if status == hipErrorNotReady
+                # don't make a new task wait for work that the previous owner left behind
+                true
+            elseif status != hipSuccess
+                # the stream is in an error state
+                false
+            elseif candidate === nothing
+                candidate = entry
+                true
+            else
+                (idle += 1) <= STREAM_POOL_IDLE
+            end
+        end
+        keep ? (i += 1) : deleteat!(pool, i)
+    end
+    candidate === nothing && return nothing
+
+    candidate.owner.value = task
+    generation = Base.@atomic :monotonic candidate.stream.generation
+    Base.@atomic :release candidate.stream.generation = generation + 1
+    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
+    mode = Ref(hipStreamCaptureModeRelaxed)
+    hipThreadExchangeStreamCaptureMode(mode)
+    try
+        return unchecked_hipStreamQuery(s)
+    finally
+        hipThreadExchangeStreamCaptureMode(mode)
+    end
+end
+
+# only bumped while holding `STREAM_POOL_LOCK`, but read without it
+generation(s::HIPStream) = Base.@atomic :acquire s.generation
+
 isvalid(s::HIPStream) = s.valid
 
-default_stream() = HIPStream(C_NULL, :normal, device(), HIPContext(), true)
+default_stream() = HIPStream(C_NULL, :normal, device(), HIPContext(), true, 0)
 
 """
     HIPStream(stream::hipStream_t)
@@ -46,8 +138,18 @@ Create HIPStream from `hipStream_t` handle.
 Device is the default device that's currently in use.
 """
 function HIPStream(stream::hipStream_t)
+    # the streams of tasks get recycled, which only the pool's objects keep track of
+    if !GC.in_finalizer()
+        Base.@lock STREAM_POOL_LOCK begin
+            for pool in values(STREAM_POOLS), entry in pool
+                s = entry.stream
+                s.stream == stream && isvalid(s) && return s
+            end
+        end
+    end
+
     d = device()
-    HIPStream(stream, priority(stream), d, HIPContext(d), true)
+    HIPStream(stream, priority(stream), d, HIPContext(d), true, 0)
 end
 
 function isdone(stream::HIPStream)

src/memory.jl

@@ -409,19 +409,31 @@ end
 mutable struct Managed{M}
     const mem::M
     const lock::ReentrantLock
+    # which stream is currently using the memory, and the generation of that stream
     stream::HIPStream
+    generation::Int
     dirty::Bool
     captured::Bool
 
     function Managed(mem; stream=AMDGPU.stream(), dirty=true, captured=false)
-        new{typeof(mem)}(mem, ReentrantLock(), stream, dirty, captured)
+        new{typeof(mem)}(mem, ReentrantLock(), stream, HIP.generation(stream),
+                         dirty, captured)
     end
 end
 
+# if the stream has been handed to another task since it last used the memory, that work
+# has finished, and waiting for the stream would only wait for the new task's work.
+recycled(m::Managed) = m.generation != HIP.generation(m.stream)
+
 function synchronize(m::Managed)
     Base.@lock m.lock begin
         m.dirty || return
-        synchronize(m.stream)
+        if recycled(m)
+            # the work has finished, but may have raised an exception
+            throw_if_exception(m.stream.device)
+        else
+            synchronize(m.stream)
+        end
         m.dirty = false
         return
     end
@@ -441,14 +453,16 @@ function take_ownership!(managed::Managed; stream::HIPStream=AMDGPU.stream())
         synchronize(managed)
         managed.stream = stream
     end
+    managed.generation = HIP.generation(managed.stream)
 
     managed.dirty = true
     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
+    (managed.stream === stream && managed.dirty &&
+     managed.generation == HIP.generation(stream)) && return
     Base.@lock managed.lock take_ownership!(managed; stream)
     return
 end
@@ -522,7 +536,7 @@ function pool_free(managed::Managed{M}) where M
 
     try
         time = Base.@elapsed Base.@lock managed.lock begin
-            _pool_free(managed.mem, managed.stream)
+            _pool_free(managed)
         end
         Base.@atomic alloc_stats.free_count += 1
         Base.@atomic alloc_stats.free_bytes += sz
@@ -536,9 +550,19 @@ function pool_free(managed::Managed{M}) where M
     return
 end
 
-function _pool_free(buf, stream::HIPStream)
-    if !HIP.isvalid(stream)
-        stream = AMDGPU.default_stream()
+function _pool_free(managed::Managed)
+    buf = managed.mem
+    AMDGPU.context!(() -> Mem.free(buf; stream=free_stream(managed)), buf.ctx)
+end
+
+function free_stream(managed::Managed)
+    if !HIP.isvalid(managed.stream)
+        return AMDGPU.default_stream()
+    elseif recycled(managed) && !GC.in_finalizer()
+        # the stream now belongs to another task, which may be capturing it.
+        # finalizers can keep using it, because capturing disables the GC.
+        return AMDGPU.stream()
+    else
+        return managed.stream
     end
-    AMDGPU.context!(() -> Mem.free(buf; stream), buf.ctx)
 end

src/tls.jl

@@ -96,7 +96,7 @@ end
 function stream(state::TaskLocalState)::HIPStream
     i = device_id(state.device)
     if state.streams[i] ≡ nothing
-        state.streams[i] = HIPStream(:normal)
+        state.streams[i] = HIP.task_stream(:normal)
     else
         state.streams[i]
     end
@@ -107,6 +107,11 @@ end
 
 Get the HIP stream that should be used as the default one
 for the currently executing task.
+
+Each task gets its own stream, which may be handed to another task once the task has
+finished and all work on the stream has completed. If you need a stream that outlives the
+task, or that is never shared, create one with `HIPStream()` and activate it using
+[`AMDGPU.stream!`](@ref).
 """
 stream()::HIPStream = stream(task_local_state!())
 
@@ -180,7 +185,7 @@ function priority!(p::Symbol)
     state = task_local_state!()
     state.stream.priority == p && return p
 
-    state.streams[device_id(state.device)] = HIPStream(p)
+    state.streams[device_id(state.device)] = HIP.task_stream(p)
     return p
 end
 
@@ -201,7 +206,7 @@ function priority!(f::Function, p::Symbol)
 
     old_s = state.stream
     swap = p != old_s.priority
-    swap && (state.streams[idx] = HIPStream(p);)
+    swap && (state.streams[idx] = HIP.task_stream(p);)
 
     return try
         f()