diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index a0230f7d6..49b282ef0 100644 --- a/src/datadeps/queue.jl +++ b/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 diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 1ea158550..1b047373d 100644 --- a/src/datadeps/scheduling.jl +++ b/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}() diff --git a/test/datadeps.jl b/test/datadeps.jl index 41fd9e800..cf3dac38d 100644 --- a/test/datadeps.jl +++ b/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)