Datadeps: Build DAG specs faster

src/datadeps/queue.jl

@@ -138,6 +138,122 @@ function _spec_has_region_dtask_arg(spec::DTaskSpec, region_uids::Set{UInt})
     return false
 end
 
+"""
+    prefetch_chunk_aliasing!(pairs::Vector{DTaskPair})
+
+Resolve the aliasing of every remote `Chunk` argument in `pairs` with one
+`remotecall_fetch` per owning worker, issued concurrently, and seed the region's
+aliasing memo with the answers.
+
+Without this, the flat path asks each owner one chunk at a time, the first time
+`dag_add_task!` or the launch loop (`aliasing!`) needs a chunk: the memo only
+deduplicates a chunk within the region, so a 512-task matmul over 128 remote
+tiles paid ~128 sequential round trips, whatever the scheduler.
+`build_aliasing_parallel` already batches this way for the hierarchical path.
+
+Only each chunk's `DRef` travels, and only its aliasing comes back: sending the
+`Chunk` (and getting it back in a `Pair`, as `_compute_aliasing_batch` does)
+made the exchange ~3x slower, since every `Chunk` carries its type, domain,
+processor and scope, and each `DRef` crossing the wire is refcounted again.
+
+Skips arguments already in the memo, `DTask`s that have not finished (in-region
+producers, which fetching would block on), and uniform execution (MPI), whose
+`aliasing` is a collective that must stay in submission order on every rank.
+"""
+function prefetch_chunk_aliasing!(pairs::Vector{DTaskPair})
+    memo = CHUNK_AINFO_MEMO[]
+    memo === nothing && return
+    uniform_execution(current_acceleration()) && return
+    candidates = Tuple{Int,DRef,Any,UInt}[]  # (owner, handle, dep_mod, memo key)
+    seen = Set{UInt}()
+    for pair in pairs
+        _collect_spec_remote_chunks!(candidates, seen, pair.spec)
+    end
+    isempty(candidates) && return
+    batches = Dict{Int,ChunkAliasingBatch}()
+    @lock memo.lock for (wid, handle, dep_mod, key) in candidates
+        haskey(memo.entries, key) && continue
+        _batch_push!(get!(ChunkAliasingBatch, batches, wid), handle, dep_mod, key)
+    end
+    isempty(batches) && return
+    wids = collect(keys(batches))
+    results = Vector{Any}(undef, length(wids))
+    @sync for (j, wid) in enumerate(wids)
+        b = batches[wid]
+        Threads.@spawn results[j] = remotecall_fetch(_compute_chunk_aliasing_batch, wid,
+                                                     b.handles, b.which, b.dep_mods)
+    end
+    # Seed on the region's own task: the spawned tasks above do not see its
+    # task-local acceleration, which is part of every memo key.
+    for (j, wid) in enumerate(wids)
+        ks = batches[wid].keys
+        rs = results[j]
+        @assert length(rs) == length(ks) "prefetch_chunk_aliasing!: $(length(rs)) results for $(length(ks)) args (wid=$wid)"
+        for i in eachindex(ks)
+            memoize_ainfo!(ks[i], rs[i])
+        end
+    end
+    return
+end
+
+# One owner's share of a prefetch. Each `DRef` is listed once, even when the
+# region needs it under several `dep_mod`s (`which` points into `handles`):
+# MemPool's `DRef` serializer ignores `serialize_cycle_header`'s answer, so the
+# same `DRef` twice in one message corrupts it, and the owner errors or hangs.
+struct ChunkAliasingBatch
+    handles::Vector{DRef}
+    handle_idx::Dict{DRef,Int}
+    which::Vector{Int}
+    dep_mods::Vector{Any}
+    keys::Vector{UInt}
+end
+ChunkAliasingBatch() = ChunkAliasingBatch(DRef[], Dict{DRef,Int}(), Int[], Any[], UInt[])
+function _batch_push!(b::ChunkAliasingBatch, handle::DRef, dep_mod, key::UInt)
+    i = get!(b.handle_idx, handle) do
+        push!(b.handles, handle)
+        return length(b.handles)
+    end
+    push!(b.which, i)
+    push!(b.dep_mods, dep_mod)
+    push!(b.keys, key)
+    return
+end
+
+# Runs on the owner of every handle: what `aliasing(::Chunk, dep_mod)` computes
+# there, for a whole batch. `map` rather than an `AbstractAliasing[...]`: when
+# every answer has the same type, as for the tiles of an array, the result is
+# concretely typed and goes back as one block instead of element by element.
+function _compute_chunk_aliasing_batch(handles::Vector{DRef}, which::Vector{Int}, dep_mods::Vector{Any})
+    values = map(unwrap, handles)
+    return map(i -> aliasing_unwrapped(values[which[i]], dep_mods[i]), eachindex(which))
+end
+
+# A function barrier: `pair.spec` is abstractly typed, and reading `fargs` off
+# it dynamically cost more than the rest of collecting the candidates.
+_collect_spec_remote_chunks!(candidates, seen, spec::DTaskSpec) =
+    _foreach_arg(_arg -> _collect_remote_chunk_args!(candidates, seen, value(_arg)), spec.fargs)
+
+# A function barrier, like `_push_argspecs!`: `value(_arg)` is not always inferrable.
+function _collect_remote_chunk_args!(candidates, seen, v)
+    arg, deps = unwrap_inout(v)
+    if arg isa DTask
+        isready(arg) || return
+        arg = fetch(arg; raw=true)
+    end
+    arg isa Chunk || return
+    # Any other handle is left to `aliasing(::Chunk, dep_mod)`, as before.
+    arg.handle isa DRef || return
+    wid = root_worker_id(arg.processor)
+    wid == myid() && return
+    for (dep_mod, _, _) in deps
+        key = ainfo_memo_key(arg, dep_mod)
+        key in seen && continue
+        push!(seen, key)
+        push!(candidates, (wid, arg.handle, dep_mod, key))
+    end
+    return
+end
+
 """
     datadeps_build_schedule!(scheduler, pairs, all_procs, all_scope;
                              region_uids=nothing, cache=true) -> (dag_spec, schedule)
@@ -197,11 +313,10 @@ function datadeps_build_schedule!(scheduler::DataDepsScheduler,
         found, stale = _schedule_cache_find(scheduler, schedule_cache, dag_spec)
         if found !== nothing
             @dagdebug nothing :spawn_datadeps "Found matching DAG spec!"
+            # `dag_add_task!` recorded each id's task; looking it up there
+            # rather than scanning `pairs` keeps a cache hit linear in tasks.
             for (id, proc) in found.id_to_proc
-                uid = dag_spec.id_to_uid[id]
-                task_idx = findfirst(spec_task -> spec_task.task.uid == uid, pairs)
-                task = pairs[task_idx].task
-                schedule[task] = proc
+                schedule[dag_spec.id_to_task[id]] = proc
             end
             cache_hit = true
         end
@@ -332,6 +447,10 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
 
     upper_queue = get_options(:task_queue)
 
+    # Both planning and launching need every argument's aliasing; ask for the
+    # remote chunks' in one batch per owner instead of one round trip each.
+    prefetch_chunk_aliasing!(queue.seen_tasks)
+
     # Plan the whole region ahead of time, if the scheduler supports it.
     # Schedulers that don't define `datadeps_schedule_dag_aot!` leave `schedule`
     # empty and every task falls through to JIT placement below. A planner is

src/datadeps/scheduling.jl

@@ -7,48 +7,16 @@ Base.isempty(spec::DAGSpec) = length(spec) == 0
 
 function dag_add_task!(dspec::DAGSpec, tspec::DTaskSpec, task::DTask)
     # Check if this task depends on any other tasks within the DAG,
-    # which we are not yet ready to handle
-    for (idx, _arg) in enumerate(tspec.fargs)
-        arg, deps = unwrap_inout(value(_arg))
-        for (dep_mod, readdep, writedep) in deps
-            if arg isa DTask
-                if arg.uid in keys(dspec.uid_to_id)
-                    # Within-DAG dependency, bail out
-                    return false
-                end
-            end
-        end
-    end
+    # which we are not yet ready to handle (a within-DAG dependency: bail out)
+    _any_arg(_arg -> _depends_on_dag_task(dspec, value(_arg)), tspec.fargs) && return false
 
     add_vertex!(dspec.g)
     id = nv(dspec.g)
 
     # Record function signature
     dspec.id_to_functype[id] = chunktype(tspec.fargs[1])
-    argtypes = DatadepsArgSpec[]
-    for (idx, _arg) in enumerate(tspec.fargs)
-        arg, deps = unwrap_inout(value(_arg))
-        pos = raw_position(_arg)
-        for (dep_mod, readdep, writedep) in deps
-            if arg isa DTask
-                #= TODO: Re-enable this when we can handle within-DAG dependencies
-                if arg.uid in keys(dspec.uid_to_id)
-                    # Within-DAG dependency
-                    arg_id = dspec.uid_to_id[arg.uid]
-                    push!(dspec.id_to_argtypes[arg_id], DatadepsArgSpec(pos, DTaskDAGID{arg_id}, dep_mod, UnknownAliasing()))
-                    add_edge!(dspec.g, arg_id, id)
-                    continue
-                end
-                =#
-
-                # External DTask, so fetch this and track it as a raw value
-                arg = fetch(arg; raw=true)
-            end
-            ainfo = aliasing(arg, dep_mod)
-            # FIXME: Generate syncdeps and add edges
-            push!(argtypes, DatadepsArgSpec(pos, typeof(arg), dep_mod, ainfo))
-        end
-    end
+    argtypes = sizehint!(DatadepsArgSpec[], length(tspec.fargs))
+    _foreach_arg(_arg -> _push_argspecs!(argtypes, raw_position(_arg), value(_arg)), tspec.fargs)
     dspec.id_to_argtypes[id] = argtypes
     dspec.id_to_scope[id] = @something(tspec.options.compute_scope,
                                        tspec.options.scope,
@@ -62,6 +30,55 @@ function dag_add_task!(dspec::DAGSpec, tspec::DTaskSpec, task::DTask)
 
     return true
 end
+
+# Visit a task's arguments. A typed spec's arguments are a heterogeneous `Tuple`,
+# whose elements a `for` loop sees as `Any`; `map` unrolls it, so each call is
+# specialized on that argument's type (cf. `map_or_ntuple` in `distribute_task!`).
+# `f::F` forces specializing on the closure: Julia does not specialize on a
+# function argument that is only passed along, so `map` was a dynamic call whose
+# result was boxed.
+_foreach_arg(f::F, fargs::Vector{Argument}) where {F} = foreach(f, fargs)
+_foreach_arg(f::F, fargs::Tuple) where {F} = (map(f, fargs); nothing)
+_any_arg(f::F, fargs::Vector{Argument}) where {F} = any(f, fargs)
+_any_arg(f::F, fargs::Tuple) where {F} = any(map(f, fargs))
+
+# Function barriers for `dag_add_task!`: `value(_arg)` is not inferrable for an
+# untyped spec, so each argument is handled by a method specialized on its
+# concrete type. Destructuring `unwrap_inout`'s result and iterating its deps on
+# an `Any` cost a generic `indexed_iterate`/`getfield` dispatch per dependency,
+# and was most of what was left of building a `DAGSpec` once aliasing was
+# prefetched.
+function _depends_on_dag_task(dspec::DAGSpec, v)
+    arg, deps = unwrap_inout(v)
+    arg isa DTask || return false
+    isempty(deps) && return false
+    return haskey(dspec.uid_to_id, arg.uid)
+end
+function _push_argspecs!(argtypes::Vector{DatadepsArgSpec}, pos, v)
+    arg, deps = unwrap_inout(v)
+    for (dep_mod, readdep, writedep) in deps
+        if arg isa DTask
+            #= TODO: Re-enable this when we can handle within-DAG dependencies
+            if arg.uid in keys(dspec.uid_to_id)
+                # Within-DAG dependency
+                arg_id = dspec.uid_to_id[arg.uid]
+                push!(dspec.id_to_argtypes[arg_id], DatadepsArgSpec(pos, DTaskDAGID{arg_id}, dep_mod, UnknownAliasing()))
+                add_edge!(dspec.g, arg_id, id)
+                continue
+            end
+            =#
+
+            # External DTask, so fetch this and track it as a raw value
+            arg = fetch(arg; raw=true)
+        end
+        # FIXME: Generate syncdeps and add edges
+        _push_argspec!(argtypes, pos, arg, dep_mod)
+    end
+    return
+end
+# A second barrier: the result of `fetch(...; raw=true)` is not inferrable either.
+_push_argspec!(argtypes::Vector{DatadepsArgSpec}, pos, arg, dep_mod) =
+    push!(argtypes, DatadepsArgSpec(pos, typeof(arg), dep_mod, aliasing(arg, dep_mod)))
 function dag_build_edges!(dag_spec::DAGSpec)
     # Naively build edges based on exact argument comparisons (no aliasing)
     arg_writes = Dict{Any,Int}()

test/datadeps.jl

@@ -81,6 +81,48 @@ end
     end
 end
 
+@testset "Flat prefetch of chunk aliasing" begin
+    # Every remote chunk the tasks name, directly or as a finished DTask's
+    # result, is resolved in one batch per owner and memoized under the key
+    # `aliasing` itself would look up. Local chunks are left to `aliasing`, and
+    # unfinished DTasks too: fetching one would block. Each remote chunk is
+    # named under two dep_mods, so its owner's batch must carry its `DRef` once
+    # (the same `DRef` twice in one message used to hang the owner).
+    remote = Chunk[remotecall_fetch(Dagger.tochunk, wid, zeros(4, 4))
+                   for wid in procs() if wid != myid()]
+    done = Dagger.@spawn scope=Dagger.scope(worker=last(procs())) zeros(4, 4)
+    wait(done)
+    done_chunk = fetch(done; raw=true)
+    local_chunk = Dagger.tochunk(zeros(4, 4))
+    unlaunched = Dagger.DTask(rand(UInt), Dagger.ThunkFuture(), Dagger.DTaskMetadata(Any))
+    task = Dagger.DTask(rand(UInt), Dagger.ThunkFuture(), Dagger.DTaskMetadata(Any))
+    pair(fargs) = Dagger.DTaskPair(Dagger.DTaskSpec(fargs, Dagger.Options()), task)
+    pairs = [pair(Dagger.Argument[Dagger.Argument(1, identity), Dagger.Argument(2, In(c)),
+                                  Dagger.Argument(3, Deps(c, InOut(UpperTriangular)))])
+             for c in remote]
+    # A typed spec, as `@spawn` builds for concretely-typed arguments
+    push!(pairs, pair((Dagger.TypedArgument(1, identity), Dagger.TypedArgument(2, InOut(done)),
+                       Dagger.TypedArgument(3, In(local_chunk)), Dagger.TypedArgument(4, In(unlaunched)))))
+    expected = Dict{Tuple{Chunk,Any},AbstractAliasing}((c, dep_mod) => aliasing(c, dep_mod)
+                                                       for c in remote for dep_mod in (identity, UpperTriangular))
+    # With no other workers (`-p 0`) the finished task ran here, so its result
+    # is local and left alone, like every other chunk
+    if Dagger.root_worker_id(done_chunk.processor) != myid()
+        expected[(done_chunk, identity)] = aliasing(done_chunk, identity)
+    end
+    memo = Dagger.ChunkAinfoMemo()
+    Dagger.with(Dagger.CHUNK_AINFO_MEMO => memo) do
+        Dagger.prefetch_chunk_aliasing!(pairs)
+        @test length(memo.entries) == length(expected)
+        for ((c, dep_mod), ainfo) in expected
+            key = Dagger.ainfo_memo_key(c, dep_mod)
+            # A cache miss must fail, not silently recompute the same answer.
+            cached = Dagger.memoized_ainfo(() -> error("prefetch did not memoize"), key)
+            @test cached == ainfo
+        end
+    end
+end
+
 @testset "Memory Aliasing" begin
     A = rand(4)
     a = Dagger.aliasing(A)