Make CUDA stream from cache match CUDA context (#294)

Co-authored-by: Christopher <christophernhill@gmail.com>

Author: christophernhill

lib/CUDAKernels/src/CUDAKernels.jl

@@ -31,6 +31,7 @@ end
 # - Add a background task that occasionally scans all streams
 # - Add a hysterisis by checking a "since last scanned" timestamp
 const STREAM_GC_LOCK = Threads.ReentrantLock()
+#=
 function next_stream()
     lock(STREAM_GC_LOCK) do
         if !isempty(FREE_STREAMS)
@@ -53,6 +54,43 @@ function next_stream()
         return stream
     end
 end
+=#
+const FREE_STREAMS_D = Dict{CUDA.CuContext,Array{CUDA.CuStream,1}}()
+const STREAMS_D      = Dict{CUDA.CuContext,Array{CUDA.CuStream,1}}()
+function next_stream()
+    ctx = CUDA.current_context()
+    lock(STREAM_GC_LOCK) do
+        # see if there is a compatible free stream
+        FREE_STREAMS_CT  = get!(FREE_STREAMS_D, ctx) do
+           CUDA.CuStream[]
+        end
+        if !isempty(FREE_STREAMS_CT)
+           return pop!(FREE_STREAMS_CT)
+        end
+
+        # GC to recover streams that are not busy
+        STREAMS_CT  = get!(STREAMS_D, ctx) do
+            CUDA.CuStream[]
+        end
+        if length(STREAMS_CT) > STREAM_GC_THRESHOLD[]
+            for stream in STREAMS_CT
+                if CUDA.query(stream)
+                    push!(FREE_STREAMS_CT, stream)
+                end
+            end
+        end
+
+        # if there is a compatible free stream after GC, return that stream
+        if !isempty(FREE_STREAMS_CT)
+            return pop!(FREE_STREAMS_CT)
+        end
+
+        # no compatible free stream available so create a new one
+        stream = CUDA.CuStream(flags = CUDA.STREAM_NON_BLOCKING)
+        push!(STREAMS_CT, stream)
+        return stream
+    end
+end
 
 import KernelAbstractions: Event, CPUEvent, NoneEvent, MultiEvent, CPU, GPU, isdone, failed