From 13f8436405b04137ee6845530c94bb93c33f4b88 Mon Sep 17 00:00:00 2001 From: antoniocaarvalh Date: Tue, 13 Oct 2026 17:16:23 -0300 Subject: [PATCH 1/2] Datadeps: Build DAG specs faster With the schedule cache on, building a region's DAGSpec was most of its planning time, and most of that went to dag_add_task! asking the owner of each remote chunk for its aliasing one round trip at a time, plus dynamic dispatch on arguments whose type isn't inferrable for an untyped spec. Resolve the aliasing of every remote chunk up front, with one remotecall_fetch per owning worker issued concurrently, and seed the region's aliasing memo with the answers, the same way build_aliasing_parallel already does in hierarchical mode, sending only each chunk's DRef and bringing back only its aliasing. Handle each argument in dag_add_task! behind a function barrier, mapping over typed argument tuples so every call is specialized on its argument's type, and look up a cached plan's tasks by id instead of scanning the region for each one. A 4096x4096 matmul with 512x512 tiles on 4 workers now builds its DAGSpec in 4.7 ms instead of 75 ms with the plan already cached, and a Cholesky in 2.3 ms instead of 23 ms, with about a third of the allocations. Each DRef goes once per worker even when the region uses it under several dep_mods, since MemPool's DRef serializer ignores what serialize_cycle_header returns and the same DRef twice in one message corrupts the stream. The prefetch is skipped under MPI, where aliasing is a collective that has to run in the same order on every rank. --- src/datadeps/queue.jl | 127 +++++++++++++++++++++++++++++++++++-- src/datadeps/scheduling.jl | 89 +++++++++++++++----------- test/datadeps.jl | 38 +++++++++++ 3 files changed, 214 insertions(+), 40 deletions(-) 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..a3113f0ef 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -81,6 +81,44 @@ 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)) + expected[(done_chunk, identity)] = aliasing(done_chunk, identity) + 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) From f789a3f46cebb8e622efe67dddd416fce9084573 Mon Sep 17 00:00:00 2001 From: antoniocaarvalh Date: Tue, 13 Oct 2026 19:01:53 -0300 Subject: [PATCH 2/2] Datadeps: Fix the prefetch test without extra workers The prefetch test spawned its finished task on the last worker and always expected that task's result in the memo, but in the CI jobs that run with no workers besides the first (-p 0) the task runs locally, so its result is a local chunk that prefetch_chunk_aliasing! is right to leave alone. Only expect it when it ran on another worker. --- test/datadeps.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/test/datadeps.jl b/test/datadeps.jl index a3113f0ef..cf3dac38d 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -105,7 +105,11 @@ end 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)) - expected[(done_chunk, identity)] = aliasing(done_chunk, identity) + # 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)