Add streaming API

Author: jpsamaroo

src/Dagger.jl

@@ -41,6 +41,9 @@ include("utils/locked-object.jl")
 include("utils/caching.jl")
 include("sch/Sch.jl"); using .Sch
 
+# Streaming
+include("stream.jl")
+
 # Array computations
 include("array/darray.jl")
 include("array/alloc.jl")

src/stream.jl

@@ -0,0 +1,185 @@
+struct StreamingTaskQueue <: AbstractTaskQueue
+    tasks::Vector{Pair{EagerTaskSpec,EagerThunk}}
+    StreamingTaskQueue() = new(Pair{EagerTaskSpec,EagerThunk}[])
+end
+
+function enqueue!(queue::StreamingTaskQueue, spec::Pair{EagerTaskSpec,EagerThunk})
+    push!(queue.tasks, spec)
+end
+function enqueue!(queue::StreamingTaskQueue, specs::Vector{Pair{EagerTaskSpec,EagerThunk}})
+    append!(queue.tasks, specs)
+end
+
+function spawn_streaming(f::Base.Callable)
+    queue = StreamingTaskQueue()
+    result = with_options(;task_queue=queue) do
+        spawn_bulk(f)
+    end
+    if length(queue.tasks) > 0
+        setup_streaming!(queue.tasks)
+        enqueue!(queue.tasks)
+    end
+    return result
+end
+
+mutable struct StreamElement
+    value
+    remaining::Int
+end
+function consume!(elem::StreamElement)
+    elem.remaining -= 1
+    return (elem.value, elem.remaining)
+end
+mutable struct StreamStore
+    waiters::Int
+    buffer::Vector{StreamElement}
+    open::Bool
+    lock::Threads.Condition
+    StreamStore(waiters) = new(waiters, StreamElement[], true, Threads.Condition())
+end
+function Base.put!(store::StreamStore, @nospecialize(value))
+    elem = StreamElement(value, store.waiters)
+    @lock store.lock begin
+        @debug "[$(myid())] adding $value"
+        push!(store.buffer, elem)
+        notify(store.lock)
+    end
+end
+function Base.take!(store::StreamStore)
+    @lock store.lock begin
+        while true
+            @debug "[$(myid())] take loop"
+            if length(store.buffer) == 0
+                wait(store.lock)
+            end
+            @debug "[$(myid())] wait finished"
+            if !isopen(store)
+                @debug "[$(myid())] closed!"
+                throw(InvalidStateException("Stream is closed", :closed))
+            end
+            value, remaining = consume!(first(store.buffer))
+            @debug "[$(myid())] took $value ($remaining remaining)"
+            if remaining < 0
+                throw(ConcurrencyViolationError("StoreElement consumed too many times"))
+            end
+            if remaining == 0
+                @debug "[$(myid())] removing from buffer"
+                popfirst!(store.buffer)
+            end
+            if remaining >= 0
+                @debug "[$(myid())] value accepted"
+                return value
+            end
+            @debug "[$(myid())] value ignored"
+        end
+    end
+end
+Base.isopen(store::StreamStore) = store.open
+function Base.close(store::StreamStore)
+    store.open = false
+end
+
+mutable struct Stream{T} <: AbstractChannel{T}
+    ref::Union{Chunk,Nothing}
+    waiters::Int
+end
+Stream() = Stream{Any}(nothing, 0)
+function initialize!(stream::Stream)
+    @assert stream.ref === nothing "Stream already initialized"
+    stream.ref = tochunk(StreamStore(stream.waiters))
+end
+function Base.put!(stream::Stream, @nospecialize(value))
+    remotecall_wait(stream.ref.handle.owner, stream.ref.handle, value) do ref, value
+        @nospecialize value
+        store = MemPool.poolget(ref)::StreamStore
+        put!(store, value)
+    end
+end
+function Base.take!(stream::Stream{T}) where T
+    return remotecall_fetch(stream.ref.handle.owner, stream.ref.handle) do ref
+        store = MemPool.poolget(ref)::StreamStore
+        return take!(store)::T
+    end
+end
+function Base.isopen(stream::Stream)::Bool
+    return remotecall_fetch(stream.ref.handle.owner, stream.ref.handle) do ref
+        return isopen(MemPool.poolget(ref)::StreamStore)
+    end
+end
+function Base.close(stream::Stream)
+    remotecall_wait(stream.ref.handle.owner, stream.ref.handle) do ref
+        close(MemPool.poolget(ref)::StreamStore)
+    end
+end
+
+struct FinishedStreaming end
+finish_streaming() = FinishedStreaming()
+
+struct StreamingFunction{F, T}
+    f::F
+    stream::Stream{T}
+end
+function (sf::StreamingFunction)(args...; kwargs...)
+    @nospecialize sf args kwargs
+    result = nothing
+    stream_args = Base.mapany(identity, args)
+    stream_kwargs = Base.mapany(identity, kwargs)
+    try
+        while true
+            # Get values from Stream args/kwargs
+            for (idx, arg) in enumerate(args)
+                if arg isa Stream
+                    stream_args[idx] = take!(arg)
+                end
+            end
+            for (idx, (pos, arg)) in enumerate(kwargs)
+                if arg isa Stream
+                    stream_kwargs[idx] = pos => take!(arg)
+                end
+            end
+
+            # Run a single cycle of f
+            stream_result = sf.f(stream_args...; stream_kwargs...)
+
+            # Exit streaming on graceful request
+            if stream_result === FinishedStreaming()
+                return
+            end
+
+            # Put the result into the output stream
+            put!(sf.stream, stream_result)
+        end
+    finally
+        # Ensure downstream consumers also terminate
+        close(sf.stream)
+    end
+end
+
+function setup_streaming!(tasks::Vector{Pair{EagerTaskSpec,EagerThunk}})
+    # Adapt called function for streaming and generate output Streams
+    self_streams = Dict{UInt,Stream}()
+    for (spec, task) in tasks
+        # FIXME: Infer type
+        stream = Stream()
+        self_streams[task.uid] = stream
+        spec.f = StreamingFunction(spec.f, stream)
+        # FIXME: Generalize to other processors
+        spec.options = merge(spec.options, (;occupancy=Dict(ThreadProc=>0)))
+    end
+
+    # Adapt args to accept Stream output of other streaming tasks
+    for (spec, task) in tasks
+        for (idx, (pos, arg)) in reverse(collect(enumerate(spec.args)))
+            if arg isa EagerThunk && haskey(self_streams, arg.uid)
+                other_stream = self_streams[arg.uid]
+                spec.args[idx] = pos => other_stream
+                other_stream.waiters += 1
+            end
+        end
+    end
+
+    # Initialize all Streams
+    for stream in values(self_streams)
+        initialize!(stream)
+    end
+end