spawn: Add configurable task queueing

To allow Dagger to fulfill more use cases, it would be helpful if we
could perform some basic transformations on tasks before they get
submitted to the scheduler. For example, we would like to enable
in-order execution semantics for regions of code that execute GPU
kernels, as this matches GPU synchronization semantics, and thus makes
it easier to upgrade GPU code to use Dagger.

Separately, to enable reliable DAG optimizations, we need to be able to
guarantee that a region of user code can be seen as a whole within the
scheduler. The one-at-a-time task submission semantics that we currently
have are insufficient to achieve this, as some tasks in the DAG region
of interest may already have been launched before the optimization can
see enough of the DAG to be useful.

To support these and other use cases, this commit adds a flexible
pre-submit task queueing system, as well as making it possible to add
additional tasks as synchronization dependencies (instead of the default
set from task arguments).

The task queueing system allows a custom task queue to be set in TLS,
which will be used by `@spawn`/`spawn` when submitting tasks. The task
queue is provided one or more task specifications and `EagerThunk`
handles, and is free to delay and/or batch task submission, as well as
to modify the task specification arbitrarily to match the desired
semantics. Task queues are nestable, and tasks submitted within sets of
nested task queues should inherit the semantics of the queue they are
contained within most directly (with further transformations occuring as
tasks move upwards within the nest of queues). The most upstream task
queue submits tasks to the worker 1 eager scheduler, but this is also
expected to be flexible to allow unique task submission semantics.

To support the goal of predictable optimizations, a `LazyTaskQueue` is
added (available via `spawn_bulk`) which batches up multiple task
submissions into just one, and locks the scheduler until all tasks have
been submitted, allowing the scheduler to see the entire DAG structure
all at once. Nesting of `spawn_bulk` queues allows multiple DAG regions
to be combined into a single total region which is submitted all at
once.

The ability to specify additional task synchronization dependencies is also a
key piece that is orthogonal to task queues. This feature enables the
goal of in-order execution semantics by enabling the creation of an
`InOrderTaskQueue` (available via `spawn_sequential`), which tracks the
last-submitted task or set of tasks, and adds those tasks as additional
synchronization dependencies to the next submitted task or set of tasks,
effectively causing serializing behavior. Nesting of `spawn_sequential`
queues allows separate sequential chains of tasks to be specified, with
deeper-nested chains sequencing after previously-submitted tasks or
chains in shallower-nested queues.

Interestingly, the nesting of `spawn_bulk` within `spawn_sequential`
allows entire DAG regions to explicitly synchronize against each other
(such that one region executes before another), while allowing tasks
within each region to still expose parallelism. The inverse nesting of
`spawn_sequential` within `spawn_bulk` allows a chain of sequential
tasks to be submitted all at once, adding interesting optimization
opportunities.

Alongside these enhancements, the eager submission pipeline is optimized
by removing the `eager_thunk` submission pathway (which submitted all
tasks into a `Channel`), and allows tasks to be directly submitted into
the scheduler without redirection. It is expected that this will improve
task submission performance and reduce memory usage.

Author: jpsamaroo

docs/make.jl

