From 02dfd4b346328b75ab90d11bf82ba30f4bfc1794 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 15 Aug 2026 03:18:09 -0700 Subject: [PATCH 01/26] datadeps/hierarchical: Submit off the planning thread, and measure planning Turning a prepared spec into scheduler thunks needs nothing from the planner, but it was running inline and accounted for ~40% of per-task planning cost in a multi-worker region. `AsyncEnqueueQueue` hands each batch to a submitter task instead, preserving FIFO order (the syncdeps recorded during planning rely on it) and keeping a synchronous drain for the two points that need a task to really exist: a value dependency's `fetch`, and the end of the region. It is used only when there is a spare thread to submit on, and never under uniform execution, where a rank's submission runs collectives that must stay ordered against planning's own. `DATADEPS_BATCH_LIMIT` goes from 4 to 16, which is where the scheduler round-trip stops amortizing (256 independent `InOut` tasks over 4 workers: 69 us/task unbatched, 55 at 16, 53.5 unbounded) while still bounding how far planning runs ahead of execution. Together these take that region from 31.2 ms to 16.9 ms. Attributing planning cost is hard from a profile, because the expensive parts are blocking waits inside communication rather than hot loops, so this also adds per-phase timing behind `JULIA_DAGGER_HIER_TIMING=1` (off by default, one `Ref` read per phase) and records what it found in the module header. The MPI hang warning now carries a backtrace, since which call site is waiting is the whole diagnosis for a wait cycle. Co-authored-by: Cursor --- ext/MPIExt.jl | 15 +- src/Dagger.jl | 2 + src/datadeps/aliasing.jl | 7 +- src/datadeps/hierarchical.jl | 392 ++++++++++++++++++++++++++++++++--- 4 files changed, 382 insertions(+), 34 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index e8831f8f3..b5e96d065 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -743,8 +743,13 @@ function tochunk_pset(x, space::MPIMemorySpace; device=nothing, force_nonlocal=f end const DEADLOCK_DETECT = TaskLocalValue{Bool}(()->true) + const DEADLOCK_WARN_PERIOD = TaskLocalValue{Float64}(()->10.0) const DEADLOCK_TIMEOUT_PERIOD = TaskLocalValue{Float64}(()->120.0) +# Off by default: a full stacktrace buries the one-line diagnosis (which call +# site is waiting) in noise. Opt in with JULIA_DAGGER_MPI_WARN_BACKTRACE=1 when +# that call site itself is the thing in question. +const DEADLOCK_WARN_BACKTRACE = TaskLocalValue{Bool}(()->parse(Bool, get(ENV, "JULIA_DAGGER_MPI_WARN_BACKTRACE", "0"))) const RECV_WAITING = LockedObject(Dict{Tuple{MPI.Comm, Int, Int}, Base.Event}()) # Envelope for the out-of-place raw-bytes MPI path: serialize a small @@ -1120,7 +1125,15 @@ end function mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, rank, tag, kind, srcdest) time_elapsed = (time_ns() - time_start) if detect && time_elapsed > warn_period - @warn "[rank $rank][tag $tag] Hit probable hang on $kind (dest: $srcdest)" + # A hang here is a wait cycle across ranks, so which call site is waiting + # (and on whose behalf) is the whole diagnosis; a bare tag is not enough. + # The backtrace that pins down the call site is off by default, though -- + # see DEADLOCK_WARN_BACKTRACE. + if DEADLOCK_WARN_BACKTRACE[] + @warn "[rank $rank][tag $tag] Hit probable hang on $kind (dest: $srcdest)" stacktrace=sprint(Base.show_backtrace, stacktrace()) + else + @warn "[rank $rank][tag $tag] Hit probable hang on $kind (dest: $srcdest)" + end return typemax(UInt64) end if detect && time_elapsed > timeout_period diff --git a/src/Dagger.jl b/src/Dagger.jl index 3c19ac4a1..2b988a956 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -241,6 +241,8 @@ function __init__() catch err @warn "Error parsing JULIA_DAGGER_DEBUG" exception=err end + + HIER_TIMING[] = parse(Bool, get(ENV, "JULIA_DAGGER_HIER_TIMING", "0")) end end # module diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 13a3ea61f..ef27070f1 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1227,6 +1227,7 @@ function generate_slot!(state::DataDepsState, dest_space, data) id = logging ? rand(Int) : 0 logging && @logstart ctx LogMove LogMoveId(0, ArgPosition(), to_proc, id) data tid = something(DATADEPS_CURRENT_TASK[], (;uid=0)).uid + t0 = time_ns() data_chunk = if slot_is_already_in_place(data, orig_space, dest_space) # Nothing to move: the slot for data already in `dest_space` is the data # itself. Going through `move_rewrap` here would allocate a second Chunk @@ -1236,10 +1237,14 @@ 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 - with(DATADEPS_THUNK_ID=>tid) do + 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) end + hier_stat_add!(:slot_moved_ns, time_ns() - t0) + orig_space == dest_space && hier_stat_add!(:slot_samespace_ns, time_ns() - t0) + moved end + hier_stat_add!(:slot_ns, time_ns() - t0) logging && @logfinish ctx LogMove LogMoveId(0, ArgPosition(), to_proc, id) data_chunk @assert memory_space(data_chunk) == dest_space "space mismatch! $dest_space (dest) != $(memory_space(data_chunk)) (actual) ($(typeof(data)) (data) vs. $(typeof(data_chunk)) (chunk)), spaces ($orig_space -> $dest_space)" dest_space_args[data] = data_chunk diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 9c558d7af..2ce45d7be 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -20,28 +20,59 @@ # serialized via `LockedEnqueueQueue`, but the (expensive) per-task # `distribute_task!` preparation runs concurrently across partitions. # +# ### Where planning time actually goes +# +# Measured with `HIER_TIMING`; Phase 4 dominates in every configuration, so the +# notes below are about it. Two configurations behave quite differently: +# +# * Distributed. Per-task planning is pure local CPU work on the calling +# process: ~55 us/task for a trivial region (256 independent `InOut` tasks over +# 4 workers). Roughly 40% of that was the scheduler round-trip of submitting +# the task, which `AsyncEnqueueQueue` now overlaps with planning; the rest is +# `distribute_task!`'s own bookkeeping, of which `populate_task_info!` (its +# aliasing lookups and slot generation) is the largest part. Slot *transfers* +# are usually free here -- `slot_is_already_in_place` keeps them at zero for a +# stencil sweep -- so there is little latency left to hide, and the remaining +# cost is CPU that a coarse lock would simply re-serialize (see 2). +# +# * MPI. Per-task planning is an order of magnitude more expensive (~1 ms/task +# for a 2-rank stencil sweep, 8 args and 14 slots per task) and it is dominated +# by *collectives*, not CPU: about 70% of `distribute_task!` is slot +# generation, and that splits roughly evenly between the point-to-point +# transfer of the data and a second collective that broadcasts the resulting +# destination chunk's aliasing info to every rank (needed because the object +# cache is replicated and keyed by pointer spans, which only the destination +# rank can compute). So MPI planning cost tracks the number of *new slot +# chunks* a region creates, at ~2 rendezvous each. +# # ### Not-yet-parallelized work (performance only; results are unaffected) # # These are the known gaps between what this pipeline does and what it could do. # Each is marked with a `# PERF(hier-N)` comment at the relevant site. # -# 1. MPI plans entirely sequentially. Uniform execution needs rank-identical -# ordering for tag / `MPIRefID` allocation, and aliasing may run collectives -# that must be ordered identically everywhere, so Phase 1 uses `nchunks == 1`, -# `_compute_aliasing_batch` refuses to thread, and Phase 4 uses -# `schedule_partitions_sequential!`. Under MPI this pipeline therefore buys -# rank affinity but *not* parallel planning. Lifting this needs a -# deterministic parallel order (e.g. pre-allocating tag ranges per partition) -# rather than simply enabling the threaded paths. +# 1. MPI plans entirely sequentially, and every rank plans every task. Uniform +# execution needs rank-identical ordering for tag / `MPIRefID` allocation, and +# aliasing may run collectives that must be ordered identically everywhere, so +# Phase 1 uses `nchunks == 1`, `_compute_aliasing_batch` refuses to thread, +# and Phase 4 uses `schedule_partitions_sequential!`. Replication makes +# per-rank planning O(all tasks) while execution is O(tasks/rank), so planning +# grows linearly with rank count: a 4096^2 stencil sweep plans in 0.4 ms at 1 +# rank, 2.5 ms at 2, and 5.6 ms at 4. Lifting this needs tags every rank can +# compute independently and deterministically (e.g. hashed from the region id +# plus the vertex and argument index) instead of a shared counter; only then +# can a rank skip preparing a task it neither runs nor sources data for. # 2. Multi-owner regions plan Phase 4 sequentially, not just MPI ones. Partitions # still carry worker/rank affinity, but they are planned in global topological # order on one shared `DataDepsState` (`use_shared_state` below). The parallel # per-partition path is correct only when every partition shares one memory # space, because `DataDepsState` keys its slot / ownership / currency # bookkeeping by memory *space* and cannot represent two partitions' distinct -# slots for one chunk. Fixing that -- tracking slot identity rather than space -# -- is what would re-enable parallel planning across workers, and is the -# single highest-value follow-up here. +# slots for one chunk. Note that simply sharing one state behind one lock +# would fix the correctness half and buy almost nothing: the measurements +# above show this path's per-task cost is state bookkeeping, so a coarse lock +# re-serializes exactly what it was meant to spread. Making it pay needs +# either per-argument striped locking of the state, or the transfers to be the +# thing being overlapped (which is the case for regions that do move data). # 3. Planning is centralized on the calling process. Phase 1's aliasing is # genuinely distributed (`remotecall` per worker), but Phase 4 runs every # partition's `distribute_task!` locally. Workers never plan their own @@ -50,6 +81,16 @@ # order-dependent (interval-tree insertion; sequential owner/reader state), # so they resist naive parallelization. Cheap relative to Phases 1 and 4 # today, but they become the ceiling once those scale. +# 5. A region's slots are allocated and freed every time it runs. An iterative +# workload (a stencil loop, a solver) re-creates the same per-space buffers +# each sweep, paying allocation, a populating transfer, an aliasing collective +# under MPI, and a free task per buffer -- all for data whose contents the +# region overwrites anyway (`generate_slot!` deliberately does not sync with +# the owner). Reusing a slot across regions is safe for that reason, and would +# remove most of MPI's planning cost for iterative code. It needs a cache +# keyed on (origin chunk identity, destination space) whose entries die with +# the origin chunk, plus an epilogue that does not free what it did not +# allocate this time round. struct HierarchicalTaskInfo arg_w::ArgumentWrapper @@ -104,8 +145,228 @@ function flush_batch!(beq::BatchedEnqueueQueue) return end -"Maximum number of tasks a hierarchical partition buffers before submitting." -const DATADEPS_BATCH_LIMIT = Ref(4) +""" + AsyncEnqueueQueue(inner, lock; limit=DATADEPS_BATCH_LIMIT[]) + +`BatchedEnqueueQueue` that hands each full batch to a submitter task rather than +submitting it on the caller's thread. + +Submission is about 40% of a region's per-task planning cost, and none of it +touches the planning state: it turns already-prepared specs into scheduler +thunks. Running it inline makes planning stop dead every `limit` tasks for work +that has nothing left to learn from the planner. Handing the batch off lets the +next batch be prepared while the previous one is submitted, which is worth +roughly that 40% whenever the caller has a spare thread. + +Ordering is preserved (one submitter, FIFO channel), which the syncdeps +recorded during planning rely on. `flush_batch!` is a *synchronous* drain, so +the points that genuinely need a task to exist -- a value dependency's `fetch`, +and the end of the region -- still get it. + +Not used under uniform execution: an MPI rank's submission runs collectives, and +overlapping those with planning's own collectives makes the message order +rank-dependent. +""" +mutable struct AsyncEnqueueQueue <: AbstractTaskQueue + const inner::AbstractTaskQueue + const lock::ReentrantLock + const pending::Vector{DTaskPair} + const limit::Int + # `Event`s are drain barriers: the submitter notifies one once everything + # queued ahead of it has been submitted. + const chan::Channel{Union{Vector{DTaskPair},Base.Event}} + submitter::Union{Task,Nothing} + # Set by the submitter, re-thrown on the planning task so a submission + # failure surfaces as the region's error rather than a silent hang. + failure::Any +end +function AsyncEnqueueQueue(inner::AbstractTaskQueue, lock::ReentrantLock; + limit::Int=DATADEPS_BATCH_LIMIT[]) + chan = Channel{Union{Vector{DTaskPair},Base.Event}}(Inf) + queue = AsyncEnqueueQueue(inner, lock, DTaskPair[], limit, chan, nothing, nothing) + queue.submitter = Threads.@spawn _async_submit_loop(queue) + return queue +end +function _async_submit_loop(queue::AsyncEnqueueQueue) + for item in queue.chan + # Nothing in here may escape: this task is the only thing that notifies + # the drain barriers, so dying on an error would hang the planner instead + # of reporting to it. Record the failure and keep draining; whoever waits + # next re-throws it. + try + if item isa Base.Event + notify(item) + else + @lock queue.lock enqueue!(queue.inner, item) + end + catch err + queue.failure === nothing && (queue.failure = err) + item isa Base.Event && notify(item) + end + end + return +end +_async_check_failure(queue::AsyncEnqueueQueue) = + queue.failure === nothing || throw(queue.failure) +function enqueue!(aeq::AsyncEnqueueQueue, pair::DTaskPair) + push!(aeq.pending, pair) + length(aeq.pending) >= aeq.limit && _async_hand_off!(aeq) + return +end +function enqueue!(aeq::AsyncEnqueueQueue, pairs::Vector{DTaskPair}) + append!(aeq.pending, pairs) + length(aeq.pending) >= aeq.limit && _async_hand_off!(aeq) + return +end +function _async_hand_off!(aeq::AsyncEnqueueQueue) + _async_check_failure(aeq) + isempty(aeq.pending) && return + # A fresh vector, since the submitter reads this one after we return. + put!(aeq.chan, copy(aeq.pending)) + empty!(aeq.pending) + return +end +function flush_batch!(aeq::AsyncEnqueueQueue) + _async_hand_off!(aeq) + drained = Base.Event() + put!(aeq.chan, drained) + wait(drained) + _async_check_failure(aeq) + return +end +"Stop the submitter task and wait for it to finish draining." +function close_submitter!(aeq::AsyncEnqueueQueue) + try + flush_batch!(aeq) + finally + close(aeq.chan) + wait(aeq.submitter) + end + _async_check_failure(aeq) + return +end +close_submitter!(::AbstractTaskQueue) = nothing + +"`flush_batch!` for queues that may or may not batch." +maybe_flush_batch!(beq::BatchedEnqueueQueue) = flush_batch!(beq) +maybe_flush_batch!(aeq::AsyncEnqueueQueue) = flush_batch!(aeq) +maybe_flush_batch!(::AbstractTaskQueue) = nothing + +""" +Maximum number of tasks a hierarchical partition buffers before submitting. + +Sized from the two costs it trades off. Submitting one task at a time makes the +scheduler round-trip about 40% of this path's per-task planning cost; batching +amortizes it down, and measurably stops paying off past ~16 (a 256-task region +over 4 workers plans at 69 us/task unbatched, 55 at 16, 53.5 unbounded). Against +that, buffered tasks cannot start running, so the batch is what planning gets +ahead of execution -- bounded here at 16 tasks' worth of planning. +""" +const DATADEPS_BATCH_LIMIT = Ref(16) + +""" +Whether uniform (SPMD) planning withholds a region's tasks until it has finished +planning them (see `schedule_partitions_sequential!`). For A/B measurement; the +unbuffered behaviour serializes execution across ranks. +""" +const DATADEPS_UNIFORM_DEFER = Ref(true) + +### Planning instrumentation ### +# +# Where a region's planning time went, per phase, printed once per region when +# `JULIA_DAGGER_HIER_TIMING=1` (or `Dagger.HIER_TIMING[] = true`). Planning cost +# is the thing that limits datadeps at scale -- especially under MPI, where every +# rank replays the whole plan -- and it is hard to attribute from a profile, +# because the expensive parts are blocking waits inside communication rather than +# hot loops. Off by default and costs one `Ref` read per phase. + +const HIER_TIMING = Ref(false) +"Whether `HIER_TIMING` also logs each region's report (vs. only recording it)." +const HIER_TIMING_REPORT = Ref(true) + +mutable struct HierPlanStats + ntasks::Int + nargs::Int + npartitions::Int + prescan_ns::UInt64 + aliasing_ns::UInt64 + dag_ns::UInt64 + partition_ns::UInt64 + schedule_ns::UInt64 + epilogue_ns::UInt64 + # Touched from `generate_slot!` / `aliasing`, which run under parallel + # partition planning, so these are atomic while the phase fields are not. + slot_ns::Threads.Atomic{UInt64} + slot_count::Threads.Atomic{Int} + slot_moved_ns::Threads.Atomic{UInt64} + slot_moved_count::Threads.Atomic{Int} + # Slots whose data was moved even though it was already in the destination + # space: pure overhead, and the reason `slot_is_already_in_place` exists. + slot_samespace_ns::Threads.Atomic{UInt64} + slot_samespace_count::Threads.Atomic{Int} + ainfo_ns::Threads.Atomic{UInt64} + ainfo_count::Threads.Atomic{Int} + # Per-event durations (ns), for distribution rather than just totals. + samples::Dict{Symbol,Vector{UInt64}} + samples_lock::ReentrantLock +end +HierPlanStats() = HierPlanStats(0, 0, 0, 0, 0, 0, 0, 0, 0, + Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), + Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), + Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), + Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), + Dict{Symbol,Vector{UInt64}}(), ReentrantLock()) + +# A `ScopedValue`, not a task-local: the parallel partition path calls +# `generate_slot!` from tasks spawned inside the region, which inherit scoped +# values but not task-locals. +const HIER_STATS = ScopedValue{Union{HierPlanStats,Nothing}}(nothing) + +"Stats for the most recently planned region, for programmatic inspection." +const LAST_HIER_STATS = Ref{Union{HierPlanStats,Nothing}}(nothing) + +"Accumulate `ns` (and one event) into field `f` of the enclosing region's stats, if enabled." +@inline function hier_stat_add!(f::Symbol, ns::Integer, count::Integer=1) + HIER_TIMING[] || return + stats = HIER_STATS[] + stats === nothing && return + Threads.atomic_add!(getfield(stats, f), UInt64(ns)) + Threads.atomic_add!(getfield(stats, Symbol(String(f)[1:end-3] * "_count")), Int(count)) + @lock stats.samples_lock push!(get!(Vector{UInt64}, stats.samples, f), UInt64(ns)) + return +end + +macro hier_phase(stats, field, ex) + quote + local _stats = $(esc(stats)) + if _stats === nothing + $(esc(ex)) + else + local _t0 = time_ns() + local _res = $(esc(ex)) + setfield!(_stats, $(QuoteNode(field)), getfield(_stats, $(QuoteNode(field))) + (time_ns() - _t0)) + _res + end + end +end + +hier_stats_total_ns(stats::HierPlanStats) = + stats.prescan_ns + stats.aliasing_ns + stats.dag_ns + stats.partition_ns + + stats.schedule_ns + stats.epilogue_ns + +function report_hier_stats(stats::HierPlanStats) + total = hier_stats_total_ns(stats) + ms(x) = round(x / 1e6; digits=2) + @info """ + datadeps plan: $(stats.ntasks) tasks, $(stats.nargs) args, $(stats.npartitions) partitions, $(ms(total)) ms total + phase 1 prescan $(ms(stats.prescan_ns)) ms + phase 1 aliasing $(ms(stats.aliasing_ns)) ms ($(stats.ainfo_count[]) ainfos, $(ms(stats.ainfo_ns[])) ms in aliasing()) + phase 2 dag $(ms(stats.dag_ns)) ms + phase 3 partition $(ms(stats.partition_ns)) ms + phase 4 schedule $(ms(stats.schedule_ns)) ms ($(stats.slot_count[]) slots, $(ms(stats.slot_ns[])) ms, of which $(stats.slot_moved_count[]) moved data in $(ms(stats.slot_moved_ns[])) ms, $(stats.slot_samespace_count[]) of those within one space in $(ms(stats.slot_samespace_ns[])) ms) + epilogue $(ms(stats.epilogue_ns)) ms""" + return +end struct HierarchicalTaskMeta pair::DTaskPair @@ -628,7 +889,9 @@ function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) else for i in 1:n arg_w = arg_ws[i] + t0 = time_ns() ainfo = AliasingWrapper(aliasing(accel, arg_w.arg, arg_w.dep_mod)) + hier_stat_add!(:ainfo_ns, time_ns() - t0) results[i] = arg_w => ainfo end end @@ -1138,13 +1401,47 @@ function schedule_partitions_sequential!(queue::DataDepsTaskQueue, partition_procs::Vector{<:Vector{<:Processor}}, vertex_to_partition::Vector{Int}, registry::Union{SharedChunkRegistry,Nothing}, - wait_all_queue) + wait_all_queue, + value_dep_verts::Set{Int}) n_partitions = length(partitions) temp_queues = Vector{DataDepsTaskQueue}(undef, n_partitions) local_scopes = Vector{AbstractScope}(undef, n_partitions) proc_to_scope_lfus = [BasicLFUCache{Processor,AbstractScope}(1024) for _ in 1:n_partitions] shared_state = DataDepsState() write_num = 1 + + # Uniform execution inverts `BatchedEnqueueQueue`'s latency/throughput + # tradeoff, so there we buffer the whole region rather than a few tasks. + # + # Planning is SPMD: every rank replays it, and its slot transfers are + # rendezvous points that all ranks must reach. Letting a task start as soon + # as it is submitted means a rank that picks one up stops planning for the + # length of that task (planning and execution share the thread), and every + # other rank waits in the next slot transfer for it. The ranks then advance + # in lockstep through one task at a time, and the region costs the *sum* of + # its tasks instead of the maximum -- measurably so: a 4-rank stencil sweep + # spent 66 of 74 ms of planning parked in slot transfers, exactly the time + # its peers were computing. + # + # Withholding submission until planning is done keeps every rank computing + # nothing while it plans, so the rendezvous chain runs at metadata speed and + # the tasks then execute concurrently across ranks. + submit_queue = if uniform_execution() && DATADEPS_UNIFORM_DEFER[] + BatchedEnqueueQueue(wait_all_queue, queue_lock; limit=typemax(Int)) + else + # Non-uniform (Distributed) still wants its tasks running while the rest + # of the region is planned, so it keeps a small batch rather than + # withholding everything, and submits each batch off the planning thread + # when there is a thread to submit on (submission is otherwise the single + # largest item in this path's per-task cost, and it needs nothing from the + # planner). With one thread the handoff has nowhere to run and only adds + # a reschedule per batch, so batch in place instead. + if Threads.nthreads() > 1 + AsyncEnqueueQueue(wait_all_queue, queue_lock) + else + BatchedEnqueueQueue(wait_all_queue, queue_lock) + end + end # Shared state already tracks global ownership/history like flat # `distribute_tasks!`. Do not pass `registry` as `ownership`: sync/commit # would fight the single-state history, and the epilogue must then ignore @@ -1160,8 +1457,7 @@ function schedule_partitions_sequential!(queue::DataDepsTaskQueue, end local_scope = UnionScope(map(ExactScope, local_procs)) local_scopes[pid] = local_scope - locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) - temp_queues[pid] = DataDepsTaskQueue(locked_queue; scheduler=similar(queue.scheduler)) + temp_queues[pid] = DataDepsTaskQueue(submit_queue; scheduler=similar(queue.scheduler)) end topo = try @@ -1170,15 +1466,26 @@ function schedule_partitions_sequential!(queue::DataDepsTaskQueue, collect(vertices(dag)) end - with_options(; task_queue=LockedEnqueueQueue(wait_all_queue, queue_lock)) do - for v in topo - pid = vertex_to_partition[v] - local_procs = partition_procs[pid] - isempty(local_procs) && continue - write_num = _schedule_vertex!( - v, pid, temp_queues[pid], shared_state, local_procs, - local_scopes[pid], dag, seen_tasks, vertex_to_partition, - proc_to_scope_lfus[pid], write_num, ownership) + # Copy tasks spawned from within `distribute_task!` go through the same queue + # as the tasks they serve, so that they are held (or not) alongside them. + with_options(; task_queue=submit_queue) do + try + for v in topo + pid = vertex_to_partition[v] + local_procs = partition_procs[pid] + isempty(local_procs) && continue + # A task taking an in-region task's *value* as an argument has + # `distribute_task!` `fetch` that producer, so it must really + # have been submitted; its own turn is topologically earlier. + v in value_dep_verts && maybe_flush_batch!(submit_queue) + write_num = _schedule_vertex!( + v, pid, temp_queues[pid], shared_state, local_procs, + local_scopes[pid], dag, seen_tasks, vertex_to_partition, + proc_to_scope_lfus[pid], write_num, ownership) + end + finally + maybe_flush_batch!(submit_queue) + close_submitter!(submit_queue) end end @@ -1269,16 +1576,28 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) return distribute_tasks!(queue) end + stats = HIER_TIMING[] ? HierPlanStats() : nothing + return with(HIER_STATS => stats) do + _distribute_tasks_hierarchical!(queue, seen_tasks, accel, all_procs, stats) + end +end + +function _distribute_tasks_hierarchical!(queue::DataDepsTaskQueue, + seen_tasks::Vector{DTaskPair}, + accel::Acceleration, + all_procs::Vector{<:Processor}, + stats::Union{HierPlanStats,Nothing}) # Phase 1: Collect arguments and compute aliasing in parallel - task_metas, unique_arg_ws = collect_aliased_args(seen_tasks) - _lookup, ainfos_overlaps, arg_to_ainfo = build_aliasing_parallel(unique_arg_ws) + task_metas, unique_arg_ws = @hier_phase stats prescan_ns collect_aliased_args(seen_tasks) + _lookup, ainfos_overlaps, arg_to_ainfo = + @hier_phase stats aliasing_ns build_aliasing_parallel(unique_arg_ws) # Phase 2: Build dependency DAG - dag = build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) + dag = @hier_phase stats dag_ns build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) # Phase 3: Partition the DAG vertex_to_partition, n_partitions, partition_procs, multi_owner = - partition_dag(dag, task_metas, all_procs) + @hier_phase stats partition_ns partition_dag(dag, task_metas, all_procs) # Detect backing chunks shared across partitions in different memory spaces. # These need runtime ownership transfer to avoid split-brain concurrent @@ -1350,12 +1669,12 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) # See PERF(hier-2)/(hier-3). exec_spaces = unique(Iterators.flatten(memory_spaces(proc) for proc in all_procs)) use_shared_state = uniform_execution(accel) || length(exec_spaces) > 1 - partition_states = try + partition_states = @hier_phase stats schedule_ns try if use_shared_state schedule_partitions_sequential!( queue, queue_lock, partitions, dag, seen_tasks, partition_procs, vertex_to_partition, registry, - wait_all_queue) + wait_all_queue, value_dep_verts) else states = Vector{DataDepsState}(undef, n_partitions) @sync for pid in 1:n_partitions @@ -1392,7 +1711,16 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) # Distributed keeps per-partition states and needs the registry for coherent # cross-partition write-back. epilogue_registry = use_shared_state ? nothing : registry - _hierarchical_copy_from_and_free!(partition_states, length(partition_states), epilogue_registry) + @hier_phase stats epilogue_ns _hierarchical_copy_from_and_free!( + partition_states, length(partition_states), epilogue_registry) + if stats !== nothing + stats.ntasks = length(seen_tasks) + stats.nargs = length(unique_arg_ws) + stats.npartitions = n_partitions + LAST_HIER_STATS[] = stats + HIER_TIMING_REPORT[] && report_hier_stats(stats) + end + return end function _hierarchical_max_write_num(state::DataDepsState, arg_w::ArgumentWrapper) From 84786b3b68a2ccb685c53cdb801e1e00e406827a Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 15 Aug 2026 03:18:58 -0700 Subject: [PATCH 02/26] Datadeps: Memoize, batch and defer aliasing round-trips A chunk's aliasing info cannot be computed locally: under Distributed it is a `remotecall_fetch` to the owner, and under MPI a broadcast from the owner that every rank must join. Planning asks the same questions repeatedly -- once per unique argument to build the DAG, again for every slot, again for the write-back epilogue -- so a region spent hundreds of round-trips re-deriving a handful of distinct answers. Under MPI each is a global synchronization point, which is what made replicated planning scale so poorly with rank count. Three changes, all invisible to the user: * `ChunkAinfoMemo`, a per-region memo keyed on argument identity, the dependency modifier and the acceleration. Per-region because aliasing info describes where a value's memory currently is: stable while one region plans, but a later region's chunk may reuse a freed address. Keys are rank-uniform, so every rank hits and misses on exactly the same calls and the remaining broadcasts are still collective. * `batch_aliasing` / `batch_ainfos`, which resolve a whole uniform list of arguments in one exchange per owning rank rather than one broadcast each, and seed the memo with the results. Phase 1 now goes through these, so the rest of planning finds its answers already computed. * Copies made by Datadeps recorded their own destination-side ainfo eagerly, costing a second rendezvous per slot on top of the transfer. That ainfo only matters when the copy is itself the source of a later move, which most regions never do, so copies are now recorded unresolved and resolved lazily on the first `derived` miss, as one batch. Safe under SPMD because the trigger is uniform. Slot generation halved for a 4-rank stencil sweep (1.97 -> 0.96 ms/sweep), and the 2-rank MPI test suite went from 12m20s to 10m07s. Co-authored-by: Cursor --- ext/MPIExt.jl | 228 ++++++++++++++++++++++++++--------- src/datadeps/aliasing.jl | 206 ++++++++++++++++++++++++++++--- src/datadeps/chunkview.jl | 16 +-- src/datadeps/hierarchical.jl | 49 +++++++- src/datadeps/queue.jl | 13 +- src/memory-spaces.jl | 42 +++---- 6 files changed, 437 insertions(+), 117 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index b5e96d065..47555058f 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -159,34 +159,154 @@ struct MPIAcceleration <: Dagger.Acceleration end MPIAcceleration() = MPIAcceleration(MPI.COMM_WORLD) -function aliasing(accel::MPIAcceleration, x::Chunk, T) - handle = x.handle - # Chunks created under a temporary DistributedAcceleration (or a worker - # thread that did not inherit MPI TLS) carry a DRef; fall back to the - # Distributed unwrap path rather than hard-failing the typeassert. - if !(handle isa MPIRef) +aliasing(accel::MPIAcceleration, x::Chunk, T) = + Dagger.memoized_chunk_aliasing(() -> _aliasing_bcast(accel, x, T), x, T) + +""" + mpi_ainfo_owner(x) -> Int + +The rank whose copy of `x` is the one that can be inspected to derive its +aliasing info, or `-1` when every rank can answer for itself. + +`-1` covers data that no single rank owns: a plain value, or a `Chunk` created +under a temporary `DistributedAcceleration` (or on a worker thread that did not +inherit the MPI TLS) and therefore carrying a `DRef` rather than an `MPIRef`. +""" +mpi_ainfo_owner(x::Chunk) = x.handle isa MPIRef ? x.handle.rank : -1 +mpi_ainfo_owner(@nospecialize(x)) = -1 + +# `x`'s aliasing info as its owner sees it, stamped with the owning rank so that +# spans from different ranks -- whose SPMD heaps have very similar address +# layouts -- can never appear to alias each other. +function mpi_owner_ainfo(@nospecialize(x), dep_mod, owner::Int) + ainfo = _with_default_acceleration() do + mpi_raw_aliasing(x, dep_mod) + end + return owner < 0 ? ainfo : mpi_remap_ainfo(ainfo, owner) +end + +# How to read the local aliasing info out of each kind of handle. Runs under the +# default acceleration (see `mpi_owner_ainfo`), on the owning rank only. +mpi_raw_aliasing(@nospecialize(x), dep_mod) = aliasing(x, dep_mod) + +""" + _aliasing_bcast(accel::MPIAcceleration, x, dep_mod) + +The owner computes `x`'s aliasing info and broadcasts it; every other rank +receives it. All ranks must call this at the same logical point, since it is a +collective. + +Prefer letting Phase 1 resolve a whole region's arguments in one exchange +(`batch_aliasing`) and reading the results back out of the region memo; this +per-argument form is the fallback for arguments the batch did not cover. +""" +function _aliasing_bcast(accel::MPIAcceleration, @nospecialize(x), dep_mod) + owner = mpi_ainfo_owner(x) + if owner < 0 return _with_default_acceleration() do - aliasing(x, T) + mpi_raw_aliasing(x, dep_mod) end end - @assert accel.comm == handle.comm "MPIAcceleration comm mismatch" tag = to_tag() check_uniform(tag) rank = MPI.Comm_rank(accel.comm) - if handle.rank == rank - ainfo = _with_default_acceleration() do - aliasing(x, T) - end - ainfo = mpi_remap_ainfo(ainfo, handle.rank) + ainfo = if owner == rank @opcounter :aliasing_bcast_send_yield - ainfo = bcast_yield(accel.comm, handle.rank, tag, ainfo) + bcast_yield(accel.comm, owner, tag, mpi_owner_ainfo(x, dep_mod, owner)) else - ainfo = bcast_yield(accel.comm, handle.rank, tag) + bcast_yield(accel.comm, owner, tag) end check_uniform(ainfo) return ainfo end +""" + batch_aliasing(accel::MPIAcceleration, arg_ws) + +Resolve a whole batch of arguments' aliasing info in one exchange per owning +rank, instead of one broadcast per argument. + +Replicated planning makes every per-argument `aliasing` call a rendezvous: a rank +that reaches argument `i` waits there until all its peers do, so a region's +planning becomes a chain of `nargs` round-trips whose length grows with both the +region and the rank count. The batch is a *uniform* list -- every rank walks the +same arguments in the same order -- so an index into it names the same argument +everywhere, which is all that is needed to exchange the answers wholesale: each +owning rank broadcasts the ainfos for its own arguments keyed by index. + +The results also seed the region's aliasing memo, so the rest of planning (slot +generation, remainder computation, the write-back epilogue) finds them already +answered and communicates no further. +""" +function Dagger.batch_aliasing(accel::MPIAcceleration, arg_ws::Vector{Dagger.ArgumentWrapper}) + n = length(arg_ws) + n == 0 && return Pair{Dagger.ArgumentWrapper,Dagger.AliasingWrapper}[] + objs = Any[arg_w.arg for arg_w in arg_ws] + dep_mods = Any[arg_w.dep_mod for arg_w in arg_ws] + ainfos = Dagger.batch_ainfos(accel, objs, dep_mods) + return Pair{Dagger.ArgumentWrapper,Dagger.AliasingWrapper}[ + arg_ws[i] => Dagger.AliasingWrapper(ainfos[i]) for i in 1:n] +end + +""" + batch_ainfos(accel::MPIAcceleration, objs, dep_mods) + +One exchange per owning rank for a whole batch of objects, instead of one +broadcast each. See `batch_aliasing` above for why this shape matters. +""" +function Dagger.batch_ainfos(accel::MPIAcceleration, objs::Vector, dep_mods::Vector) + n = length(objs) + n == 0 && return AbstractAliasing[] + comm = accel.comm + rank = MPI.Comm_rank(comm) + check_uniform(UInt64(n)) + t0 = time_ns() + + owners = Vector{Int}(undef, n) + ainfos = Vector{Union{Nothing,AbstractAliasing}}(nothing, n) + mine = Dict{Int,AbstractAliasing}() + for i in 1:n + owner = mpi_ainfo_owner(objs[i]) + owners[i] = owner + if owner < 0 + # Answerable on every rank, so nothing to exchange + ainfos[i] = mpi_owner_ainfo(objs[i], dep_mods[i], owner) + elseif owner == rank + mine[i] = mpi_owner_ainfo(objs[i], dep_mods[i], owner) + end + end + + # `owners` is itself uniform, so every rank knows which ranks have something + # to say and skips the rest: one broadcast per *contributing* rank. + roots = sort!(unique!(filter(>=(0), copy(owners)))) + for root in roots + tag = to_tag() + check_uniform(tag) + received = if root == rank + @opcounter :aliasing_bcast_send_yield + bcast_yield(comm, root, tag, mine) + else + bcast_yield(comm, root, tag) + end::Dict{Int,AbstractAliasing} + for (i, ainfo) in received + ainfos[i] = ainfo + end + end + + results = Vector{AbstractAliasing}(undef, n) + for i in 1:n + ainfo = ainfos[i] + if ainfo === nothing + error("batch_ainfos: rank $rank received no aliasing info for object $i (owner $(owners[i]))") + end + # Take back whatever the memo now holds, so this batch and every later + # lookup of the same object agree even if an entry was already there. + results[i] = Dagger.memoize_ainfo!(Dagger.ainfo_memo_key(objs[i], dep_mods[i]), ainfo) + end + Dagger.hier_stat_add!(:ainfo_ns, time_ns() - t0, n) + return results +end + default_processor(accel::MPIAcceleration) = MPIOSProc(accel.comm, 0) default_processor(accel::MPIAcceleration, x) = MPIOSProc(accel.comm, 0) default_processor(accel::MPIAcceleration, x::Chunk) = MPIOSProc(x.handle.comm, x.handle.rank) @@ -1657,7 +1777,15 @@ end struct MPIWireValue{T} value::Union{Some{T},Nothing} space::MPIMemorySpace -end + # Aliasing-memo identity of the `Chunk` this value stands in for, or 0 for a + # sub-value produced by wrapper recursion (which has no chunk of its own). + # Asking for a wire value's aliasing info means broadcasting from the owner -- + # a global synchronization point during planning -- so when the answer is + # already memoized for the originating chunk, reuse it (`Dagger.memoized_ainfo`). + origin_key::UInt +end +MPIWireValue{T}(value::Union{Some{T},Nothing}, space::MPIMemorySpace) where T = + MPIWireValue{T}(value, space, UInt(0)) wire_type(::MPIWireValue{T}) where T = T has_value(w::MPIWireValue) = w.value !== nothing wire_value(w::MPIWireValue) = something(w.value) @@ -1666,9 +1794,12 @@ memory_space(w::MPIWireValue) = w.space default_memory_space(accel::MPIAcceleration, w::MPIWireValue) = w.space function check_uniform(w::MPIWireValue{T}, original=w) where T - # Compare logical metadata only: the value itself is rank-local + # Compare logical metadata only: the value itself is rank-local. The origin + # key gates a collective (see `MPIWireValue`), so a rank disagreeing on it + # would hang in `aliasing` rather than fail loudly. return check_uniform(hash(T), original) && - check_uniform(w.space, original) + check_uniform(w.space, original) && + check_uniform(w.origin_key, original) end function tochunk(w::MPIWireValue{T}, proc::P, scope::S=Dagger.AnyScope(); kwargs...) where {T,P<:Processor,S} @@ -1714,36 +1845,30 @@ mpi_remap_ainfo(a::Dagger.AliasingWrapper, owner::Int) = Dagger.AliasingWrapper(mpi_remap_ainfo(a.inner, owner)) mpi_remap_ainfo(a::Dagger.AbstractAliasing, owner::Int) = a -# Owner computes the aliasing info for its local value and broadcasts it; -# all ranks must call this collectively at the same logical point. +# A wire value stands in for the chunk it was unwrapped from, so its aliasing +# info is the chunk's: reuse the memo entry rather than repeating the collective. function Dagger.aliasing(accel::MPIAcceleration, w::MPIWireValue, dep_mod) - tag = to_tag() - check_uniform(tag) - rank = MPI.Comm_rank(accel.comm) - if w.space.rank == rank - ainfo = Dagger._with_default_acceleration() do - Dagger.aliasing(wire_value(w), dep_mod) - end - ainfo = mpi_remap_ainfo(ainfo, w.space.rank) - @opcounter :aliasing_bcast_send_yield - ainfo = bcast_yield(accel.comm, w.space.rank, tag, ainfo) - else - ainfo = bcast_yield(accel.comm, w.space.rank, tag) - end - check_uniform(ainfo) - return ainfo + w.origin_key == 0 && return _aliasing_bcast(accel, w, dep_mod) + return Dagger.memoized_ainfo(() -> _aliasing_bcast(accel, w, dep_mod), + Dagger.ainfo_memo_key(w.origin_key, dep_mod)) end +mpi_ainfo_owner(w::MPIWireValue) = w.space.rank +mpi_raw_aliasing(w::MPIWireValue, dep_mod) = Dagger.aliasing(wire_value(w), dep_mod) + # All ranks enter collectively; the owner unwraps its local value, all other # ranks construct a wire proxy carrying only type and origin space. The # `move_rewrap` recursion below then walks the same path on every rank. function remotecall_endpoint_toplevel(f, accel::MPIAcceleration, cache::AliasedObjectCache, from_proc, to_proc, from_space, to_space, data::Chunk) local_rank = MPI.Comm_rank(accel.comm) T = chunktype(data) + # The chunk's identity alone: `aliasing` combines it with the `dep_mod` to + # form the same memo key `ainfo_memo_key` would. + origin_key = Dagger._identity_hash(data) w = if local_rank == from_space.rank - MPIWireValue{T}(Some{T}(unwrap(data)), from_space) + MPIWireValue{T}(Some{T}(unwrap(data)), from_space, origin_key) else - MPIWireValue{T}(nothing, from_space) + MPIWireValue{T}(nothing, from_space, origin_key) end return f(accel, cache, from_proc, to_proc, from_space, to_space, w)::Chunk end @@ -2011,27 +2136,16 @@ function move_rewrap(accel::MPIAcceleration, cache::AliasedObjectCache, from_pro end # Owner computes the view's aliasing info locally and broadcasts it -function aliasing(accel::MPIAcceleration, x::ChunkView, dep_mod) +aliasing(accel::MPIAcceleration, x::ChunkView, dep_mod) = + Dagger.memoized_chunk_aliasing(() -> _aliasing_bcast(accel, x, dep_mod), x, dep_mod) + +mpi_ainfo_owner(x::ChunkView) = mpi_ainfo_owner(x.chunk) +function mpi_raw_aliasing(x::ChunkView, dep_mod) @assert dep_mod === identity "Dependency modifiers not yet supported for ChunkView: $dep_mod" - handle = x.chunk.handle::MPIRef - tag = to_tag() - check_uniform(tag) - rank = MPI.Comm_rank(accel.comm) - if handle.rank == rank - ainfo = _with_default_acceleration() do - v = view(unwrap(x.chunk), x.slices...) - # Resolve whole-object containers (e.g. `DSparseArray`) where `v` - # lives; see `aliasing_unwrapped`. - aliasing_unwrapped(v, dep_mod) - end - ainfo = mpi_remap_ainfo(ainfo, handle.rank) - @opcounter :aliasing_bcast_send_yield - ainfo = bcast_yield(accel.comm, handle.rank, tag, ainfo) - else - ainfo = bcast_yield(accel.comm, handle.rank, tag) - end - check_uniform(ainfo) - return ainfo + v = view(unwrap(x.chunk), x.slices...) + # Resolve whole-object containers (e.g. `DSparseArray`) where `v` + # lives; see `aliasing_unwrapped`. + return aliasing_unwrapped(v, dep_mod) end # The aliased-object cache is a per-rank replicated store; every rank updates diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index ef27070f1..31644bd64 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -225,6 +225,106 @@ _identity_hash(arg::Chunk, h::UInt=UInt(0)) = hash(arg.handle, hash(Chunk, h)) _identity_hash(arg::SubArray, h::UInt=UInt(0)) = hash(arg.indices, hash(arg.offset1, hash(arg.stride1, _identity_hash(arg.parent, h)))) _identity_hash(arg::CartesianIndices, h::UInt=UInt(0)) = hash(arg.indices, hash(typeof(arg), h)) +""" + ChunkAinfoMemo + +Per-region memo for the aliasing info of a `Chunk` / `ChunkView`. + +A chunk's aliasing info cannot be computed locally: under Distributed it takes a +`remotecall_fetch` to the owner, and under MPI a broadcast from the owner that +every rank has to join. Datadeps asks the same question repeatedly while planning +one region -- once per unique argument to build the dependency DAG, again for +every slot (`aliased_object!`, `aliasing!`), and again for the write-back +epilogue -- so left unmemoized a single region spends hundreds of round-trips +re-deriving a handful of distinct answers. Under MPI each of those is a +*global synchronization point*, which is what makes replicated planning scale +poorly with rank count. + +Memoization is per-region because aliasing info describes where a value's memory +currently is: stable while one region plans (datadeps copies into existing +buffers, it never relocates a live argument), but a chunk allocated by a later +region may well reuse a freed address. + +Keys (`ainfo_memo_key`) combine the argument's identity, the dependency modifier +and the acceleration, and are identical on every rank (chunk handles hash by rank ++ id). Uniform keys are what makes the memo safe under SPMD replay: every rank +hits and misses on exactly the same calls, so the broadcasts that remain are +still reached collectively by all ranks. +""" +struct ChunkAinfoMemo + entries::Dict{UInt,AbstractAliasing} + lock::ReentrantLock +end +ChunkAinfoMemo() = ChunkAinfoMemo(Dict{UInt,AbstractAliasing}(), ReentrantLock()) + +const CHUNK_AINFO_MEMO = ScopedValue{Union{ChunkAinfoMemo,Nothing}}(nothing) + +"The memo key for `arg`'s aliasing info under `dep_mod`; see `ChunkAinfoMemo`." +ainfo_memo_key(arg, dep_mod) = ainfo_memo_key(_identity_hash(arg), dep_mod) + +""" + ainfo_memo_key(idhash::UInt, dep_mod) -> UInt + +The memo key for an argument whose identity hash is already known (an indirect +handle, e.g. an MPI wire value, keys on the chunk it stands in for). + +The current acceleration is part of the key because the memo answers "what does +*this* acceleration say the aliasing is". Under MPI the acceleration-level answer +is the owner's local answer rank-stamped and broadcast, and it is derived by +re-asking the *default* acceleration for the raw local one -- a nested query with +the same argument. Keying both alike lets the outer, uniform answer be displaced +by the inner, rank-local one: the owner then returns an unstamped span while +every other rank holds the stamped one, and the region fails its uniformity check +(or, with checks off, mistakes two ranks' buffers for the same memory). +""" +ainfo_memo_key(idhash::UInt, dep_mod) = + hash(accel_kind(current_acceleration()), hash(dep_mod, idhash)) + +"Run `f` (a chunk-aliasing computation) at most once per region per `(arg, dep_mod)`." +memoized_chunk_aliasing(f, arg, dep_mod) = + memoized_ainfo(f, ainfo_memo_key(arg, dep_mod)) + +""" + memoized_ainfo(f, key::UInt) + +`memoized_chunk_aliasing` for callers that hold a key rather than the argument +itself, which is how an indirect handle (e.g. an MPI wire value standing in for a +chunk) reuses the entry already computed for the chunk it came from. +""" +function memoized_ainfo(f, key::UInt) + memo = CHUNK_AINFO_MEMO[] + memo === nothing && return f() + @lock memo.lock begin + cached = get(memo.entries, key, nothing) + cached === nothing || return cached + end + # Computed outside the lock: it blocks on the owner, and under parallel + # partition planning holding the lock across that would serialize planners. + # Two planners racing on one key both compute; the loser's result is dropped. + ainfo = f() + @lock memo.lock begin + return get!(memo.entries, key, ainfo) + end +end + +""" + memoize_ainfo!(key::UInt, ainfo) -> AbstractAliasing + +Record `ainfo` as the region's answer for `key`, for a caller that obtained it +some other way than by asking `memoized_ainfo` -- a batch exchange +(`batch_aliasing`) resolves many arguments at once, and seeding its results here +is what keeps the rest of planning from re-deriving them one collective at a +time. Returns the entry in force, which is the existing one if any (a batch never +contradicts what has already been computed). +""" +function memoize_ainfo!(key::UInt, ainfo::AbstractAliasing) + memo = CHUNK_AINFO_MEMO[] + memo === nothing && return ainfo + @lock memo.lock begin + return get!(memo.entries, key, ainfo) + end +end + struct ArgumentWrapper arg dep_mod @@ -316,6 +416,9 @@ struct AliasedObjectCacheStore # which is exactly where its value is the user's own object. Every *other* # space holding that key is a copy we allocated and may free. originals::Set{Tuple{MemorySpace,AbstractAliasing}} + # Copies whose own (destination-side) ainfo has not been computed yet, as + # `key => copy`. See `resolve_pending!`. + pending::Vector{Pair{AbstractAliasing,Chunk}} end AliasedObjectCacheStore(accel::Acceleration) = AliasedObjectCacheStore(accel, @@ -324,7 +427,69 @@ AliasedObjectCacheStore(accel::Acceleration) = Dict{MemorySpace,Set{AbstractAliasing}}(), Dict{MemorySpace,Dict{AbstractAliasing,Chunk}}(), Dict{MemorySpace,Dict{AbstractAliasing,AbstractAliasing}}(), - Set{Tuple{MemorySpace,AbstractAliasing}}()) + Set{Tuple{MemorySpace,AbstractAliasing}}(), + Vector{Pair{AbstractAliasing,Chunk}}()) + +""" + resolve_pending!(cache) -> Bool + +Give every copy recorded by `set_stored!` its `derived` entry, returning whether +there was anything to do. + +`set_stored!` defers this because obtaining a copy's own ainfo is expensive +exactly where it is least likely to be needed. The ainfo describes pointer spans +in the destination space, so under MPI only the destination rank can compute it +and it has to be broadcast to every other rank -- a second rendezvous per slot, +on top of the transfer, which measured as about a quarter of `distribute_task!` +for a 2-rank stencil sweep. What it buys is the ability to recognize a copy when +that copy is itself the *source* of a later move, so that both hops share one +cache key. Regions that never take a second hop -- the common case -- never need +it at all. + +Deferring is safe under SPMD because the trigger is uniform: `derived` and the +ainfo being looked up are rank-uniform, so every rank misses, and resolves, at +the same point and in the same order. +""" +function resolve_pending!(cache::AliasedObjectCacheStore) + isempty(cache.pending) && return false + # Task-local scratch: snapshotting into a fresh `Vector` every call (here, + # and for `values`/`dep_mods` below) added up under an iterative workload + # that resolves a pending batch every sweep. Safe to reuse because + # `cache.pending` is cleared before any call that could re-enter this + # function (see the module docstring on deferred resolution being + # SPMD-safe), so there is nothing left to alias by the time that could + # happen. + entries = @reusable_vector :resolve_pending_entries Pair{AbstractAliasing,Chunk} nothing 32 + append!(entries, cache.pending) + empty!(cache.pending) + @opcounter :aliasing_resolve_pending + @opcounter :aliasing_resolve_pending_entries length(entries) + # Resolved as a batch: the whole point of deferring was to not pay a + # rendezvous per copy, which asking one at a time here would reintroduce. + values = @reusable_vector :resolve_pending_values Chunk nothing 32 + dep_mods = @reusable_vector :resolve_pending_dep_mods Any nothing 32 + for (_, value) in entries + push!(values, value) + push!(dep_mods, identity) + end + for (i, value_ainfo) in enumerate(batch_ainfos(cache.accel, values, dep_mods)) + key, value = entries[i] + cache.derived[value_ainfo] = key + # Keep `value_ainfo` around: this is the one point where the buffer's own + # aliasing is computed on every rank, and the free loop needs it later. + ainfos_dict = get!(Dict{AbstractAliasing,AbstractAliasing}, cache.value_ainfos, memory_space(value)) + ainfos_dict[key] = value_ainfo + end + return true +end + +"`cache.derived[ainfo]`, resolving deferred copies first, or `nothing` if absent." +function derived_key(cache::AliasedObjectCacheStore, ainfo::AbstractAliasing) + key = get(cache.derived, ainfo, nothing) + key === nothing || return key + resolve_pending!(cache) || return nothing + return get(cache.derived, ainfo, nothing) +end """ is_original(cache, space, ainfo) -> Bool @@ -340,33 +505,26 @@ function is_stored(cache::AliasedObjectCacheStore, space::MemorySpace, ainfo::Ab if !haskey(cache.stored, space) return false end - if !haskey(cache.derived, ainfo) - return false - end - key = cache.derived[ainfo] + key = derived_key(cache, ainfo) + key === nothing && return false return key in cache.stored[space] end function is_key_present(cache::AliasedObjectCacheStore, space::MemorySpace, ainfo::AbstractAliasing) - return haskey(cache.derived, ainfo) + return derived_key(cache, ainfo) !== nothing end function get_stored(cache::AliasedObjectCacheStore, space::MemorySpace, ainfo::AbstractAliasing) - @assert is_stored(cache, space, ainfo) "Cache does not have derived ainfo $ainfo" - key = cache.derived[ainfo] + key = derived_key(cache, ainfo) + @assert key !== nothing "Cache does not have derived ainfo $ainfo" return cache.values[space][key] end function set_stored!(cache::AliasedObjectCacheStore, dest_space::MemorySpace, value::Chunk, ainfo::AbstractAliasing) @assert !is_stored(cache, dest_space, ainfo) "Cache already has derived ainfo $ainfo" @check_uniform(value) key = cache.derived[ainfo] - value_ainfo = aliasing(cache.accel, value, identity) - cache.derived[value_ainfo] = key + push!(cache.pending, key => value) push!(get!(Set{AbstractAliasing}, cache.stored, dest_space), key) values_dict = get!(Dict{AbstractAliasing,Chunk}, cache.values, dest_space) values_dict[key] = value - # Keep `value_ainfo` around: this is the one point where the buffer's own - # aliasing is computed on every rank, and the free loop needs it later. - ainfos_dict = get!(Dict{AbstractAliasing,AbstractAliasing}, cache.value_ainfos, dest_space) - ainfos_dict[key] = value_ainfo return end @@ -374,14 +532,21 @@ end stored_value_ainfo(cache, space, key) -> Union{AbstractAliasing,Nothing} The aliasing of the buffer `cache.values[space][key]` *in `space`*, as recorded -by `set_stored!` when the buffer was allocated. `nothing` for the user's -original data, which never gets one (and is never freed). +by `set_stored!` when the buffer was allocated (possibly deferred; see +`resolve_pending!`). `nothing` for the user's original data, which never gets +one (and is never freed). This is deliberately a lookup rather than a fresh `aliasing` call: under uniform (SPMD) execution `aliasing` is a collective, so the buffer's extent cannot be -recomputed once planning is over. +recomputed once planning is over. A miss instead resolves any batch of copies +still pending, exactly like `derived_key` -- the free loop may be the first +thing to ask about a copy that no other lookup happened to flush already. """ function stored_value_ainfo(cache::AliasedObjectCacheStore, space::MemorySpace, key::AbstractAliasing) + ainfos = get(cache.value_ainfos, space, nothing) + value = ainfos === nothing ? nothing : get(ainfos, key, nothing) + value === nothing || return value + resolve_pending!(cache) || return nothing ainfos = get(cache.value_ainfos, space, nothing) ainfos === nothing && return nothing return get(ainfos, key, nothing) @@ -462,10 +627,11 @@ function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(cache.a else y = f(x) @assert y isa Chunk "Didn't get a Chunk from functor" + # N.B. Deliberately not also checking that `y`'s ainfo differs from `x`'s + # when the spaces differ: distinct memory spaces hold distinct memory, so + # the space assertion above already implies it, and asking for `y`'s ainfo + # here would cost a collective per slot (see `resolve_pending!`). @assert memory_space(y) == cache.space "Space mismatch! $(memory_space(y)) != $(cache.space)" - if memory_space(x) != cache.space - @assert ainfo != aliasing(cache.accel, y, identity) "Aliasing mismatch! $ainfo == $(aliasing(cache.accel, y, identity))" - end set_stored!(cache.accel, cache, y, ainfo) return y end diff --git a/src/datadeps/chunkview.jl b/src/datadeps/chunkview.jl index 1593a6ddc..27899f38c 100644 --- a/src/datadeps/chunkview.jl +++ b/src/datadeps/chunkview.jl @@ -66,13 +66,15 @@ Base.view(c::DTask, slices...) = view(fetch(c; raw=true), slices...) function aliasing(accel::Acceleration, x::ChunkView{N}, dep_mod) where N @assert dep_mod === identity "Dependency modifiers not yet supported for ChunkView: $dep_mod" - return remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices - x = unwrap(x) - v = view(x, slices...) - # A view of a whole-object container (e.g. `DSparseArray`) must alias the - # entire container; `aliasing_unwrapped` resolves that (otherwise it just - # aliases the view), and crucially does so here where `v` lives. - return aliasing_unwrapped(v) + return memoized_chunk_aliasing(x, dep_mod) do + remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices + x = unwrap(x) + v = view(x, slices...) + # A view of a whole-object container (e.g. `DSparseArray`) must alias the + # entire container; `aliasing_unwrapped` resolves that (otherwise it just + # aliases the view), and crucially does so here where `v` lives. + return aliasing_unwrapped(v) + end end end aliasing(x::ChunkView) = aliasing(current_acceleration(), x, identity) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 2ce45d7be..5e45150c9 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -797,7 +797,7 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr # `_compute_aliasing_batch` still uses threads internally when there # are enough args to make it worthwhile. wid, worker_args = only(by_worker) - results = wid == myid() ? _compute_aliasing_batch(worker_args) : + results = wid == myid() ? batch_aliasing(current_acceleration(), worker_args) : remotecall_fetch(_compute_aliasing_batch, wid, worker_args) # Key by the *local* `arg_w`, not the pair's: for a remote worker the # returned `ArgumentWrapper` is a deserialized copy that need not be @@ -823,7 +823,7 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr # insidiously, as silently wrong aliasing when it is longer). local results results = if wid == myid() - _compute_aliasing_batch(worker_args) + batch_aliasing(current_acceleration(), worker_args) else remotecall_fetch(_compute_aliasing_batch, wid, worker_args) end @@ -863,13 +863,56 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr return lookup, ainfos_overlaps, arg_to_ainfo end +""" + batch_aliasing(accel, arg_ws) -> Vector{Pair{ArgumentWrapper,AliasingWrapper}} + +Phase 1's aliasing computation for a whole batch of arguments, in input order. + +The default walks the batch one argument at a time. An acceleration that has to +*communicate* to answer (MPI, where the answer is the owner's and every rank +needs it) should override this to exchange the batch in one go: per-argument +collectives make a region's planning a chain of `nargs` rendezvous, which is +latency the batch shape can pay once instead. + +Distributed does not need its own override for this reason: `aliasing` there +is a plain `remotecall_fetch` to one owning worker rather than a collective, so +there is no per-argument rendezvous *chain* to collapse in the first place -- +only a rendezvous *count*, and `build_aliasing_parallel` already collapses that +by grouping `arg_ws` by owning worker before calling this function once per +worker (see its `by_worker` grouping). Overriding this function too would just +re-batch already-grouped, single-worker input; the default's own +`Threads.@threads` split (`_compute_aliasing_batch`) is what actually +parallelizes each worker's share once its data has landed. +""" +batch_aliasing(::Acceleration, arg_ws::Vector{ArgumentWrapper}) = + _compute_aliasing_batch(arg_ws) + +""" + batch_ainfos(accel, objs, dep_mods) -> Vector{AbstractAliasing} + +`aliasing(accel, objs[i], dep_mods[i])` for a whole uniform list at once, in +input order. + +The list is what `batch_aliasing` exchanges underneath, exposed separately for +the other places that end up holding a batch of objects needing ainfos -- +notably `resolve_pending!`, which is resolving copies rather than the region's +arguments. Same reasoning: an acceleration that must communicate to answer +should override this and pay one rendezvous rather than `length(objs)`. +""" +batch_ainfos(accel::Acceleration, objs::Vector, dep_mods::Vector) = + AbstractAliasing[aliasing(accel, objs[i], dep_mods[i]) for i in eachindex(objs)] + # Below this many args, the fixed cost of forking/joining `Threads.@threads` # outweighs the benefit of parallelizing the (typically cheap) `aliasing()` calls. const COMPUTE_ALIASING_BATCH_MIN_PARALLEL = 8 function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) n = length(arg_ws) - results = Vector{Pair{ArgumentWrapper, AliasingWrapper}}(undef, n) + # Task-local: safe to reuse across calls because the result is fully + # consumed by the caller (copied into `arg_to_ainfo`) before this task + # could re-enter this function for another region. + results = @reusable_vector :compute_aliasing_batch_results Pair{ArgumentWrapper,AliasingWrapper} nothing 32 + resize!(results, n) accel = current_acceleration() # Under uniform execution (MPI), aliasing may perform collectives that must # run in the same sequential order on every rank -- never Threads.@threads. diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 0dab5ddd7..9e89d2d6a 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -84,10 +84,15 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, # local) so it shows up by name in stacktraces and profiles, which is # the boundary between region setup and the whole planning pipeline. function run_distribute(queue) - if hierarchical - distribute_tasks_hierarchical!(queue) - else - distribute_tasks!(queue) + # One aliasing memo per region: planning asks for the same chunks' + # aliasing info from Phase 1, from every slot, and from the write-back + # epilogue, and each answer costs a round-trip to the owner. + with(CHUNK_AINFO_MEMO => ChunkAinfoMemo()) do + if hierarchical + distribute_tasks_hierarchical!(queue) + else + distribute_tasks!(queue) + end end end if launch_wait diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 2796b6636..63d18cf05 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -661,34 +661,24 @@ aliasing_unwrapped(x) = aliasing(aliasing_root(x)) aliasing_unwrapped(x, dep_mod) = aliasing(aliasing_root(x), dep_mod) function aliasing(x::Chunk, T) - # Under uniform execution (MPI), `root_worker_id` is always `myid()` and is - # not a valid owner key -- defer to the acceleration so non-owning ranks - # take the owner-broadcast path instead of a local unwrap. - accel = current_acceleration() - if uniform_execution(accel) - return aliasing(accel, x, T) - end - if root_worker_id(x.processor) == myid() - return aliasing_unwrapped(unwrap(x), T) - end - @assert x.handle isa DRef - return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T - aliasing_unwrapped(unwrap(x), T) - end -end -function aliasing(x::Chunk) - accel = current_acceleration() - if uniform_execution(accel) - return aliasing(accel, x, identity) - end - if root_worker_id(x.processor) == myid() - return aliasing_unwrapped(unwrap(x)) - end - @assert x.handle isa DRef - return remotecall_fetch(root_worker_id(x.processor), x) do x - aliasing_unwrapped(unwrap(x)) + return memoized_chunk_aliasing(x, T) do + # Under uniform execution (MPI), `root_worker_id` is always `myid()` and is + # not a valid owner key -- defer to the acceleration so non-owning ranks + # take the owner-broadcast path instead of a local unwrap. + accel = current_acceleration() + if uniform_execution(accel) + return aliasing(accel, x, T) + end + if root_worker_id(x.processor) == myid() + return aliasing_unwrapped(unwrap(x), T) + end + @assert x.handle isa DRef + return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T + aliasing_unwrapped(unwrap(x), T) + end end end +aliasing(x::Chunk) = aliasing(x, identity) aliasing(x::DTask, T) = aliasing(fetch(x; move_value=false, unwrap=false), T) aliasing(x::DTask) = aliasing(fetch(x; move_value=false, unwrap=false)) From 3f846fba38b1f288aab1ade9b957d723edcd0e85 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 7 Sep 2026 14:08:08 -0700 Subject: [PATCH 03/26] Benchmark CI: Add per-job regression/improvement summary, bump default samples The regressions/improvements lists in the CI report require opening the report to see, and the "within noise" breakdown was buried behind a
section entirely -- there was no way to tell at a glance which suite/method jobs actually moved. Add a summary table (job, regression count, improvement count, within-noise count) at the top of both the Markdown report and the stdout summary, grouped by the `suite/method` prefix that the `BENCHMARK` spec already uses to identify one job. Also bump the default BENCHMARK_SAMPLES from 5 to 7 for less noisy CI comparisons. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01NHPKqRaGTGLt7CPKyZGuPe --- benchmark/benchmarks.jl | 2 +- benchmark/ci.jl | 48 +++++++++++++++++++++++++++++++++++++++++ benchmark/common.jl | 2 +- 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 1f5f45273..6f18a6492 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -99,7 +99,7 @@ # request/response protocol -- see `run_all_mpi` vs `run_all_external`). # Defaults to "0" (disabled; plain subprocess, worker.jl). # - BENCHMARK_SECONDS: Time budget (seconds) per benchmark. Defaults to "30". -# - BENCHMARK_SAMPLES: Max samples per benchmark. Defaults to "5". +# - BENCHMARK_SAMPLES: Max samples per benchmark. Defaults to "7". # - BENCHMARK_PROC_TIMEOUT: Wall-clock seconds to wait for a single benchmark # before assuming the worker is wedged, killing it, and treating it like an # OOM (skip this and larger scales, restart). Defaults to "3600". diff --git a/benchmark/ci.jl b/benchmark/ci.jl index 6af535478..944acb8bf 100644 --- a/benchmark/ci.jl +++ b/benchmark/ci.jl @@ -194,11 +194,56 @@ sort!(within_noise; by = last, rev = true) pct(r) = string(round((r - 1) * 100; digits = 1), "%") +# --- Per-job summary (counts at a glance, no need to open the detail lists) - +# +# A "job" is a benchmark's `suite/method[+accels]` prefix -- e.g. +# `stencil/dagger` -- which is exactly one entry of the `BENCHMARK` spec +# (`suite:method+accel,...;...`). Grouping at that level keeps the summary to +# one row per suite/method combination actually run, regardless of how many +# individual benchmarks or metrics it contains. +job_of(name) = join(split(name, '/')[1:min(2, end)], "/") + +job_names = Set{String}() +for (name, _) in cur + name == "time_to_load" && continue + push!(job_names, job_of(name)) +end + +function job_tally(entries) + counts = Dict{String,Int}() + for (name, _metric, _ratio) in entries + j = job_of(name) + counts[j] = get(counts, j, 0) + 1 + end + return counts +end +job_regression_counts = job_tally(regressions) +job_improvement_counts = job_tally(improvements) +job_noise_counts = job_tally(within_noise) + +# Worst jobs (most regressions, then fewest improvements) sort to the top. +job_summary_rows = [(j, get(job_regression_counts, j, 0), get(job_improvement_counts, j, 0), + get(job_noise_counts, j, 0)) for j in job_names] +sort!(job_summary_rows; by = r -> (-r[2], -r[3], r[1])) + +function job_summary_table(io, rows) + println(io, "| Job | Regressions | Improvements | Within noise |") + println(io, "|:---|---:|---:|---:|") + for (job, nreg, nimp, nnoise) in rows + marker = nreg > 0 ? " ⚠️" : (nimp > 0 ? " ✅" : "") + println(io, "| `", job, "`", marker, " | ", nreg, " | ", nimp, " | ", nnoise, " |") + end +end + # --- Markdown report (for the Buildkite annotation / optional PR comment) --- open(joinpath(OUTPUT_DIR, "report.md"), "w") do io println(io, "### Dagger benchmarks: `$CUR_REV` vs `$BASE_REV`") println(io) + println(io, "#### Summary by job") + println(io) + job_summary_table(io, job_summary_rows) + println(io) println(io, "#### Median time") println(io) println(io, table) @@ -256,6 +301,9 @@ end # --- Summary + exit status ------------------------------------------------- +println("\nSummary by job:\n") +job_summary_table(stdout, job_summary_rows) + if !isempty(within_noise) println("\n$(length(within_noise)) metric(s) moved past their threshold but stayed within the measured spread (not counted):") for (name, metric, r) in within_noise diff --git a/benchmark/common.jl b/benchmark/common.jl index e988be090..50ba1df03 100644 --- a/benchmark/common.jl +++ b/benchmark/common.jl @@ -127,7 +127,7 @@ function supported(f, label) end const bench_seconds = parse(Float64, get(ENV, "BENCHMARK_SECONDS", "30")) -const bench_samples = parse(Int, get(ENV, "BENCHMARK_SAMPLES", "5")) +const bench_samples = parse(Int, get(ENV, "BENCHMARK_SAMPLES", "7")) # Rendering/logging are not used under AirspeedVelocity; these globals are kept # defined because the suite files reference them. From 410695b3af4117002fa5fe5e23a24c8cee60b036 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 8 Sep 2026 14:19:38 -0700 Subject: [PATCH 04/26] MPIExt: Judge cross-rank waits by progress, not by elapsed time `mpi_deadlock_detect` timed each wait on its own and threw once one exceeded 120s. That reads a long wait as a deadlock, but how long a rank waits for a peer is bounded by the peer's *backlog*, not by anything about the wait itself -- and `DATADEPS_UNIFORM_DEFER` deliberately lets that backlog grow, submitting a region's whole task set in one burst so planning is not interleaved with execution. A non-owner then reaches a task's `execute!` metadata wait long before the owner starts it, and if the owner still has to JIT the task body first, one legitimate wait runs past any fixed timeout. That is the MPI CPU CI failure: the `@stencil` 4D case (81 blocks x 80 `Wrap()` neighbors, 82 inputs) takes minutes to compile on first use, so rank 1 aborted the run on a wait that was making perfectly good progress. The suite passed with the detector disabled, and every stall warning fell in the first repetition -- the compiling one -- with later repetitions of the same region clean at ~3s. Count the cross-rank operations each rank completes (finished requests, delivered broadcast payloads) and restart a wait's clock whenever that counter moves. The thresholds then measure a stall rather than a wait. A real cycle still trips them, once the work that does not depend on it has drained; verified by a rank waiting on a message that never comes, which still errors at exactly the timeout. Wait-loop state moves into an isbits `DeadlockTimer` so the loops keep allocating nothing, and the two periods pick up env-var overrides, which is what made the false positive falsifiable in the first place. 2 ranks x 2 threads, test/mpi.jl: 420/446 pass, 10m04s, no warnings (master: 18m03s; this branch before the fix: exit 1 at 14m42s). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01NHPKqRaGTGLt7CPKyZGuPe --- ext/MPIExt.jl | 116 ++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 84 insertions(+), 32 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 47555058f..6c198e6f3 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -864,12 +864,56 @@ end const DEADLOCK_DETECT = TaskLocalValue{Bool}(()->true) -const DEADLOCK_WARN_PERIOD = TaskLocalValue{Float64}(()->10.0) -const DEADLOCK_TIMEOUT_PERIOD = TaskLocalValue{Float64}(()->120.0) +const DEADLOCK_WARN_PERIOD = TaskLocalValue{Float64}(()->parse(Float64, get(ENV, "JULIA_DAGGER_MPI_WARN_PERIOD", "10.0"))) +const DEADLOCK_TIMEOUT_PERIOD = TaskLocalValue{Float64}(()->parse(Float64, get(ENV, "JULIA_DAGGER_MPI_TIMEOUT_PERIOD", "120.0"))) # Off by default: a full stacktrace buries the one-line diagnosis (which call # site is waiting) in noise. Opt in with JULIA_DAGGER_MPI_WARN_BACKTRACE=1 when # that call site itself is the thing in question. const DEADLOCK_WARN_BACKTRACE = TaskLocalValue{Bool}(()->parse(Bool, get(ENV, "JULIA_DAGGER_MPI_WARN_BACKTRACE", "0"))) + +# Monotonic count of cross-rank operations this rank has *completed*: every +# finished `Isend`/`Irecv` request and every broadcast payload the relay hands +# to a slot. Only the volume matters, never the value, so relaxed increments +# from any thread are fine. +# +# This is what lets a wait tell "the job is moving and I am behind a peer" from +# "nothing is moving anywhere". A per-wait stopwatch cannot: how long one rank +# waits for another is bounded by the peer's *backlog*, not by any property of +# the wait itself, and datadeps deliberately lets that backlog grow (a region's +# tasks are submitted in one burst, so a non-owner can reach a task's metadata +# wait long before the owner starts it). Add a first-call JIT of the task body +# -- minutes, for a wide `@stencil` -- and a single legitimate wait runs well +# past any fixed timeout. Timing that wait alone reports a deadlock that is not +# there; timing it against this counter reports one only when the rank has +# genuinely stopped doing cross-rank work. +const MPI_PROGRESS = Threads.Atomic{UInt64}(0) +@inline note_mpi_progress!() = (Threads.atomic_add!(MPI_PROGRESS, UInt64(1)); nothing) + +""" + DeadlockTimer() + +The stopwatch one cross-rank wait carries, rebuilt (not mutated) on each poll so +it stays isbits and costs no allocation on the wait loops. + +`warn_period` is the live threshold and goes to `typemax` once this wait has +warned, so it warns at most once; `base_warn` is the value to restore when the +clock is reset. `last_progress` is the `MPI_PROGRESS` reading this wait was last +judged against. See `mpi_deadlock_detect`. +""" +struct DeadlockTimer + detect::Bool + time_start::UInt64 + warn_period::UInt64 + base_warn::UInt64 + timeout_period::UInt64 + last_progress::UInt64 +end +function DeadlockTimer() + warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) + return DeadlockTimer(DEADLOCK_DETECT[], time_ns(), warn_period, warn_period, + round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9), + MPI_PROGRESS[]) +end const RECV_WAITING = LockedObject(Dict{Tuple{MPI.Comm, Int, Int}, Base.Event}()) # Envelope for the out-of-place raw-bytes MPI path: serialize a small @@ -1056,10 +1100,7 @@ end recv_yield_inplace!(array::Array, comm, my_rank, their_rank, tag) = _recv_yield_inplace_raw!(array, comm, my_rank, their_rank, tag) function _recv_yield_inplace_raw!(array, comm, my_rank, their_rank, tag) - time_start = time_ns() - detect = DEADLOCK_DETECT[] - warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) - timeout_period = round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9) + timer = DeadlockTimer() while true (got, msg, stat) = MPI.Improbe(their_rank, tag, comm, MPI.Status) @@ -1074,7 +1115,7 @@ function _recv_yield_inplace_raw!(array, comm, my_rank, their_rank, tag) __wait_for_request(req, comm, my_rank, their_rank, tag, "recv_yield", "recv") return array end - warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, my_rank, tag, "recv", their_rank) + timer = mpi_deadlock_detect(timer, my_rank, tag, "recv", their_rank) yield() end end @@ -1087,10 +1128,7 @@ function recv_yield_inplace(_value::InplaceInfo, comm, my_rank, their_rank, tag) end function recv_yield_serialized(comm, my_rank, their_rank, tag) - time_start = time_ns() - detect = DEADLOCK_DETECT[] - warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) - timeout_period = round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9) + timer = DeadlockTimer() while true (got, msg, stat) = MPI.Improbe(their_rank, tag, comm, MPI.Status) @@ -1104,7 +1142,7 @@ function recv_yield_serialized(comm, my_rank, their_rank, tag) __wait_for_request(req, comm, my_rank, their_rank, tag, "recv_yield", "recv") return MPI.deserialize(buf) end - warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, my_rank, tag, "recv", their_rank) + timer = mpi_deadlock_detect(timer, my_rank, tag, "recv", their_rank) yield() end end @@ -1185,19 +1223,19 @@ function send_yield_serialized(value, comm, my_rank, their_rank, tag) end function __wait_for_request(req, comm, my_rank, their_rank, tag, fn::String, kind::String) - time_start = time_ns() - detect = DEADLOCK_DETECT[] - warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) - timeout_period = round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9) + timer = DeadlockTimer() while true finish, status = MPI.Test(req, MPI.Status) if finish if MPI.Get_error(status) != MPI.SUCCESS error("$fn failed with error $(MPI.Get_error(status))") end + # Every completed request is this rank's proof of life for the + # waits running concurrently with it (see `MPI_PROGRESS`). + note_mpi_progress!() return end - warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, my_rank, tag, kind, their_rank) + timer = mpi_deadlock_detect(timer, my_rank, tag, kind, their_rank) yield() end end @@ -1242,9 +1280,23 @@ function bcast_yield(comm, root::Integer, tag, value=nothing) return value end -function mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, rank, tag, kind, srcdest) - time_elapsed = (time_ns() - time_start) - if detect && time_elapsed > warn_period +# Escalate a wait only while this rank is doing no cross-rank work at all. +# +# Waiting a long time is not evidence of a deadlock (see `MPI_PROGRESS`): it is +# the normal cost of being behind a peer. What *is* evidence is waiting while +# nothing else on this rank completes, so every poll that sees the counter move +# restarts the clock, and the warn/timeout thresholds then measure a stall +# rather than a wait. A real cycle still trips them, just once the work that +# does not depend on it has drained. +function mpi_deadlock_detect(timer::DeadlockTimer, rank, tag, kind, srcdest) + timer.detect || return timer + progress = MPI_PROGRESS[] + if progress != timer.last_progress + return DeadlockTimer(true, time_ns(), timer.base_warn, timer.base_warn, + timer.timeout_period, progress) + end + time_elapsed = (time_ns() - timer.time_start) + if time_elapsed > timer.warn_period # A hang here is a wait cycle across ranks, so which call site is waiting # (and on whose behalf) is the whole diagnosis; a bare tag is not enough. # The backtrace that pins down the call site is off by default, though -- @@ -1254,12 +1306,13 @@ function mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, ra else @warn "[rank $rank][tag $tag] Hit probable hang on $kind (dest: $srcdest)" end - return typemax(UInt64) + return DeadlockTimer(true, timer.time_start, typemax(UInt64), timer.base_warn, + timer.timeout_period, progress) end - if detect && time_elapsed > timeout_period + if time_elapsed > timer.timeout_period error("[rank $rank][tag $tag] Hit hang on $kind (dest: $srcdest)") end - return warn_period + return timer end # --------------------------------------------------------------------------- @@ -1360,6 +1413,9 @@ function bcast_deliver!(state::BcastState, root::Int, tag::UInt32, value) push!(get!(BcastSlot, state.slots, (root, tag)).values, value) notify(state.cond) end + # A delivered payload is cross-rank progress even when it is not the one a + # given consumer is waiting for (see `MPI_PROGRESS`). + note_mpi_progress!() return end @@ -1384,16 +1440,13 @@ function bcast_slot_wait(state::BcastState, root::Int, tag::UInt32) # here forever with nothing to report -- a CI job exhausting its timeout # with no error to point at. Re-checks are driven by `bcast_heartbeat!` # rather than by polling, so a blocked consumer costs no CPU. - time_start = time_ns() - detect = DEADLOCK_DETECT[] - warn_period = round(UInt64, DEADLOCK_WARN_PERIOD[] * 1e9) - timeout_period = round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9) + timer = DeadlockTimer() rank = MPI.Comm_rank(state.bcast_comm) # `@lock`, not `lock(...) do`: the loop closes over enough of the enclosing # frame (`key`, the deadlock-timer state, `rank`) that the closure boxes - # them, and `warn_period` -- reassigned across iterations and captured -- - # becomes a heap `Box`. Inlining the body keeps all of it on the stack; the - # `return` below still releases the lock, via `@lock`'s `finally`. + # them, and `timer` -- reassigned across iterations and captured -- becomes + # a heap `Box`. Inlining the body keeps all of it on the stack; the `return` + # below still releases the lock, via `@lock`'s `finally`. @lock state.cond begin while true slot = get(state.slots, key, nothing) @@ -1402,8 +1455,7 @@ function bcast_slot_wait(state::BcastState, root::Int, tag::UInt32) isempty(slot.values) && delete!(state.slots, key) return value end - warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, - rank, tag, "bcast_meta delivery", root) + timer = mpi_deadlock_detect(timer, rank, tag, "bcast_meta delivery", root) wait(state.cond) end end From e0f813a50f273a8b4638fc81376629a956dba296 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 11 Sep 2026 15:12:57 -0700 Subject: [PATCH 05/26] datadeps/hierarchical: Replace HIER_TIMING with TimespanLogging categories TimespanLogging's typed-category storage (LogCategory/steal_typed, from "Make TimespanLogging cheap enough to leave on") is a lock-free per-thread alternative to the custom HierPlanStats struct HIER_TIMING built for itself because the old TimespanLogging was too slow to measure planning with. It is now cheap enough to reuse directly. HierPlanStats (mutable struct with Atomics, a locked samples Dict, and a ScopedValue to thread it through parallel partition planning) is replaced by three LogCategory declarations (LogHierPhase/LogHierSlot/LogHierAinfo) and hier_log!, a thin wrapper around TimespanLogging's internal `_emit` gated by `HIER_TIMING[]` directly rather than the shared `enable!` bits -- so this diagnostic stays independent of whatever else `enable_logging!`/ `disable_logging!` is doing. report_hier_stats steals this region's events right after planning finishes and builds the same breakdown the old code printed. This assumes one region plans at a time (true today), so it's documented rather than solved with a region id. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01GaPaEAUCJMwpQjFFPbU7Ck --- ext/MPIExt.jl | 2 +- src/datadeps/aliasing.jl | 7 +- src/datadeps/hierarchical.jl | 189 +++++++++++++++++--------------- src/utils/logging-categories.jl | 8 ++ 4 files changed, 116 insertions(+), 90 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 6c198e6f3..38c970c4a 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -303,7 +303,7 @@ function Dagger.batch_ainfos(accel::MPIAcceleration, objs::Vector, dep_mods::Vec # lookup of the same object agree even if an entry was already there. results[i] = Dagger.memoize_ainfo!(Dagger.ainfo_memo_key(objs[i], dep_mods[i]), ainfo) end - Dagger.hier_stat_add!(:ainfo_ns, time_ns() - t0, n) + Dagger.hier_log!(Dagger.LogHierAinfo, 0x01, Dagger.LogHierAinfoId(), (time_ns() - t0, n)) return results end diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 31644bd64..151f0476c 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1406,11 +1406,12 @@ function generate_slot!(state::DataDepsState, dest_space, data) 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) end - hier_stat_add!(:slot_moved_ns, time_ns() - t0) - orig_space == dest_space && hier_stat_add!(:slot_samespace_ns, time_ns() - t0) + move_ns = time_ns() - t0 + hier_log!(LogHierSlot, 0x01, LogHierSlotId(:moved), move_ns) + orig_space == dest_space && hier_log!(LogHierSlot, 0x01, LogHierSlotId(:samespace), move_ns) moved end - hier_stat_add!(:slot_ns, time_ns() - t0) + hier_log!(LogHierSlot, 0x01, LogHierSlotId(:total), time_ns() - t0) logging && @logfinish ctx LogMove LogMoveId(0, ArgPosition(), to_proc, id) data_chunk @assert memory_space(data_chunk) == dest_space "space mismatch! $dest_space (dest) != $(memory_space(data_chunk)) (actual) ($(typeof(data)) (data) vs. $(typeof(data_chunk)) (chunk)), spaces ($orig_space -> $dest_space)" dest_space_args[data] = data_chunk diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 5e45150c9..bbf54a871 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -278,94 +278,121 @@ const DATADEPS_UNIFORM_DEFER = Ref(true) # is the thing that limits datadeps at scale -- especially under MPI, where every # rank replays the whole plan -- and it is hard to attribute from a profile, # because the expensive parts are blocking waits inside communication rather than -# hot loops. Off by default and costs one `Ref` read per phase. +# hot loops. +# +# Built on TimespanLogging's typed categories (`LogHierPhase`/`LogHierSlot`/ +# `LogHierAinfo` in `logging-categories.jl`) for their lock-free per-thread +# storage, which is what makes leaving this on cheap -- but gated by +# `HIER_TIMING` rather than the shared `enable!`/`enable_logging!` bits: +# planning cost is a standing question independent of whatever else logging is +# on for, and `Dagger.disable_logging!()` should not silently turn it off. +# `TimespanLogging.steal_typed` drains exactly the events produced since the +# last steal, which is exactly this region's as long as one region plans at a +# time (true today -- `distribute_tasks_hierarchical!` runs synchronously on +# its caller); concurrent regions would interleave into one report. const HIER_TIMING = Ref(false) "Whether `HIER_TIMING` also logs each region's report (vs. only recording it)." const HIER_TIMING_REPORT = Ref(true) -mutable struct HierPlanStats - ntasks::Int - nargs::Int - npartitions::Int - prescan_ns::UInt64 - aliasing_ns::UInt64 - dag_ns::UInt64 - partition_ns::UInt64 - schedule_ns::UInt64 - epilogue_ns::UInt64 - # Touched from `generate_slot!` / `aliasing`, which run under parallel - # partition planning, so these are atomic while the phase fields are not. - slot_ns::Threads.Atomic{UInt64} - slot_count::Threads.Atomic{Int} - slot_moved_ns::Threads.Atomic{UInt64} - slot_moved_count::Threads.Atomic{Int} - # Slots whose data was moved even though it was already in the destination - # space: pure overhead, and the reason `slot_is_already_in_place` exists. - slot_samespace_ns::Threads.Atomic{UInt64} - slot_samespace_count::Threads.Atomic{Int} - ainfo_ns::Threads.Atomic{UInt64} - ainfo_count::Threads.Atomic{Int} - # Per-event durations (ns), for distribution rather than just totals. - samples::Dict{Symbol,Vector{UInt64}} - samples_lock::ReentrantLock +"Record `id`/`data` under category `C` when `HIER_TIMING[]` is set." +@inline function hier_log!(::Type{C}, phase::UInt8, id, data) where C <: TimespanLogging.LogCategory + HIER_TIMING[] || return nothing + TimespanLogging._emit(C, phase, id, data) + return nothing end -HierPlanStats() = HierPlanStats(0, 0, 0, 0, 0, 0, 0, 0, 0, - Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), - Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), - Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), - Threads.Atomic{UInt64}(0), Threads.Atomic{Int}(0), - Dict{Symbol,Vector{UInt64}}(), ReentrantLock()) - -# A `ScopedValue`, not a task-local: the parallel partition path calls -# `generate_slot!` from tasks spawned inside the region, which inherit scoped -# values but not task-locals. -const HIER_STATS = ScopedValue{Union{HierPlanStats,Nothing}}(nothing) -"Stats for the most recently planned region, for programmatic inspection." -const LAST_HIER_STATS = Ref{Union{HierPlanStats,Nothing}}(nothing) - -"Accumulate `ns` (and one event) into field `f` of the enclosing region's stats, if enabled." -@inline function hier_stat_add!(f::Symbol, ns::Integer, count::Integer=1) - HIER_TIMING[] || return - stats = HIER_STATS[] - stats === nothing && return - Threads.atomic_add!(getfield(stats, f), UInt64(ns)) - Threads.atomic_add!(getfield(stats, Symbol(String(f)[1:end-3] * "_count")), Int(count)) - @lock stats.samples_lock push!(get!(Vector{UInt64}, stats.samples, f), UInt64(ns)) - return +"Time `ex` as phase `field` of the enclosing region's `HIER_TIMING` report." +macro hier_phase(field, ex) + id = :(LogHierPhaseId($(QuoteNode(field)))) + quote + hier_log!(LogHierPhase, 0x00, $id, nothing) + local _res = $(esc(ex)) + hier_log!(LogHierPhase, 0x01, $id, nothing) + _res + end end -macro hier_phase(stats, field, ex) - quote - local _stats = $(esc(stats)) - if _stats === nothing - $(esc(ex)) +"Pair up `LogHierPhase` start/finish events (one region, single-threaded) into total ns per phase." +function _hier_phase_ns(events) + open = Dict{Symbol,UInt64}() + ns = Dict{Symbol,UInt64}() + for ev in events + phase = ev.id.phase + if ev.phase == 0x00 + open[phase] = ev.timestamp else - local _t0 = time_ns() - local _res = $(esc(ex)) - setfield!(_stats, $(QuoteNode(field)), getfield(_stats, $(QuoteNode(field))) + (time_ns() - _t0)) - _res + ns[phase] = get(ns, phase, UInt64(0)) + (ev.timestamp - get(open, phase, ev.timestamp)) end end + return ns end -hier_stats_total_ns(stats::HierPlanStats) = - stats.prescan_ns + stats.aliasing_ns + stats.dag_ns + stats.partition_ns + - stats.schedule_ns + stats.epilogue_ns +"Stats for the most recently planned region, for programmatic inspection." +const LAST_HIER_STATS = Ref{Union{NamedTuple,Nothing}}(nothing) + +""" + report_hier_stats(ntasks, nargs, npartitions) + +Steal this region's `LogHier*` events (recorded only when `HIER_TIMING[]`), +print a breakdown (unless `HIER_TIMING_REPORT[]` is false), and record it in +`LAST_HIER_STATS[]` for programmatic inspection. No-op unless `HIER_TIMING[]`. +""" +function report_hier_stats(ntasks::Int, nargs::Int, npartitions::Int) + HIER_TIMING[] || return nothing -function report_hier_stats(stats::HierPlanStats) - total = hier_stats_total_ns(stats) + phase_ns = _hier_phase_ns(TimespanLogging.steal_typed(LogHierPhase)) + ns(phase) = get(phase_ns, phase, UInt64(0)) + + slot_ns = UInt64(0); slot_count = 0 + moved_ns = UInt64(0); moved_count = 0 + samespace_ns = UInt64(0); samespace_count = 0 + # Per-event durations (ns), for distribution rather than just totals. + slot_samples = Dict{Symbol,Vector{UInt64}}() + for ev in TimespanLogging.steal_typed(LogHierSlot) + kind, d = ev.id.kind, ev.data + push!(get!(Vector{UInt64}, slot_samples, kind), d) + if kind === :total + slot_ns += d; slot_count += 1 + elseif kind === :moved + moved_ns += d; moved_count += 1 + elseif kind === :samespace + samespace_ns += d; samespace_count += 1 + end + end + + ainfo_ns = UInt64(0); ainfo_count = 0 + ainfo_samples = UInt64[] + for ev in TimespanLogging.steal_typed(LogHierAinfo) + this_ns, this_count = ev.data + ainfo_ns += this_ns; ainfo_count += this_count + push!(ainfo_samples, this_ns) + end + slot_samples[:ainfo] = ainfo_samples + + stats = (;ntasks, nargs, npartitions, + prescan_ns=ns(:prescan), aliasing_ns=ns(:aliasing), dag_ns=ns(:dag), + partition_ns=ns(:partition), schedule_ns=ns(:schedule), epilogue_ns=ns(:epilogue), + slot_ns, slot_count, moved_ns, moved_count, samespace_ns, samespace_count, + ainfo_ns, ainfo_count, samples=slot_samples) + LAST_HIER_STATS[] = stats + HIER_TIMING_REPORT[] && _print_hier_stats(stats) + return stats +end + +function _print_hier_stats(stats) + total = stats.prescan_ns + stats.aliasing_ns + stats.dag_ns + + stats.partition_ns + stats.schedule_ns + stats.epilogue_ns ms(x) = round(x / 1e6; digits=2) @info """ datadeps plan: $(stats.ntasks) tasks, $(stats.nargs) args, $(stats.npartitions) partitions, $(ms(total)) ms total phase 1 prescan $(ms(stats.prescan_ns)) ms - phase 1 aliasing $(ms(stats.aliasing_ns)) ms ($(stats.ainfo_count[]) ainfos, $(ms(stats.ainfo_ns[])) ms in aliasing()) + phase 1 aliasing $(ms(stats.aliasing_ns)) ms ($(stats.ainfo_count) ainfos, $(ms(stats.ainfo_ns)) ms in aliasing()) phase 2 dag $(ms(stats.dag_ns)) ms phase 3 partition $(ms(stats.partition_ns)) ms - phase 4 schedule $(ms(stats.schedule_ns)) ms ($(stats.slot_count[]) slots, $(ms(stats.slot_ns[])) ms, of which $(stats.slot_moved_count[]) moved data in $(ms(stats.slot_moved_ns[])) ms, $(stats.slot_samespace_count[]) of those within one space in $(ms(stats.slot_samespace_ns[])) ms) + phase 4 schedule $(ms(stats.schedule_ns)) ms ($(stats.slot_count) slots, $(ms(stats.slot_ns)) ms, of which $(stats.moved_count) moved data in $(ms(stats.moved_ns)) ms, $(stats.samespace_count) of those within one space in $(ms(stats.samespace_ns)) ms) epilogue $(ms(stats.epilogue_ns)) ms""" - return + return nothing end struct HierarchicalTaskMeta @@ -934,7 +961,7 @@ function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) arg_w = arg_ws[i] t0 = time_ns() ainfo = AliasingWrapper(aliasing(accel, arg_w.arg, arg_w.dep_mod)) - hier_stat_add!(:ainfo_ns, time_ns() - t0) + hier_log!(LogHierAinfo, 0x01, LogHierAinfoId(), (time_ns() - t0, 1)) results[i] = arg_w => ainfo end end @@ -1619,28 +1646,24 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) return distribute_tasks!(queue) end - stats = HIER_TIMING[] ? HierPlanStats() : nothing - return with(HIER_STATS => stats) do - _distribute_tasks_hierarchical!(queue, seen_tasks, accel, all_procs, stats) - end + return _distribute_tasks_hierarchical!(queue, seen_tasks, accel, all_procs) end function _distribute_tasks_hierarchical!(queue::DataDepsTaskQueue, seen_tasks::Vector{DTaskPair}, accel::Acceleration, - all_procs::Vector{<:Processor}, - stats::Union{HierPlanStats,Nothing}) + all_procs::Vector{<:Processor}) # Phase 1: Collect arguments and compute aliasing in parallel - task_metas, unique_arg_ws = @hier_phase stats prescan_ns collect_aliased_args(seen_tasks) + task_metas, unique_arg_ws = @hier_phase prescan collect_aliased_args(seen_tasks) _lookup, ainfos_overlaps, arg_to_ainfo = - @hier_phase stats aliasing_ns build_aliasing_parallel(unique_arg_ws) + @hier_phase aliasing build_aliasing_parallel(unique_arg_ws) # Phase 2: Build dependency DAG - dag = @hier_phase stats dag_ns build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) + dag = @hier_phase dag build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) # Phase 3: Partition the DAG vertex_to_partition, n_partitions, partition_procs, multi_owner = - @hier_phase stats partition_ns partition_dag(dag, task_metas, all_procs) + @hier_phase partition partition_dag(dag, task_metas, all_procs) # Detect backing chunks shared across partitions in different memory spaces. # These need runtime ownership transfer to avoid split-brain concurrent @@ -1712,7 +1735,7 @@ function _distribute_tasks_hierarchical!(queue::DataDepsTaskQueue, # See PERF(hier-2)/(hier-3). exec_spaces = unique(Iterators.flatten(memory_spaces(proc) for proc in all_procs)) use_shared_state = uniform_execution(accel) || length(exec_spaces) > 1 - partition_states = @hier_phase stats schedule_ns try + partition_states = @hier_phase schedule try if use_shared_state schedule_partitions_sequential!( queue, queue_lock, partitions, dag, seen_tasks, @@ -1754,15 +1777,9 @@ function _distribute_tasks_hierarchical!(queue::DataDepsTaskQueue, # Distributed keeps per-partition states and needs the registry for coherent # cross-partition write-back. epilogue_registry = use_shared_state ? nothing : registry - @hier_phase stats epilogue_ns _hierarchical_copy_from_and_free!( + @hier_phase epilogue _hierarchical_copy_from_and_free!( partition_states, length(partition_states), epilogue_registry) - if stats !== nothing - stats.ntasks = length(seen_tasks) - stats.nargs = length(unique_arg_ws) - stats.npartitions = n_partitions - LAST_HIER_STATS[] = stats - HIER_TIMING_REPORT[] && report_hier_stats(stats) - end + report_hier_stats(length(seen_tasks), length(unique_arg_ws), n_partitions) return end diff --git a/src/utils/logging-categories.jl b/src/utils/logging-categories.jl index 297e42688..eff0958c6 100644 --- a/src/utils/logging-categories.jl +++ b/src/utils/logging-categories.jl @@ -17,3 +17,11 @@ TimespanLogging.@logcategory LogFinish as=:finish id=(uid::UInt, thunk_id::Int) TimespanLogging.@logcategory LogEnqueue as=:enqueue id=(uid::UInt, processor::Any, thunk_id::Int) data=Nothing TimespanLogging.@logcategory LogSchedule as=:schedule id=(uid::UInt, thunk_id::Int) TimespanLogging.@logcategory LogFire as=:fire id=(uid::UInt, worker::Int) data=Nothing + +# Hierarchical datadeps planning instrumentation -- see `hier_log!` in +# datadeps/hierarchical.jl. Gated by `Dagger.HIER_TIMING[]`, not the shared +# `enable!` bits, so these three are declared (and their ids constructed) even +# when the rest of logging is off. +TimespanLogging.@logcategory LogHierPhase as=:hier_phase id=(phase::Symbol,) data=Nothing +TimespanLogging.@logcategory LogHierSlot as=:hier_slot id=(kind::Symbol,) data=UInt64 +TimespanLogging.@logcategory LogHierAinfo as=:hier_ainfo id=() data=Tuple{UInt64,Int} From ffbbab68557e284f84c4c731a473338ad0644550 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 08:54:30 -0700 Subject: [PATCH 06/26] benchmark/stencil: Await every asynchronous fixture Stencil setup launched both input and destination DArrays but waited only for the input. The timed body therefore absorbed a nondeterministic share of destination allocation and scheduling, producing multi-fold Distributed regressions. Materialize both arrays before timing so the benchmark measures only the stencil. --- AGENTS.md | 7 +++++++ benchmark/suites/stencil.jl | 12 ++++++------ 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5f18e0547..132e17e59 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -324,3 +324,10 @@ lesson. Category IDs are assigned lazily on first `category_id` use. Keep it that way — a `const ID = register_category!(...)` at Dagger toplevel is not safe. + +30. **Benchmark fixtures must be fully awaited.** DArray construction is + asynchronous: if setup launches `A` and `B` but waits only for `A`, the + timed body inherits an arbitrary fraction of `B`'s allocation and scheduling + work. That contamination is placement- and timing-dependent, so a scheduler + improvement can look like a multi-fold regression. Wait every fixture that + the measured operation consumes. diff --git a/benchmark/suites/stencil.jl b/benchmark/suites/stencil.jl index 764c58b88..4a388ce78 100644 --- a/benchmark/suites/stencil.jl +++ b/benchmark/suites/stencil.jl @@ -118,38 +118,38 @@ function stencil_suite(ctx; method, accels) if wrap_ok sub["neighbors (Wrap)"] = @benchmarkable(stencil_neighbors_wrap!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if pad_ok sub["neighbors (Pad)"] = @benchmarkable(stencil_neighbors_pad!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if clamp_ok sub["neighbors (Clamp)"] = @benchmarkable(stencil_neighbors_clamp!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if reflect_ok sub["neighbors (Reflect)"] = @benchmarkable(stencil_neighbors_reflect!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if update_ok sub["update (+)"] = @benchmarkable(stencil_update_plus!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if multi_ok sub["multi-expr"] = @benchmarkable(stencil_multi_expr!(A, B, $T), setup = (A = zeros(Blocks($b, $b), $T, $N, $N); - B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end end From 223eb881f583276b651d68f1ae7cf8636c181ecd Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 08:56:02 -0700 Subject: [PATCH 07/26] benchmark/mpi: Synchronize SPMD sampling across ranks BenchmarkTools applies its seconds budget independently, so rank timing skew let one rank stop while peers entered another collective sample. It also let a fast rank time its wait for a peer's prior GC as part of the next operation. Drive samples in lockstep, barrier before timing, and stop from the maximum per-sample wall time across ranks. A synthetic skew test changed sample counts from [5,2,2,2] to [2,2,2,2]; a four-rank allocation sample measured 1.93 ms instead of CI's spurious 14-35 ms. Declare JSON3 explicitly because both benchmark workers import it. --- AGENTS.md | 11 ++++++++++ benchmark/Project.toml | 2 ++ benchmark/worker_mpi.jl | 45 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 57 insertions(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index 132e17e59..2e13a831f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -331,3 +331,14 @@ lesson. work. That contamination is placement- and timing-dependent, so a scheduler improvement can look like a multi-fold regression. Wait every fixture that the measured operation consumes. + +31. **SPMD benchmarks need a rank-uniform sample loop.** BenchmarkTools applies + its `seconds` cutoff independently on each process; small timing differences + can make one MPI rank stop while a peer enters another collective sample, + deadlocking the next benchmark. Run one sample per rank at a time, use a + collective maximum to make the stop decision, and use a cooperatively-waited + `MPI.Ibarrier` after per-sample GC/teardown before timing the next sample (a + blocking `MPI.Barrier` has the same progress-engine deadlock risk as lesson + 12). Otherwise a fast rank also charges its next sample for waiting on a + peer's preceding GC, creating large but fake timing regressions in short + collective operations. diff --git a/benchmark/Project.toml b/benchmark/Project.toml index 124898801..d875efd6f 100644 --- a/benchmark/Project.toml +++ b/benchmark/Project.toml @@ -10,6 +10,7 @@ DTables = "20c56dc6-594c-4682-91cf-1d46875b1eba" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94" @@ -20,4 +21,5 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [compat] BenchmarkTools = "1" +JSON3 = "1" julia = "1.9" diff --git a/benchmark/worker_mpi.jl b/benchmark/worker_mpi.jl index de0cbaead..bf615ab06 100644 --- a/benchmark/worker_mpi.jl +++ b/benchmark/worker_mpi.jl @@ -112,6 +112,49 @@ for (_, b) in BenchmarkTools.leaves(SUITE) b.params.gcsample = true end +# BenchmarkTools enforces `seconds` independently on each process. That is not +# safe for SPMD work: small timing differences can make one rank stop sampling +# while another enters one more Dagger collective, leaving the latter hung. GC +# and benchmark teardown also finish at different times, so without a barrier a +# fast rank charges its next sample for waiting on a slow peer. +# +# Keep the ordinary BenchmarkTools sampling semantics, but make the stop +# decision from the maximum elapsed time across ranks and synchronize directly +# before every timed sample. All ranks therefore execute exactly the same number +# of samples, and rank 0's trial measures the operation rather than another +# rank's preceding GC. +function run_mpi_benchmark(bench::BenchmarkTools.Benchmark) + params = bench.params + + wait(MPI.Ibarrier(comm)) + bench.samplefunc(bench.quote_vals, BenchmarkTools.Parameters(params; evals=1)) + + trial = BenchmarkTools.Trial(params) + params.gctrial && BenchmarkTools.gcscrub() + elapsed = 0.0 + for sample in 1:params.samples + sample > 1 && params.gcsample && BenchmarkTools.gcscrub() + # A blocking MPI.Barrier can park the only Julia thread while a peer + # still needs this rank's Dagger progress tasks. Base.wait on Ibarrier + # cooperatively yields instead. + wait(MPI.Ibarrier(comm)) + sample_start = Base.time() + result = bench.samplefunc(bench.quote_vals, params) + push!(trial, result[1:(end - 1)]...) + + # Do not charge this rank for time spent at the pre-sample barrier. + # Sum the slowest rank's setup/body/teardown duration for each sample, + # which is the SPMD analogue of BenchmarkTools' local wall-time budget. + local_elapsed = Base.time() - sample_start + # Align cooperatively before the blocking reduction, so no rank can be + # parked in MPI while a peer is still finishing Dagger work. + wait(MPI.Ibarrier(comm)) + elapsed += MPI.Allreduce(local_elapsed, MPI.MAX, comm) + elapsed >= params.seconds && break + end + return trial +end + # --- Run every benchmark once, identically on every rank -------------------- # Sorted by keypath (not Dict/BenchmarkGroup insertion order) so ranks agree # even if suite construction were ever to introduce nondeterministic @@ -125,7 +168,7 @@ for (keypath, bench) in leaves kp = String[string(k) for k in keypath] rank == 0 && @info "[worker_mpi] Running: $(join(kp, " / "))" try - trial = BenchmarkTools.run(bench) + trial = run_mpi_benchmark(bench) rank == 0 && push!(results, (kp, trial)) catch err if err isa OutOfMemoryError From a6d065f73fc081435fc305ea1db9a6ee8600b6a1 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 11:12:57 -0700 Subject: [PATCH 08/26] benchmark/mpi: Support BenchmarkTools 1.8 sampling BenchmarkTools 1.8 changed generated sample functions to write into a Ref. The old internal call failed every MPI leaf, and the worker then reported an empty successful manifest. Dispatch to either sampling API, abort on SPMD leaf failures, and reject missing or empty manifests so CI cannot pass without measurements.\n\nVerified with 4 local ranks: the reduced array matrix now completes 14/14 leaves (previously 0/14 with MethodError). --- AGENTS.md | 9 +++++++++ benchmark/benchmarks.jl | 8 ++++++-- benchmark/worker_mpi.jl | 35 +++++++++++++++++++++++++++-------- 3 files changed, 42 insertions(+), 10 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 2e13a831f..d0593b22d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -342,3 +342,12 @@ lesson. 12). Otherwise a fast rank also charges its next sample for waiting on a peer's preceding GC, creating large but fake timing regressions in short collective operations. + +32. **An empty benchmark suite is a benchmark failure, not a successful no-op.** + AirspeedVelocity adds its own `time_to_load` result, so an orchestrator that + swallows every leaf error still emits a plausible-looking green report with + one row. Require a nonempty manifest, and abort an SPMD run on a leaf error: + after one rank leaves a failed operation, attempting the next collective is + unsafe. Also treat BenchmarkTools' generated `samplefunc` as a versioned + internal API: 1.8 changed it from a two-argument function returning a tuple + to a four-argument function writing measurements through a `Ref`. diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 6f18a6492..ec4b2ef27 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -299,10 +299,10 @@ function run_all_mpi() proc = spawn_worker() donepath = joinpath(WORKDIR, "done") + completed = false t0 = time() while !isfile(donepath) if !process_running(proc) - @error "MPI benchmark worker exited before signaling completion; producing partial results." break end if time() - t0 > PROC_TIMEOUT @@ -312,13 +312,17 @@ function run_all_mpi() end sleep(POLL) end + completed = isfile(donepath) if process_running(proc) try; wait(proc); catch; end end + completed || error("MPI benchmark worker exited before signaling completion") + manifestpath = joinpath(WORKDIR, "results_mpi_manifest.json") - isfile(manifestpath) || return results + isfile(manifestpath) || error("MPI benchmark worker produced no result manifest") manifest = JSON3.read(read(manifestpath, String)) + isempty(manifest) && error("MPI benchmark worker produced an empty result manifest") for entry in manifest kp = String[string(k) for k in entry.keypath] resultpath = joinpath(WORKDIR, String(entry.file)) diff --git a/benchmark/worker_mpi.jl b/benchmark/worker_mpi.jl index bf615ab06..0e9b1dbd1 100644 --- a/benchmark/worker_mpi.jl +++ b/benchmark/worker_mpi.jl @@ -22,10 +22,9 @@ # caught `OutOfMemoryError` aborts the *whole* MPI job (`MPI.Abort`) rather # than exiting only the local rank and letting the orchestrator retry smaller # scales, since a partially-alive rank set can't make collective progress. -# Other exceptions are caught per-rank and just skip that one benchmark: -# since every rank runs the identical deterministic computation, ranks are -# expected to fail together at the same call, so no cross-rank coordination -# is needed for the common case. +# Any other exception aborts the MPI job. A partial SPMD result set is not a +# trustworthy benchmark run, and continuing after one rank fails can deadlock +# as soon as its peers enter their next collective. using BenchmarkTools using Distributed @@ -126,8 +125,25 @@ end function run_mpi_benchmark(bench::BenchmarkTools.Benchmark) params = bench.params + # BenchmarkTools 1.8 changed its generated sample function from returning + # `(time, gctime, memory, allocs, value)` to writing the four measurements + # through a Ref. Keep the compatibility check here, at the one unavoidable + # use of BenchmarkTools' internal sampling API. + sample_ref = Ref{Tuple{Float64,Float64,Int,Int}}((0.0, 0.0, 0, 0)) + new_sample_api = applicable(bench.samplefunc, bench.quote_vals, params, + sample_ref, nothing) + function sample!(sample_params) + if new_sample_api + bench.samplefunc(bench.quote_vals, sample_params, sample_ref, nothing) + return sample_ref[] + else + result = bench.samplefunc(bench.quote_vals, sample_params) + return result[1:(end - 1)] + end + end + wait(MPI.Ibarrier(comm)) - bench.samplefunc(bench.quote_vals, BenchmarkTools.Parameters(params; evals=1)) + sample!(BenchmarkTools.Parameters(params; evals=1)) trial = BenchmarkTools.Trial(params) params.gctrial && BenchmarkTools.gcscrub() @@ -139,8 +155,8 @@ function run_mpi_benchmark(bench::BenchmarkTools.Benchmark) # cooperatively yields instead. wait(MPI.Ibarrier(comm)) sample_start = Base.time() - result = bench.samplefunc(bench.quote_vals, params) - push!(trial, result[1:(end - 1)]...) + result = sample!(params) + push!(trial, result...) # Do not charge this rank for time spent at the pre-sample barrier. # Sum the slowest rank's setup/body/teardown duration for each sample, @@ -177,7 +193,10 @@ for (keypath, bench) in leaves MPI.Abort(comm, 137) exit(137) # unreachable unless MPI.Abort fails to terminate us else - rank == 0 && @warn "[worker_mpi] Benchmark errored (skipped)" benchmark = join(kp, " / ") exception = (err, catch_backtrace()) + @error "[worker_mpi] Benchmark errored; aborting MPI job" rank benchmark = join(kp, " / ") exception = (err, catch_backtrace()) + flush(stdout); flush(stderr) + MPI.Abort(comm, 1) + exit(1) # unreachable unless MPI.Abort fails to terminate us end end end From 8ea4c90d935df08992e774c5eb6d4a5b2fd81578 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 11:34:33 -0700 Subject: [PATCH 09/26] datadeps: Keep synchronous submission batches short A 16-task synchronous batch withholds most small Krylov regions until planning finishes, serializing planning against execution on one-thread Distributed drivers. Retain 16 for the asynchronous submitter, where it amortizes handoff overhead, but restore the synchronous batch to 4.\n\nDistributed CG (1024, block 64, 4 processes; one post-warmup run):\n batch 16: 8.22 s, 16.85M allocs, 832.9 MB\n batch 8: 7.91 s, 16.59M allocs, 818.6 MB\n batch 4: 6.92 s, 15.90M allocs, 783.1 MB\n\nThe Distributed datadeps suite passes (1462 pass, 2 broken), as does the allocation suite (19/19). --- AGENTS.md | 8 ++++++++ src/datadeps/hierarchical.jl | 31 ++++++++++++++++++------------- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index d0593b22d..74b17043d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -351,3 +351,11 @@ lesson. unsafe. Also treat BenchmarkTools' generated `samplefunc` as a versioned internal API: 1.8 changed it from a two-argument function returning a tuple to a four-argument function writing measurements through a `Ref`. + +33. **Synchronous and asynchronous submission need different batch sizes.** A + large asynchronous batch amortizes channel and scheduler overhead while its + submitter overlaps with continued planning. The same batch size on a + one-thread `BatchedEnqueueQueue` withholds every task in a short region until + planning has nearly or completely finished, serializing planning against + execution. This is especially costly for iterative solvers, which execute + many small datadeps regions. Tune the two queue types independently. diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index bbf54a871..0c09f2830 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -146,7 +146,7 @@ function flush_batch!(beq::BatchedEnqueueQueue) end """ - AsyncEnqueueQueue(inner, lock; limit=DATADEPS_BATCH_LIMIT[]) + AsyncEnqueueQueue(inner, lock; limit=DATADEPS_ASYNC_BATCH_LIMIT[]) `BatchedEnqueueQueue` that hands each full batch to a submitter task rather than submitting it on the caller's thread. @@ -181,7 +181,7 @@ mutable struct AsyncEnqueueQueue <: AbstractTaskQueue failure::Any end function AsyncEnqueueQueue(inner::AbstractTaskQueue, lock::ReentrantLock; - limit::Int=DATADEPS_BATCH_LIMIT[]) + limit::Int=DATADEPS_ASYNC_BATCH_LIMIT[]) chan = Channel{Union{Vector{DTaskPair},Base.Event}}(Inf) queue = AsyncEnqueueQueue(inner, lock, DTaskPair[], limit, chan, nothing, nothing) queue.submitter = Threads.@spawn _async_submit_loop(queue) @@ -252,17 +252,22 @@ maybe_flush_batch!(beq::BatchedEnqueueQueue) = flush_batch!(beq) maybe_flush_batch!(aeq::AsyncEnqueueQueue) = flush_batch!(aeq) maybe_flush_batch!(::AbstractTaskQueue) = nothing -""" -Maximum number of tasks a hierarchical partition buffers before submitting. - -Sized from the two costs it trades off. Submitting one task at a time makes the -scheduler round-trip about 40% of this path's per-task planning cost; batching -amortizes it down, and measurably stops paying off past ~16 (a 256-task region -over 4 workers plans at 69 us/task unbatched, 55 at 16, 53.5 unbounded). Against -that, buffered tasks cannot start running, so the batch is what planning gets -ahead of execution -- bounded here at 16 tasks' worth of planning. -""" -const DATADEPS_BATCH_LIMIT = Ref(16) +"""Maximum tasks buffered by a synchronous submitter. + +Keep this small enough to start short regions before planning reaches their +end; Krylov's vector operations are commonly only one or two batches long. +Submission blocks planning here, so a large batch serializes planning against +execution.""" +const DATADEPS_BATCH_LIMIT = Ref(4) + +"""Maximum tasks buffered by an asynchronous submitter. + +Submitting one task at a time makes the scheduler round-trip about 40% of this +path's per-task planning cost; batching measurably stops paying off past ~16 (a +256-task region over 4 workers plans at 69 us/task unbatched, 55 at 16, 53.5 +unbounded). The submitter overlaps that cost with planning, so it can use the +larger batch without serializing planning against execution.""" +const DATADEPS_ASYNC_BATCH_LIMIT = Ref(16) """ Whether uniform (SPMD) planning withholds a region's tasks until it has finished From eb8d98b7229551a43bf9e5024b549680ed95f6df Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 11:35:04 -0700 Subject: [PATCH 10/26] benchmark: Await vector fixtures before sampling DArray construction is asynchronous. The sparse SpMV/CG and dense matvec/solve setups waited for their matrix but not their vector, allowing vector creation and placement to leak into the timed operation. Await every fixture before sampling so scheduler changes cannot shift setup work into these measurements. --- benchmark/suites/linalg.jl | 4 ++-- benchmark/suites/sparse.jl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/benchmark/suites/linalg.jl b/benchmark/suites/linalg.jl index fd42a530d..5ac97f91b 100644 --- a/benchmark/suites/linalg.jl +++ b/benchmark/suites/linalg.jl @@ -46,7 +46,7 @@ function linalg_suite(ctx; method, accels) sub["solve (A\\b via lu)"] = @benchmarkable(wait(lu(A, RowMaximum()) \ b), setup = (A = rand(Blocks($b, $b), $T, $N, $N); - b = rand(Blocks($b), $T, $N); wait(A)), + b = rand(Blocks($b), $T, $N); wait(A); wait(b)), teardown = (A = nothing; b = nothing; @everywhere GC.gc())) end @@ -70,7 +70,7 @@ function linalg_suite(ctx; method, accels) if fits_budget(dense_bytes(N; nmats=1, T=T)) sub["matvec (A*x)"] = @benchmarkable(wait(A * x), setup = (A = rand(Blocks($b, $b), $T, $N, $N); - x = rand(Blocks($b), $T, $N); wait(A)), + x = rand(Blocks($b), $T, $N); wait(A); wait(x)), teardown = (A = nothing; x = nothing; @everywhere GC.gc())) end diff --git a/benchmark/suites/sparse.jl b/benchmark/suites/sparse.jl index d0e638c89..0f7daaeca 100644 --- a/benchmark/suites/sparse.jl +++ b/benchmark/suites/sparse.jl @@ -72,7 +72,7 @@ function sparse_suite(ctx; method, accels) if spmv_ok && fits_budget(sparse_bytes(N; nmats=2, density=density, T=T)) sub["spmv (S*x)"] = @benchmarkable(wait(S * x), setup = (S = distribute(sprand($T, $N, $N, $density), Blocks($b, $b)); - x = distribute(rand($T, $N), Blocks($b)); wait(S)), + x = distribute(rand($T, $N), Blocks($b)); wait(S); wait(x)), teardown = (S = nothing; x = nothing; @everywhere GC.gc())) end @@ -87,7 +87,7 @@ function sparse_suite(ctx; method, accels) if cg_ok && fits_budget(sparse_bytes(N; nmats=2, density=3 / N, T=T)) sub["cg solve (laplacian)"] = @benchmarkable(wait(first(Dagger.cg(A, rhs; atol=1e-8, rtol=1e-6, itmax=200))), setup = (A = distribute(laplacian_1d($T, $N), Blocks($b, $b)); - rhs = distribute(rand($T, $N), Blocks($b)); wait(A)), + rhs = distribute(rand($T, $N), Blocks($b)); wait(A); wait(rhs)), teardown = (A = nothing; rhs = nothing; @everywhere GC.gc())) end From 4005db6db6f49502b326a1f650e4a252b6499d8f Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 15:26:10 -0700 Subject: [PATCH 11/26] datadeps: Join disjoint remainder copies before whole use Multiple disjoint copy tasks were each registered as owner of the whole destination, so a later whole-object consumer could run after only the last-registered copy and be overwritten by a slower sibling. Make the final copy depend on its earlier siblings, preserving their parallelism while making the recorded owner a truthful batch-completion point.\n\nValidated with the four-rank delayed ChunkView race (9/9 repeated passes) and the Distributed datadeps suite (1462 passed, 2 expected broken). --- AGENTS.md | 12 ++++++++++++ src/datadeps/remainders.jl | 35 +++++++++++++++++++++++++++-------- test/mpi.jl | 8 +++++++- 3 files changed, 46 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 74b17043d..7af3e5291 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -359,3 +359,15 @@ lesson. planning has nearly or completely finished, serializing planning against execution. This is especially costly for iterative solvers, which execute many small datadeps regions. Tune the two queue types independently. + +35. **A whole-region copy owner must represent the whole copy batch.** A + remainder can be assembled by several copy tasks from disjoint source + spaces. Registering each task as the owner of the whole destination makes + the last registration replace the earlier ones; unless that last task is a + causal successor of the other copies, a whole-region consumer waits for one + piece and a late copy can overwrite its result. Fuse the join into the final + copy by making it depend on the earlier copies (which may remain mutually + parallel), then the one recorded owner truthfully means "the whole batch is + complete." Registering every exact destination span is correct too, but is a + performance cliff for halo exchange: a 64-tile stencil produced hundreds of + megabytes of interval-tree/overlap bookkeeping and regressed by 7–8x. diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 9566b6537..31683f6af 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -324,11 +324,22 @@ Enqueues a copy operation to update the remainder regions of an object before a """ function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::MultiRemainderAliasing, f, idx, dest_scope, task, write_num::Int) - for remainder in remainder_aliasing.remainders + remainders = remainder_aliasing.remainders + last_idx = length(remainders) + for (i, remainder) in enumerate(remainders) @check_uniform(remainder.space) @assert !isempty(remainder.spans) @check_uniform(remainder.spans) - enqueue_remainder_copy_to!(state, dest_space, arg_w, remainder, f, idx, dest_scope, task, write_num) + copy_task = enqueue_remainder_copy_to!(state, dest_space, arg_w, + remainder, f, idx, dest_scope, task, write_num) + if i < last_idx + # Every copy is recorded as the owner of the whole destination, so + # the final one replaces the earlier owners. Make that replacement + # truthful: once the final copy completes, all disjoint pieces have + # completed. Earlier copies remain mutually parallel, and this + # fuses the batch join into work we already had to schedule. + push!(remainders[last_idx].syncdeps, ThunkSyncdep(copy_task)) + end end end function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::RemainderAliasing, @@ -356,7 +367,6 @@ function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpac source_ainfos = copy(remainder_aliasing.ainfos) empty!(remainder_aliasing.ainfos) get_write_deps!(state, dest_space, target_ainfo, write_num, remainder_syncdeps) - @dagdebug task.uid :spawn_datadeps "($(repr(f)))[$(idx-1)][$dep_mod] Remainder copy-to has $(length(remainder_syncdeps)) syncdeps" # Launch the remainder copy task @@ -373,7 +383,9 @@ function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpac for ainfo in source_ainfos add_reader!(state, arg_w, source_space, ainfo, copy_task, write_num) end - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; + copy_src=source_space) + return copy_task end """ enqueue_remainder_copy_from!(state::DataDepsState, target_ainfo::AliasingWrapper, arg, remainder_aliasing, @@ -383,11 +395,17 @@ Enqueues a copy operation to update the remainder regions of an object back to t """ function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::MultiRemainderAliasing, dest_scope, write_num::Int) - for remainder in remainder_aliasing.remainders + remainders = remainder_aliasing.remainders + last_idx = length(remainders) + for (i, remainder) in enumerate(remainders) @check_uniform(remainder.space) @assert !isempty(remainder.spans) @check_uniform(remainder.spans) - enqueue_remainder_copy_from!(state, dest_space, arg_w, remainder, dest_scope, write_num) + copy_task = enqueue_remainder_copy_from!(state, dest_space, arg_w, + remainder, dest_scope, write_num) + if i < last_idx + push!(remainders[last_idx].syncdeps, ThunkSyncdep(copy_task)) + end end end function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::RemainderAliasing, @@ -415,7 +433,6 @@ function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySp source_ainfos = copy(remainder_aliasing.ainfos) empty!(remainder_aliasing.ainfos) get_write_deps!(state, dest_space, target_ainfo, write_num, remainder_syncdeps) - @dagdebug nothing :spawn_datadeps "($(typeof(arg_w.arg)))[$dep_mod] Remainder copy-from has $(length(remainder_syncdeps)) syncdeps" # Launch the remainder copy task @@ -432,7 +449,9 @@ function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySp for ainfo in source_ainfos add_reader!(state, arg_w, source_space, ainfo, copy_task, write_num) end - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; + copy_src=source_space) + return copy_task end # FIXME: Document me diff --git a/test/mpi.jl b/test/mpi.jl index cda70f5f4..f7e138589 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -67,6 +67,7 @@ end inc!(X) = (X .+= 1; nothing) add1!(X) = (X .+= 1; nothing) +delayed_add1!(X) = (sleep(0.25); X .+= 1; nothing) scale2!(X) = (X .*= 2; nothing) sum_into!(r, X) = (r[] = Int(sum(X)); nothing) sum_into_f!(r, X) = (r[] = Float64(sum(X)); nothing) @@ -170,6 +171,7 @@ end @test c.handle isa MPIExt.MPIRef @test Dagger.check_uniform(c.handle) end + end @testset "check_uniform" begin @@ -451,7 +453,11 @@ end r1 = min(1, nranks-1) r2 = min(2, nranks-1) Dagger.spawn_datadeps() do - Dagger.@spawn scope=rank_scope(r1) add1!(InOut(cv_top)) + # Delay the first producer so the bottom-half copy reaches rank 0 + # first. A whole-ainfo owner for both disjoint copy tasks used to let + # the following whole-chunk write run after only the faster copy, then + # the late top-half copy overwrote its result. + Dagger.@spawn scope=rank_scope(r1) delayed_add1!(InOut(cv_top)) Dagger.@spawn scope=rank_scope(r2) scale2!(InOut(cv_bot)) Dagger.@spawn scope=rank_scope(0) add1!(InOut(c)) end From ec2326674dedf6c7bdf74f4ce55983ae40ebfa5e Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 15:26:53 -0700 Subject: [PATCH 12/26] benchmark: Report the MPI leaf that aborts MPI SVD uses a Distributed-only processor grid and dies before sampling. Probe it during suite construction so unsupported revision/backend combinations are omitted, and persist failures from every rank before MPI.Abort so the orchestrator reports the real leaf and CapturedException. Preserve ordinary external-worker exception text as well.\n\nA four-rank end-to-end benchmark smoke produced a nonempty 22-leaf completion manifest with all supported Array, LinearAlgebra, and Stencil leaves. --- AGENTS.md | 12 ++++++++++++ benchmark/benchmarks.jl | 33 ++++++++++++++++++++++++++------- benchmark/suites/linalg.jl | 11 ++++++++++- benchmark/worker_mpi.jl | 15 ++++++++++++++- 4 files changed, 62 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7af3e5291..3ff26e18b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -360,6 +360,18 @@ lesson. execution. This is especially costly for iterative solvers, which execute many small datadeps regions. Tune the two queue types independently. +34. **A benchmark script shared across revisions must probe backend capability.** + Airspeed runs today's script against old Dagger code, and a generic suite can + also contain a leaf unsupported by the selected backend. MPI SVD, for + example, uses a Distributed-only processor grid and divides by zero before + sampling. Probe optional or newly backend-enabled operations while + constructing the suite, so unsupported leaves are absent rather than + aborting the comparison. When an external MPI worker does abort, persist + each failing rank's `CapturedException` before `MPI.Abort`: the first failure + may be off rank 0 and can terminate rank 0 before it writes anything, mpiexec + output may disappear from the parent runner's log, and a generic "worker + exited" message discards the only actionable diagnosis. + 35. **A whole-region copy owner must represent the whole copy batch.** A remainder can be assembled by several copy tasks from disjoint source spaces. Registering each task as the owner of the whole destination makes diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index ec4b2ef27..2b9d4729b 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -171,6 +171,10 @@ function spawn_worker() rm(joinpath(WORKDIR, "request.json.tmp"); force=true) rm(joinpath(WORKDIR, "done"); force=true) rm(joinpath(WORKDIR, "results_mpi_manifest.json"); force=true) + for name in readdir(WORKDIR) + startswith(name, "error_mpi_rank_") || continue + rm(joinpath(WORKDIR, name); force=true) + end if MPI_RANKS > 0 # `MPI.mpiexec` only sets up the environment mpiexec needs (library # paths for the bundled MPICH_jll, etc.) for the dynamic extent of the @@ -204,7 +208,8 @@ send_request(id, action, keypath) = JSON3.write((; id=id, action=action, keypath=keypath))) # Wait for a response to request `id`, or for the worker to die / time out. -# Returns the status string ("ok"/"error"/"missing"/"oom"/"died"/"timeout"). +# Returns an object with `status` and `message` fields. Preserve the worker's +# error text here: this orchestrator is the only output Airspeed reliably shows. function await_response(proc, id) resppath = joinpath(WORKDIR, "response_$(id).json") t0 = time() @@ -215,10 +220,13 @@ function await_response(proc, id) catch nothing # caught mid-write; retry end - resp === nothing || return String(resp.status) + resp === nothing || return (; + status=String(resp.status), + message=hasproperty(resp, :message) ? String(resp.message) : "", + ) end - process_running(proc) || return "died" - time() - t0 > PROC_TIMEOUT && return "timeout" + process_running(proc) || return (; status="died", message="") + time() - t0 > PROC_TIMEOUT && return (; status="timeout", message="") sleep(POLL) end end @@ -254,7 +262,8 @@ function run_all_external() req_id += 1 send_request(req_id, "run", it.keypath) - status = await_response(proc, req_id) + response = await_response(proc, req_id) + status = response.status if status == "ok" results[it.keypath] = @@ -262,7 +271,7 @@ function run_all_external() elseif status == "missing" @warn "Worker has no such benchmark (capability probe differs?): $(join(it.keypath, " / "))" elseif status == "error" - @warn "Benchmark errored (skipped): $(join(it.keypath, " / "))" + @warn "Benchmark errored (skipped): $(join(it.keypath, " / "))" message=response.message else # "oom" / "died" / "timeout" @warn "Benchmark $(status); skipping this and larger scales of $(it.suite)/$(it.method)" benchmark = join(it.keypath, " / ") if it.N >= 0 @@ -317,7 +326,17 @@ function run_all_mpi() try; wait(proc); catch; end end - completed || error("MPI benchmark worker exited before signaling completion") + if !completed + errorpaths = sort!(filter(name -> startswith(name, "error_mpi_rank_") && + endswith(name, ".json"), + readdir(WORKDIR))) + if !isempty(errorpaths) + errorpath = joinpath(WORKDIR, first(errorpaths)) + failure = JSON3.read(read(errorpath, String)) + error("MPI benchmark rank $(failure.rank) failed while running $(failure.benchmark):\n$(failure.exception)") + end + error("MPI benchmark worker exited before signaling completion") + end manifestpath = joinpath(WORKDIR, "results_mpi_manifest.json") isfile(manifestpath) || error("MPI benchmark worker produced no result manifest") diff --git a/benchmark/suites/linalg.jl b/benchmark/suites/linalg.jl index 5ac97f91b..e5a26efe0 100644 --- a/benchmark/suites/linalg.jl +++ b/benchmark/suites/linalg.jl @@ -20,6 +20,15 @@ function linalg_suite(ctx; method, accels) @assert accel == "cpu" "Linalg suite only supports CPU execution" T = Float64 + # Some older Dagger revisions use a Distributed-only processor grid for + # tiled SVD. Under MPI that grid is empty and `_tile_index` divides by zero. + # This script is shared by both Airspeed revisions, so probe once and omit + # SVD where the revision/backend combination cannot execute it. + svd_ok = supported("linalg/svd") do + A = rand(Blocks(4, 4), T, 8, 8) + wait(A) + wait(svd(A).U) + end suite = BenchmarkGroup() for N in scales @@ -60,7 +69,7 @@ function linalg_suite(ctx; method, accels) # SVD (tiled one-sided Jacobi) additionally holds the internally-copied # scratch matrix, the accumulated V factor, and (across multiple # workers) a restaged copy of A, on top of the resident input. - if fits_budget(dense_bytes(N; nmats=5, T=T)) + if svd_ok && fits_budget(dense_bytes(N; nmats=5, T=T)) sub["svd"] = @benchmarkable(wait(svd(A).U), setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), teardown = (A = nothing; @everywhere GC.gc())) diff --git a/benchmark/worker_mpi.jl b/benchmark/worker_mpi.jl index 0e9b1dbd1..e181d7af5 100644 --- a/benchmark/worker_mpi.jl +++ b/benchmark/worker_mpi.jl @@ -16,6 +16,8 @@ # that completed (in sorted-keypath order). # - `results_mpi_manifest.json`: [{keypath, file}, ...] pairing each # completed leaf's key path to its result file. +# - `error_mpi_rank_.json`: a rank-local `CapturedException`, written +# before aborting the communicator when benchmark work fails. # - `done`: written last, once every benchmark has been attempted. # # Note this worker does not have worker.jl's per-scale OOM isolation: a @@ -187,13 +189,24 @@ for (keypath, bench) in leaves trial = run_mpi_benchmark(bench) rank == 0 && push!(results, (kp, trial)) catch err + bt = catch_backtrace() + # Any rank may be the first to fail, and MPI.Abort can terminate rank 0 + # before it observes the same error. Give every failing rank its own + # sidecar so the orchestrator can report the first one that survives. + failure = (; + benchmark=join(kp, " / "), + rank, + exception=sprint(showerror, CapturedException(err, bt)), + ) + atomic_write(joinpath(WORKDIR, "error_mpi_rank_$(rank).json"), + JSON3.write(failure)) if err isa OutOfMemoryError rank == 0 && @warn "[worker_mpi] OutOfMemoryError; aborting MPI job" benchmark = join(kp, " / ") flush(stdout); flush(stderr) MPI.Abort(comm, 137) exit(137) # unreachable unless MPI.Abort fails to terminate us else - @error "[worker_mpi] Benchmark errored; aborting MPI job" rank benchmark = join(kp, " / ") exception = (err, catch_backtrace()) + @error "[worker_mpi] Benchmark errored; aborting MPI job" rank benchmark = join(kp, " / ") exception = (err, bt) flush(stdout); flush(stderr) MPI.Abort(comm, 1) exit(1) # unreachable unless MPI.Abort fails to terminate us From 8563030d2dd43682a21ea2d5592dfee9e6089eeb Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 14 Sep 2026 15:27:29 -0700 Subject: [PATCH 13/26] benchmark: Stabilize Distributed fixture placement Define generated stencil kernels and their boundary-condition bindings on every Distributed worker, then use a deterministic cyclic fixture grid only when real Distributed workers exist. This prevents remote UndefVar/world-age failures and stops driver-local BenchmarkTools allocation counts from varying with accidental tile ownership while retaining MPI-aware arbitrary placement under MPI.\n\nExact N=1024 comparisons completed all 16 stencil and 14 array leaves on both revisions: stencil improved 12/16 times and 14/16 memory samples, while array removed the visible allocation regressions and retained the large X+X improvement. --- AGENTS.md | 24 +++++++ benchmark/suites/array.jl | 24 +++++-- benchmark/suites/stencil.jl | 132 ++++++++++++++++++++---------------- 3 files changed, 115 insertions(+), 65 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 3ff26e18b..99b837dd7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -383,3 +383,27 @@ lesson. complete." Registering every exact destination span is correct too, but is a performance cliff for halo exchange: a 64-tile stencil produced hundreds of megabytes of interval-tree/overlap bookkeeping and regressed by 7–8x. + +36. **Distributed benchmark kernels must exist on every worker before + sampling.** Defining an `@stencil` wrapper only on the driver serializes its + generated closure to whichever worker happens to receive a tile. Globals + used by the closure (`Clamp`, `Reflect`, etc.) can then be missing or too new + for that worker's world age, and even successful leaves randomly pay remote + compilation during a timed sample. Import macros/globals in one + `@everywhere` statement, then define the macro-using wrappers in a second + `@everywhere` statement (the import must be evaluated before remote macro + expansion). A one-block capability probe is not a distributed warmup. + +37. **BenchmarkTools sees only the driver process's allocations.** In a + Distributed benchmark, arbitrary tile placement makes the reported bytes + include however many payload tiles happened to execute on the driver. A + 1024² `Float64` allocation consequently varied from 2–8 MiB with no global + allocation change, and task overhead varied with the driver's share too. + Give Distributed benchmark fixtures a deterministic balanced proc grid (the + array and stencil suites use `assignment=:cyclicrow` when + `length(procs()) > 1`) so both revisions measure the same local fraction of + the workload. Do not force that assignment under MPI: named cyclic grids are + built from Distributed processors, and an MPI rank's `procs()` is only `[1]`, + so the grid is empty and allocation divides by zero. Retain `:arbitrary` + there so the MPI-aware scheduler places tiles. The reported number is still + process-local; deterministic placement only makes the comparison meaningful. diff --git a/benchmark/suites/array.jl b/benchmark/suites/array.jl index 7d93c4889..a84c9c890 100644 --- a/benchmark/suites/array.jl +++ b/benchmark/suites/array.jl @@ -13,6 +13,10 @@ function array_suite(ctx; method, accels) T = Float64 suite = BenchmarkGroup() + # The named cyclic grids are currently backed by Distributed processors; + # an MPI worker has no Distributed workers and must use the MPI-aware + # arbitrary allocator instead. + fixture_assignment = length(procs()) > 1 ? :cyclicrow : :arbitrary for N in scales # Elementwise ops hold at most the input plus a same-size result. @@ -20,31 +24,37 @@ function array_suite(ctx; method, accels) for b in blocks_for(N) sub = BenchmarkGroup() - sub["alloc (rand)"] = @benchmarkable(wait(rand(Blocks($b, $b), $T, $N, $N)), + # Multi-process BenchmarkTools only observes allocations made on + # the driver process. Arbitrary placement therefore turns each + # driver-owned tile payload into benchmark noise (a 1024² result + # can vary by several MiB between otherwise-identical revisions). + # A cyclic grid keeps the Distributed workload balanced and gives + # both revisions the same share of driver-local tiles. + sub["alloc (rand)"] = @benchmarkable(wait(rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment)), teardown = (@everywhere GC.gc())) sub["broadcast (X .+ 1)"] = @benchmarkable(wait(X .+ 1), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(X)), teardown = (X = nothing; @everywhere GC.gc())) sub["add (X + X)"] = @benchmarkable(wait(X + X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(X)), teardown = (X = nothing; @everywhere GC.gc())) sub["map (sin.(X))"] = @benchmarkable(wait(sin.(X)), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(X)), teardown = (X = nothing; @everywhere GC.gc())) sub["transpose (permutedims)"] = @benchmarkable(wait(permutedims(X)), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(X)), teardown = (X = nothing; @everywhere GC.gc())) sub["reduce (sum)"] = @benchmarkable(sum(X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(X)), teardown = (X = nothing; @everywhere GC.gc())) sub["norm"] = @benchmarkable(norm(X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(X)), teardown = (X = nothing; @everywhere GC.gc())) suite["N=$N (block $b)"] = sub diff --git a/benchmark/suites/stencil.jl b/benchmark/suites/stencil.jl index 4a388ce78..413dc503c 100644 --- a/benchmark/suites/stencil.jl +++ b/benchmark/suites/stencil.jl @@ -7,50 +7,60 @@ # resident; sizes whose estimated peak allocation exceeds the memory budget are # skipped. +# Generated stencil task closures resolve boundary-condition constructors in +# `Main` on whichever process executes them. Import the bindings everywhere; +# importing only in the benchmark worker makes a remotely scheduled Clamp or +# Reflect task fail with `UndefVarError` (and made capability probes depend on +# the tiny probe's accidental placement). import Dagger: @stencil, Wrap, Pad, Reflect, Clamp +@everywhere import Dagger: @stencil, Wrap, Pad, Reflect, Clamp +@everywhere begin + # BenchmarkTools cannot parse @stencil inside @benchmarkable, so each + # kernel is wrapped in a plain function. Define the wrappers everywhere: + # their macro-generated task closures otherwise arrive at a remote worker + # as new Main bindings during a timed sample, with both world-age failures + # and compilation noise depending on accidental task placement. + function stencil_assign!(B, ::Type{T}) where {T} + @stencil B[idx] = one(T) + return B + end -# BenchmarkTools cannot parse @stencil inside @benchmarkable, so each kernel is -# wrapped in a plain function. -function stencil_assign!(B, ::Type{T}) where {T} - @stencil B[idx] = one(T) - return B -end - -function stencil_neighbors_wrap!(A, B) - @stencil B[idx] = sum(@neighbors(A[idx], 1, Wrap())) - return B -end + function stencil_neighbors_wrap!(A, B) + @stencil B[idx] = sum(@neighbors(A[idx], 1, Wrap())) + return B + end -function stencil_neighbors_pad!(A, B) - @stencil B[idx] = sum(@neighbors(A[idx], 1, Pad(0))) - return B -end + function stencil_neighbors_pad!(A, B) + @stencil B[idx] = sum(@neighbors(A[idx], 1, Pad(0))) + return B + end -function stencil_neighbors_clamp!(A, B) - @stencil B[idx] = sum(@neighbors(A[idx], 1, Clamp())) - return B -end + function stencil_neighbors_clamp!(A, B) + @stencil B[idx] = sum(@neighbors(A[idx], 1, Clamp())) + return B + end -function stencil_neighbors_reflect!(A, B) - @stencil B[idx] = sum(@neighbors(A[idx], 1, Reflect(true))) - return B -end + function stencil_neighbors_reflect!(A, B) + @stencil B[idx] = sum(@neighbors(A[idx], 1, Reflect(true))) + return B + end -function stencil_alloc_neighbors_wrap(A) - return @stencil sum(@neighbors(A[idx], 1, Wrap())) -end + function stencil_alloc_neighbors_wrap(A) + return @stencil sum(@neighbors(A[idx], 1, Wrap())) + end -function stencil_update_plus!(A, B) - @stencil B[idx] = B[idx] + A[idx] - return B -end + function stencil_update_plus!(A, B) + @stencil B[idx] = B[idx] + A[idx] + return B + end -function stencil_multi_expr!(A, B, ::Type{T}) where {T} - @stencil begin - A[idx] = one(T) - B[idx] = A[idx] * 2 + function stencil_multi_expr!(A, B, ::Type{T}) where {T} + @stencil begin + A[idx] = one(T) + B[idx] = A[idx] * 2 + end + return B end - return B end function stencil_suite(ctx; method, accels) @@ -60,47 +70,53 @@ function stencil_suite(ctx; method, accels) T = Float64 suite = BenchmarkGroup() + # Named cyclic grids use Distributed processors. MPI ranks do not populate + # `Distributed.procs()`, so retain the MPI-aware arbitrary allocator there. + fixture_assignment = length(procs()) > 1 ? :cyclicrow : :arbitrary + # Keep fixture placement identical across revisions. Besides stabilizing + # transfer work, this makes the driver's process-local allocation sample + # represent the same fraction of the distributed computation every time. # Capability probes (run once, at a tiny size). @stencil is a relatively new # Dagger feature: a baseline revision in an AirspeedVelocity comparison may # lack it entirely, in which case running the kernel would abort the whole # benchmark run. assign_ok = supported("stencil/assign (const)") do - B = zeros(Blocks(2, 2), T, 8, 8) + B = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) stencil_assign!(B, T) end wrap_ok = supported("stencil/neighbors (Wrap)") do - A = ones(Blocks(2, 2), T, 8, 8) - B = zeros(Blocks(2, 2), T, 8, 8) + A = ones(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) + B = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) stencil_neighbors_wrap!(A, B) end pad_ok = supported("stencil/neighbors (Pad)") do - A = ones(Blocks(2, 2), T, 8, 8) - B = zeros(Blocks(2, 2), T, 8, 8) + A = ones(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) + B = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) stencil_neighbors_pad!(A, B) end clamp_ok = supported("stencil/neighbors (Clamp)") do - A = ones(Blocks(2, 2), T, 8, 8) - B = zeros(Blocks(2, 2), T, 8, 8) + A = ones(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) + B = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) stencil_neighbors_clamp!(A, B) end reflect_ok = supported("stencil/neighbors (Reflect)") do - A = ones(Blocks(2, 2), T, 8, 8) - B = zeros(Blocks(2, 2), T, 8, 8) + A = ones(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) + B = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) stencil_neighbors_reflect!(A, B) end alloc_ok = supported("stencil/alloc (neighbors Wrap)") do - A = ones(Blocks(2, 2), T, 8, 8) + A = ones(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) wait(stencil_alloc_neighbors_wrap(A)) end update_ok = supported("stencil/update (+)") do - A = ones(Blocks(2, 2), T, 8, 8) - B = zeros(Blocks(2, 2), T, 8, 8) + A = ones(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) + B = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) stencil_update_plus!(A, B) end multi_ok = supported("stencil/multi-expr") do - A = zeros(Blocks(2, 2), T, 8, 8) - B = zeros(Blocks(2, 2), T, 8, 8) + A = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) + B = zeros(Blocks(2, 2), T, 8, 8; assignment=fixture_assignment) stencil_multi_expr!(A, B, T) end @@ -112,44 +128,44 @@ function stencil_suite(ctx; method, accels) if fits_budget(dense_bytes(N; nmats=2, T=T)) if assign_ok sub["assign (const)"] = @benchmarkable(stencil_assign!(B, $T), - setup = (B = zeros(Blocks($b, $b), $T, $N, $N); wait(B)), + setup = (B = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(B)), teardown = (B = nothing; @everywhere GC.gc())) end if wrap_ok sub["neighbors (Wrap)"] = @benchmarkable(stencil_neighbors_wrap!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); B = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if pad_ok sub["neighbors (Pad)"] = @benchmarkable(stencil_neighbors_pad!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); B = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if clamp_ok sub["neighbors (Clamp)"] = @benchmarkable(stencil_neighbors_clamp!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); B = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if reflect_ok sub["neighbors (Reflect)"] = @benchmarkable(stencil_neighbors_reflect!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); B = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if update_ok sub["update (+)"] = @benchmarkable(stencil_update_plus!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); B = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end if multi_ok sub["multi-expr"] = @benchmarkable(stencil_multi_expr!(A, B, $T), - setup = (A = zeros(Blocks($b, $b), $T, $N, $N); - B = zeros(Blocks($b, $b), $T, $N, $N); wait(A); wait(B)), + setup = (A = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); + B = zeros(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A); wait(B)), teardown = (A = nothing; B = nothing; @everywhere GC.gc())) end end @@ -157,7 +173,7 @@ function stencil_suite(ctx; method, accels) # Functional allocation syntax also materializes an output DArray. if alloc_ok && fits_budget(dense_bytes(N; nmats=3, T=T)) sub["alloc (neighbors Wrap)"] = @benchmarkable(wait(stencil_alloc_neighbors_wrap(A)), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A)), teardown = (A = nothing; @everywhere GC.gc())) end From c9f7430570bed54396cda1c5db8d12bc6642d712 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 15 Sep 2026 11:38:45 -0700 Subject: [PATCH 14/26] datadeps: Track every producer of a disjoint copy batch A whole-object consumer must wait for every disjoint remainder copy. Recording only the last copy as whole-object owner loses sibling producers; forcing the last copy to depend on siblings repairs correctness but serializes an otherwise parallel tail. Launch pieces independently, then record one logical whole-object writer backed by a CopyBatchOwner. Expand its physical producers for live writer, history, and free dependencies, while leaving singleton copies on the direct path. Two-worker 4x4 ChunkView plus whole read/write steady-state GC minima (10 warmups, GC, 5 runs): before e04a672f 13,833 allocations / 679,952 bytes; after 14,255 / 706,384 (+422 / +26,432), the cost of explicit fan-in. Allocation suite: 19/19; Distributed datadeps: 1,462 passed, 2 pre-existing broken; four-rank focused race: 8 repetitions passed; full four-rank MPI suite: 443-470 passed per rank. --- AGENTS.md | 29 +++++++++++----- src/datadeps/aliasing.jl | 51 +++++++++++++++++++++++----- src/datadeps/remainders.jl | 68 +++++++++++++++++++++++--------------- test/mpi.jl | 5 +++ 4 files changed, 110 insertions(+), 43 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 99b837dd7..5ee982032 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -375,14 +375,15 @@ lesson. 35. **A whole-region copy owner must represent the whole copy batch.** A remainder can be assembled by several copy tasks from disjoint source spaces. Registering each task as the owner of the whole destination makes - the last registration replace the earlier ones; unless that last task is a - causal successor of the other copies, a whole-region consumer waits for one - piece and a late copy can overwrite its result. Fuse the join into the final - copy by making it depend on the earlier copies (which may remain mutually - parallel), then the one recorded owner truthfully means "the whole batch is - complete." Registering every exact destination span is correct too, but is a - performance cliff for halo exchange: a 64-tile stencil produced hundreds of - megabytes of interval-tree/overlap bookkeeping and regressed by 7–8x. + the last registration replace the earlier ones; a whole-region consumer + then waits for one piece and a late copy can overwrite its result. Making + the last task depend on every earlier copy repairs that owner invariant, + but couples disjoint copies unnecessarily and puts the last one on the + critical path. Represent the logical write as a batch producer instead + (lesson 38). Registering every exact destination span is correct too, but + is a performance cliff for halo exchange: a 64-tile stencil produced + hundreds of megabytes of interval-tree/overlap bookkeeping and regressed + by 7–8x. 36. **Distributed benchmark kernels must exist on every worker before sampling.** Defining an `@stencil` wrapper only on the driver serializes its @@ -407,3 +408,15 @@ lesson. so the grid is empty and allocation divides by zero. Retain `:arbitrary` there so the MPI-aware scheduler places tiles. The reported number is still process-local; deterministic placement only makes the comparison meaningful. + +38. **A logical copy batch must keep every physical producer in every + dependency view.** When a `MultiRemainderAliasing` restores one whole-region + replica from disjoint pieces, launch each copy with its own source readers + and syncdeps, but do not let each copy rewrite the destination's whole + `ainfos_owner`, `arg_history`, or `arg_current`. After all copies are + launched, record one logical writer whose producer is the complete task + batch. Expand that producer in whole-object read/write dependencies, + historical remainder dependencies, and free-buffer syncdeps; widening only + the live owner silently leaves history or teardown waiting for one copy. + Keep the singleton-copy path direct so ordinary per-argument moves do not + allocate a batch vector. diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 151f0476c..e1872199b 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -385,6 +385,27 @@ logged_task_args(deps_vec::Vector{DataDepsTaskDependency}, infos::Vector{TaskArg [LoggedTaskArg(i.arg, i.pos, i.may_alias, i.inplace_move, deps_vec[arg_deps_range(i)]) for i in infos] +"A logical copy whose disjoint pieces may execute concurrently." +struct CopyBatchOwner + tasks::Vector{DTask} + function CopyBatchOwner(tasks::Vector{DTask}) + @assert length(tasks) > 1 + new(tasks) + end +end +const WriterProducer = Union{DTask,CopyBatchOwner} + +@inline function append_producer_syncdeps!(syncdeps, task::DTask) + push!(syncdeps, ThunkSyncdep(task)) + return +end +function append_producer_syncdeps!(syncdeps, batch::CopyBatchOwner) + for task in batch.tasks + push!(syncdeps, ThunkSyncdep(task)) + end + return +end + struct HistoryEntry ainfo::AliasingWrapper space::MemorySpace @@ -393,7 +414,7 @@ struct HistoryEntry # directly instead of re-resolving through live `ainfos_owner`, which may # later name a different task for the same ainfo or miss the producer when # only an overlapping ainfo is consulted. - task::DTask + task::WriterProducer end struct AliasedObjectCacheStore @@ -722,7 +743,7 @@ mutable struct DataDepsState # Track writers ("owners") and readers # Updated as new writer and reader tasks are launched # Used by task dependency tracking to calculate syncdeps and ensure correct launch ordering - ainfos_owner::Dict{AliasingWrapper,Union{Pair{DTask,Int},Nothing}} + ainfos_owner::Dict{AliasingWrapper,Union{Pair{DTask,Int},Pair{CopyBatchOwner,Int},Nothing}} ainfos_readers::Dict{AliasingWrapper,Vector{Pair{DTask,Int}}} # Per-task scratch buffers filled by `populate_task_info!` and consumed by @@ -755,7 +776,7 @@ mutable struct DataDepsState ainfos_lookup = AliasingLookup() ainfos_overlaps = Dict{AliasingWrapper,Set{AliasingWrapper}}() - ainfos_owner = Dict{AliasingWrapper,Union{Pair{DTask,Int},Nothing}}() + ainfos_owner = Dict{AliasingWrapper,Union{Pair{DTask,Int},Pair{CopyBatchOwner,Int},Nothing}}() ainfos_readers = Dict{AliasingWrapper,Vector{Pair{DTask,Int}}}() 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, @@ -1115,7 +1136,7 @@ function _get_write_deps!(state::DataDepsState, dest_space::MemorySpace, ainfo:: other_task, other_write_num = other_task_write_num write_num == other_write_num && continue @dagdebug nothing :spawn_datadeps_sync "Sync with writer via $ainfo -> $other_ainfo" - push!(syncdeps, ThunkSyncdep(other_task)) + append_producer_syncdeps!(syncdeps, other_task) end end function _get_read_deps!(state::DataDepsState, dest_space::MemorySpace, ainfo::AbstractAliasing, write_num, syncdeps) @@ -1163,7 +1184,7 @@ function gather_overlap_syncdeps!(state::DataDepsState, ainfo::AliasingWrapper, owner = get(state.ainfos_owner, other_ainfo, nothing) if owner !== nothing owner_task, owner_write_num = owner - owner_write_num != write_num && push!(syncdeps, ThunkSyncdep(owner_task)) + owner_write_num != write_num && append_producer_syncdeps!(syncdeps, owner_task) end for (reader_task, reader_write_num) in get(state.ainfos_readers, other_ainfo, ()) reader_write_num != write_num && push!(syncdeps, ThunkSyncdep(reader_task)) @@ -1192,7 +1213,13 @@ function assert_free_syncdeps!(state::DataDepsState, ainfo::AliasingWrapper, wri owner_task, owner_write_num = owner owner_write_num == write_num && continue will_alias(ainfo, other_ainfo) || continue - @assert ThunkSyncdep(owner_task) in syncdeps "gather_free_syncdeps! omitted a live writer $owner_task ($other_ainfo) for buffer overlapping $ainfo" + if owner_task isa CopyBatchOwner + for task in owner_task.tasks + @assert ThunkSyncdep(task) in syncdeps "gather_free_syncdeps! omitted a live copy $task ($other_ainfo) for buffer overlapping $ainfo" + end + else + @assert ThunkSyncdep(owner_task) in syncdeps "gather_free_syncdeps! omitted a live writer $owner_task ($other_ainfo) for buffer overlapping $ainfo" + end end for (other_ainfo, readers) in state.ainfos_readers for (reader_task, reader_write_num) in readers @@ -1255,7 +1282,7 @@ function gather_free_syncdeps!(state::DataDepsState, space::MemorySpace, buf_ain assert_free_syncdeps!(state, wrapped, write_num, syncdeps) return end -function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::MemorySpace, ainfo::AbstractAliasing, task, write_num; copy_src::Union{MemorySpace,Nothing}=nothing) +function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::MemorySpace, ainfo::AbstractAliasing, task::WriterProducer, write_num; copied::Bool=false) state.ainfos_owner[ainfo] = task=>write_num empty!(state.ainfos_readers[ainfo]) @@ -1272,7 +1299,7 @@ function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::M end # Track which spaces hold a fully-current replica of this region - if copy_src === nothing + if !copied # Task write: only the written space is current, and other spaces' # replicas of overlapping regions become stale # N.B. The `Set` is reused in place rather than replaced; nothing ever @@ -1308,7 +1335,13 @@ function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::M state.arg_owner[arg_w] = dest_space # Not necessary to assert a read, but conceptually it's true - add_reader!(state, arg_w, dest_space, ainfo, task, write_num) + if task isa CopyBatchOwner + for copy_task in task.tasks + add_reader!(state, arg_w, dest_space, ainfo, copy_task, write_num) + end + else + add_reader!(state, arg_w, dest_space, ainfo, task, write_num) + end end function add_reader!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::MemorySpace, ainfo::AbstractAliasing, task, write_num) push!(state.ainfos_readers[ainfo], task=>write_num) diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 31683f6af..28a9ed615 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -263,7 +263,7 @@ function compute_remainder_for_arg!(state::DataDepsState, # ownership of `other_ainfo` may have moved to a later copy, or # the producer may only be registered on an overlapping ainfo # missing from `ainfos_overlaps[other_ainfo]`. - push!(tracker_other_space[3], ThunkSyncdep(other_entry.task)) + append_producer_syncdeps!(tracker_other_space[3], other_entry.task) else # idx==0 owner fallback: no HistoryEntry, use live writer set. get_read_deps!(state, other_space, other_ainfo, write_num, tracker_other_space[3]) @@ -325,25 +325,32 @@ Enqueues a copy operation to update the remainder regions of an object before a function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::MultiRemainderAliasing, f, idx, dest_scope, task, write_num::Int) remainders = remainder_aliasing.remainders - last_idx = length(remainders) - for (i, remainder) in enumerate(remainders) + if length(remainders) == 1 + remainder = only(remainders) @check_uniform(remainder.space) @assert !isempty(remainder.spans) @check_uniform(remainder.spans) - copy_task = enqueue_remainder_copy_to!(state, dest_space, arg_w, + return enqueue_remainder_copy_to!(state, dest_space, arg_w, remainder, f, idx, dest_scope, task, write_num) - if i < last_idx - # Every copy is recorded as the owner of the whole destination, so - # the final one replaces the earlier owners. Make that replacement - # truthful: once the final copy completes, all disjoint pieces have - # completed. Earlier copies remain mutually parallel, and this - # fuses the batch join into work we already had to schedule. - push!(remainders[last_idx].syncdeps, ThunkSyncdep(copy_task)) - end end + copy_tasks = Vector{DTask}(undef, length(remainders)) + for (i, remainder) in enumerate(remainders) + @check_uniform(remainder.space) + @assert !isempty(remainder.spans) + @check_uniform(remainder.spans) + copy_tasks[i] = enqueue_remainder_copy_to!(state, dest_space, arg_w, + remainder, f, idx, dest_scope, task, write_num; record_dest=false) + end + # The destination becomes logically current when all disjoint copies finish. + # Record that one write event only after every piece has been launched, so + # neither copy ownership nor history can discard a sibling producer. + target_ainfo = aliasing!(state, dest_space, arg_w) + add_writer!(state, arg_w, dest_space, target_ainfo, + CopyBatchOwner(copy_tasks), write_num; copied=true) + return end function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::RemainderAliasing, - f, idx, dest_scope, task, write_num::Int) + f, idx, dest_scope, task, write_num::Int; record_dest::Bool=true) dep_mod = arg_w.dep_mod # Find the source space for the remainder data @@ -383,8 +390,8 @@ function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpac for ainfo in source_ainfos add_reader!(state, arg_w, source_space, ainfo, copy_task, write_num) end - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; - copy_src=source_space) + record_dest && add_writer!(state, arg_w, dest_space, target_ainfo, + copy_task, write_num; copied=true) return copy_task end """ @@ -396,20 +403,29 @@ Enqueues a copy operation to update the remainder regions of an object back to t function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::MultiRemainderAliasing, dest_scope, write_num::Int) remainders = remainder_aliasing.remainders - last_idx = length(remainders) - for (i, remainder) in enumerate(remainders) + if length(remainders) == 1 + remainder = only(remainders) @check_uniform(remainder.space) @assert !isempty(remainder.spans) @check_uniform(remainder.spans) - copy_task = enqueue_remainder_copy_from!(state, dest_space, arg_w, + return enqueue_remainder_copy_from!(state, dest_space, arg_w, remainder, dest_scope, write_num) - if i < last_idx - push!(remainders[last_idx].syncdeps, ThunkSyncdep(copy_task)) - end end + copy_tasks = Vector{DTask}(undef, length(remainders)) + for (i, remainder) in enumerate(remainders) + @check_uniform(remainder.space) + @assert !isempty(remainder.spans) + @check_uniform(remainder.spans) + copy_tasks[i] = enqueue_remainder_copy_from!(state, dest_space, arg_w, + remainder, dest_scope, write_num; record_dest=false) + end + target_ainfo = aliasing!(state, dest_space, arg_w) + add_writer!(state, arg_w, dest_space, target_ainfo, + CopyBatchOwner(copy_tasks), write_num; copied=true) + return end function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, remainder_aliasing::RemainderAliasing, - dest_scope, write_num::Int) + dest_scope, write_num::Int; record_dest::Bool=true) dep_mod = arg_w.dep_mod # Find the source space for the remainder data @@ -449,8 +465,8 @@ function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySp for ainfo in source_ainfos add_reader!(state, arg_w, source_space, ainfo, copy_task, write_num) end - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; - copy_src=source_space) + record_dest && add_writer!(state, arg_w, dest_space, target_ainfo, + copy_task, write_num; copied=true) return copy_task end @@ -488,7 +504,7 @@ function enqueue_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w:: # This copy task reads the source and writes to the target add_reader!(state, arg_w, source_space, source_ainfo, copy_task, write_num) - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copied=true) end function enqueue_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w::ArgumentWrapper, dest_scope, write_num::Int) @@ -523,7 +539,7 @@ function enqueue_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w # This copy task reads the source and writes to the target add_reader!(state, arg_w, source_space, source_ainfo, copy_task, write_num) - add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copy_src=source_space) + add_writer!(state, arg_w, dest_space, target_ainfo, copy_task, write_num; copied=true) end # Main copy function for RemainderAliasing diff --git a/test/mpi.jl b/test/mpi.jl index f7e138589..c10f8b75e 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -449,6 +449,7 @@ end cv_top = view(c, 1:2, 1:4) cv_bot = view(c, 3:4, 1:4) @test cv_top isa Dagger.ChunkView + whole_sum = Ref(0.0) r1 = min(1, nranks-1) r2 = min(2, nranks-1) @@ -459,12 +460,16 @@ end # the late top-half copy overwrote its result. Dagger.@spawn scope=rank_scope(r1) delayed_add1!(InOut(cv_top)) Dagger.@spawn scope=rank_scope(r2) scale2!(InOut(cv_bot)) + # A whole-object read must wait for both disjoint writeback copies, + # even though neither copy covers the whole chunk by itself. + Dagger.@spawn scope=rank_scope(0) sum_into_f!(Out(whole_sum), In(c)) Dagger.@spawn scope=rank_scope(0) add1!(InOut(c)) end ref_blk = A[1:4, 1:4] ref_blk[1:2, :] .+= 1 ref_blk[3:4, :] .*= 2 + rank == 0 && @test whole_sum[] ≈ sum(ref_blk) ref_blk .+= 1 # Collective uniform fetches: identical on every rank @test fetch(c) ≈ ref_blk From 4ad6e59b49da7b9e528f7be47e400ddc8a6acb7a Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 10:50:15 -0700 Subject: [PATCH 15/26] datadeps: Skip clock reads when hierarchical timing is disabled hier_log! gates event emission but cannot gate evaluation of its arguments. Per-argument aliasing and slot creation still read the clock even with HIER_TIMING=false; guard both start timestamps and duration construction, preserving enabled diagnostics and all scheduling/dependency behavior. Julia 1.12 GC minima, 10 warmups then GC.gc() and 5 runs, 10,000 four-argument aliasing batches: before and after 160,000 allocations / 5,760,000 bytes; minimum elapsed 11.25 -> 8.63 ms (-23.3%). Timing on/off and identity/general slots: 15 passed. Unmodified allocation bounds: 19 passed. --- AGENTS.md | 7 +++++++ ext/MPIExt.jl | 5 +++-- src/datadeps/aliasing.jl | 13 +++++++----- src/datadeps/hierarchical.jl | 5 +++-- test/datadeps.jl | 39 ++++++++++++++++++++++++++++++++++++ 5 files changed, 60 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5ee982032..fec240883 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -420,3 +420,10 @@ lesson. the live owner silently leaves history or teardown waiting for one copy. Keep the singleton-copy path direct so ordinary per-argument moves do not allocate a batch vector. + +39. **Gate diagnostic work at the call site, not only inside the logger.** + `hier_log!` checks `HIER_TIMING[]`, but Julia evaluates its arguments before + that check. Unconditional `time_ns()` calls around per-argument aliasing and + slot creation therefore survived with timing disabled. Gate both the start + timestamp and the finish/event argument evaluation; a disabled diagnostic + should not pay for the measurement it discards. diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 38c970c4a..e19bdfb0e 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -260,7 +260,8 @@ function Dagger.batch_ainfos(accel::MPIAcceleration, objs::Vector, dep_mods::Vec comm = accel.comm rank = MPI.Comm_rank(comm) check_uniform(UInt64(n)) - t0 = time_ns() + timing = Dagger.HIER_TIMING[] + t0 = timing ? time_ns() : UInt64(0) owners = Vector{Int}(undef, n) ainfos = Vector{Union{Nothing,AbstractAliasing}}(nothing, n) @@ -303,7 +304,7 @@ function Dagger.batch_ainfos(accel::MPIAcceleration, objs::Vector, dep_mods::Vec # lookup of the same object agree even if an entry was already there. results[i] = Dagger.memoize_ainfo!(Dagger.ainfo_memo_key(objs[i], dep_mods[i]), ainfo) end - Dagger.hier_log!(Dagger.LogHierAinfo, 0x01, Dagger.LogHierAinfoId(), (time_ns() - t0, n)) + timing && Dagger.hier_log!(Dagger.LogHierAinfo, 0x01, Dagger.LogHierAinfoId(), (time_ns() - t0, n)) return results end diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index e1872199b..e020d864b 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1426,7 +1426,8 @@ function generate_slot!(state::DataDepsState, dest_space, data) id = logging ? rand(Int) : 0 logging && @logstart ctx LogMove LogMoveId(0, ArgPosition(), to_proc, id) data tid = something(DATADEPS_CURRENT_TASK[], (;uid=0)).uid - t0 = time_ns() + timing = HIER_TIMING[] + t0 = timing ? time_ns() : UInt64(0) data_chunk = if slot_is_already_in_place(data, orig_space, dest_space) # Nothing to move: the slot for data already in `dest_space` is the data # itself. Going through `move_rewrap` here would allocate a second Chunk @@ -1439,12 +1440,14 @@ function generate_slot!(state::DataDepsState, dest_space, data) 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) end - move_ns = time_ns() - t0 - hier_log!(LogHierSlot, 0x01, LogHierSlotId(:moved), move_ns) - orig_space == dest_space && hier_log!(LogHierSlot, 0x01, LogHierSlotId(:samespace), move_ns) + if timing + move_ns = time_ns() - t0 + hier_log!(LogHierSlot, 0x01, LogHierSlotId(:moved), move_ns) + orig_space == dest_space && hier_log!(LogHierSlot, 0x01, LogHierSlotId(:samespace), move_ns) + end moved end - hier_log!(LogHierSlot, 0x01, LogHierSlotId(:total), time_ns() - t0) + timing && hier_log!(LogHierSlot, 0x01, LogHierSlotId(:total), time_ns() - t0) logging && @logfinish ctx LogMove LogMoveId(0, ArgPosition(), to_proc, id) data_chunk @assert memory_space(data_chunk) == dest_space "space mismatch! $dest_space (dest) != $(memory_space(data_chunk)) (actual) ($(typeof(data)) (data) vs. $(typeof(data_chunk)) (chunk)), spaces ($orig_space -> $dest_space)" dest_space_args[data] = data_chunk diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 0c09f2830..2c4bcb015 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -964,9 +964,10 @@ function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) else for i in 1:n arg_w = arg_ws[i] - t0 = time_ns() + timing = HIER_TIMING[] + t0 = timing ? time_ns() : UInt64(0) ainfo = AliasingWrapper(aliasing(accel, arg_w.arg, arg_w.dep_mod)) - hier_log!(LogHierAinfo, 0x01, LogHierAinfoId(), (time_ns() - t0, 1)) + timing && hier_log!(LogHierAinfo, 0x01, LogHierAinfoId(), (time_ns() - t0, 1)) results[i] = arg_w => ainfo end end diff --git a/test/datadeps.jl b/test/datadeps.jl index edf99ad07..86d1641a6 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -2,6 +2,45 @@ import Dagger: ChunkView, Chunk, AbstractAliasing, MemorySpace, ArgumentWrapper import Dagger: aliasing, memory_space using LinearAlgebra, Graphs +@testset "Hierarchical timing gates" begin + # Exercise both the identity slot and the general move_rewrap path. Timing + # must remain available when enabled, independently of scheduler logging. + tsl = Dagger.TimespanLogging + old_timing = Dagger.HIER_TIMING[] + try + for timing in (false, true) + tsl.steal_typed(Dagger.LogHierAinfo) + tsl.steal_typed(Dagger.LogHierSlot) + Dagger.HIER_TIMING[] = timing + A = zeros(4) + chunk = Dagger.tochunk(A) + arg_ws = [ArgumentWrapper(chunk, identity)] + result = Dagger._compute_aliasing_batch(arg_ws) + @test only(result).second == Dagger.AliasingWrapper(aliasing(chunk)) + state = Dagger.DataDepsState() + space = memory_space(chunk) + @test Dagger.generate_slot!(state, space, chunk) === chunk + @test fetch(Dagger.generate_slot!(state, space, view(chunk, 1:4))) == A + ainfo_events = tsl.steal_typed(Dagger.LogHierAinfo) + slot_events = tsl.steal_typed(Dagger.LogHierSlot) + if timing + @test length(ainfo_events) == 1 + @test only(ainfo_events).data[1] > 0 + @test only(ainfo_events).data[2] == 1 + @test count(ev -> ev.id.kind === :total, slot_events) == 2 + @test count(ev -> ev.id.kind === :moved, slot_events) == 1 + @test count(ev -> ev.id.kind === :samespace, slot_events) == 1 + @test all(ev -> ev.data > 0, slot_events) + else + @test isempty(ainfo_events) + @test isempty(slot_events) + end + end + finally + Dagger.HIER_TIMING[] = old_timing + end +end + @testset "Memory Aliasing" begin A = rand(4) a = Dagger.aliasing(A) From e2b9d3b50e97b5d3135c57201053946328c55e90 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 10:51:26 -0700 Subject: [PATCH 16/26] datadeps: Preserve concrete types for reusable aliasing scratch The task-local vector macro reads a mutable global and infers Any. The newly pooled aliasing-result buffer therefore made every element store a dynamically dispatched boxed Pair; deferred copy-info scratch had the same problem for iteration and batch dispatch. Assert the declared concrete vector types at the use sites, retaining pooling and exactly the same contents and ownership semantics. Check inferred results on sequential and threaded paths. Julia 1.12 GC minima, 10 warmups then GC.gc() and 5 runs, 10,000 four-argument aliasing batches: 160,000 -> 120,000 allocations (-25%), 5,760,000 -> 3,840,000 bytes (-33.3%), minimum elapsed after gating clocks 8.63 -> 7.56 ms. Scratch/timing tests: 17 passed; allocation suite: 19 passed without changing bounds. --- AGENTS.md | 9 +++++++++ src/datadeps/aliasing.jl | 6 +++--- src/datadeps/hierarchical.jl | 2 +- test/datadeps.jl | 11 ++++++++++- 4 files changed, 23 insertions(+), 5 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fec240883..c0565e9e5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -427,3 +427,12 @@ lesson. slot creation therefore survived with timing disabled. Gate both the start timestamp and the finish/event argument evaluation; a disabled diagnostic should not pay for the measurement it discards. + +40. **Assert the concrete type of reusable scratch at the point of use.** + `@reusable_vector` expands through a non-const global `TaskLocalValue`, so + the expression infers as `Any` even though its runtime vector is typed. + Without a `::Vector{...}` assertion, iteration, element stores, and closures + capturing it lose that type throughout the hot loop. The aliasing-result + scratch introduced one boxed pair allocation per argument this way. Keep + reuse, but recover the concrete type before entering the loop; check the + return type with `@inferred` as well as the container's runtime type. diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index e020d864b..5b2475137 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -480,15 +480,15 @@ function resolve_pending!(cache::AliasedObjectCacheStore) # function (see the module docstring on deferred resolution being # SPMD-safe), so there is nothing left to alias by the time that could # happen. - entries = @reusable_vector :resolve_pending_entries Pair{AbstractAliasing,Chunk} nothing 32 + entries = (@reusable_vector :resolve_pending_entries Pair{AbstractAliasing,Chunk} nothing 32)::Vector{Pair{AbstractAliasing,Chunk}} append!(entries, cache.pending) empty!(cache.pending) @opcounter :aliasing_resolve_pending @opcounter :aliasing_resolve_pending_entries length(entries) # Resolved as a batch: the whole point of deferring was to not pay a # rendezvous per copy, which asking one at a time here would reintroduce. - values = @reusable_vector :resolve_pending_values Chunk nothing 32 - dep_mods = @reusable_vector :resolve_pending_dep_mods Any nothing 32 + values = (@reusable_vector :resolve_pending_values Chunk nothing 32)::Vector{Chunk} + dep_mods = (@reusable_vector :resolve_pending_dep_mods Any nothing 32)::Vector{Any} for (_, value) in entries push!(values, value) push!(dep_mods, identity) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 2c4bcb015..420b36689 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -943,7 +943,7 @@ function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) # Task-local: safe to reuse across calls because the result is fully # consumed by the caller (copied into `arg_to_ainfo`) before this task # could re-enter this function for another region. - results = @reusable_vector :compute_aliasing_batch_results Pair{ArgumentWrapper,AliasingWrapper} nothing 32 + results = (@reusable_vector :compute_aliasing_batch_results Pair{ArgumentWrapper,AliasingWrapper} nothing 32)::Vector{Pair{ArgumentWrapper,AliasingWrapper}} resize!(results, n) accel = current_acceleration() # Under uniform execution (MPI), aliasing may perform collectives that must diff --git a/test/datadeps.jl b/test/datadeps.jl index 86d1641a6..7401941f1 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -8,6 +8,15 @@ using LinearAlgebra, Graphs tsl = Dagger.TimespanLogging old_timing = Dagger.HIER_TIMING[] try + # Also cover the threaded branch when the driver has multiple threads. + # The captured result buffer must infer concretely there as well. + Dagger.HIER_TIMING[] = false + parallel_args = [ArgumentWrapper(Dagger.tochunk(zeros(4)), identity) + for _ in 1:8] + parallel_result = @inferred Dagger._compute_aliasing_batch(parallel_args) + @test length(parallel_result) == length(parallel_args) + @test all(p -> p.second == Dagger.AliasingWrapper(aliasing(p.first.arg)), + parallel_result) for timing in (false, true) tsl.steal_typed(Dagger.LogHierAinfo) tsl.steal_typed(Dagger.LogHierSlot) @@ -15,7 +24,7 @@ using LinearAlgebra, Graphs A = zeros(4) chunk = Dagger.tochunk(A) arg_ws = [ArgumentWrapper(chunk, identity)] - result = Dagger._compute_aliasing_batch(arg_ws) + result = @inferred Dagger._compute_aliasing_batch(arg_ws) @test only(result).second == Dagger.AliasingWrapper(aliasing(chunk)) state = Dagger.DataDepsState() space = memory_space(chunk) From a3b80f0e6db1677cad1da2acc7af033b2567fdd3 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 10:56:57 -0700 Subject: [PATCH 17/26] mpi: Serialize each integer uniformity payload only once Replicated planning checks the same integer against every peer. MPI.isend serialized that identical value separately for each peer, allocating an IOBuffer and serialization state each time. Reuse one serialized byte vector for all completed nonblocking sends, preserving the reserved tag, peer order, receive wire format, cooperative waits, and deadlock detection. Retain the general raw-parts transport for non-integer values and the no-communication single-rank path. Four ranks, Julia 1.12, 10 warmups then GC.gc() and 5 measurements of 1,000 checks: 93,000 -> 69,000 allocations (-25.8%), 5,280,000 -> 3,744,000 bytes (-29.1%), minimum elapsed 7.64 -> 7.35 ms. Actual MPI chunk/uniformity tests: 126 passed per rank with two threads; focused mixed-width/BigInt/array/mismatch recovery tests passed with four ranks and with one rank. Both regressed SpGEMM sizes completed 10 warmups and 5 measurements per rank with orderly backend teardown. --- ext/MPIExt.jl | 14 +++++++++++++- test/mpi.jl | 5 +++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index e19bdfb0e..f810c141b 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -105,10 +105,22 @@ compare_all_mpi_tag() = UInt32(MPI.tag_ub()) function compare_all(value, comm) rank = MPI.Comm_rank(comm) size = MPI.Comm_size(comm) + size == 1 && return true tag = compare_all_mpi_tag() + # Integer checks dominate replicated planning. Every peer receives the + # exact same serialized value, so build its payload once rather than have + # MPI.isend serialize it anew for each peer. Keep the general transport for + # values which may use the raw-parts protocol (e.g. arrays). + serialized = value isa Integer ? MPI.serialize(value) : nothing for i in 0:(size-1) if i != rank - send_yield(value, comm, i, tag) + if serialized === nothing + send_yield(value, comm, i, tag) + else + @opcounter :send_yield_serialized + req = MPI.Isend(serialized, comm; dest=i, tag) + __wait_for_request(req, comm, rank, i, tag, "send_yield", "send") + end end end match = true diff --git a/test/mpi.jl b/test/mpi.jl index c10f8b75e..c66bc5fa2 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -177,6 +177,11 @@ end @testset "check_uniform" begin @test Dagger.check_uniform(42) @test Dagger.check_uniform(hash((1, :a, "x"))) + # Reused serialized payloads retain the general Integer wire format, + # including arbitrary precision and equality across integer types. + @test Dagger.check_uniform(big(2)^256) + @test Dagger.check_uniform(rank == 0 ? UInt8(42) : Int64(42)) + @test MPIExt.compare_all([1, 2, 3], comm) # Rank-dependent values must be detected on every rank @test_throws ArgumentError Dagger.check_uniform(rank) # The compare stream stays aligned after a detected failure From 1977cd5fee8a987c2fa676b86cc85b2ce0034a35 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 11:17:20 -0700 Subject: [PATCH 18/26] reuse: Assert scratch container types inside the macros Non-const task-local globals erase inference even though reusable_vector and reusable_dict already declare concrete container types. Assert those types before empty! in the macro expansions so every caller retains inferred scratch, and escape supplied type expressions to preserve caller-side name resolution. Remove redundant datadeps and scheduler call-site assertions. Add inference checks for caller-defined types, aliases, unions, and abstract elements, plus reuse, clearing, and task isolation coverage. Update the scratch inference lesson. Profile: Julia 1.12.6, 4 threads, 10 warmups, GC.gc(), minimum of 5 Base.gc_num() measurements; each run fills and iterates 1,000 batches of 32 pairs. Vector: 127,979 allocations / 3,583,664 bytes -> 0 / 0. Dict: 128,986 allocations / 3,599,776 bytes -> 0 / 0. code_warntype shows concrete scratch and entry locals with Int return types. Validation: test/reuse.jl passes (14 new macro checks); test/allocations.jl passes all 19 checks with unchanged bounds; Distributed datadeps passes 2,083 checks with 2 existing expected-broken tests; four-rank, two-thread MPI passes 126 uniformity/chunk checks plus 27 sparse/dense operation checks per rank. --- AGENTS.md | 18 +++++++------ src/datadeps/aliasing.jl | 6 ++--- src/datadeps/hierarchical.jl | 2 +- src/sch/Sch.jl | 7 +---- src/utils/reuse.jl | 6 +++-- test/reuse.jl | 52 +++++++++++++++++++++++++++++++++++- 6 files changed, 70 insertions(+), 21 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index c0565e9e5..24a601d62 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -428,11 +428,13 @@ lesson. timestamp and the finish/event argument evaluation; a disabled diagnostic should not pay for the measurement it discards. -40. **Assert the concrete type of reusable scratch at the point of use.** - `@reusable_vector` expands through a non-const global `TaskLocalValue`, so - the expression infers as `Any` even though its runtime vector is typed. - Without a `::Vector{...}` assertion, iteration, element stores, and closures - capturing it lose that type throughout the hot loop. The aliasing-result - scratch introduced one boxed pair allocation per argument this way. Keep - reuse, but recover the concrete type before entering the loop; check the - return type with `@inferred` as well as the container's runtime type. +40. **Reusable scratch macros must return their declared container types.** + `@reusable_vector` and `@reusable_dict` expand through non-const global + `TaskLocalValue`s. Without a type assertion, their expressions infer as + `Any` even though the runtime containers are typed: iteration, element + stores, and closures capturing them lose that type throughout the hot loop. + The aliasing-result scratch introduced one boxed pair allocation per + argument this way. Assert the container type inside the macros, before + `empty!`, rather than relying on each caller to do so. Escape the supplied + type expressions to resolve caller-defined types correctly, and check the + return types with `@inferred` as well as the containers' runtime types. diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 5b2475137..e020d864b 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -480,15 +480,15 @@ function resolve_pending!(cache::AliasedObjectCacheStore) # function (see the module docstring on deferred resolution being # SPMD-safe), so there is nothing left to alias by the time that could # happen. - entries = (@reusable_vector :resolve_pending_entries Pair{AbstractAliasing,Chunk} nothing 32)::Vector{Pair{AbstractAliasing,Chunk}} + entries = @reusable_vector :resolve_pending_entries Pair{AbstractAliasing,Chunk} nothing 32 append!(entries, cache.pending) empty!(cache.pending) @opcounter :aliasing_resolve_pending @opcounter :aliasing_resolve_pending_entries length(entries) # Resolved as a batch: the whole point of deferring was to not pay a # rendezvous per copy, which asking one at a time here would reintroduce. - values = (@reusable_vector :resolve_pending_values Chunk nothing 32)::Vector{Chunk} - dep_mods = (@reusable_vector :resolve_pending_dep_mods Any nothing 32)::Vector{Any} + values = @reusable_vector :resolve_pending_values Chunk nothing 32 + dep_mods = @reusable_vector :resolve_pending_dep_mods Any nothing 32 for (_, value) in entries push!(values, value) push!(dep_mods, identity) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 420b36689..2c4bcb015 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -943,7 +943,7 @@ function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) # Task-local: safe to reuse across calls because the result is fully # consumed by the caller (copied into `arg_to_ainfo`) before this task # could re-enter this function for another region. - results = (@reusable_vector :compute_aliasing_batch_results Pair{ArgumentWrapper,AliasingWrapper} nothing 32)::Vector{Pair{ArgumentWrapper,AliasingWrapper}} + results = @reusable_vector :compute_aliasing_batch_results Pair{ArgumentWrapper,AliasingWrapper} nothing 32 resize!(results, n) accel = current_acceleration() # Under uniform execution (MPI), aliasing may perform collectives that must diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 7878c92e7..299462f9f 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -1306,12 +1306,7 @@ end # to keep the local copy alive/untouched for the duration of the send. # The pool lives under `istate.queue`'s lock, and the whole batch is # borrowed in a single acquisition (rather than one per task). - # - # N.B. The `::Vector{TaskSpec}` assertion matters: `@reusable_vector` - # expands to a *non-const* global, so without it `spec` infers as `Any` and - # every field store below becomes a boxing dynamic `setproperty!` (which - # cost more than the pooling saved). - pooled = (@reusable_vector :fire_tasks!_spec_pool TaskSpec nothing 32)::Vector{TaskSpec} + pooled = @reusable_vector :fire_tasks!_spec_pool TaskSpec nothing 32 fire_istate = nothing if root_worker_id(gproc) == myid() pstate = maybe_proc_state(state.uid, proc) diff --git a/src/utils/reuse.jl b/src/utils/reuse.jl index 5eecdb96b..34a4186e9 100644 --- a/src/utils/reuse.jl +++ b/src/utils/reuse.jl @@ -565,19 +565,21 @@ end # N.B. Emptied on take (cheap when already empty), so users need no cleanup # registration for correctness — explicit `empty!` at the end of use remains # good hygiene to release references promptly. +# The task-local globals are non-const, so assert the declared container type +# before empty! to keep both that call and the caller's scratch use inferred. macro reusable_vector(name, T, null, N) vec_name = Symbol("__$(name)_TLV_ReusableVector") if !hasproperty(__module__, vec_name) __module__.eval(:(#=const=# $vec_name = $TaskLocalValue{$Vector{$T}}(()->$Vector{$T}()))) end - return :(empty!($(esc(vec_name))[])) + return :(empty!($(esc(vec_name))[]::$Vector{$(esc(T))})) end macro reusable_dict(name, K, V, null_key, null_value, N) dict_name = Symbol("__$(name)_TLV_ReusableDict") if !hasproperty(__module__, dict_name) __module__.eval(:(#=const=# $dict_name = $TaskLocalValue{$Dict{$K,$V}}(()->$Dict{$K,$V}()))) end - return :(empty!($(esc(dict_name))[])) + return :(empty!($(esc(dict_name))[]::$Dict{$(esc(K)),$(esc(V))})) end mutable struct ReusableTaskCache diff --git a/test/reuse.jl b/test/reuse.jl index 5a80df7c9..82b52d5a1 100644 --- a/test/reuse.jl +++ b/test/reuse.jl @@ -1,6 +1,56 @@ using Test import Dagger: ReusableLinkedList, ReusableDict, ReusableCache import Dagger: take_or_alloc!, maybe_take_or_alloc!, maybetake!, putback! +import Dagger: @reusable_vector, @reusable_dict + +struct ReusableScratchTestEntry + value::Int +end +const ReusableScratchTestPair = Pair{Symbol,ReusableScratchTestEntry} +const ReusableScratchTestKey = Int + +reuse_test_vector() = @reusable_vector :test_reusable_vector ReusableScratchTestPair nothing 32 +reuse_test_dict() = @reusable_dict :test_reusable_dict ReusableScratchTestKey Union{Nothing,ReusableScratchTestEntry} 0 nothing 32 +reuse_test_any_vector() = @reusable_vector :test_reusable_any_vector Any nothing 32 +reuse_test_abstract_dict() = @reusable_dict :test_reusable_abstract_dict Number Any 0 nothing 32 + +@testset "Reusable scratch macros" begin + # Types and aliases belong to the caller, not the macro's defining module. + # @inferred checks the expansion itself, without a call-site assertion. + vector = @inferred reuse_test_vector() + @test vector isa Vector{ReusableScratchTestPair} + push!(vector, :entry => ReusableScratchTestEntry(1)) + @test (@inferred reuse_test_vector()) === vector + @test isempty(vector) + + dict = @inferred reuse_test_dict() + @test dict isa Dict{Int,Union{Nothing,ReusableScratchTestEntry}} + dict[1] = ReusableScratchTestEntry(1) + dict[2] = nothing + @test (@inferred reuse_test_dict()) === dict + @test isempty(dict) + + # The container type remains concrete with abstract element/key/value types. + @test (@inferred reuse_test_any_vector()) isa Vector{Any} + @test (@inferred reuse_test_abstract_dict()) isa Dict{Number,Any} + + # A different task gets independent scratch; taking it must not clear ours. + push!(vector, :entry => ReusableScratchTestEntry(2)) + dict[1] = ReusableScratchTestEntry(2) + other_vector, other_dict = fetch(Threads.@spawn begin + task_vector = @inferred reuse_test_vector() + task_dict = @inferred reuse_test_dict() + return task_vector, task_dict + end) + @test isempty(other_vector) + @test isempty(other_dict) + @test other_vector !== vector + @test other_dict !== dict + @test length(vector) == 1 + @test length(dict) == 1 + empty!(vector) + empty!(dict) +end @testset "ReusableCache Tests" begin @testset "Construction and Basic Properties" begin @@ -1119,4 +1169,4 @@ end @test dict["apple"] == 1 # Dict unchanged @test !haskey(dict, "apricot") # New value not in dict end -end \ No newline at end of file +end From ec27fc47038c53425b4f0ee813a312ec69a8ccda Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 12:55:35 -0700 Subject: [PATCH 19/26] Reuse remote aliasing batches throughout Distributed planning Distributed RPCs do not inherit the driver region memo. Merge remote Phase 1 answers into that memo with the original local chunk/view and dependency modifier keys, so Phase 4 does not ask each owner again. Keep local and MPI batching unchanged; a typed merge barrier restores the RPC result type. Julia 1.13, one driver thread plus three one-thread workers: a 64-chunk batch/plan probe, after 10 warmups and GC, minimum of 5 Base.gc_num deltas: 18,397 allocations / 822,992 bytes before, 9,118 / 393,128 after. Redundant remote queries drop from 48 to zero; minimum elapsed time drops from 6.262 ms to 1.024 ms. Validation: Distributed datadeps 2,137 passed, 2 existing broken; allocation suite 19 passed with all bounds unchanged. New memo tests cover local, single-remote, and multiple-owner batches with modifiers and views. --- AGENTS.md | 10 +++++++++ src/datadeps/hierarchical.jl | 40 +++++++++++++++++++++--------------- test/datadeps.jl | 31 ++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 16 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 24a601d62..e039e018e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -438,3 +438,13 @@ lesson. `empty!`, rather than relying on each caller to do so. Escape the supplied type expressions to resolve caller-defined types correctly, and check the return types with `@inferred` as well as the containers' runtime types. + +41. **Remote aliasing batches must seed the driver's region memo.** Phase 1 + groups arguments by owner and computes aliasing remotely, but Distributed + RPCs do not inherit the caller's `CHUNK_AINFO_MEMO` scoped value. Keeping + those answers only in `arg_to_ainfo` makes Phase 4 ask the owners again. + Seed the memo when merging each remote batch, keyed by the original local + `Chunk`/`ChunkView` and dependency modifier, not the deserialized argument + wrapper. Keep the local/MPI path unchanged: it already fills that memo. + A typed merge barrier also prevents `remotecall_fetch`'s `Any` result from + erasing the result container's type in the per-argument loop. diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 2c4bcb015..b195a879e 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -831,15 +831,7 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr wid, worker_args = only(by_worker) results = wid == myid() ? batch_aliasing(current_acceleration(), worker_args) : remotecall_fetch(_compute_aliasing_batch, wid, worker_args) - # Key by the *local* `arg_w`, not the pair's: for a remote worker the - # returned `ArgumentWrapper` is a deserialized copy that need not be - # identity/hash-equal to the entry in `arg_ws_vec` we later look up - # (which would raise a `KeyError`). `_compute_aliasing_batch` preserves - # input order, so pair by index. - @assert length(results) == length(worker_args) "build_aliasing_parallel: _compute_aliasing_batch returned $(length(results)) results for $(length(worker_args)) args (wid=$wid)" - for i in eachindex(worker_args) - arg_to_ainfo[worker_args[i]] = results[i].second - end + _record_aliasing_batch!(arg_to_ainfo, worker_args, results, wid) else all_results_lock = ReentrantLock() @sync for (wid, worker_args) in by_worker @@ -859,14 +851,8 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr else remotecall_fetch(_compute_aliasing_batch, wid, worker_args) end - # Key by the *local* `arg_w` (see single-worker note above): a - # remote worker returns deserialized `ArgumentWrapper` copies - # that may not compare equal to our `arg_ws_vec` lookup keys. - @assert length(results) == length(worker_args) "build_aliasing_parallel: _compute_aliasing_batch returned $(length(results)) results for $(length(worker_args)) args (wid=$wid)" @lock all_results_lock begin - for i in eachindex(worker_args) - arg_to_ainfo[worker_args[i]] = results[i].second - end + _record_aliasing_batch!(arg_to_ainfo, worker_args, results, wid) end end end @@ -895,6 +881,28 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr return lookup, ainfos_overlaps, arg_to_ainfo end +function _record_aliasing_batch!(arg_to_ainfo::Dict{ArgumentWrapper,AliasingWrapper}, + worker_args::Vector{ArgumentWrapper}, + results::Vector{Pair{ArgumentWrapper,AliasingWrapper}}, + wid::Int) + @assert length(results) == length(worker_args) "build_aliasing_parallel: _compute_aliasing_batch returned $(length(results)) results for $(length(worker_args)) args (wid=$wid)" + # Pair with the local arguments by index: the returned wrappers have crossed + # serialization and need not be identity/hash-equal to the caller's keys. + # This barrier also restores the known batch type after remotecall_fetch. + for i in eachindex(worker_args) + arg_w = worker_args[i] + ainfo = results[i].second + arg_to_ainfo[arg_w] = ainfo + # Remote calls do not inherit the driver's ScopedValues memo. Phase 1 + # already got the owner's answer, so retain it rather than asking again + # in Phase 4. Local batches (including MPI's override) seed it themselves. + if wid != myid() && (arg_w.arg isa Chunk || arg_w.arg isa ChunkView) + memoize_ainfo!(ainfo_memo_key(arg_w.arg, arg_w.dep_mod), ainfo.inner) + end + end + return +end + """ batch_aliasing(accel, arg_ws) -> Vector{Pair{ArgumentWrapper,AliasingWrapper}} diff --git a/test/datadeps.jl b/test/datadeps.jl index 7401941f1..461221f5b 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -50,6 +50,37 @@ using LinearAlgebra, Graphs end end +@testset "Distributed aliasing batch memo" begin + chunks = Chunk[wid == myid() ? Dagger.tochunk(zeros(4,4)) : + remotecall_fetch(Dagger.tochunk, wid, zeros(4,4)) + for wid in procs()] + # Cover the direct local batch, a single remote owner, and parallel owner + # batches. Dependency modifiers and views must each keep their own memo key. + groups = [[first(chunks)], [last(chunks)], chunks] + for group in groups + arg_ws = Dict{ArgumentWrapper,ArgumentWrapper}() + for chunk in group + for arg_w in (ArgumentWrapper(chunk, identity), + ArgumentWrapper(chunk, UpperTriangular), + ArgumentWrapper(view(chunk, 1:2, :), identity)) + arg_ws[arg_w] = arg_w + end + end + memo = Dagger.ChunkAinfoMemo() + Dagger.with(Dagger.CHUNK_AINFO_MEMO => memo) do + _, _, arg_to_ainfo = Dagger.build_aliasing_parallel(arg_ws) + for arg_w in keys(arg_ws) + key = Dagger.ainfo_memo_key(arg_w.arg, arg_w.dep_mod) + @test haskey(memo.entries, key) + # A cache miss must fail, not silently recompute the same answer. + cached = Dagger.memoized_ainfo(() -> error("aliasing batch was not memoized"), key) + @test cached === arg_to_ainfo[arg_w].inner + @test aliasing(Dagger.current_acceleration(), arg_w.arg, arg_w.dep_mod) === cached + end + end + end +end + @testset "Memory Aliasing" begin A = rand(4) a = Dagger.aliasing(A) From b37ad6db39dfdb5597f5d70123b18650a866471d Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 12:57:21 -0700 Subject: [PATCH 20/26] Show benchmark change lists immediately after the count summary Keep each backend report readable in the combined GitHub comment: summarize the outcomes, then list the individual regressions and improvements, before the full timing/allocation tables and plots. Classification, thresholds, and exit behavior are unchanged. Validation: 14 rendering assertions pass across reports with and without regressions/improvements, including within-noise entries. --- benchmark/ci.jl | 35 ++++++++++++++++++----------------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/benchmark/ci.jl b/benchmark/ci.jl index 944acb8bf..20383eaa3 100644 --- a/benchmark/ci.jl +++ b/benchmark/ci.jl @@ -244,23 +244,6 @@ open(joinpath(OUTPUT_DIR, "report.md"), "w") do io println(io) job_summary_table(io, job_summary_rows) println(io) - println(io, "#### Median time") - println(io) - println(io, table) - println(io) - println(io, "#### Allocations / memory") - println(io) - println(io, alloc_table) - println(io) - if !isempty(plot_files) - println(io, "#### Plots") - println(io) - for f in plot_files - # `artifact://` references render inline in Buildkite annotations. - println(io, "![", basename(f), "](artifact://", relpath(f, dirname(OUTPUT_DIR)), ")") - end - println(io) - end if isempty(regressions) println(io, "No time regressions beyond ", pct(1 + THRESHOLD), " or allocation regressions beyond ", pct(1 + ALLOC_THRESHOLD), @@ -297,6 +280,24 @@ open(joinpath(OUTPUT_DIR, "report.md"), "w") do io println(io) println(io, "
") end + println(io) + println(io, "#### Median time") + println(io) + println(io, table) + println(io) + println(io, "#### Allocations / memory") + println(io) + println(io, alloc_table) + println(io) + if !isempty(plot_files) + println(io, "#### Plots") + println(io) + for f in plot_files + # `artifact://` references render inline in Buildkite annotations. + println(io, "![", basename(f), "](artifact://", relpath(f, dirname(OUTPUT_DIR)), ")") + end + println(io) + end end # --- Summary + exit status ------------------------------------------------- From 2be0122cdbe18c338abb927a511d28c5cbe6ce72 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 13:00:01 -0700 Subject: [PATCH 21/26] Compare Distributed linalg with consistently placed input fixtures Arbitrary input placement changes data movement and the driver allocation share between samples and revisions. Match the existing array/stencil fixture policy with cyclic rows under Distributed; keep MPI on its native allocator. Pin the SPD destination too: assigning G does not propagate placement through similar(G), so form the SPD fixture with an explicitly assigned destination and mul! in untimed setup. Measured operations and normal output allocation remain unchanged. No scheduler heuristics, BLAS settings, classification thresholds, allocation bounds, or sample budgets change. Matched-layout Julia 1.13 check, driver 1 thread plus 3 one-thread workers, BLAS 1: 10 warmups then GC and minimum of 5 Base.gc_num deltas. Master -> branch: SYRK128 302,929 allocs / 30,448,824 bytes -> 274,930 / 29,212,920; matmul512 16,321 / 13,295,376 -> 14,393 / 13,211,776; Cholesky512 19,695 / 15,596,896 -> 19,903 / 15,612,016. Minimum times: 127.59 -> 103.11 ms, 50.36 -> 46.84 ms, 24.01 -> 25.71 ms respectively. Local results are not a guarantee of CI timing. Validation: 9 Distributed fixture checks, including exact SPD layout and numerical products/factorization; 8 fixture checks on each of 4 MPI ranks, including execution of all three affected benchmark entries. --- AGENTS.md | 10 +++++++++ benchmark/suites/linalg.jl | 43 +++++++++++++++++++++++++++----------- 2 files changed, 41 insertions(+), 12 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e039e018e..8d474eb7b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -448,3 +448,13 @@ lesson. wrapper. Keep the local/MPI path unchanged: it already fills that memo. A typed merge barrier also prevents `remotecall_fetch`'s `Any` result from erasing the result container's type in the per-argument loop. + +42. **Fix the SPD fixture's destination placement, not just its input's.** + Giving `G` a cyclic assignment does not give `G * G'` that assignment: + `similar(G)` allocates an arbitrarily placed result. A Cholesky benchmark + must explicitly allocate its SPD destination with the same assignment and + use `mul!` in untimed setup, or its input layout still varies between + samples/revisions. Keep named grids Distributed-only and leave MPI on its + native allocator. This controls the fixture, not the measured operation: + out-of-place multiplication and Cholesky still allocate/place their own + outputs normally, and neither scheduling heuristics nor thresholds change. diff --git a/benchmark/suites/linalg.jl b/benchmark/suites/linalg.jl index e5a26efe0..4e460a513 100644 --- a/benchmark/suites/linalg.jl +++ b/benchmark/suites/linalg.jl @@ -12,7 +12,22 @@ # Build a distributed symmetric positive-definite matrix (for Cholesky). `G*G'` # is PD almost surely for a full-rank square `G`. -_spd(T, N, b) = (G = rand(Blocks(b, b), T, N, N); A = G * G'; wait(A); A) +function _spd(T, N, b; assignment=:arbitrary) + G = rand(Blocks(b, b), T, N, N; assignment) + wait(G) + A = if assignment === :arbitrary + G * G' + else + # G's assignment does not propagate through similar(G). Allocate the + # SPD destination explicitly too, so Cholesky's input layout is fixed. + result = DArray{T}(undef, Blocks(b, b), N, N; assignment) + wait(result) + mul!(result, G, G') + result + end + wait(A) + return A +end function linalg_suite(ctx; method, accels) @assert method == "dagger" "Linalg suite only supports `dagger` execution" @@ -20,12 +35,16 @@ function linalg_suite(ctx; method, accels) @assert accel == "cpu" "Linalg suite only supports CPU execution" T = Float64 + # Match the array/stencil suites: arbitrary fixture placement changes both + # data movement and the driver's allocation share across revisions/samples. + # Named cyclic grids are Distributed-only; leave MPI on its native allocator. + fixture_assignment = length(procs()) > 1 ? :cyclicrow : :arbitrary # Some older Dagger revisions use a Distributed-only processor grid for # tiled SVD. Under MPI that grid is empty and `_tile_index` divides by zero. # This script is shared by both Airspeed revisions, so probe once and omit # SVD where the revision/backend combination cannot execute it. svd_ok = supported("linalg/svd") do - A = rand(Blocks(4, 4), T, 8, 8) + A = rand(Blocks(4, 4), T, 8, 8; assignment=fixture_assignment) wait(A) wait(svd(A).U) end @@ -38,31 +57,31 @@ function linalg_suite(ctx; method, accels) # gemm needs A and the result resident; factorizations copy internally. if fits_budget(dense_bytes(N; nmats=3, T=T)) sub["matmul (A*A)"] = @benchmarkable(wait(A * A), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A)), teardown = (A = nothing; @everywhere GC.gc())) sub["syrk (A'*A)"] = @benchmarkable(wait(A' * A), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A)), teardown = (A = nothing; @everywhere GC.gc())) sub["lu"] = @benchmarkable(wait(lu(A, RowMaximum()).factors), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A)), teardown = (A = nothing; @everywhere GC.gc())) sub["qr"] = @benchmarkable(wait(qr(A).factors), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A)), teardown = (A = nothing; @everywhere GC.gc())) sub["solve (A\\b via lu)"] = @benchmarkable(wait(lu(A, RowMaximum()) \ b), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); - b = rand(Blocks($b), $T, $N); wait(A); wait(b)), + setup = (A = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); + b = rand(Blocks($b), $T, $N; assignment=$fixture_assignment); wait(A); wait(b)), teardown = (A = nothing; b = nothing; @everywhere GC.gc())) end # Cholesky additionally holds the SPD-construction temporary. if fits_budget(dense_bytes(N; nmats=4, T=T)) sub["cholesky"] = @benchmarkable(wait(cholesky(A).factors), - setup = (A = _spd($T, $N, $b)), + setup = (A = _spd($T, $N, $b; assignment=$fixture_assignment)), teardown = (A = nothing; @everywhere GC.gc())) end @@ -71,15 +90,15 @@ function linalg_suite(ctx; method, accels) # workers) a restaged copy of A, on top of the resident input. if svd_ok && fits_budget(dense_bytes(N; nmats=5, T=T)) sub["svd"] = @benchmarkable(wait(svd(A).U), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + setup = (A = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); wait(A)), teardown = (A = nothing; @everywhere GC.gc())) end # gemv is cheap (one matrix + two vectors). if fits_budget(dense_bytes(N; nmats=1, T=T)) sub["matvec (A*x)"] = @benchmarkable(wait(A * x), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); - x = rand(Blocks($b), $T, $N); wait(A); wait(x)), + setup = (A = rand(Blocks($b, $b), $T, $N, $N; assignment=$fixture_assignment); + x = rand(Blocks($b), $T, $N; assignment=$fixture_assignment); wait(A); wait(x)), teardown = (A = nothing; x = nothing; @everywhere GC.gc())) end From 8bddfe7bf77ac15d8ad8a13456d6ba7eedc36191 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 16:00:16 -0700 Subject: [PATCH 22/26] mpi: Wake only consumers of the delivered broadcast tag A shared condition notified every unrelated waiter for each metadata delivery, making a large submitted region pay a quadratic wakeup cost. Give each FIFO a lazily allocated condition sharing the registry lock; retain slots until every consumer leaves, and preserve heartbeat/shutdown notification, including already-woken consumers. Julia 1.13, one thread: ten warmups, GC.gc(), minimum of five runs delivering 1024 tags: 395.105 ms -> 2.849 ms. Allocation count/bytes: 15375 / 877448 -> 17442 / 1095872 (blocked-only conditions; queued delivery avoids them). Deterministic FIFO, targeted wakeup, heartbeat, multiple-consumer and shutdown tests: 64 pass with four threads. Full SpGEMM timings remain noisy; no batching, scheduling, sampling or thresholds changed. --- AGENTS.md | 11 +++++ ext/MPIExt.jl | 77 ++++++++++++++++++++++--------- test/mpi-broadcast.jl | 102 ++++++++++++++++++++++++++++++++++++++++++ test/mpi.jl | 2 + 4 files changed, 172 insertions(+), 20 deletions(-) create mode 100644 test/mpi-broadcast.jl diff --git a/AGENTS.md b/AGENTS.md index 8d474eb7b..4fb8e2de8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -458,3 +458,14 @@ lesson. native allocator. This controls the fixture, not the measured operation: out-of-place multiplication and Cholesky still allocate/place their own outputs normally, and neither scheduling heuristics nor thresholds change. + +43. **Broadcast delivery must wake only its own consumers.** One shared + condition for all `(root, tag)` slots makes a delivery wake every unrelated + waiter: draining 1024 tags one at a time took 395 ms, versus 3 ms with + per-slot conditions sharing the registry lock. Create conditions only when + a consumer actually blocks. Keep a slot alive until all its consumers have + left, including consumers already notified but still reacquiring the lock; + an empty FIFO alone does not mean its condition can be replaced. Heartbeats + and shutdown still notify every slot, and an empty consumer must also check + `running`: a consumer already woken by a heartbeat is outside the shutdown + notification queues and otherwise goes back to sleep after the relay stops. diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index f810c141b..2f7f76be6 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -1389,17 +1389,17 @@ end # tag from interleaving into a single queue, where no ordering rule applies. mutable struct BcastSlot values::Vector{Any} - BcastSlot() = new(Any[]) + cond::Union{Threads.Condition,Nothing} + waiters::Int + BcastSlot() = new(Any[], nothing, 0) end mutable struct BcastState bcast_comm::MPI.Comm slots::Dict{Tuple{Int,UInt32},BcastSlot} - # Guards `slots` and wakes consumers. A condition rather than a plain lock - # so a consumer waiting on an undelivered payload sleeps instead of - # spinning: a datadeps region can have hundreds of tasks blocked on a - # broadcast at once, and polling them all would burn exactly the CPU the - # ranks need to *produce* those broadcasts. + # Guards `slots`; each slot's condition shares this lock but has its own + # wait queue. A delivery must only wake that (root, tag)'s consumers, not + # every task blocked on an unrelated broadcast in the region. cond::Threads.Condition running::Threads.Atomic{Bool} relay::Union{Task,Nothing} @@ -1418,13 +1418,24 @@ bcast_state_for(comm::MPI.Comm) = bcast_serialize(x) = (io = IOBuffer(); Serialization.serialize(io, x); take!(io)) +function bcast_slot!(state::BcastState, key::Tuple{Int,UInt32}) + slot = get(state.slots, key, nothing) + if slot === nothing + slot = BcastSlot() + state.slots[key] = slot + end + return slot +end + function bcast_deliver!(state::BcastState, root::Int, tag::UInt32, value) # `@lock`, not `lock(...) do`: the `do` block is a closure over `state`, # `root`, `tag` and `value`, which boxes them and allocates on a path the # relay runs for every delivered payload. @lock state.cond begin - push!(get!(BcastSlot, state.slots, (root, tag)).values, value) - notify(state.cond) + slot = bcast_slot!(state, (root, tag)) + push!(slot.values, value) + cond = slot.cond + cond === nothing || notify(cond) end # A delivered payload is cross-rank progress even when it is not the one a # given consumer is waiting for (see `MPI_PROGRESS`). @@ -1437,7 +1448,12 @@ end # deadlock detection -- no `Timer` involved, and so nothing that depends on a # thread reaching the scheduler's idle loop to service libuv. function bcast_heartbeat!(state::BcastState) - @lock state.cond notify(state.cond) + @lock state.cond begin + for slot in values(state.slots) + cond = slot.cond + cond === nothing || notify(cond) + end + end return end @@ -1461,15 +1477,29 @@ function bcast_slot_wait(state::BcastState, root::Int, tag::UInt32) # a heap `Box`. Inlining the body keeps all of it on the stack; the `return` # below still releases the lock, via `@lock`'s `finally`. @lock state.cond begin - while true - slot = get(state.slots, key, nothing) - if slot !== nothing && !isempty(slot.values) - value = popfirst!(slot.values) - isempty(slot.values) && delete!(state.slots, key) - return value + slot = bcast_slot!(state, key) + slot.waiters += 1 + try + while true + if !isempty(slot.values) + return popfirst!(slot.values) + end + # A heartbeat/delivery may have already woken this consumer + # when teardown runs, leaving it outside the notified queues. + state.running[] || throw(ConcurrencyViolationError("MPI broadcast relay stopped while a consumer was waiting")) + timer = mpi_deadlock_detect(timer, rank, tag, "bcast_meta delivery", root) + # Already-delivered payloads need no condition allocation. + if slot.cond === nothing + slot.cond = Threads.Condition(state.cond.lock) + end + wait(slot.cond::Threads.Condition) end - timer = mpi_deadlock_detect(timer, rank, tag, "bcast_meta delivery", root) - wait(state.cond) + finally + slot.waiters -= 1 + # A woken consumer may not have reacquired the lock yet. Retain + # its slot even if another consumer drains the last payload, or + # the next delivery would notify a different condition. + isempty(slot.values) && slot.waiters == 0 && delete!(state.slots, key) end end end @@ -1602,9 +1632,16 @@ function stop_bcast_relay!(accel_comm::MPI.Comm) # deadlock timer, so once it is gone nothing would ever wake a consumer # still waiting on a payload that is now never coming. Fail them instead of # letting teardown hang. - @lock state.cond notify(state.cond, - ConcurrencyViolationError("MPI broadcast relay stopped while a consumer was waiting"); - error=true) + @lock state.cond begin + for slot in values(state.slots) + cond = slot.cond + if cond !== nothing + notify(cond, + ConcurrencyViolationError("MPI broadcast relay stopped while a consumer was waiting"); + error=true) + end + end + end return nothing end diff --git a/test/mpi-broadcast.jl b/test/mpi-broadcast.jl new file mode 100644 index 000000000..3368fac56 --- /dev/null +++ b/test/mpi-broadcast.jl @@ -0,0 +1,102 @@ +# Rank-local tests of the relay's FIFO and wait queues. No relay is started: +# delivery, heartbeat, and teardown are driven explicitly for determinism. +function test_bcast_state(comm) + MPIExt.BcastState(comm, Dict{Tuple{Int,UInt32},MPIExt.BcastSlot}(), + Threads.Condition(), Threads.Atomic{Bool}(true), nothing) +end + +function wait_bcast_consumers(state, n) + status = timedwait(10.0; pollint=0.001) do + @lock state.cond begin + sum(slot -> slot.waiters, values(state.slots); init=0) == n && + all(slot -> !isempty(slot.cond), values(state.slots)) + end + end + status === :ok || error("Broadcast consumers did not reach their wait queues") +end + +function bcast_test_consumer(state, tag, threaded) + if threaded + return Threads.@spawn MPIExt.bcast_slot_wait(state, 0, tag) + else + return @async MPIExt.bcast_slot_wait(state, 0, tag) + end +end + +@testset "Broadcast wait queues" begin + @testset "Queued delivery preserves FIFO and root identity" begin + state = test_bcast_state(MPI.COMM_SELF) + tag = UInt32(1) + MPIExt.bcast_deliver!(state, 0, tag, :first) + MPIExt.bcast_deliver!(state, 0, tag, :second) + MPIExt.bcast_deliver!(state, 1, tag, :other_root) + @test MPIExt.bcast_slot_wait(state, 0, tag) === :first + @test MPIExt.bcast_slot_wait(state, 1, tag) === :other_root + @test MPIExt.bcast_slot_wait(state, 0, tag) === :second + @test isempty(state.slots) + end + + @testset "Delivery wakes only the matching tag (threaded=$threaded)" for threaded in (false, true) + state = test_bcast_state(MPI.COMM_SELF) + n = 32 + tasks = [bcast_test_consumer(state, UInt32(i), threaded) for i in 1:n] + wait_bcast_consumers(state, n) + @lock state.cond begin + MPIExt.bcast_deliver!(state, 0, UInt32(1), 1) + @test isempty(state.slots[(0, UInt32(1))].cond) + @test all(i -> !isempty(state.slots[(0, UInt32(i))].cond), 2:n) + end + @test fetch(first(tasks)) == 1 + + # Heartbeats still wake every blocked consumer for deadlock checks, + # without deleting its empty slot or dropping a subsequent delivery. + @lock state.cond begin + MPIExt.bcast_heartbeat!(state) + @test all(slot -> isempty(slot.cond), values(state.slots)) + @test length(state.slots) == n - 1 + end + wait_bcast_consumers(state, n - 1) + for i in 2:n + MPIExt.bcast_deliver!(state, 0, UInt32(i), i) + end + @test fetch.(tasks) == collect(1:n) + @test isempty(state.slots) + end + + @testset "Keep a slot while another consumer is still waiting (threaded=$threaded)" for threaded in (false, true) + state = test_bcast_state(MPI.COMM_SELF) + tag = UInt32(2) + tasks = [bcast_test_consumer(state, tag, threaded) for _ in 1:2] + wait_bcast_consumers(state, 2) + MPIExt.bcast_deliver!(state, 0, tag, 1) + @test timedwait(() -> count(istaskdone, tasks) == 1, 10.0; pollint=0.001) === :ok + @lock state.cond begin + @test state.slots[(0, tag)].waiters == 1 + end + MPIExt.bcast_deliver!(state, 0, tag, 2) + @test sort(fetch.(tasks)) == [1, 2] + @test isempty(state.slots) + end + + @testset "Teardown fails every blocked tag (already woken=$woken, threaded=$threaded)" for woken in (false, true), threaded in (false, true) + state = test_bcast_state(MPI.COMM_SELF) + tasks = [bcast_test_consumer(state, UInt32(i), threaded) for i in 1:4] + wait_bcast_consumers(state, 4) + lock(MPIExt.BCAST_STATES) do states + @assert !haskey(states, MPI.COMM_SELF) + states[MPI.COMM_SELF] = state + end + @lock state.cond begin + if woken + MPIExt.bcast_heartbeat!(state) + @test all(slot -> isempty(slot.cond), values(state.slots)) + end + MPIExt.stop_bcast_relay!(MPI.COMM_SELF) + end + for task in tasks + @test_throws TaskFailedException fetch(task) + @test task.exception isa ConcurrencyViolationError + end + @test isempty(state.slots) + end +end diff --git a/test/mpi.jl b/test/mpi.jl index c66bc5fa2..39a3505ff 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -51,6 +51,8 @@ const rank = MPI.Comm_rank(comm) const nranks = MPI.Comm_size(comm) const accel = Dagger.current_acceleration() +include("mpi-broadcast.jl") + mpi_procs() = sort(collect(Dagger.get_processors(MPIExt.MPIClusterProc(comm))); by=p->(p.rank, Dagger.short_name(p))) proc_for_rank(r) = first(filter(p->p.rank == r, mpi_procs())) From e92d6c1b9ff8ab8db48b7913ef81bb42086f3739 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 16:18:51 -0700 Subject: [PATCH 23/26] mpi: Allocate receive wait events only for contended streams Uncontended receives now register a lease with nothing; competitors attach a one-shot event under the registry lock. Explicit locked blocks remove captured event boxing. Both receive paths release in finally, including failures. Julia 1.13, one thread, ten warmups, GC.gc(), min of five 10000-guard runs: 50000 allocations / 1760000 bytes / 0.998 ms -> zero allocations / zero bytes / 0.544 ms. No Core.Box in recv_yield lowering. Four-thread lifecycle tests pass on Julia 1.12 and 1.13. Full four-rank one-thread MPI suite passes (446/473/458/470 tests). Warm SpGEMM loses 2-3 million allocations; concurrent timings are not speedup claims. No scheduling, sampling or thresholds changed. --- AGENTS.md | 11 +++++ ext/MPIExt.jl | 99 +++++++++++++++++++++------------------------ test/mpi-receive.jl | 63 +++++++++++++++++++++++++++++ test/mpi.jl | 1 + 4 files changed, 120 insertions(+), 54 deletions(-) create mode 100644 test/mpi-receive.jl diff --git a/AGENTS.md b/AGENTS.md index 4fb8e2de8..53ab1cfe8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -469,3 +469,14 @@ lesson. and shutdown still notify every slot, and an empty consumer must also check `running`: a consumer already woken by a heartbeat is outside the shutdown notification queues and otherwise goes back to sleep after the relay stops. + +44. **An uncontended MPI receive needs a lease, not an event.** Uniformity + checks made the old receive guard create hundreds of thousands of + `Base.Event`s per sparse product, although its `(comm, source, tag)` stream + had only one receiver. Register `nothing` for an owned stream and create + a one-shot event only when a competitor arrives; delete the stream and + notify that event when the owner leaves. Never reset/recycle the event: + notified competitors can still be inside `wait`. Use explicit locked + blocks, since the old captured `our_event` lowered to a `Core.Box`, and + release the lease in `finally` on both serialized and in-place paths so a + failed receive cannot strand every subsequent receiver of that stream. diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 2f7f76be6..f76540c2e 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -927,7 +927,36 @@ function DeadlockTimer() round(UInt64, DEADLOCK_TIMEOUT_PERIOD[] * 1e9), MPI_PROGRESS[]) end -const RECV_WAITING = LockedObject(Dict{Tuple{MPI.Comm, Int, Int}, Base.Event}()) +# A present key with `nothing` denotes an active receiver with no competitors. +# Allocate an event only when a second receiver actually needs to wait. +const RECV_WAITING = LockedObject(Dict{Tuple{MPI.Comm, Int, Int}, Union{Base.Event,Nothing}}()) + +function acquire_recv!(key::Tuple{MPI.Comm,Int,Int}) + while true + local other_event + @lock RECV_WAITING begin + waiting = Dagger.payload(RECV_WAITING) + if !haskey(waiting, key) + waiting[key] = nothing + return + end + other_event = waiting[key] + if other_event === nothing + other_event = Base.Event() + waiting[key] = other_event + end + end + wait(other_event::Base.Event) + end +end + +function release_recv!(key::Tuple{MPI.Comm,Int,Int}) + @lock RECV_WAITING begin + other_event = pop!(Dagger.payload(RECV_WAITING), key) + other_event === nothing || notify(other_event) + end + return +end # Envelope for the out-of-place raw-bytes MPI path: serialize a small # descriptor, then send contiguous bitstype buffers. Types register via @@ -1025,68 +1054,30 @@ function recv_yield!(buffer, comm, src, tag) end #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Starting recv! from [$src]") - # Ensure no other receiver is waiting - our_event = Base.Event() - @label retry - other_event = lock(RECV_WAITING) do waiting - if haskey(waiting, (comm, src, tag)) - waiting[(comm, src, tag)] - else - waiting[(comm, src, tag)] = our_event - nothing - end - end - if other_event !== nothing - #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Waiting for other receiver...") - wait(other_event) - @goto retry - end - - buffer = recv_yield_inplace!(buffer, comm, rank, src, tag) - - lock(RECV_WAITING) do waiting - delete!(waiting, (comm, src, tag)) - notify(our_event) + key = (comm, Int(src), Int(tag)) + acquire_recv!(key) + try + return recv_yield_inplace!(buffer, comm, rank, src, tag), true + finally + release_recv!(key) end - - return buffer, true - end function recv_yield(comm, src, tag) rank = MPI.Comm_rank(comm) #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Starting recv from [$src]") - # Ensure no other receiver is waiting - our_event = Base.Event() - @label retry - other_event = lock(RECV_WAITING) do waiting - if haskey(waiting, (comm, src, tag)) - waiting[(comm, src, tag)] - else - waiting[(comm, src, tag)] = our_event - nothing + key = (comm, Int(src), Int(tag)) + acquire_recv!(key) + try + value = recv_yield_serialized(comm, rank, src, tag) + if value isa InplaceInfo + return recv_yield_inplace(value, comm, rank, src, tag) end + return value + finally + release_recv!(key) end - if other_event !== nothing - #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Waiting for other receiver...") - wait(other_event) - @goto retry - end - #Core.println("[rank $(MPI.Comm_rank(comm))][tag $tag] Receiving...") - - type = nothing - @label receive - value = recv_yield_serialized(comm, rank, src, tag) - if value isa InplaceInfo - value = recv_yield_inplace(value, comm, rank, src, tag) - end - - lock(RECV_WAITING) do waiting - delete!(waiting, (comm, src, tag)) - notify(our_event) - end - return value end # Device-resident dense array: receive directly when the MPI library is diff --git a/test/mpi-receive.jl b/test/mpi-receive.jl new file mode 100644 index 000000000..773746d26 --- /dev/null +++ b/test/mpi-receive.jl @@ -0,0 +1,63 @@ +# Fail source conversion before MPI posts a receive, after the stream guard +# has been acquired. This exercises error cleanup without corrupting the wire. +struct InvalidMPIReceiveSource <: Integer end +struct MPIReceiveConversionError <: Exception end +Base.Int(::InvalidMPIReceiveSource) = 0 +Base.cconvert(::Type{Cint}, ::InvalidMPIReceiveSource) = throw(MPIReceiveConversionError()) + +@testset "MPI receive stream ownership" begin + key = (MPI.COMM_SELF, 0, 1) + + @testset "Uncontended receives need no event" begin + @test (@inferred MPIExt.acquire_recv!(key)) === nothing + @lock MPIExt.RECV_WAITING begin + @test haskey(Dagger.payload(MPIExt.RECV_WAITING), key) + @test Dagger.payload(MPIExt.RECV_WAITING)[key] === nothing + end + @test (@inferred MPIExt.release_recv!(key)) === nothing + @lock MPIExt.RECV_WAITING begin + @test !haskey(Dagger.payload(MPIExt.RECV_WAITING), key) + end + end + + @testset "Contenders serialize without losing wakeups" begin + MPIExt.acquire_recv!(key) + active = Threads.Atomic{Int}(0) + tasks = [Threads.@spawn begin + MPIExt.acquire_recv!(key) + sole_owner = Threads.atomic_add!(active, 1) == 0 + try + sleep(0.001) + sole_owner + finally + Threads.atomic_sub!(active, 1) + MPIExt.release_recv!(key) + end + end for _ in 1:8] + status = timedwait(10.0; pollint=0.001) do + @lock MPIExt.RECV_WAITING begin + Dagger.payload(MPIExt.RECV_WAITING)[key] isa Base.Event + end + end + @test status === :ok + MPIExt.release_recv!(key) + @test timedwait(() -> all(istaskdone, tasks), 10.0; pollint=0.001) === :ok + @test all(fetch, tasks) + @test active[] == 0 + @lock MPIExt.RECV_WAITING begin + @test !haskey(Dagger.payload(MPIExt.RECV_WAITING), key) + end + end + + @testset "A failed receive releases its stream" begin + source = InvalidMPIReceiveSource() + @test_throws MPIReceiveConversionError MPIExt.recv_yield(MPI.COMM_SELF, source, UInt32(1)) + @lock MPIExt.RECV_WAITING begin + @test !haskey(Dagger.payload(MPIExt.RECV_WAITING), key) + end + @test_throws MPIReceiveConversionError MPIExt.recv_yield!(zeros(1), MPI.COMM_SELF, source, UInt32(1)) + @lock MPIExt.RECV_WAITING begin + @test !haskey(Dagger.payload(MPIExt.RECV_WAITING), key) + end + end +end diff --git a/test/mpi.jl b/test/mpi.jl index 39a3505ff..6e0767956 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -52,6 +52,7 @@ const nranks = MPI.Comm_size(comm) const accel = Dagger.current_acceleration() include("mpi-broadcast.jl") +include("mpi-receive.jl") mpi_procs() = sort(collect(Dagger.get_processors(MPIExt.MPIClusterProc(comm))); by=p->(p.rank, Dagger.short_name(p))) From 974f201c5bd9ec41452f194a7effaa44a29884f7 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 16 Sep 2026 16:33:32 -0700 Subject: [PATCH 24/26] reuse: Create pooled task slots only when they are needed Temporary completion and placement callers often dispatch once, yet initialized all 32 task-local slots, channels and monitors. Initialize each slot on demand without changing capacity or overflow behavior. Preserve cleared dynamic scope, setup-before-schedule and registration-before-publication; skip unassigned slots during finalization. Julia 1.13, ten warmups, GC.gc(), minimum of five runs creating 100 single-use 32-slot caches: 75100 allocations / 4004800 bytes / 5.744 ms -> 6800 allocations / 296000 bytes / 0.318 ms. Slot initialization and cached dispatch infer Task returns. Reuse suites pass on Julia 1.12 and 1.13 with four threads; all 19 unchanged allocation-bound/scope checks pass. Final combined fixes pass lifecycle and sparse/dense product, Cholesky and LU reconstruction checks on each of four MPI ranks with two threads. No pool sizing, scheduling, sampling or thresholds changed. --- AGENTS.md | 10 +++++++ src/utils/reuse.jl | 41 ++++++++++++++------------- test/reuse-task-cache.jl | 60 ++++++++++++++++++++++++++++++++++++++++ test/reuse.jl | 2 ++ 4 files changed, 94 insertions(+), 19 deletions(-) create mode 100644 test/reuse-task-cache.jl diff --git a/AGENTS.md b/AGENTS.md index 53ab1cfe8..207e7a46d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -480,3 +480,13 @@ lesson. blocks, since the old captured `our_event` lowered to a `Core.Box`, and release the lease in `finally` on both serialized and in-place paths so a failed receive cannot strand every subsequent receiver of that stream. + +45. **Task-pool capacity is not a reason to create every slot upfront.** + Completion and placement tasks can be short-lived, but each owns a + task-local fire cache with capacity 32. Creating all 32 tasks, channels, + and monitors on its first dispatch pays for 31 idle slots when that caller + dispatches only once. Initialize slots on demand without changing capacity + or the overflow policy; 100 single-use caches fell from 75,100 allocations + / 4,004,800 bytes to 6,800 / 296,000. Clear the dynamic scope at actual slot + creation, run its setup before scheduling, and register each dispatch + before publishing the payload. Finalization must skip unassigned slots. diff --git a/src/utils/reuse.jl b/src/utils/reuse.jl index 34a4186e9..34f80c968 100644 --- a/src/utils/reuse.jl +++ b/src/utils/reuse.jl @@ -593,21 +593,12 @@ mutable struct ReusableTaskCache tasks = Vector{Task}(undef, N) chans = Vector{Channel{Any}}(undef, N) ready = [Threads.Atomic{Bool}(true) for _ in 1:N] - for idx in 1:N - chans[idx] = Channel{Any}(1) - chan, r = chans[idx], ready[idx] - # N.B. These tasks are created on whichever call first touches this - # (task-local) cache and then serve every later payload, so they - # must not inherit that call's dynamic scope (see - # `clear_task_scope!`). - tasks[idx] = clear_task_scope!(@task reusable_task_loop(chan, r)) - end cache = new(tasks, chans, ready, t->nothing, N, false) finalizer(cache) do cache # Ask tasks to shut down for idx in 1:N Threads.atomic_xchg!(cache.ready[idx], false) - close(cache.chans[idx]) + isassigned(cache.chans, idx) && close(cache.chans[idx]) end end return cache @@ -616,15 +607,27 @@ end function reusable_task_cache_init!(setup_f::Function, cache::ReusableTaskCache) cache.init && return cache.setup_f = setup_f - for idx in 1:cache.N - task = cache.tasks[idx] - setup_f(task) - schedule(task) - Sch.errormonitor_tracked("reusable_task_$idx", task) - end cache.init = true return end + +# A temporary completion/placement task often dispatches only one payload +# through its task-local cache. Keep the full capacity without creating N +# channels, tasks and monitors which that caller will never use. +function initialize_reusable_task_slot!(cache::ReusableTaskCache, idx::Int) + chan = Channel{Any}(1) + ready = cache.ready[idx] + # Slots are created by whichever call first needs them, and then serve + # unrelated callers. Clear dynamic scope before scheduling, just as for + # the overflow path below. + task = clear_task_scope!(@task reusable_task_loop(chan, ready)) + cache.setup_f(task) + cache.chans[idx] = chan + cache.tasks[idx] = task + schedule(task) + Sch.errormonitor_tracked("reusable_task_$idx", task) + return task +end function reusable_task_loop(chan::Channel{Any}, ready::Threads.Atomic{Bool}) r = rand(1:128) while true @@ -655,12 +658,12 @@ function (cache::ReusableTaskCache)(f, name::String, register=nothing) idx = findfirst(getindex, cache.ready) if idx !== nothing @assert Threads.atomic_xchg!(cache.ready[idx], false) - t = cache.tasks[idx] + t = isassigned(cache.tasks, idx) ? cache.tasks[idx] : initialize_reusable_task_slot!(cache, idx) register === nothing || register(t) put!(cache.chans[idx], f) # N.B. No errormonitor_tracked_set! here: pooled tasks are registered - # once at init, and renaming the tracked entry per dispatch was an - # O(n) locked scan on the hot path for a debugging-only list + # once at slot creation; renaming the tracked entry per dispatch was + # an O(n) locked scan on the hot path for a debugging-only list. return t else t = @task try diff --git a/test/reuse-task-cache.jl b/test/reuse-task-cache.jl new file mode 100644 index 000000000..2ffe70e8e --- /dev/null +++ b/test/reuse-task-cache.jl @@ -0,0 +1,60 @@ +using Test + +@testset "Lazy reusable task slots" begin + setup_calls = Ref(0) + setup(task) = (setup_calls[] += 1; Dagger.set_task_migratable!(task)) + cache = Dagger.ReusableTaskCache(32) + Dagger.reusable_task_cache_init!(setup, cache) + @test setup_calls[] == 0 + @test all(i -> !isassigned(cache.tasks, i), 1:cache.N) + @test all(i -> !isassigned(cache.chans, i), 1:cache.N) + + # First use allocates just one slot, and registration still precedes the + # payload even if the new task has already started on another thread. + registered = Ref(false) + result = Channel{Tuple{Any,Bool}}(1) + task = Dagger.with_options(; scope=Dagger.scope(worker=1, thread=1)) do + cache(() -> put!(result, (Dagger.get_options(:scope, nothing), registered[])), + "lazy task slot", task -> (registered[] = true)) + end + @test take!(result) == (nothing, true) + @test task === cache.tasks[1] + @test setup_calls[] == 1 + @test count(i -> isassigned(cache.tasks, i), 1:cache.N) == 1 + @test count(i -> isassigned(cache.chans, i), 1:cache.N) == 1 + @test timedwait(() -> cache.ready[1][], 10.0; pollint=0.001) === :ok + done = Base.Event() + @test (@inferred cache(() -> notify(done), "reuse task slot")) === task + wait(done) + @test setup_calls[] == 1 + @test timedwait(() -> cache.ready[1][], 10.0; pollint=0.001) === :ok + finalize(cache) + @test timedwait(() -> istaskdone(task), 10.0; pollint=0.001) === :ok + @test !istaskfailed(task) + + # Saturation still uses the original overflow path; capacity is unchanged. + cache = Dagger.ReusableTaskCache(2) + setup_calls[] = 0 + Dagger.reusable_task_cache_init!(setup, cache) + entered = Channel{Task}(3) + release = Base.Event() + f() = (put!(entered, current_task()); wait(release)) + tasks = [cache(f, "saturated task slots") for _ in 1:3] + @test Set(take!(entered) for _ in 1:3) == Set(tasks) + @test tasks[1] === cache.tasks[1] + @test tasks[2] === cache.tasks[2] + @test tasks[3] ∉ cache.tasks + @test setup_calls[] == 3 + notify(release) + wait(tasks[3]) + @test timedwait(() -> all(getindex, cache.ready), 10.0; pollint=0.001) === :ok + finalize(cache) + @test timedwait(() -> all(istaskdone, tasks), 10.0; pollint=0.001) === :ok + @test !any(istaskfailed, tasks) + + # No payload ever used this cache: finalization must skip its empty cells. + empty_cache = Dagger.ReusableTaskCache(32) + Dagger.reusable_task_cache_init!(setup, empty_cache) + finalize(empty_cache) + @test all(i -> !isassigned(empty_cache.chans, i), 1:empty_cache.N) +end diff --git a/test/reuse.jl b/test/reuse.jl index 82b52d5a1..d13a6c93f 100644 --- a/test/reuse.jl +++ b/test/reuse.jl @@ -3,6 +3,8 @@ import Dagger: ReusableLinkedList, ReusableDict, ReusableCache import Dagger: take_or_alloc!, maybe_take_or_alloc!, maybetake!, putback! import Dagger: @reusable_vector, @reusable_dict +include("reuse-task-cache.jl") + struct ReusableScratchTestEntry value::Int end From 34a01f75c9132beba21e725fad69867be8b74fc4 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 17 Sep 2026 12:45:17 -0700 Subject: [PATCH 25/26] Avoid boxed signature iteration and unused default specialization The defaults LFU can immediately evict a new frequency-1 entry when its resident entries are hotter. Those signatures keep taking the dynamically typed fallback, even after deep warmup. Splat the existing full type vector for non-keyword signatures rather than its identical view, and avoid specializing the generic fallback on arguments it never inspects. Keep custom type-based defaults, keyword filtering, explicit options, signature representation, cache capacity and eviction policy unchanged. In particular, do not add a keyword-signature copy on the cached path. Add custom/default dispatch coverage and a repeated-miss allocation guard using a fresh task-local cache. Leave existing allocation bounds unchanged. Julia 1.13, 10 warmups + GC, minimum of 5 Base.gc_num measurements: - Cold four-type default population: 217 allocs / 9,088 bytes -> 49 / 1,792. - Cached population: unchanged at 1 alloc / 256 bytes. - Distributed LU 1024/128: 405,166 / 36,317,096 -> 396,400 / 35,955,080. - Distributed LU solve 1024/128: 511,038 / 40,974,680 -> 484,016 / 39,796,616. The Distributed probes use 3 one-thread workers and matched cyclic fixtures. The actual CI worker setup additionally shows lower allocations for all four rand cases and LU/solve, with LU/solve memory lower by about 13-14%. Julia 1.12 also measures 49 / 1,792 for repeated cold default population. Validated unchanged allocation-suite bounds, option/default and reuse tests, and Distributed LU/solve correctness on Julia 1.12 and 1.13. --- AGENTS.md | 13 +++++++++++++ src/options.jl | 18 ++++++++++++++++-- test/allocations.jl | 25 +++++++++++++++++++++++++ test/options.jl | 39 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 93 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 207e7a46d..79d960b58 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -490,3 +490,16 @@ lesson. / 4,004,800 bytes to 6,800 / 296,000. Clear the dynamic scope at actual slot creation, run its setup before scheduling, and register each dispatch before publishing the payload. Finalization must skip unassigned slots. + +46. **Repeated calls do not necessarily warm the option-default cache.** + `BasicLFUCache` can immediately evict a newly inserted frequency-1 key + when every resident key is more frequent, so a new signature keeps taking + the default fallback after arbitrarily many warmups. Test that path with + a full cache in a fresh task, not by changing capacity or eviction policy. + Splat `Signature.sig` directly for non-keyword calls: splatting its + identical `sig_nokw` view boxes the view and an iteration pair per type. + The generic fallback also needs no specialization of its unused arguments; + leave type-dependent user overrides untouched. A four-type cold default + population fell from 217 allocations / 9,088 bytes to 49 / 1,792, while + cache hits stayed at 1 / 256. Keep keyword views instead of adding a copy + to every signature, especially signatures whose defaults are already cached. diff --git a/src/options.jl b/src/options.jl index f53401812..bf01364c3 100644 --- a/src/options.jl +++ b/src/options.jl @@ -208,7 +208,15 @@ end function maybe_default!(opts::Options, ::Val{opt}, sig::Signature) where opt if getfield(opts, opt) === nothing default_opt = get!(SIGNATURE_DEFAULT_CACHE[], (sig.hash_nokw, opt)) do - Dagger.default_option(Val{opt}(), sig.sig_nokw...) + # Julia can splat a Vector{Any} directly, whereas splatting its + # view boxes the view and allocates an iteration pair per type. + # Most calls have no kwargs: use the original vector in that case. + # Keep the kwarg view, without adding a copy on cache hits. + if length(sig.sig_nokw) == length(sig.sig) + Dagger.default_option(Val{opt}(), sig.sig...) + else + Dagger.default_option(Val{opt}(), sig.sig_nokw...) + end end setfield!(opts, opt, default_opt) end @@ -307,7 +315,13 @@ This function may be executed within the scheduler, so it should generally be made very cheap to execute. If the function throws an error, the scheduler will use whatever the global default value is for that option instead. """ -default_option(::Val{name}, Tf, Targs...) where name = nothing +function default_option(::Val{name}, Tf, Targs...) where name + # This fallback never inspects the types; specializing on them constructs + # otherwise-unused dispatch types on every dynamically-typed cache miss. + # User overrides still dispatch on the original, concrete argument types. + @nospecialize Tf Targs + return nothing +end default_option(::Val) = throw(ArgumentError("default_option requires a function type and any argument types")) """ diff --git a/test/allocations.jl b/test/allocations.jl index 2ec03ad39..4b8cfbf39 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -56,6 +56,31 @@ function measure_steady_state_allocs(f; warmup=10, runs=5) return (allocs=best_allocs, bytes=best_bytes) end +# A full LFU cache can immediately evict a newly inserted signature when all +# existing entries are more frequent. Warmup alone therefore does not cover +# the default-option fallback; exercise repeated misses in a fresh task so +# this test cannot disturb the scheduler's task-local defaults cache. +@testset "Uncached option defaults" begin + measured = fetch(Threads.@spawn begin + cache = Dagger.SIGNATURE_DEFAULT_CACHE[] + empty!(cache) + for i in 1:cache.max_size + key = (UInt(i), :meta) + cache.cache[key] = nothing + cache.freq[key] = 2 + end + sig = Dagger.Signature(Any[typeof(Dagger.allocate_array), typeof(rand), + Type{Float64}, Tuple{Int,Int}]) + result = measure_steady_state_allocs() do + Dagger.populate_defaults!(Dagger.Options(), sig) + end + @test !haskey(cache.cache, (sig.hash_nokw, :meta)) + result + end) + @test measured.allocs <= 100 * ALLOC_BOUND_MULTIPLIER + @test measured.bytes <= 4_000 * ALLOC_BOUND_MULTIPLIER +end + # name => (; allocs, bytes) upper bounds (see header for how these are set). # Measured steady-state values at the time of writing are noted inline. const ALLOC_BOUNDS = Dict( diff --git a/test/options.jl b/test/options.jl index b7fd7c818..0b2e737ef 100644 --- a/test/options.jl +++ b/test/options.jl @@ -22,6 +22,45 @@ (to.owner == 1 && to.tid == 1) ? SpecialFunc(1) : sf end +option_default_target(x; ignored=nothing) = x +Dagger.@option :meta option_default_target(Integer) = true +Dagger.@option :get_result option_default_target(AbstractFloat) = true +Dagger.default_option(::Val{:name}, ::Type{typeof(option_default_target)}, + ::Type{T}) where {T<:AbstractFloat} = string(T) + +@testset "Signature option defaults" begin + positional = Dagger.Signature(Any[typeof(option_default_target), Int]) + keyword = Dagger.Signature(Any[typeof(Core.kwcall), @NamedTuple{ignored::String}, + typeof(option_default_target), Int]) + other_keyword = Dagger.Signature(Any[typeof(Core.kwcall), @NamedTuple{other::Float64}, + typeof(option_default_target), Int]) + @test positional.hash_nokw == keyword.hash_nokw == other_keyword.hash_nokw + @test positional.hash != keyword.hash + + # Exercise both the vector and kwarg-view cache-miss paths, independently + # of the cache entries left by other tests in this task. + cache = Dagger.SIGNATURE_DEFAULT_CACHE[] + for sig in (positional, keyword, other_keyword) + for option in (:meta, :get_result) + key = (sig.hash_nokw, option) + delete!(cache.cache, key) + delete!(cache.freq, key) + end + opts = Dagger.populate_defaults!(Dagger.Options(), sig) + @test opts.meta === true + @test opts.get_result === nothing + @test Dagger.populate_defaults!(Dagger.Options(; meta=false), sig).meta === false + end + + floating = Dagger.Signature(Any[typeof(option_default_target), Float64]) + opts = Dagger.populate_defaults!(Dagger.Options(), floating) + @test opts.meta === nothing + @test opts.get_result === true + @test opts.name == "Float64" + @test Dagger.default_option(Val(:meta), typeof(option_default_target), String) === nothing + @test_throws ArgumentError Dagger.default_option(Val(:meta)) +end + @testset "Scope propagation" begin first_wid = first(workers()) last_wid = last(workers()) From 6421ac55d3ec75f1fbe6172ac4bc7d4cf38b06c5 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 17 Sep 2026 12:49:07 -0700 Subject: [PATCH 26/26] Scan LFU frequencies without allocating an eviction result The typed defaults cache still boxed findmin's (frequency, key) result on every eviction. This is particularly costly for signatures immediately evicted after insertion: warmup never eliminates those repeated misses. Scan frequencies directly, initializing from the first entry and updating only for a strictly smaller frequency. This preserves the original first- minimum tie breaking, capacity, admission and eviction decisions exactly. Do not resize caches or change their retention policy. Add an allocation guard and compare exact values, contents and frequencies against the original findmin algorithm after mixed hits/misses at capacities 0, 1, 8 and 256. Check empty! too. Leave existing allocation bounds unchanged. Julia 1.13, 10 warmups + GC, minimum of 5 Base.gc_num measurements: - 10,000 repeated evictions: 10,000 allocs / 320,000 bytes -> 0 / 0. - Cold four-type default population: 49 / 1,792 -> 25 / 1,024. - Cached default population: unchanged at 1 / 256. Julia 1.12 also measures zero allocations for the eviction probe and 25 / 1,024 for cold default population. A separate 40,000-operation randomized probe verified exact equivalence to the original eviction code on both Julia versions. Validated the complete allocation, option/default and reuse suites with Distributed workers on Julia 1.13; all allocation bounds remain green. --- AGENTS.md | 12 ++++++++++++ src/utils/lfucache.jl | 12 ++++++++++-- test/allocations.jl | 22 ++++++++++++++++++++++ test/reuse.jl | 34 ++++++++++++++++++++++++++++++++++ 4 files changed, 78 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 79d960b58..e34336c96 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -503,3 +503,15 @@ lesson. population fell from 217 allocations / 9,088 bytes to 49 / 1,792, while cache hits stayed at 1 / 256. Keep keyword views instead of adding a copy to every signature, especially signatures whose defaults are already cached. + +47. **Even a typed `findmin(::Dict)` can box its result in an eviction loop.** + The defaults LFU allocated one `(frequency, key)` tuple per eviction + despite its concrete key/frequency types. A direct minimum-frequency scan + removes that allocation without changing capacity, admission or eviction + decisions. Initialize from the first dictionary entry and update only on + strict `<` so equal-frequency ties preserve the original iteration order. + Compare exact cache contents and frequencies against the old algorithm + after mixed hits/misses, including zero capacity. Ten thousand repeated + evictions fell from 10,000 allocations / 320,000 bytes to zero; cache-hit + behavior is unchanged. Together with lesson 46, cold default population + is 25 allocations / 1,024 bytes rather than 217 / 9,088. diff --git a/src/utils/lfucache.jl b/src/utils/lfucache.jl index c9ca69aa8..2af1900b5 100644 --- a/src/utils/lfucache.jl +++ b/src/utils/lfucache.jl @@ -20,9 +20,17 @@ function Base.get!(f, cache::BasicLFUCache{K,V}, key::K) where {K,V} cache.freq[key] = 1 if length(cache.cache) > cache.max_size # Find the least frequently used key - _, lfu_key::K = findmin(cache.freq) + # `findmin(::Dict)` boxes its (frequency, key) result. Scan directly + # and keep the first minimum, preserving its iteration-order ties. + lfu_key, min_freq = first(cache.freq) + for (candidate, frequency) in cache.freq + if frequency < min_freq + lfu_key = candidate + min_freq = frequency + end + end delete!(cache.cache, lfu_key) delete!(cache.freq, lfu_key) end return val -end \ No newline at end of file +end diff --git a/test/allocations.jl b/test/allocations.jl index 4b8cfbf39..5b70e402b 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -81,6 +81,28 @@ end @test measured.bytes <= 4_000 * ALLOC_BOUND_MULTIPLIER end +function exercise_lfu_evictions!(cache, key) + for _ in 1:1_000 + get!(() -> nothing, cache, key) + end +end + +@testset "LFU eviction allocations" begin + cache = Dagger.BasicLFUCache{Tuple{UInt,Symbol},Any}(16) + for i in 1:cache.max_size + key = (UInt(i), :meta) + cache.cache[key] = nothing + cache.freq[key] = 2 + end + key = (UInt(1024), :meta) + measured = measure_steady_state_allocs() do + exercise_lfu_evictions!(cache, key) + end + @test !haskey(cache.cache, key) + @test measured.allocs <= 100 * ALLOC_BOUND_MULTIPLIER + @test measured.bytes <= 4_000 * ALLOC_BOUND_MULTIPLIER +end + # name => (; allocs, bytes) upper bounds (see header for how these are set). # Measured steady-state values at the time of writing are noted inline. const ALLOC_BOUNDS = Dict( diff --git a/test/reuse.jl b/test/reuse.jl index d13a6c93f..db07242a9 100644 --- a/test/reuse.jl +++ b/test/reuse.jl @@ -1172,3 +1172,37 @@ end @test !haskey(dict, "apricot") # New value not in dict end end + +@testset "LFU eviction equivalence" begin + function reference_lfu_get!(cache, key) + if haskey(cache.cache, key) + cache.freq[key] += 1 + return cache.cache[key] + end + value = key[1] + cache.cache[key] = value + cache.freq[key] = 1 + if length(cache.cache) > cache.max_size + _, lfu_key = findmin(cache.freq) + delete!(cache.cache, lfu_key) + delete!(cache.freq, lfu_key) + end + return value + end + + # Include zero capacity, ties, retained hot entries and immediate eviction + # of new entries. Compare the exact contents/frequencies after every call. + for capacity in (0, 1, 8, 256) + current = Dagger.BasicLFUCache{Tuple{UInt,Symbol},Any}(capacity) + original = Dagger.BasicLFUCache{Tuple{UInt,Symbol},Any}(capacity) + rng = MersenneTwister(1729) + for _ in 1:1_000 + key = (rand(rng, UInt(1):UInt(512)), :meta) + @test get!(() -> key[1], current, key) == reference_lfu_get!(original, key) + @test current.cache == original.cache + @test current.freq == original.freq + end + @test empty!(current) === current + @test isempty(current.cache) && isempty(current.freq) + end +end