From 0f47f2a8aa0415fe6ef3bc1ea0965d325383a9c8 Mon Sep 17 00:00:00 2001 From: antoniocaarvalh Date: Tue, 13 Oct 2026 17:16:23 -0300 Subject: [PATCH 1/4] 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 1d878f2edd19333c8d5cecca232f2b3b02828c88 Mon Sep 17 00:00:00 2001 From: antoniocaarvalh Date: Tue, 13 Oct 2026 19:01:53 -0300 Subject: [PATCH 2/4] 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) From c5273c6e28d03cdee74e1aeeaed4440b58b55c92 Mon Sep 17 00:00:00 2001 From: antoniocaarvalh Date: Wed, 14 Oct 2026 10:49:14 -0300 Subject: [PATCH 3/4] Datadeps: Allocate uninitialized slots with one round trip Creating a slot went through move_rewrap on the source worker, which calls back to the planning process for every lookup in the aliased object cache and then on to the destination to allocate, about seven round trips in series per slot, and in flat mode that was most of the launch loop: a 4096x4096 matmul with 512x512 tiles on 4 workers spent 2.8 s of a 3.4 s loop creating 461 slots under RoundRobin. When the slot is only going to be allocated, because slot_may_be_uninit lets the copy-to phase fill it, the source worker has nothing to add, so generate_slot! now does the cache lookups itself and allocates with a single round trip to the destination that also brings back the new slot's aliasing. This only applies under Distributed and to dense, non-wrapper arrays whose type and size the chunk already tells, and everything else keeps going through move_rewrap. The RoundRobin launch loop on that matmul gets about 30-40% faster. --- src/datadeps/aliasing.jl | 74 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 72 insertions(+), 2 deletions(-) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index f8de25146..99a0a074b 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1437,8 +1437,20 @@ function generate_slot!(state::DataDepsState, dest_space, data) # original for its aliasing key, exactly as the general path does. aliased_object!(Returns(data), aliased_object_cache, data)::Chunk else - moved = with(DATADEPS_THUNK_ID=>tid) do - remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + direct = direct_uninit_slot(current_acceleration(), data, from_proc, to_proc, orig_space, dest_space) + moved = if direct !== nothing + # What `move_rewrap` would do on the source worker, done from here: + # the object cache lives with the planner, and the slot is only + # allocated, so the source worker has nothing to contribute. + T_dest, header = direct + slot = aliased_object!(aliased_object_cache, data) do _ + alloc_slot_with_aliasing!(to_proc, dest_space, T_dest, header) + end + slot::Chunk + else + with(DATADEPS_THUNK_ID=>tid) do + remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + end end if timing move_ns = time_ns() - t0 @@ -1497,6 +1509,64 @@ function remotecall_endpoint_alloc(accel::DistributedAcceleration, to_proc, to_s end end +""" + direct_uninit_slot(accel, data, from_proc, to_proc, from_space, to_space) + -> Union{Tuple{Type,Any},Nothing} + +The type and `alloc_header` of the slot `move_rewrap` would conjure +uninitialized for `data` in `to_space`, when the planning process can tell both +from `data`'s handle alone; `nothing` when only the source worker can. + +Creating a slot through `move_rewrap` runs on the source worker, which then +calls back to the planner for each object cache lookup (the cache lives with +the planner) and on to the destination to allocate: ~6 round trips in series +per slot. When the slot is only allocated, none of that needs the source +worker, so `generate_slot!` does the lookups locally and allocates with one +round trip (`alloc_slot_with_aliasing!`). That was ~90% of a flat-mode launch +loop: a 4096^2 matmul with 512^2 tiles on 4 Distributed workers spent 2.8 s of +a 3.4 s loop creating 461 slots under RoundRobin. + +Only Distributed qualifies: under MPI the cache is replicated and every rank +takes part in each step. +""" +direct_uninit_slot(::Acceleration, data, from_proc, to_proc, from_space, to_space) = nothing +function direct_uninit_slot(::DistributedAcceleration, data::Chunk, from_proc, to_proc, from_space, to_space) + T = chunktype(data) + isconcretetype(T) && T <: DenseArray || return nothing + # A wrapper is rebuilt from its children on the destination + move_rewrap_child_types(T) === nothing || return nothing + T_dest = move_type(from_proc, to_proc, T) + slot_may_be_uninit(from_space, to_space, T_dest) || return nothing + # `tochunk` records `domain(x)`, whose ranges span `size(x)`: what + # `alloc_header` would compute from the array itself. + dom = domain(data) + dom isa ArrayDomain || return nothing + header = map(length, dom.indexes) + header isa alloc_header_type(T_dest) || return nothing + return T_dest, header +end + +""" + alloc_slot_with_aliasing!(to_proc, to_space, ::Type{T}, header) -> Chunk + +`remotecall_endpoint_alloc`, but bringing the new slot's aliasing back with it +and memoizing it, so that planning does not pay another round trip to ask. +""" +function alloc_slot_with_aliasing!(to_proc, to_space, ::Type{T}, header) where T + wid = root_worker_id(to_proc) + slot, ainfo = if wid == myid() + _alloc_slot_with_aliasing(to_proc, to_space, T, header) + else + remotecall_fetch(_alloc_slot_with_aliasing, wid, to_proc, to_space, T, header) + end + memoize_ainfo!(ainfo_memo_key(slot, identity), ainfo) + return slot +end +function _alloc_slot_with_aliasing(to_proc, to_space, ::Type{T}, header) where T + x = alloc_uninit(to_space, T, header) + return tochunk(x, to_proc, to_space), aliasing_unwrapped(x, identity) +end + function remotecall_endpoint_transfer(f, accel::DistributedAcceleration, from_proc, to_proc, from_space, to_space, data) wid = root_worker_id(to_proc) if wid == myid() From cd284631b26537f37cafdd1da9ac7f199a34fbf6 Mon Sep 17 00:00:00 2001 From: antoniocaarvalh Date: Mon, 19 Oct 2026 19:20:08 -0300 Subject: [PATCH 4/4] Datadeps: Carve slots out of per-region arenas Even with one round trip per slot, a flat launch loop still waited on an allocation for every slot, queued behind the region's copies on the destination. Instead, before the launch loop, reserve one arena per destination worker, sized exactly from the plan for AOT schedulers and grown on demand for the rest, and hand slots out of it from the planning process as base + offset, so each slot's ArenaRef and aliasing are known without asking anyone. The arenas are freed whole when the region ends, right away when it succeeds, and in the background once its tasks are done when one of them fails, since the scheduler fails a task's dependents outright and a per-slot free task would never run. Only dense CPU arrays of isbits elements use arenas; everything else allocates as before. A 4096x4096 matmul with 512x512 tiles on 4 workers now launches its RoundRobin region in 0.17 s instead of 3 s, and the whole region takes 1.97 s instead of 3.3 s. --- src/datadeps/aliasing.jl | 293 ++++++++++++++++++++++++++++++++++- src/datadeps/hierarchical.jl | 7 + src/datadeps/queue.jl | 26 ++++ src/memory-spaces.jl | 2 +- test/datadeps.jl | 25 +++ 5 files changed, 348 insertions(+), 5 deletions(-) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 99a0a074b..997eaa943 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -658,6 +658,79 @@ function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(cache.a end end +""" + ArenaRef{T,N} + +The handle of a Datadeps slot carved out of a `SlotArena`: an `Array{T,N}` of +size `dims` at address `addr` on worker `owner`, inside arena `arena`. Unlike a +`DRef`, the planner makes one without asking `owner`, since it already knows the +arena's base address: so it knows the slot's aliasing before the slot exists. +""" +struct ArenaRef{T,N} + owner::Int + arena::UInt + addr::UInt + dims::NTuple{N,Int} +end + +""" + SlotArena + +One block of memory a region reserved on a worker for its slots, which the +planner hands out by bumping `used`. It is freed whole, once, when the region +ends (see `RegionSlotArenas`). +""" +mutable struct SlotArena + const space::CPURAMMemorySpace + const id::UInt + const base::UInt + const capacity::Int + used::Int +end + +""" + RegionSlotArenas + +Every slot arena a `spawn_datadeps` region reserved, as `(owner, id)`, which it +frees once all of its tasks have finished, failed or not: right away when the +region succeeds, or in the background when a task fails +(`free_region_arenas_when_done!`). Freeing an arena from a task that depends on +its slots' users would leak it whenever one of them fails, since the scheduler +fails a task's dependents outright; the per-slot buffers this replaces were +reclaimed by their finalizers in that case, but an arena has none. +""" +mutable struct RegionSlotArenas + const lock::ReentrantLock + const arenas::Vector{Tuple{Int,UInt}} + # The region's `wait_all` queue's tasks, which every task it launches joins + tasks::Union{Vector{DTask},Nothing} +end +RegionSlotArenas() = RegionSlotArenas(ReentrantLock(), Tuple{Int,UInt}[], nothing) +const DATADEPS_REGION_ARENAS = ScopedValue{Union{RegionSlotArenas,Nothing}}(nothing) + +""" + SlotPool + +The memory a region allocates its slots from, ahead of need. `generate_slot!`'s +direct path (see `direct_uninit_slot`) takes each dense CPU slot from an arena +(`SlotArena`) on its destination, computing its address and aliasing itself, +instead of paying a round trip per slot -- most of what was left of a flat +launch loop, where one allocation costs ~0.3 ms on an idle worker but ~5 ms in +the middle of a region, queued behind the region's copies. + +For AOT plans `prefill_slot_pool!` reserves exactly what the plan needs, with +one call per destination worker. Anything else reserves arenas on first need, +growing from `SLOT_ARENA_MIN_BYTES` to `SLOT_ARENA_BYTES`. Each arena is freed +whole when the region ends (`RegionSlotArenas`), so a region no longer spawns a +free task per slot. The pool is only enabled inside a `spawn_datadeps` region, +which is what frees its arenas. +""" +mutable struct SlotPool + enabled::Bool + arenas::Dict{CPURAMMemorySpace,Vector{SlotArena}} +end +SlotPool() = SlotPool(false, Dict{CPURAMMemorySpace,Vector{SlotArena}}()) + # N.B. Declared `mutable` (despite no field ever being rebound) purely for # identity: an immutable struct this large is re-boxed on every dynamic call # that passes it, and planning threads it through hundreds of such calls per @@ -753,6 +826,9 @@ mutable struct DataDepsState scratch_deps::Vector{DataDepsTaskDependency} scratch_remote::Vector{Any} + # Slot buffers allocated ahead of need (see `SlotPool`) + slot_pool::SlotPool + function DataDepsState() arg_to_chunk = IdDict{Any,Any}() sparse_arg_wrap = IdDict{Any,Any}() @@ -781,7 +857,7 @@ mutable struct DataDepsState return new(arg_to_chunk, sparse_arg_wrap, arg_origin, remote_args, remote_arg_to_original, remote_arg_w, ainfo_arg, arg_history, arg_owner, arg_current, arg_overlaps, ainfo_backing_chunk, supports_inplace_cache, ainfo_cache, ainfos_lookup, ainfos_overlaps, ainfos_owner, ainfos_readers, - TaskArgInfo[], DataDepsTaskDependency[], Any[]) + TaskArgInfo[], DataDepsTaskDependency[], Any[], SlotPool()) end end @@ -1444,7 +1520,7 @@ function generate_slot!(state::DataDepsState, dest_space, data) # allocated, so the source worker has nothing to contribute. T_dest, header = direct slot = aliased_object!(aliased_object_cache, data) do _ - alloc_slot_with_aliasing!(to_proc, dest_space, T_dest, header) + alloc_slot_with_aliasing!(state.slot_pool, to_proc, dest_space, T_dest, header) end slot::Chunk else @@ -1547,12 +1623,17 @@ function direct_uninit_slot(::DistributedAcceleration, data::Chunk, from_proc, t end """ - alloc_slot_with_aliasing!(to_proc, to_space, ::Type{T}, header) -> Chunk + alloc_slot_with_aliasing!(pool::SlotPool, to_proc, to_space, ::Type{T}, header) -> Chunk `remotecall_endpoint_alloc`, but bringing the new slot's aliasing back with it and memoizing it, so that planning does not pay another round trip to ask. +With `pool` enabled, a dense CPU slot instead comes out of an arena on its +destination, with no round trip at all once the arena is reserved. """ -function alloc_slot_with_aliasing!(to_proc, to_space, ::Type{T}, header) where T +function alloc_slot_with_aliasing!(pool::SlotPool, to_proc, to_space, ::Type{T}, header) where T + if pool.enabled && _arena_slot_type(T, to_space) + return _arena_slot!(pool, to_proc, to_space, T, header) + end wid = root_worker_id(to_proc) slot, ainfo = if wid == myid() _alloc_slot_with_aliasing(to_proc, to_space, T, header) @@ -1567,6 +1648,210 @@ function _alloc_slot_with_aliasing(to_proc, to_space, ::Type{T}, header) where T return tochunk(x, to_proc, to_space), aliasing_unwrapped(x, identity) end +_arena_slot_type(::Type{<:Array{E}}, ::CPURAMMemorySpace) where {E} = isbitstype(E) +_arena_slot_type(::Type, ::MemorySpace) = false + +# Arenas hand out slots at this alignment. When no plan says how much a region +# needs, they are reserved between these sizes (see `_arena_with_room!`). +const SLOT_ARENA_ALIGN = 64 +const SLOT_ARENA_MIN_BYTES = 64 * 2^10 +const SLOT_ARENA_BYTES = 64 * 2^20 +_arena_aligned(n::Int) = cld(n, SLOT_ARENA_ALIGN) * SLOT_ARENA_ALIGN +_arena_slot_bytes(::Type{<:Array{E}}, header) where {E} = sizeof(E) * prod(header) + +function _arena_slot!(pool::SlotPool, to_proc, to_space::CPURAMMemorySpace, ::Type{T}, header::Dims{N}) where {T<:Array,N} + E = eltype(T) + nbytes = _arena_slot_bytes(T, header) + arena = _arena_with_room!(pool, to_space, nbytes) + offset = arena.used + arena.used += _arena_aligned(nbytes) + ref = ArenaRef{E,N}(root_worker_id(to_space), arena.id, arena.base + offset, header) + # What `tochunk(x, to_proc, to_space)` would make of the array `unwrap(ref)` + slot = Chunk{T,typeof(ref),typeof(to_proc),AnyScope,typeof(to_space)}( + T, ArrayDomain([1:l for l in header]), ref, to_proc, AnyScope(), to_space) + # What `aliasing(::Array)` would compute on the owner + ainfo = ContiguousAliasing(MemorySpan{CPURAMMemorySpace}(RemotePtr{Cvoid}(arena.base + offset, to_space), nbytes)) + memoize_ainfo!(ainfo_memo_key(slot, identity), ainfo) + return slot +end + +function _arena_with_room!(pool::SlotPool, space::CPURAMMemorySpace, nbytes::Int) + arenas = get!(Vector{SlotArena}, pool.arenas, space) + prev = 0 + if !isempty(arenas) + arena = last(arenas) + arena.capacity - arena.used >= nbytes && return arena + prev = arena.capacity + end + # On demand, start small and double: an arena reserved for a few small + # slots would be mostly waste. + want = clamp(max(2prev, 8 * _arena_aligned(nbytes)), SLOT_ARENA_MIN_BYTES, SLOT_ARENA_BYTES) + reserve_slot_arenas!(pool, Dict(space => max(want, _arena_aligned(nbytes)))) + return last(pool.arenas[space]) +end + +"Reserve an arena of `bytes[space]` on each `space`'s worker, concurrently." +function reserve_slot_arenas!(pool::SlotPool, bytes::Dict{CPURAMMemorySpace,Int}) + region = DATADEPS_REGION_ARENAS[]::RegionSlotArenas + spaces = collect(keys(bytes)) + results = Vector{Tuple{UInt,UInt}}(undef, length(spaces)) + @sync for (j, space) in enumerate(spaces) + Threads.@spawn begin + wid = root_worker_id(space) + id, base = wid == myid() ? _reserve_slot_arena(bytes[space]) : + remotecall_fetch(_reserve_slot_arena, wid, bytes[space]) + # Registered as soon as it exists, so that the region frees it even + # if reserving another worker's arena fails + @lock region.lock push!(region.arenas, (wid, id)) + results[j] = (id, base) + end + end + for (j, space) in enumerate(spaces) + id, base = results[j] + push!(get!(Vector{SlotArena}, pool.arenas, space), SlotArena(space, id, base, bytes[space], 0)) + end + return +end + +# Worker side: the arenas this process holds, by id +const SLOT_ARENAS = Dict{UInt,Ptr{Cvoid}}() +const SLOT_ARENAS_LOCK = ReentrantLock() +const SLOT_ARENA_NEXT_ID = Threads.Atomic{UInt}(1) +function _reserve_slot_arena(nbytes::Int) + p = Libc.malloc(nbytes) + p == C_NULL && throw(OutOfMemoryError()) + id = Threads.atomic_add!(SLOT_ARENA_NEXT_ID, UInt(1)) + @lock SLOT_ARENAS_LOCK SLOT_ARENAS[id] = p + return id, UInt(p) +end +function _free_slot_arenas!(ids::Vector{UInt}) + for id in ids + p = @lock SLOT_ARENAS_LOCK pop!(SLOT_ARENAS, id) + Libc.free(p) + end + return +end + +root_worker_id(h::ArenaRef) = h.owner +datasize(h::ArenaRef{T}) where {T} = sizeof(T) * prod(h.dims) +function unwrap(h::ArenaRef{T,N}) where {T,N} + @assert h.owner == myid() "ArenaRef $h is not owned by this process: $(h.owner) != $(myid())" + # Not registered with `_libc_register!`: the arena owns this memory, so + # `unsafe_free!` of the slot alone is a no-op and the arena is freed whole. + # Wrapping needs only the address, not the arena, so a slot outliving its + # region still answers `aliasing` (logs hold slots), like a freed + # `alloc_libc_array` buffer does; reading it then is a use-after-free either way. + return unsafe_wrap(Array{T,N}, Ptr{T}(h.addr), h.dims; own=false) +end +fetch_handle(h::ArenaRef; uniform::Bool=false) = + h.owner == myid() ? unwrap(h) : remotecall_fetch(unwrap, h.owner, h) +MemPool.poolget(h::ArenaRef) = fetch_handle(h) +unwrappable(::ArenaRef) = true +move(from_proc::Processor, to_proc::Processor, x::ArenaRef) = move(from_proc, to_proc, fetch_handle(x)) +move(to_proc::Processor, d::ArenaRef) = move(OSProc(d.owner), to_proc, d) + +""" + prefill_slot_pool!(state, pairs, schedule) + +Enable `state.slot_pool` and reserve on each destination worker an arena with +room for every slot that `schedule` (an AOT plan) will make `generate_slot!` +create through its direct path: one per argument per destination space it is +placed in, counted the way `direct_uninit_slot` decides them. Arguments it +cannot see yet (`DTask`s still running) get arenas reserved on first need. +""" +function prefill_slot_pool!(state::DataDepsState, pairs::Vector{DTaskPair}, schedule::Dict{DTask,Processor}) + accel = current_acceleration() + accel isa DistributedAcceleration || return + # Arenas are only safe where a region will free them + DATADEPS_REGION_ARENAS[] === nothing && return + state.slot_pool.enabled = true + bytes = Dict{CPURAMMemorySpace,Int}() + seen = Set{Tuple{UInt,MemorySpace}}() + for pair in pairs + proc = get(schedule, pair.task, nothing) + proc === nothing && continue + _count_spec_slots!(bytes, seen, accel, pair.spec, only(memory_spaces(proc))) + end + isempty(bytes) || reserve_slot_arenas!(state.slot_pool, bytes) + return +end +# A function barrier: `pair.spec` is abstractly typed (cf. `_collect_spec_remote_chunks!`) +_count_spec_slots!(bytes, seen, accel, spec::DTaskSpec, to_space) = + _foreach_arg(_arg -> _count_arg_slot!(bytes, seen, accel, value(_arg), to_space), spec.fargs) +function _count_arg_slot!(bytes, seen, accel, v, to_space) + arg, deps = unwrap_inout(v) + if arg isa DTask + isready(arg) || return + arg = fetch(arg; raw=true) + end + arg isa Chunk || return + from_space = memory_space(arg) + direct = direct_uninit_slot(accel, arg, first(processors(from_space)), + first(processors(to_space)), from_space, to_space) + direct === nothing && return + T, header = direct + _arena_slot_type(T, to_space) || return + id = (_identity_hash(arg), to_space) + id in seen && return + push!(seen, id) + bytes[to_space] = get(bytes, to_space, 0) + _arena_aligned(_arena_slot_bytes(T, header)) + return +end + +"Whether `slot` lives in a slot arena, which its region frees whole (no free task of its own)." +freed_with_its_arena(slot) = slot isa Chunk && slot.handle isa ArenaRef + +""" + free_region_arenas!(region::RegionSlotArenas) + +Free every arena `region` reserved, with one call per worker. Only valid once +all of the region's tasks have finished (see `spawn_datadeps`). Never throws: +it runs in a `finally`, where an error would hide the region's own; an arena +that cannot be freed (its worker died, say) is reported and dropped. +""" +function free_region_arenas!(region::RegionSlotArenas) + arenas = @lock region.lock copy(region.arenas) + isempty(arenas) && return + by_owner = Dict{Int,Vector{UInt}}() + for (wid, id) in arenas + push!(get!(Vector{UInt}, by_owner, wid), id) + end + @sync for (wid, ids) in by_owner + Threads.@spawn try + wid == myid() ? _free_slot_arenas!(ids) : remotecall_wait(_free_slot_arenas!, wid, ids) + catch err + @warn "Could not free $(length(ids)) Datadeps slot arena(s) on worker $wid" exception=err + end + end + @lock region.lock empty!(region.arenas) + return +end + +""" + free_region_arenas_when_done!(region::RegionSlotArenas) + +`free_region_arenas!`, from a background task, once every task `region` +launched has finished: for a region that failed, whose `wait_all` rethrew the +first task error without waiting for the rest, which may still be using their +slots. The error reaches the caller as soon as it did without arenas. +""" +function free_region_arenas_when_done!(region::RegionSlotArenas) + @lock(region.lock, isempty(region.arenas)) && return + tasks = region.tasks === nothing ? DTask[] : copy(region.tasks) + errormonitor(Threads.@spawn begin + for task in tasks + # A task that was never launched has nothing to wait for + istaskstarted(task) || continue + try + wait(task) + catch + end + end + free_region_arenas!(region) + end) + return +end + function remotecall_endpoint_transfer(f, accel::DistributedAcceleration, from_proc, to_proc, from_space, to_space, data) wid = root_worker_id(to_proc) if wid == myid() diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index b0f364294..bfa065d2f 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -1446,6 +1446,9 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, end filtered end + # Arenas for this partition's slots, in the spaces its plan uses (see + # `SlotPool`); freed by the epilogue of `distribute_tasks_hierarchical!`. + prefill_slot_pool!(state, partition_pairs, schedule) # N.B. If this partition throws partway through (e.g. from # `distribute_task!`), any of our vertices that haven't yet been @@ -1546,6 +1549,8 @@ function schedule_partitions_sequential!(queue::DataDepsTaskQueue, local_scopes = Vector{AbstractScope}(undef, n_partitions) proc_to_scope_lfus = [BasicLFUCache{Processor,AbstractScope}(1024) for _ in 1:n_partitions] shared_state = DataDepsState() + # Arenas for the region's slots (see `SlotPool`); freed by the epilogue + prefill_slot_pool!(shared_state, seen_tasks, schedule) write_num = 1 # Uniform execution inverts `BatchedEnqueueQueue`'s latency/throughput @@ -2148,6 +2153,8 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat is_original(obj_cache, remote_space, ainfo) && continue haskey(freed, remote_arg) && continue freed[remote_arg] = nothing + # A slot in an arena is freed with it, when the region ends + freed_with_its_arena(remote_arg) && continue free_syncdeps = Set{ThunkSyncdep}() buf_ainfo = stored_value_ainfo(obj_cache, remote_space, ainfo) gather_free_syncdeps!(state, remote_space, buf_ainfo, remote_arg, diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 49b282ef0..54d521126 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -76,7 +76,27 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, if !aliasing throw(ArgumentError("Aliasing analysis is no longer optional")) end + # The region's slot arenas (see `SlotPool`) are freed once every task it + # launched has finished, since a failed task would never let a dependent + # free task run. When the region succeeds, `wait_all` has waited for them + # all; when a task fails, it rethrows right away, possibly with others + # still running, so the arenas are freed in the background once those end. + arenas = RegionSlotArenas() + ok = false + try + result = with(DATADEPS_REGION_ARENAS => arenas) do + _spawn_datadeps_region(f, scheduler, launch_wait, hierarchical) + end + ok = true + return result + finally + ok ? free_region_arenas!(arenas) : free_region_arenas_when_done!(arenas) + end +end +function _spawn_datadeps_region(f, scheduler, launch_wait, hierarchical) wait_all(; check_errors=true) do + # Every task the region launches joins this queue + DATADEPS_REGION_ARENAS[].tasks = (get_options(:task_queue)::WaitAllQueue).tasks scheduler = something(scheduler, DATADEPS_SCHEDULER[], RoundRobinScheduler()) launch_wait = something(launch_wait, DATADEPS_LAUNCH_WAIT[], false)::Bool hierarchical = something(hierarchical, DATADEPS_HIERARCHICAL[], true)::Bool @@ -470,6 +490,10 @@ function distribute_tasks!(queue::DataDepsTaskQueue) write_num = 1 proc_to_scope_lfu = BasicLFUCache{Processor,AbstractScope}(1024) launch_start = time_ns() + # Allocate the slots the plan will need up front, one call per worker, so + # the loop below does not wait on one allocation per slot. Counted in the + # launch time, which is what it replaces. + prefill_slot_pool!(state, queue.seen_tasks, schedule) for pair in queue.seen_tasks spec = pair.spec task = pair.task @@ -550,6 +574,8 @@ function distribute_tasks!(queue::DataDepsTaskQueue) is_original(obj_cache, remote_space, ainfo) && continue haskey(freed, remote_arg) && continue freed[remote_arg] = nothing + # A slot in an arena is freed with it, when the region ends + freed_with_its_arena(remote_arg) && continue free_syncdeps = Set{ThunkSyncdep}() buf_ainfo = stored_value_ainfo(obj_cache, remote_space, ainfo) gather_free_syncdeps!(state, remote_space, buf_ainfo, remote_arg, write_num, chunk_to_ainfos, free_syncdeps) diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 5cd345b06..5a8c60112 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -712,7 +712,7 @@ function aliasing(x::Chunk, T) if root_worker_id(x.processor) == myid() return aliasing_unwrapped(unwrap(x), T) end - @assert x.handle isa DRef + @assert x.handle isa Union{DRef,ArenaRef} return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T aliasing_unwrapped(unwrap(x), T) end diff --git a/test/datadeps.jl b/test/datadeps.jl index cf3dac38d..23b5215ce 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -123,6 +123,31 @@ end end end +@testset "Slot arenas are freed when a region fails" begin + # A region frees its slot arenas once all of its tasks are done, in the + # background when one fails. Freeing them from a task depending on the + # slots' users leaked them whenever one failed, since the scheduler fails a + # task's dependents outright. + if nprocs() >= 3 + arenas_alive() = sum(w -> remotecall_fetch(() -> length(Dagger.SLOT_ARENAS), w), procs()) + arenas_made() = sum(w -> remotecall_fetch(() -> Int(Dagger.SLOT_ARENA_NEXT_ID[]), w), procs()) + chunks = [remotecall_fetch(Dagger.tochunk, w, rand(16, 16)) for w in procs()] + made = arenas_made() + @test_throws Exception Dagger.with(Dagger.DATADEPS_SCHEDULER => Dagger.RoundRobinScheduler(), + Dagger.DATADEPS_HIERARCHICAL => false) do + Dagger.spawn_datadeps() do + for _ in 1:2, c in chunks + Dagger.@spawn (x -> (x .+= 1; nothing))(InOut(c)) + end + Dagger.@spawn (x -> error("failing on purpose"))(InOut(first(chunks))) + end + end + # RoundRobin moved tasks off their chunks' owners, so arenas were used + @test arenas_made() > made + @test timedwait(() -> arenas_alive() == 0, 60; pollint=0.5) === :ok + end +end + @testset "Memory Aliasing" begin A = rand(4) a = Dagger.aliasing(A)