@@ -17,12 +17,13 @@ makedocs(;
         "Data Management" => "data-management.md",
         "Scopes" => "scopes.md",
         "Processors" => "processors.md",
-        "Dynamic Scheduler Control" => "dynamic.md",
+        "Task Queues" => "task-queues.md",
         "Option Propagation" => "propagation.md",
         "Logging and Graphing" => "logging.md",
         "Checkpointing" => "checkpointing.md",
         "Scheduler Visualization" => "scheduler-visualization.md",
         "Benchmarking" => "benchmarking.md",
+        "Dynamic Scheduler Control" => "dynamic.md",
         "Scheduler Internals" => "scheduler-internals.md",
         "Dagger API" => [
             "Types" => "api-dagger/types.md",

docs/src/task-queues.md

@@ -0,0 +1,74 @@
+# Task Queues
+
+By default, `@spawn`/`spawn` submit tasks immediately and directly into
+Dagger's scheduler without modifications. However, sometimes you want to be
+able to tweak this behavior for a region of code; for example, when working
+with GPUs or other operations which operate in-place, you might want to emulate
+CUDA's stream semantics by ensuring that tasks execute sequentially (to avoid
+one kernel reading from an array while another kernel is actively writing to
+it). Or, you might want to ensure that a set of Dagger tasks are submitted into
+the scheduler all at once for benchmarking purposes or to emulate the behavior
+of `delayed`. This and more is possible through a mechanism called "task
+queues".
+
+A task queue in Dagger is an object that can be configured to accept unlaunched
+tasks from `@spawn`/`spawn` and either modify them or delay their launching
+arbitrarily. By default, Dagger tasks are enqueued through the
+`EagerTaskQueue`, which submits tasks directly into the scheduler before
+`@spawn`/`spawn` returns. However, Dagger also has an `InOrderTaskQueue`, which
+ensures that tasks enqueued through it execute sequentially with respect to
+each other. This queue can be allocated with `Dagger.spawn_sequential`:
+
+```julia
+A = rand(16)
+B = zeros(16)
+C = zeros(16)
+function vcopy!(B, A)
+    B .= A .+ 1.0
+    return
+end
+function vadd!(C, A, B)
+    C .+= A .+ B
+    return
+end
+wait(Dagger.spawn_sequential() do
+    Dagger.@spawn vcopy!(B, A)
+    Dagger.@spawn vadd!(C, A, B)
+end)
+```
+
+In the above example, `vadd!` is guaranteed to wait until `vcopy!` is
+completed, even though `vadd!` isn't taking the result of `vcopy!` as an
+argument (which is how tasks are normally ordered).
+
+What if we wanted to launch multiple `vcopy!` calls within a `spawn_sequential`
+region and allow them to execute in parallel, but still ensure that the `vadd!`
+happens after they all finish? In this case, we want to switch to another kind
+of task queue: the `LazyTaskQueue`. This task queue batches up task submissions
+into groups, so that all tasks enqueued with it are placed in the scheduler all
+at once. But what would happen if we used this task queue (via `spawn_bulk`)
+within a region using `spawn_sequential`:
+
+```julia
+A = rand(16)
+B1 = zeros(16)
+B2 = zeros(16)
+C = zeros(16)
+wait(Dagger.spawn_sequential() do
+    Dagger.spawn_bulk() do
+        Dagger.@spawn vcopy!(B1, A)
+        Dagger.@spawn vcopy!(B2, A)
+    end
+    Dagger.@spawn vadd!(C, B1, B2)
+end)
+```
+
+Conveniently, Dagger's task queues can be nested to get the expected behavior;
+the above example will submit the two `vcopy!` tasks as a group (and they can
+execute concurrently), while still ensuring that those two tasks finish before
+the `vadd!` task executes.
+
+!!! warn
+    Task queues do not propagate to nested tasks; if a Dagger task launches
+    another task internally, the child task doesn't inherit the task queue that
+    the parent task was enqueued in.

src/Dagger.jl

@@ -19,13 +19,16 @@ using MacroTools
 using TimespanLogging
 
 include("lib/util.jl")
+include("utils/dagdebug.jl")
 
 # Distributed data
 include("options.jl")
 include("processor.jl")
 include("scopes.jl")
+include("eager_thunk.jl")
+include("queue.jl")
 include("thunk.jl")
-include("utils/dagdebug.jl")
+include("submission.jl")
 include("chunks.jl")
 
 # Task scheduling

src/chunks.jl

@@ -62,6 +62,8 @@ shouldpersist(p::Chunk) = t.persist
 processor(c::Chunk) = c.processor
 affinity(c::Chunk) = affinity(c.handle)
 
+is_task_or_chunk(c::Chunk) = true
+
 Base.:(==)(c1::Chunk, c2::Chunk) = c1.handle == c2.handle
 Base.hash(c::Chunk, x::UInt64) = hash(c.handle, x)
 
@@ -272,6 +274,7 @@ function unwrap_weak_checked(c::WeakChunk)
     @assert c !== nothing
     return c
 end
+is_task_or_chunk(c::WeakChunk) = true
 
 Base.@deprecate_binding AbstractPart Union{Chunk, Thunk}
 Base.@deprecate_binding Part Chunk

src/compute.jl

@@ -85,7 +85,7 @@ function dependents(node::Thunk)
         if !haskey(deps, next)
             deps[next] = Set{Thunk}()
         end
-        for (_, inp) in next.inputs
+        for inp in next.syncdeps
             if istask(inp) || (inp isa Chunk)
                 s = get!(()->Set{Thunk}(), deps, inp)
                 push!(s, next)
@@ -96,7 +96,7 @@ function dependents(node::Thunk)
         end
         push!(visited, next)
     end
-    deps
+    return deps
 end
 
 """
@@ -126,7 +126,7 @@ function noffspring(dpents::Dict{Union{Thunk,Chunk}, Set{Thunk}})
         has_all || continue
         noff[next] = off
     end
-    noff
+    return noff
 end
 
 """
@@ -153,7 +153,7 @@ function order(node::Thunk, ndeps)
         haskey(output, next) && continue
         s += 1
         output[next] = s
-        parents = filter(istask, map(last, next.inputs))
+        parents = collect(filter(istask, next.syncdeps))
         if !isempty(parents)
             # If parents is empty, sort! should be a no-op, but raises an ambiguity error
             # when InlineStrings.jl is loaded (at least, version 1.1.0), because InlineStrings

src/eager_thunk.jl

@@ -0,0 +1,87 @@
+"A future holding the result of a `Thunk`."
+struct ThunkFuture
+    future::Future
+end
+ThunkFuture(x::Integer) = ThunkFuture(Future(x))
+ThunkFuture() = ThunkFuture(Future())
+Base.isready(t::ThunkFuture) = isready(t.future)
+Base.wait(t::ThunkFuture) = Dagger.Sch.thunk_yield() do
+    wait(t.future)
+end
+function Base.fetch(t::ThunkFuture; proc=OSProc(), raw=false)
+    error, value = Dagger.Sch.thunk_yield() do
+        fetch(t.future)
+    end
+    if error
+        throw(value)
+    end
+    if raw
+        return value
+    else
+        return move(proc, value)
+    end
+end
+Base.put!(t::ThunkFuture, x; error=false) = put!(t.future, (error, x))
+
+struct Options
+    options::NamedTuple
+end
+Options(;options...) = Options((;options...))
+Options(options...) = Options((;options...))
+
+"""
+    EagerThunk
+
+Returned from `spawn`/`@spawn` calls. Represents a task that is in the
+scheduler, potentially ready to execute, executing, or finished executing. May
+be `fetch`'d or `wait`'d on at any time.
+"""
+mutable struct EagerThunk
+    uid::UInt
+    future::ThunkFuture
+    finalizer_ref::DRef
+    thunk_ref::DRef
+    EagerThunk(uid, future, finalizer_ref) = new(uid, future, finalizer_ref)
+end
+
+Base.isready(t::EagerThunk) = isready(t.future)
+function Base.wait(t::EagerThunk)
+    if !isdefined(t, :thunk_ref)
+        throw(ConcurrencyViolationError("Cannot `wait` on an unlaunched `EagerThunk`"))
+    end
+    wait(t.future)
+end
+function Base.fetch(t::EagerThunk; raw=false)
+    if !isdefined(t, :thunk_ref)
+        throw(ConcurrencyViolationError("Cannot `fetch` an unlaunched `EagerThunk`"))
+    end
+    return fetch(t.future; raw)
+end
+function Base.show(io::IO, t::EagerThunk)
+    status = if isdefined(t, :thunk_ref)
+        isready(t) ? "finished" : "running"
+    else
+        "not launched"
+    end
+    print(io, "EagerThunk ($status)")
+end
+istask(t::EagerThunk) = true
+
+"When finalized, cleans-up the associated `EagerThunk`."
+mutable struct EagerThunkFinalizer
+    uid::UInt
+    function EagerThunkFinalizer(uid)
+        x = new(uid)
+        finalizer(Sch.eager_cleanup, x)
+        x
+    end
+end
+
+const EAGER_ID_COUNTER = Threads.Atomic{UInt64}(1)
+function eager_next_id()
+    if myid() == 1
+        Threads.atomic_add!(EAGER_ID_COUNTER, one(UInt64))
+    else
+        remotecall_fetch(eager_next_id, 1)
+    end
+end

src/queue.jl

@@ -0,0 +1,79 @@
+mutable struct EagerTaskSpec
+    f
+    args::Vector{Pair{Union{Symbol,Nothing},Any}}
+    options::NamedTuple
+end
+
+abstract type AbstractTaskQueue end
+
+function enqueue! end
+
+struct EagerTaskQueue <: AbstractTaskQueue end
+enqueue!(::EagerTaskQueue, spec::Pair{EagerTaskSpec,EagerThunk}) =
+    eager_launch!(spec)
+enqueue!(::EagerTaskQueue, specs::Vector{Pair{EagerTaskSpec,EagerThunk}}) =
+    eager_launch!(specs)
+
+enqueue!(spec::Pair{EagerTaskSpec,EagerThunk}) =
+    enqueue!(get_options(:task_queue, EagerTaskQueue()), spec)
+enqueue!(specs::Vector{Pair{EagerTaskSpec,EagerThunk}}) =
+    enqueue!(get_options(:task_queue, EagerTaskQueue()), specs)
+
+struct LazyTaskQueue <: AbstractTaskQueue
+    tasks::Vector{Pair{EagerTaskSpec,EagerThunk}}
+    LazyTaskQueue() = new(Pair{EagerTaskSpec,EagerThunk}[])
+end
+function enqueue!(queue::LazyTaskQueue, spec::Pair{EagerTaskSpec,EagerThunk})
+    push!(queue.tasks, spec)
+end
+function enqueue!(queue::LazyTaskQueue, specs::Vector{Pair{EagerTaskSpec,EagerThunk}})
+    append!(queue.tasks, specs)
+end
+function spawn_bulk(f::Base.Callable)
+    queue = LazyTaskQueue()
+    result = with_options(f; task_queue=queue)
+    if length(queue.tasks) > 0
+        enqueue!(queue.tasks)
+    end
+    return result
+end
+
+struct InOrderTaskQueue <: AbstractTaskQueue
+    upper_queue::AbstractTaskQueue
+    prev_tasks::Set{EagerThunk}
+    InOrderTaskQueue(upper_queue) = new(upper_queue,
+                                        Set{EagerThunk}())
+end
+function _add_prev_deps!(queue::InOrderTaskQueue, spec::EagerTaskSpec)
+    # Add previously-enqueued task(s) to this task's syncdeps
+    opts = spec.options
+    syncdeps = get(Set{Any}, opts, :syncdeps)
+    for task in queue.prev_tasks
+        push!(syncdeps, task)
+    end
+    spec.options = merge(opts, (;syncdeps,))
+end
+function enqueue!(queue::InOrderTaskQueue, spec::Pair{EagerTaskSpec,EagerThunk})
+    if length(queue.prev_tasks) > 0
+        _add_prev_deps!(queue, first(spec))
+        empty!(queue.prev_tasks)
+    end
+    push!(queue.prev_tasks, last(spec))
+    enqueue!(queue.upper_queue, spec)
+end
+function enqueue!(queue::InOrderTaskQueue, specs::Vector{Pair{EagerTaskSpec,EagerThunk}})
+    if length(queue.prev_tasks) > 0
+        for (spec, task) in specs
+            _add_prev_deps!(queue, spec)
+        end
+        empty!(queue.prev_tasks)
+    end
+    for (spec, task) in specs
+        push!(queue.prev_tasks, task)
+    end
+    enqueue!(queue.upper_queue, specs)
+end
+function spawn_sequential(f::Base.Callable)
+    queue = InOrderTaskQueue(get_options(:task_queue, EagerTaskQueue()))
+    return with_options(f; task_queue=queue)
+end

src/sch/Sch.jl

@@ -11,6 +11,7 @@ import Base: @invokelatest
 import ..Dagger
 import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, ThunkFailedException, Chunk, WeakChunk, OSProc, AnyScope, DefaultScope, LockedObject
 import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak_checked, affinity, tochunk, timespan_start, timespan_finish, procs, move, chunktype, processor, default_enabled, get_processors, get_parent, execute!, rmprocs!, addprocs!, thunk_processor, constrain, cputhreadtime
+import ..Dagger: @dagdebug
 import DataStructures: PriorityQueue, enqueue!, dequeue_pair!, peek
 
 import ..Dagger
@@ -130,7 +131,7 @@ function start_state(deps::Dict, node_order, chan)
 
     for k in sort(collect(keys(deps)), by=node_order)
         if istask(k)
-            waiting = Set{Thunk}(Iterators.filter(istask, map(last, inputs(k))))
+            waiting = Set{Thunk}(Iterators.filter(istask, k.syncdeps))
             if isempty(waiting)
                 push!(state.ready, k)
             else
@@ -883,7 +884,7 @@ function finish_task!(ctx, state, node, thunk_failed)
     schedule_dependents!(state, node, thunk_failed)
     fill_registered_futures!(state, node, thunk_failed)
 
-    to_evict = cleanup_inputs!(state, node)
+    to_evict = cleanup_syncdeps!(state, node)
     if node.f isa Chunk
         # FIXME: Check the graph for matching chunks
         push!(to_evict, node.f)