From 00e67b0cc8b036ff5d78206dffd5fb509b4e60eb Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Fri, 17 Jul 2026 16:19:49 -0500 Subject: [PATCH 01/44] datadeps: fix _eft_runtime_ns view() on Tuple fargs `_eft_runtime_ns` did `@view spec.fargs[2:end]`, which breaks under typed tasks (DTaskSpec{true}) because `spec.fargs` is a `Tuple` and `Base.view(::Tuple, ::UnitRange)` isn't defined. Every AOT scheduling call under hierarchical scheduling hit this because typed tasks are the code path hierarchical routes through. Dispatch on argument type: `Base.tail` for tuples (zero allocation), `@view` for vectors. Both cases produce a signature-compatible view of the arguments-after-function. --- src/datadeps/scheduling.jl | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index c8d200968..085c241ea 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -740,8 +740,16 @@ Mutates and returns `state`. The caller is responsible for providing a freshly-`copy(prev_state)` if the previous state must be preserved on rejection. """ +# Take all arguments after the function. `spec.fargs` is a `Tuple` for typed +# tasks (`DTaskSpec{true}`) and a `Vector{Argument}` otherwise. `Base.view` +# is not defined for `Tuple`, so `@view fargs[2:end]` breaks under typed +# tasks; use `Base.tail` for the tuple case (zero allocation) and a view +# for the vector case. +_eft_tail_args(fargs::Tuple) = Base.tail(fargs) +_eft_tail_args(fargs::AbstractVector) = @view fargs[2:end] + function _eft_runtime_ns(snap::MT.MetricsSnapshot, spec, proc::Processor) - sig = Sch.signature(spec.fargs[1], @view spec.fargs[2:end]).sig + sig = Sch.signature(spec.fargs[1], _eft_tail_args(spec.fargs)).sig worker_id = root_worker_id(proc) runtime_lookup = metrics_lookup_runtime_median(snap, sig, proc, worker_id) return runtime_lookup === nothing ? Float64(GREEDY_DEFAULT_RUNTIME_NS) : Float64(runtime_lookup) From 80df355fff2cdafd9658df54525c9afc15be9a3f Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Fri, 17 Jul 2026 18:15:47 -0500 Subject: [PATCH 02/44] datadeps: cache hierarchical schedules; fix similar; fix JuMPExt import Adds cross-`spawn_datadeps` reuse of hierarchical AOT schedules, so iterative workflows (typical for solvers and stencil codes) pay the AOT scheduling cost once per DAG shape instead of every call. Together with two related bug fixes surfaced along the way, this makes the full scheduler test suite pass on `jps/riteshsc26` (486 pass / 1 pre-existing broken). ## Cache under hierarchical scheduling `distribute_tasks_hierarchical!` now: * Builds a full-region `DAGSpec` up front (using the same `dag_add_task!` logic as the flat path's `datadeps_build_schedule!`) and consults `queue.scheduler`'s task-local schedule cache before spinning up any partitioning or per-partition scheduling. * On a cache hit, extracts the recovered task-to-processor mapping from the matching `DAGSpecSchedule` and threads it through the parallel per-partition pipeline via a new `precomputed_schedule` keyword on `schedule_partition_full!`. Each partition filters the full-region schedule to its own tasks and skips the `datadeps_build_schedule!` call entirely; tasks whose cached processor is not in `local_procs` (possible if `all_procs` changed since the schedule was cached) fall back to JIT via the existing guard in `_schedule_vertex!`. * On a cache miss, runs the existing pipeline unchanged, collects each partition's task-to-processor assignments alongside its `DataDepsState`, and persists them to `queue.scheduler.cache` keyed by the full-region `DAGSpec`. Julian's per-partition `similar(queue.scheduler)` shard remains untouched on the miss path (its concurrency and per-partition processor-list reasoning is unchanged); the cache lookup / persistence sits strictly above that layer so no shared mutable state or cross-thread synchronisation is introduced. ## Bug fixes surfaced by the cache work 1. `_eft_runtime_ns` did `@view spec.fargs[2:end]`, which breaks under typed tasks (`DTaskSpec{true}`) because `spec.fargs` is a `Tuple` and `Base.view(::Tuple, ::UnitRange)` isn't defined. Dispatch Tuple -> `Base.tail`, Vector -> `@view` (zero allocation on both paths). This is the same fix as PR #722; carried here so the branch is testable end-to-end. 2. The parameterised schedulers -- `IteratedGreedyScheduler`, `SimulatedAnnealingScheduler`, `OptimizingScheduler`, `JuMPScheduler` -- had no `Base.similar` methods, so hierarchical partition-shard creation fell through to the default `similar(::DataDepsScheduler) = typeof(s)()` and raised a `MethodError` (no zero-arg constructor for the parameterised types). Added type-specific methods that preserve configuration and deep-copy the RNG so parallel partition shards produce independent, race-free samples without behavioural drift. 3. `ext/JuMPExt.jl` imported `MetricsTracker` as a top-level package, but `jps/riteshsc26` folds MetricsTracker into the `Dagger` module via a direct `include` (see `src/Dagger.jl`). Changed to `import Dagger.MetricsTracker as MT` so the JuMP extension loads under this branch. ## Verified Full `test/datadeps/scheduling.jl` on `jps/riteshsc26 @ 0c9aff81` (Julia 1.12.3, macOS aarch64): * Before this commit: 30 pass / 4 fail / 7 error in GreedyScheduler alone, halted before the other schedulers ran. * After this commit: 486 pass / 0 fail / 0 error / 1 pre-existing broken across `equivalent_structure`, `DAGSpec`, `PtrStrict`, `Greedy`, `IteratedGreedy`, `SimulatedAnnealing`, `OptimizingScheduler` (adaptive and MILP paths), and `JuMPScheduler`. Determinism spot check (matmul in `spawn_datadeps` with an `IteratedGreedyScheduler`): first call is a cache miss (cache size 0 -> 1); second call with the same DAG shape hits (cache size stays at 1, no crash). --- ext/JuMPExt.jl | 5 +- src/datadeps/hierarchical.jl | 193 ++++++++++++++++++++++++++++++----- src/datadeps/scheduling.jl | 45 ++++++++ 3 files changed, 219 insertions(+), 24 deletions(-) diff --git a/ext/JuMPExt.jl b/ext/JuMPExt.jl index a139cdfeb..179ca76b9 100644 --- a/ext/JuMPExt.jl +++ b/ext/JuMPExt.jl @@ -13,7 +13,10 @@ import Dagger: JuMPScheduler, DAGSpec, datadeps_schedule_dag_aot!, GREEDY_DEFAULT_RUNTIME_NS, ScheduleState, greedy_schedule! import Dagger.Sch -import MetricsTracker as MT +# `jps/riteshsc26` folds `MetricsTracker` into the `Dagger` module via a +# direct `include` (see `src/Dagger.jl`), so it is `Dagger.MetricsTracker` +# rather than a top-level package the extension can `import` directly. +import Dagger.MetricsTracker as MT import Graphs: nv, outdegree, edges, src, dst function _milp_compatible_procs(spec, all_procs::Vector{Dagger.Processor}) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index b4e2f0734..1570d1f9b 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -754,19 +754,32 @@ end schedule_partition_full!(queue, queue_lock, partition_id, partition_verts, dag, seen_tasks, local_procs, vertex_to_partition, task_submitted, - region_uids) -> DataDepsState + region_uids, registry; + precomputed_schedule=nothing) + -> (DataDepsState, Dict{DTask,Processor}) Per-partition scheduling for both single-worker and multi-worker hierarchical paths. Uses existing `distribute_task!` logic with per-partition `DataDepsState`, `all_procs` limited to this partition's processors, and -cross-partition syncdeps from the precomputed DAG. - -AOT scheduling is run *locally per partition*: a partition-local `DAGSpec` is -built from just this partition's tasks and scheduled over just `local_procs` via -`datadeps_build_schedule!`. `region_uids` (all in-region task uids) lets that -builder bail cleanly when a task depends on a same-region producer in another -partition (which it must not `fetch`). Tasks without an AOT assignment fall back -to JIT scheduling in `distribute_task!`. +cross-partition syncdeps from the precomputed DAG. Returns this partition's +`DataDepsState` along with the task-to-processor assignments it computed (or +reused), so the top-level `distribute_tasks_hierarchical!` can persist the +combined per-region schedule to `queue.scheduler`'s cache. + +AOT scheduling defaults to running *locally per partition*: a partition-local +`DAGSpec` is built from just this partition's tasks and scheduled over just +`local_procs` via `datadeps_build_schedule!`. `region_uids` (all in-region task +uids) lets that builder bail cleanly when a task depends on a same-region +producer in another partition (which it must not `fetch`). Tasks without an AOT +assignment fall back to JIT scheduling in `distribute_task!`. + +When `precomputed_schedule` is supplied (cache hit at the top of +`distribute_tasks_hierarchical!`), the per-partition AOT step is skipped: this +partition's task-to-processor assignments are filtered from +`precomputed_schedule` directly. Tasks whose cached processor is not in +`local_procs` (possible if `all_procs` changed since the schedule was cached) +fall back to JIT scheduling in `distribute_task!` via the guard in +`_schedule_vertex!`. """ function schedule_partition_full!(queue::DataDepsTaskQueue, queue_lock::ReentrantLock, @@ -778,9 +791,10 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, vertex_to_partition::Vector{Int}, task_submitted::Vector{Base.Event}, region_uids::Set{UInt}, - registry::Union{SharedChunkRegistry,Nothing}) + registry::Union{SharedChunkRegistry,Nothing}; + precomputed_schedule::Union{Dict{DTask,Processor},Nothing}=nothing) if isempty(partition_verts) || isempty(local_procs) - return DataDepsState(DAGSpec()) + return DataDepsState(DAGSpec()), Dict{DTask,Processor}() end local_scope = UnionScope(map(ExactScope, local_procs)) @@ -824,8 +838,27 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, # Per-partition (local) AOT scheduling over just this partition's procs. # `partition_verts` is in ascending vertex (= submission) order. partition_pairs = DTaskPair[seen_tasks[v] for v in partition_verts] - _pdag, schedule = datadeps_build_schedule!(temp_queue.scheduler, partition_pairs, - local_procs, local_scope; region_uids) + + if precomputed_schedule === nothing + # Cache miss: run the per-partition AOT scheduler over just this + # partition's tasks and processors. + _pdag, schedule = datadeps_build_schedule!(temp_queue.scheduler, partition_pairs, + local_procs, local_scope; region_uids) + else + # Cache hit: `distribute_tasks_hierarchical!` already resolved the whole + # region's task-to-processor mapping from `queue.scheduler.cache`. Filter + # it to just this partition's tasks. Tasks that fell through the + # top-level `dag_add_task!` bail (missing from `precomputed_schedule`) + # still get JIT-scheduled inside `distribute_task!` via + # `temp_queue.scheduler`, matching the miss-path semantics. + schedule = Dict{DTask,Processor}() + for pair in partition_pairs + proc = get(precomputed_schedule, pair.task, nothing) + if proc !== nothing + schedule[pair.task] = proc + end + end + end # N.B. If this partition throws partway through (e.g. from # `distribute_task!`), any of our vertices that haven't yet been @@ -866,7 +899,7 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, end end - return state + return state, schedule end struct LockedEnqueueQueue <: AbstractTaskQueue @@ -880,6 +913,92 @@ function enqueue!(leq::LockedEnqueueQueue, pairs::Vector{DTaskPair}) @lock leq.lock enqueue!(leq.inner, pairs) end +""" + _hierarchical_dag_spec_and_cache_lookup(scheduler, seen_tasks) + -> (dag_spec::DAGSpec, precomputed::Union{Dict{DTask,Processor},Nothing}) + +Build a full-region `DAGSpec` by walking `seen_tasks` in submission order, then +look it up in `scheduler`'s task-local schedule cache. When a structurally +equivalent DAG is found (per `datadeps_dag_equivalent`), return the recovered +task-to-processor mapping in the second slot; otherwise return `nothing`. + +Mirrors the DAGSpec-building and cache-matching logic of the flat path's +`datadeps_build_schedule!` in `queue.jl`, but stops before running the AOT +scheduler so that `distribute_tasks_hierarchical!` retains its parallel +per-partition scheduling pipeline on a cache miss. The DAGSpec built here is +also returned so the miss path can persist a freshly-computed schedule at the +end of the hierarchical run using the same key. + +`dag_add_task!` may return `false` (bail) on tasks whose arguments include a +within-DAG `DTask` -- the returned `DAGSpec` then only covers the prefix that +was successfully added, matching the flat path's behaviour. Tasks past the +bail point are not in the cache key; on a hit they simply won't be in +`precomputed`, and per-partition JIT scheduling handles them exactly as in the +miss path. +""" +function _hierarchical_dag_spec_and_cache_lookup(scheduler::DataDepsScheduler, + seen_tasks::Vector{DTaskPair}) + dag_spec = DAGSpec() + dummy_state = DataDepsState(dag_spec) + for (spec, task) in seen_tasks + if !dag_add_task!(dag_spec, dummy_state, spec, task) + break + end + end + isempty(dag_spec) && return dag_spec, nothing + + for (other_spec, spec_schedule) in datadeps_schedule_cache(scheduler) + if datadeps_dag_equivalent(scheduler, dag_spec, other_spec) + @dagdebug nothing :spawn_datadeps "Hierarchical DAG cache hit" + schedule = Dict{DTask,Processor}() + uid_to_task = Dict{UInt,DTask}() + for pair in seen_tasks + uid_to_task[pair.task.uid] = pair.task + end + for (id, proc) in spec_schedule.id_to_proc + uid = dag_spec.id_to_uid[id] + task = get(uid_to_task, uid, nothing) + task === nothing && continue + schedule[task] = proc + end + return dag_spec, schedule + end + end + return dag_spec, nothing +end + +""" + _hierarchical_persist_schedule!(scheduler, dag_spec, partition_schedules) + +After a cache-miss hierarchical run, collect the per-partition +task-to-processor assignments produced by `schedule_partition_full!` and store +them in `scheduler`'s schedule cache keyed by `dag_spec`. This allows a later +`distribute_tasks_hierarchical!` call with a structurally equivalent DAG to +reuse the schedule via `_hierarchical_dag_spec_and_cache_lookup`. + +No-op when `dag_spec` is empty (the top-level DAG build bailed on the first +task) or when no partition produced any assignments (every task fell back to +JIT scheduling). Only tasks captured in `dag_spec` are cached, so partial-DAG +cases match flat-path semantics: tasks past the bail point are re-JIT'd on +subsequent runs. +""" +function _hierarchical_persist_schedule!(scheduler::DataDepsScheduler, + dag_spec::DAGSpec, + partition_schedules::Vector{Dict{DTask,Processor}}) + isempty(dag_spec) && return + spec_schedule = DAGSpecSchedule() + for pschedule in partition_schedules + for (task, proc) in pschedule + haskey(dag_spec.uid_to_id, task.uid) || continue + id = dag_spec.uid_to_id[task.uid] + spec_schedule.id_to_proc[id] = proc + end + end + isempty(spec_schedule.id_to_proc) && return + push!(datadeps_schedule_cache(scheduler), dag_spec => spec_schedule) + return +end + # Parallel partition scheduling wraps errors in TaskFailedException / # CompositeException via `@sync`/`Threads.@spawn`. Unwrap so callers and tests # see the root `SchedulingException` / scheduler error. @@ -919,6 +1038,19 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) return end + # Top-level schedule cache lookup. Build a full-region DAGSpec and consult + # `queue.scheduler`'s task-local cache before spinning up any partitioning + # or per-partition scheduling. Under hierarchical, each partition otherwise + # runs with its own `similar(queue.scheduler)` shard (see + # `schedule_partition_full!`) whose cache is discarded after the call, so + # this top-level lookup is the only place where cross-`spawn_datadeps` + # amortisation of AOT scheduling cost can occur. On a hit + # (`precomputed_schedule !== nothing`) the per-partition scheduler is + # skipped; on a miss we run the normal pipeline and persist the collected + # schedule at the end for the next call to reuse. + dag_spec, precomputed_schedule = + _hierarchical_dag_spec_and_cache_lookup(queue.scheduler, seen_tasks) + # Get the set of all processors all_procs = Processor[] scope = get_compute_scope() @@ -966,21 +1098,26 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) # Phase 4: parallel per-partition scheduling. Each partition runs on its own # task, prepares its tasks concurrently, computes its own partition-local - # AOT schedule, and coordinates cross-partition dependencies via the - # `task_submitted` events. Concurrent submission is safe now that task - # futures are backed by `MemPool.DFuture` (see the header note). + # AOT schedule (or reuses `precomputed_schedule` on a cache hit, in which + # case the AOT step is skipped and only the task submission side runs), and + # coordinates cross-partition dependencies via the `task_submitted` events. + # Concurrent submission is safe now that task futures are backed by + # `MemPool.DFuture` (see the header note). partition_states = Vector{DataDepsState}(undef, n_partitions) + partition_schedules = Vector{Dict{DTask,Processor}}(undef, n_partitions) try @sync for pid in 1:n_partitions Threads.@spawn begin locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) with_options(; task_queue=locked_queue) do - partition_states[pid] = schedule_partition_full!( - queue, queue_lock, pid, partitions[pid], - dag, seen_tasks, - partition_procs[pid], vertex_to_partition, - task_submitted, region_uids, registry - ) + partition_states[pid], partition_schedules[pid] = + schedule_partition_full!( + queue, queue_lock, pid, partitions[pid], + dag, seen_tasks, + partition_procs[pid], vertex_to_partition, + task_submitted, region_uids, registry; + precomputed_schedule + ) end end end @@ -989,6 +1126,16 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) end _hierarchical_copy_from_and_free!(partition_states, n_partitions, registry) + + # Persist the freshly-computed schedule for reuse by a later + # `distribute_tasks_hierarchical!` call on a structurally equivalent DAG. + # Only runs on a cache miss (`precomputed_schedule === nothing`) to avoid + # both duplicating the cache entry we just consumed and appending it under + # a potentially different DAGSpec key that would fail subsequent + # `datadeps_dag_equivalent` matches. + if precomputed_schedule === nothing + _hierarchical_persist_schedule!(queue.scheduler, dag_spec, partition_schedules) + end end function _hierarchical_max_write_num(state::DataDepsState, arg_w::ArgumentWrapper) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 085c241ea..9d3488331 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -715,6 +715,18 @@ end IteratedGreedyScheduler(; kwargs...) = IteratedGreedyScheduler(GreedyScheduler(); kwargs...) +# Hierarchical scheduling calls `similar` per partition to avoid sharing mutable +# state (esp. the RNG) across the parallel per-partition scheduling tasks. Deep- +# copy the RNG so each partition-shard produces independent, race-free samples; +# recursively call `similar` on the inner scheduler so its own mutable state is +# refreshed too. Static configuration (n_iters, destroy_frac) carries through +# unchanged so the shard's behaviour matches the original's. +Base.similar(s::IteratedGreedyScheduler) = + IteratedGreedyScheduler(similar(s.inner); + n_iters=s.n_iters, + destroy_frac=s.destroy_frac, + rng=copy(s.rng)) + # Forward equivalence / cache / argspec hooks so cache reuse is partitioned by # the inner scheduler's type, not by IG's own wrapper type. Same idea as the # bench's TimedScheduler wrapper. @@ -977,6 +989,15 @@ end SimulatedAnnealingScheduler(; kwargs...) = SimulatedAnnealingScheduler(IteratedGreedyScheduler(); kwargs...) +# See the note on `similar(::IteratedGreedyScheduler)`: hierarchical scheduling +# hands each partition its own scheduler shard, so we deep-copy the RNG and +# recursively refresh the inner scheduler, while preserving the parameterisation +# (q, k, n_restarts). +Base.similar(s::SimulatedAnnealingScheduler) = + SimulatedAnnealingScheduler(similar(s.inner); + q=s.q, k=s.k, n_restarts=s.n_restarts, + rng=copy(s.rng)) + datadeps_schedule_cache(sched::SimulatedAnnealingScheduler) = datadeps_schedule_cache(sched.inner) datadeps_dag_equivalent(sched::SimulatedAnnealingScheduler, dspec1::DAGSpec, dspec2::DAGSpec) = @@ -1308,6 +1329,13 @@ struct JuMPScheduler <: DataDepsScheduler end end +# See the note on `similar(::IteratedGreedyScheduler)`. `JuMPScheduler` has no +# mutable state and only immutable config fields; return a fresh instance with +# the same configuration so hierarchical partition shards behave identically +# to the original. +Base.similar(s::JuMPScheduler) = + JuMPScheduler(s.optimizer; Z=s.Z, time_limit_sec=s.time_limit_sec) + # `datadeps_schedule_dag_aot!(::JuMPScheduler, ...)` lives in ext/JuMPExt.jl; # the constructor prevents instantiation without the extension. const OPT_DEFAULT_MILP_THRESHOLD = 12 @@ -1388,6 +1416,23 @@ struct OptimizingScheduler{R<:Random.AbstractRNG} <: DataDepsScheduler end end +# See the note on `similar(::IteratedGreedyScheduler)`. `OptimizingScheduler` +# doesn't own an inner scheduler instance (it constructs the underlying +# JuMP/SA/IG/Greedy pipeline fresh per invocation), so only the RNG needs a +# deep copy; every other field is immutable configuration and is forwarded +# verbatim. +Base.similar(s::OptimizingScheduler) = + OptimizingScheduler(; optimizer=s.optimizer, + milp_threshold=s.milp_threshold, + milp_time_limit_sec=s.milp_time_limit_sec, + milp_Z=s.milp_Z, + ig_n_iters=s.ig_n_iters, + ig_destroy_frac=s.ig_destroy_frac, + sa_q=s.sa_q, + sa_k=s.sa_k, + sa_n_restarts=s.sa_n_restarts, + rng=copy(s.rng)) + """ opt_uses_milp(sched::OptimizingScheduler, n_tasks::Integer) -> Bool From 53b61aef35a2756f7cbad69f6ef93c51150e1c15 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sat, 18 Jul 2026 02:22:22 -0500 Subject: [PATCH 03/44] datadeps: add wall-clock time_limit_sec to IG, SA, OptimizingScheduler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per Przemek's & Julian's SC26 workshop paper feedback: give the metaheuristic schedulers a "quick-and-dirty 1 minute timeout" so long-K benchmarks stay within a reasonable everyday-user budget instead of exhausting the full Orsila cooling schedule × n_restarts safety cap. MILP already had `time_limit_sec` (forwarded to HiGHS); IG and SA were purely iteration-bounded until now. Changes: - `IteratedGreedyScheduler` and `SimulatedAnnealingScheduler` grow a `time_limit_sec::Float64` field (default `Inf`, opt-in), validated strictly-positive in the inner constructor. - `iterated_greedy_schedule!` / `simulated_annealing_schedule!` compute `start_ns = time_ns()` once and check `(time_ns() - start_ns) >= budget_ns` at the top of the outer loop (IG) or every restart and every SA inner-loop iteration. `Inf` is represented as `typemax(UInt64)` so the check never trips. - The check placement returns best-so-far without spending another neighbour proposal, preserving the "non-worsening vs seed" invariant covered by the existing test. - `Base.similar` for both schedulers carries the field so hierarchical partition shards inherit the same budget. - `OptimizingScheduler` grows `ig_time_limit_sec` and `sa_time_limit_sec` and forwards them into the constructed IG / SA when routed to the heuristic path. When routed to MILP, only `milp_time_limit_sec` applies (unchanged). - Bench driver's `default_scheduler_factories` sets `heuristic_time_limit_sec=60.0` for IG, SA, and the heuristic path of OptimizingScheduler; MILP retains its 120 s default. Both budgets are kwarg-configurable so callers can widen or disable. - New tests: constructor validation (accept finite, reject 0/negative, default is `Inf`), `similar` propagation, budget-triggered early exit under a nanosecond limit on a real Cholesky DAG, `Inf` opt-out reproduces the pre-budget deterministic result, and end-to-end `OptimizingScheduler` heuristic-path forwarding. All 516 scheduler tests pass with and without JuMP loaded, plus the one pre-existing broken test unchanged. --- bench/datadeps_schedulers/driver.jl | 23 +++-- src/datadeps/scheduling.jl | 101 +++++++++++++++--- test/datadeps/scheduling.jl | 152 ++++++++++++++++++++++++++++ 3 files changed, 258 insertions(+), 18 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 3d86125b0..be7ba0a5a 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -186,13 +186,20 @@ const DEFAULT_WARMUP = 1 # Factories — not instances — because RoundRobin holds mutable state and each # trial needs a fresh copy. Default MILP budget is set generously since a -# K~64 solve can exceed a minute; callers override as needed. -function default_scheduler_factories(; milp_time_limit_sec::Real=120.0) +# K~64 solve can exceed a minute; callers override as needed. Heuristic +# (IG, SA) wall-clock budgets default to 60 s per Przemek's & Julian's ask: +# "quick-and-dirty 1 minute timeout" so metaheuristic runs stay within a +# reasonable everyday-user budget instead of exhausting the full iteration +# schedule at K~10000. Pass `Inf` to reproduce the pre-budget behaviour. +function default_scheduler_factories(; milp_time_limit_sec::Real=120.0, + heuristic_time_limit_sec::Real=60.0) factories = [ "RoundRobinScheduler" => () -> Dagger.RoundRobinScheduler(), "GreedyScheduler" => () -> Dagger.GreedyScheduler(), - "IteratedGreedyScheduler" => () -> Dagger.IteratedGreedyScheduler(), - "SimulatedAnnealingScheduler" => () -> Dagger.SimulatedAnnealingScheduler(), + "IteratedGreedyScheduler" => () -> Dagger.IteratedGreedyScheduler(; + time_limit_sec=heuristic_time_limit_sec), + "SimulatedAnnealingScheduler" => () -> Dagger.SimulatedAnnealingScheduler(; + time_limit_sec=heuristic_time_limit_sec), ] if HAS_MILP push!(factories, @@ -201,10 +208,14 @@ function default_scheduler_factories(; milp_time_limit_sec::Real=120.0) push!(factories, "OptimizingScheduler" => () -> Dagger.OptimizingScheduler(; optimizer=HiGHS.Optimizer, - milp_time_limit_sec=milp_time_limit_sec)) + milp_time_limit_sec=milp_time_limit_sec, + ig_time_limit_sec=heuristic_time_limit_sec, + sa_time_limit_sec=heuristic_time_limit_sec)) else push!(factories, - "OptimizingScheduler" => () -> Dagger.OptimizingScheduler()) + "OptimizingScheduler" => () -> Dagger.OptimizingScheduler(; + ig_time_limit_sec=heuristic_time_limit_sec, + sa_time_limit_sec=heuristic_time_limit_sec)) end return factories end diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 9d3488331..854ab3d92 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -670,6 +670,10 @@ _greedy_arg_ready_time_ns_cached(::Any, ::MT.MetricsSnapshot, ::DAGSpec, const IG_DEFAULT_N_ITERS = 32 const IG_DEFAULT_DESTROY_FRAC = 0.30 +# `Inf` disables the wall-clock stop; callers (e.g. benchmarks) opt in by +# passing a finite budget. Requested by Przemek & Julian: enable "1-minute +# metaheuristic budget" as the operational default for the paper. +const IG_DEFAULT_TIME_LIMIT_SEC = Inf """ IteratedGreedyScheduler{S<:DataDepsScheduler} <: DataDepsScheduler @@ -691,6 +695,11 @@ Fields: - `n_iters::Int` — destroy/reinsert cycles per call. - `destroy_frac::Float64` — fraction of tasks destroyed per cycle, in (0, 1]. - `rng::Random.AbstractRNG` — RNG used for destroy-set sampling. +- `time_limit_sec::Float64` — wall-clock stop, checked between iterations. + `Inf` (default) disables the stop; a finite + value returns the best-so-far the moment the + budget is exceeded, so the scheduler always + terminates within one iteration of the limit. The schedule cache is delegated to `inner` so entries are partitioned by inner-scheduler type, matching the `TimedScheduler` wrapper convention used @@ -701,15 +710,18 @@ struct IteratedGreedyScheduler{S<:DataDepsScheduler, R<:Random.AbstractRNG} <: D n_iters::Int destroy_frac::Float64 rng::R + time_limit_sec::Float64 function IteratedGreedyScheduler(inner::S; n_iters::Integer=IG_DEFAULT_N_ITERS, destroy_frac::Real=IG_DEFAULT_DESTROY_FRAC, - rng::R=Random.default_rng()) where {S<:DataDepsScheduler, R<:Random.AbstractRNG} + rng::R=Random.default_rng(), + time_limit_sec::Real=IG_DEFAULT_TIME_LIMIT_SEC) where {S<:DataDepsScheduler, R<:Random.AbstractRNG} n_iters >= 0 || throw(ArgumentError("IteratedGreedyScheduler: n_iters must be ≥ 0, got $n_iters")) (destroy_frac > 0 && destroy_frac <= 1) || throw(ArgumentError("IteratedGreedyScheduler: destroy_frac must be in (0, 1], got $destroy_frac")) - return new{S, R}(inner, Int(n_iters), Float64(destroy_frac), rng) + time_limit_sec > 0 || throw(ArgumentError("IteratedGreedyScheduler: time_limit_sec must be > 0, got $time_limit_sec")) + return new{S, R}(inner, Int(n_iters), Float64(destroy_frac), rng, Float64(time_limit_sec)) end end @@ -725,7 +737,8 @@ Base.similar(s::IteratedGreedyScheduler) = IteratedGreedyScheduler(similar(s.inner); n_iters=s.n_iters, destroy_frac=s.destroy_frac, - rng=copy(s.rng)) + rng=copy(s.rng), + time_limit_sec=s.time_limit_sec) # Forward equivalence / cache / argspec hooks so cache reuse is partitioned by # the inner scheduler's type, not by IG's own wrapper type. Same idea as the @@ -851,10 +864,12 @@ function iterated_greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapsho n_iters::Integer=IG_DEFAULT_N_ITERS, destroy_frac::Real=IG_DEFAULT_DESTROY_FRAC, rng::Random.AbstractRNG=Random.default_rng(), - cache::Union{EFTCostCache, Nothing}=nothing) + cache::Union{EFTCostCache, Nothing}=nothing, + time_limit_sec::Real=IG_DEFAULT_TIME_LIMIT_SEC) n_iters >= 0 || throw(ArgumentError("iterated_greedy_schedule!: n_iters must be ≥ 0")) (destroy_frac > 0 && destroy_frac <= 1) || throw(ArgumentError("iterated_greedy_schedule!: destroy_frac must be in (0, 1]")) + time_limit_sec > 0 || throw(ArgumentError("iterated_greedy_schedule!: time_limit_sec must be > 0")) n_tasks = nv(dag_spec.g) (n_tasks == 0 || n_iters == 0) && return state if cache === nothing @@ -874,7 +889,17 @@ function iterated_greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapsho destroyed = Set{Int}() sizehint!(destroyed, n_destroy) + # Wall-clock guard checked once per outer iteration. Elapsed nanoseconds + # are computed as `time_ns() - start_ns` (unsigned subtraction, so no + # overflow arithmetic on the deadline itself). `Inf` disables the guard + # by setting `budget_ns` to the maximum representable `UInt64`, ensuring + # the comparison never trips. + start_ns = time_ns() + budget_ns = isinf(time_limit_sec) ? typemax(UInt64) : + round(UInt64, time_limit_sec * 1e9) + for _ in 1:n_iters + (time_ns() - start_ns) >= budget_ns && break # Sample `n_destroy` distinct task IDs uniformly at random. Doing a # partial Fisher–Yates on perm_buf is O(n_destroy) and avoids # allocating a fresh randperm vector each iter. @@ -916,6 +941,7 @@ function datadeps_schedule_dag_aot!(scheduler::IteratedGreedyScheduler, greedy_schedule!(state, snap, dag_spec, all_procs; cache=cache) state = iterated_greedy_schedule!(state, snap, dag_spec, all_procs; + time_limit_sec=scheduler.time_limit_sec, n_iters=scheduler.n_iters, destroy_frac=scheduler.destroy_frac, rng=scheduler.rng, @@ -935,6 +961,12 @@ const SA_DEFAULT_Q = 0.95 const SA_DEFAULT_K = 1.0 const SA_DEFAULT_N_RESTARTS = 1 const SA_TF_FLOOR = 1e-12 +# See `IG_DEFAULT_TIME_LIMIT_SEC`. When SA wraps an IG seed inside +# `datadeps_schedule_dag_aot!(::SimulatedAnnealingScheduler, ...)`, both the +# seed and the SA refinement share the same clock but are governed by their +# own `time_limit_sec` fields, so setting both to 60 s does NOT double the +# effective budget of the SA pipeline. See the docstring for details. +const SA_DEFAULT_TIME_LIMIT_SEC = Inf """ SimulatedAnnealingScheduler{S<:DataDepsScheduler, R<:Random.AbstractRNG} <: DataDepsScheduler @@ -958,6 +990,16 @@ Fields: - `k::Float64` — coefficient `k > 0` in Orsila §4.3 Eqs. 18, 19. - `n_restarts::Int` — independent SA runs, each seeded from the best-known solution so far (Orsila §5.3 rule 7). +- `time_limit_sec::Float64` — wall-clock stop for the SA pipeline in + `datadeps_schedule_dag_aot!`, checked between + restarts and once per cooling level. `Inf` + (default) disables the stop; a finite value + returns the best-so-far state once the budget + is exceeded. When an `IteratedGreedyScheduler` + seed is used, the seed observes its own + `time_limit_sec`; the two budgets are summed + (worst-case), matching Julian's requested + "quick-and-dirty" semantics. - `rng::R` — RNG used for moves and acceptance sampling. The `rng` is mutated during scheduling, so a single `SimulatedAnnealingScheduler` instance @@ -973,16 +1015,19 @@ struct SimulatedAnnealingScheduler{S<:DataDepsScheduler, R<:Random.AbstractRNG} k::Float64 n_restarts::Int rng::R + time_limit_sec::Float64 function SimulatedAnnealingScheduler(inner::S; q::Real=SA_DEFAULT_Q, k::Real=SA_DEFAULT_K, n_restarts::Integer=SA_DEFAULT_N_RESTARTS, - rng::R=Random.default_rng()) where {S<:DataDepsScheduler, R<:Random.AbstractRNG} + rng::R=Random.default_rng(), + time_limit_sec::Real=SA_DEFAULT_TIME_LIMIT_SEC) where {S<:DataDepsScheduler, R<:Random.AbstractRNG} (0 < q < 1) || throw(ArgumentError("SimulatedAnnealingScheduler: q must be in (0, 1), got $q")) k > 0 || throw(ArgumentError("SimulatedAnnealingScheduler: k must be > 0, got $k")) n_restarts >= 1 || throw(ArgumentError("SimulatedAnnealingScheduler: n_restarts must be ≥ 1, got $n_restarts")) - return new{S, R}(inner, Float64(q), Float64(k), Int(n_restarts), rng) + time_limit_sec > 0 || throw(ArgumentError("SimulatedAnnealingScheduler: time_limit_sec must be > 0, got $time_limit_sec")) + return new{S, R}(inner, Float64(q), Float64(k), Int(n_restarts), rng, Float64(time_limit_sec)) end end @@ -996,7 +1041,8 @@ SimulatedAnnealingScheduler(; kwargs...) = Base.similar(s::SimulatedAnnealingScheduler) = SimulatedAnnealingScheduler(similar(s.inner); q=s.q, k=s.k, n_restarts=s.n_restarts, - rng=copy(s.rng)) + rng=copy(s.rng), + time_limit_sec=s.time_limit_sec) datadeps_schedule_cache(sched::SimulatedAnnealingScheduler) = datadeps_schedule_cache(sched.inner) @@ -1143,10 +1189,12 @@ function simulated_annealing_schedule!(state::ScheduleState, snap::MT.MetricsSna q::Real=SA_DEFAULT_Q, k::Real=SA_DEFAULT_K, n_restarts::Integer=SA_DEFAULT_N_RESTARTS, rng::Random.AbstractRNG=Random.default_rng(), - cache::Union{EFTCostCache, Nothing}=nothing) + cache::Union{EFTCostCache, Nothing}=nothing, + time_limit_sec::Real=SA_DEFAULT_TIME_LIMIT_SEC) (0 < q < 1) || throw(ArgumentError("simulated_annealing_schedule!: q must be in (0, 1)")) k > 0 || throw(ArgumentError("simulated_annealing_schedule!: k must be > 0")) n_restarts >= 1 || throw(ArgumentError("simulated_annealing_schedule!: n_restarts must be ≥ 1")) + time_limit_sec > 0 || throw(ArgumentError("simulated_annealing_schedule!: time_limit_sec must be > 0")) n_tasks = nv(dag_spec.g) n_tasks == 0 && return state @@ -1180,7 +1228,17 @@ function simulated_annealing_schedule!(state::ScheduleState, snap::MT.MetricsSna expected_levels = max(1, ceil(Int, log(Tf / T0) / log(Float64(q)))) max_iters_per_restart = 16 * L * expected_levels + # See the note on the same construction in `iterated_greedy_schedule!`: + # single `time_ns()` origin, unsigned elapsed subtraction, `Inf` disables + # via `typemax(UInt64)` so the check never trips. Inner check is placed + # at the top of the SA loop so a hit returns the current best-so-far + # without spending another neighbour proposal. + start_ns = time_ns() + budget_ns = isinf(time_limit_sec) ? typemax(UInt64) : + round(UInt64, time_limit_sec * 1e9) + for _ in 1:n_restarts + (time_ns() - start_ns) >= budget_ns && break _copy_state!(current, overall_best) current_cost = overall_best_cost _copy_state!(best_in_run, current) @@ -1194,6 +1252,7 @@ function simulated_annealing_schedule!(state::ScheduleState, snap::MT.MetricsSna while true (T <= Tf && R >= Rmax) && break total_iters >= max_iters_per_restart && break + (time_ns() - start_ns) >= budget_ns && break _copy_state!(candidate, current) _sa_propose_neighbor!(candidate, snap, dag_spec, all_procs, rng, cache) @@ -1251,14 +1310,16 @@ function datadeps_schedule_dag_aot!(scheduler::SimulatedAnnealingScheduler, n_iters=ig.n_iters, destroy_frac=ig.destroy_frac, rng=ig.rng, - cache=cache) + cache=cache, + time_limit_sec=ig.time_limit_sec) end final_state = simulated_annealing_schedule!(seed_state, snap, dag_spec, all_procs; q=scheduler.q, k=scheduler.k, n_restarts=scheduler.n_restarts, rng=scheduler.rng, - cache=cache) + cache=cache, + time_limit_sec=scheduler.time_limit_sec) @inbounds for idx in 1:n_tasks task = dag_spec.id_to_task[idx] @@ -1362,9 +1423,13 @@ Fields: - `milp_Z` — forwarded to `JuMPScheduler`. - `ig_n_iters` — forwarded to `IteratedGreedyScheduler`. - `ig_destroy_frac` — forwarded to `IteratedGreedyScheduler`. +- `ig_time_limit_sec` — forwarded to `IteratedGreedyScheduler` + (heuristic path). `Inf` disables the stop. - `sa_q` — forwarded to `SimulatedAnnealingScheduler`. - `sa_k` — forwarded to `SimulatedAnnealingScheduler`. - `sa_n_restarts` — forwarded to `SimulatedAnnealingScheduler`. +- `sa_time_limit_sec` — forwarded to `SimulatedAnnealingScheduler` + (heuristic path). `Inf` disables the stop. - `rng::R` — RNG shared by IG and SA on the heuristic path; unused on the MILP path. """ @@ -1375,9 +1440,11 @@ struct OptimizingScheduler{R<:Random.AbstractRNG} <: DataDepsScheduler milp_Z::Float64 ig_n_iters::Int ig_destroy_frac::Float64 + ig_time_limit_sec::Float64 sa_q::Float64 sa_k::Float64 sa_n_restarts::Int + sa_time_limit_sec::Float64 rng::R function OptimizingScheduler(; @@ -1387,9 +1454,11 @@ struct OptimizingScheduler{R<:Random.AbstractRNG} <: DataDepsScheduler milp_Z::Real=10.0, ig_n_iters::Integer=IG_DEFAULT_N_ITERS, ig_destroy_frac::Real=IG_DEFAULT_DESTROY_FRAC, + ig_time_limit_sec::Real=IG_DEFAULT_TIME_LIMIT_SEC, sa_q::Real=SA_DEFAULT_Q, sa_k::Real=SA_DEFAULT_K, sa_n_restarts::Integer=SA_DEFAULT_N_RESTARTS, + sa_time_limit_sec::Real=SA_DEFAULT_TIME_LIMIT_SEC, rng::R=Random.default_rng(), ) where {R<:Random.AbstractRNG} milp_threshold >= 0 || throw(ArgumentError("OptimizingScheduler: milp_threshold must be ≥ 0, got $milp_threshold")) @@ -1397,9 +1466,11 @@ struct OptimizingScheduler{R<:Random.AbstractRNG} <: DataDepsScheduler ig_n_iters >= 0 || throw(ArgumentError("OptimizingScheduler: ig_n_iters must be ≥ 0, got $ig_n_iters")) (0 < ig_destroy_frac && ig_destroy_frac <= 1) || throw(ArgumentError("OptimizingScheduler: ig_destroy_frac must be in (0, 1], got $ig_destroy_frac")) + ig_time_limit_sec > 0 || throw(ArgumentError("OptimizingScheduler: ig_time_limit_sec must be > 0, got $ig_time_limit_sec")) (0 < sa_q && sa_q < 1) || throw(ArgumentError("OptimizingScheduler: sa_q must be in (0, 1), got $sa_q")) sa_k > 0 || throw(ArgumentError("OptimizingScheduler: sa_k must be > 0, got $sa_k")) sa_n_restarts >= 1 || throw(ArgumentError("OptimizingScheduler: sa_n_restarts must be ≥ 1, got $sa_n_restarts")) + sa_time_limit_sec > 0 || throw(ArgumentError("OptimizingScheduler: sa_time_limit_sec must be > 0, got $sa_time_limit_sec")) return new{R}( optimizer, @@ -1408,9 +1479,11 @@ struct OptimizingScheduler{R<:Random.AbstractRNG} <: DataDepsScheduler Float64(milp_Z), Int(ig_n_iters), Float64(ig_destroy_frac), + Float64(ig_time_limit_sec), Float64(sa_q), Float64(sa_k), Int(sa_n_restarts), + Float64(sa_time_limit_sec), rng, ) end @@ -1428,9 +1501,11 @@ Base.similar(s::OptimizingScheduler) = milp_Z=s.milp_Z, ig_n_iters=s.ig_n_iters, ig_destroy_frac=s.ig_destroy_frac, + ig_time_limit_sec=s.ig_time_limit_sec, sa_q=s.sa_q, sa_k=s.sa_k, sa_n_restarts=s.sa_n_restarts, + sa_time_limit_sec=s.sa_time_limit_sec, rng=copy(s.rng)) """ @@ -1466,12 +1541,14 @@ function datadeps_schedule_dag_aot!(sched::OptimizingScheduler, schedule, dag_sp ig = IteratedGreedyScheduler(GreedyScheduler(); n_iters=sched.ig_n_iters, destroy_frac=sched.ig_destroy_frac, - rng=sched.rng) + rng=sched.rng, + time_limit_sec=sched.ig_time_limit_sec) sub = SimulatedAnnealingScheduler(ig; q=sched.sa_q, k=sched.sa_k, n_restarts=sched.sa_n_restarts, - rng=sched.rng) + rng=sched.rng, + time_limit_sec=sched.sa_time_limit_sec) end datadeps_schedule_dag_aot!(sub, schedule, dag_spec, all_procs, all_scope) diff --git a/test/datadeps/scheduling.jl b/test/datadeps/scheduling.jl index c12f2f363..a3b08c08a 100644 --- a/test/datadeps/scheduling.jl +++ b/test/datadeps/scheduling.jl @@ -1281,6 +1281,57 @@ end # Cache populated by the IG run (shares with Greedy's cache). @test !isempty(cache) end + + @testset "Wall-clock time_limit_sec early-exits" begin + # Constructor: accepts finite budget, rejects nonpositive, defaults to Inf. + s_default = IteratedGreedyScheduler() + @test s_default.time_limit_sec === Inf + s_finite = IteratedGreedyScheduler(; time_limit_sec=0.5) + @test s_finite.time_limit_sec == 0.5 + @test_throws ArgumentError IteratedGreedyScheduler(; time_limit_sec=0.0) + @test_throws ArgumentError IteratedGreedyScheduler(; time_limit_sec=-1.0) + @test_throws ArgumentError iterated_greedy_schedule!(ScheduleState(), + Dagger.MT.snapshot(Dagger.MT.global_metrics_cache()), + DAGSpec(), collect(Dagger.all_processors()); + time_limit_sec=0.0) + + # `similar` carries the field over so hierarchical partition shards + # inherit the same budget (shares the fix path with the RNG/config + # forwarding). + @test similar(s_finite).time_limit_sec == 0.5 + + # Nanosecond budget stops the loop before its natural stop condition. + # A DAG of a few hundred tasks with n_iters=10000 will otherwise + # run for well over a millisecond; the budget must dominate. + dag_spec, all_procs, snap = _capture_cholesky_dag(3, 16) + seed = ScheduleState() + greedy_schedule!(seed, snap, dag_spec, all_procs) + seed_cost = cost_of_schedule(seed) + + t0 = time_ns() + result = iterated_greedy_schedule!(copy(seed), snap, dag_spec, all_procs; + n_iters=10_000, destroy_frac=0.3, + rng=MersenneTwister(0), + time_limit_sec=1e-6) + elapsed_ms = (time_ns() - t0) / 1e6 + # Bound: even one iteration of a small DAG comfortably fits in 100 ms + # on any dev/CI box; a 10000-iter unbounded run does not. + @test elapsed_ms < 100.0 + # Correctness invariant preserved under early exit: best-so-far + # returned, never worse than the seed. + @test cost_of_schedule(result) <= seed_cost + 1e-6 + + # `Inf` is a true opt-out: behavior must match not passing the kwarg + # at all. Same seed → identical result. + r_inf = iterated_greedy_schedule!(copy(seed), snap, dag_spec, all_procs; + n_iters=4, destroy_frac=0.3, + rng=MersenneTwister(11), + time_limit_sec=Inf) + r_none = iterated_greedy_schedule!(copy(seed), snap, dag_spec, all_procs; + n_iters=4, destroy_frac=0.3, + rng=MersenneTwister(11)) + @test cost_of_schedule(r_inf) == cost_of_schedule(r_none) + end end @testset "SimulatedAnnealingScheduler" begin @@ -1649,6 +1700,51 @@ end @test length(r_greedy.task_proc) == Dagger.nv(dag_spec.g) @test length(r_ig.task_proc) == Dagger.nv(dag_spec.g) end + + @testset "Wall-clock time_limit_sec early-exits" begin + # Constructor: accepts finite budget, rejects nonpositive, defaults to Inf. + s_default = SimulatedAnnealingScheduler() + @test s_default.time_limit_sec === Inf + s_finite = SimulatedAnnealingScheduler(; time_limit_sec=0.5) + @test s_finite.time_limit_sec == 0.5 + @test_throws ArgumentError SimulatedAnnealingScheduler(; time_limit_sec=0.0) + @test_throws ArgumentError SimulatedAnnealingScheduler(; time_limit_sec=-1.0) + @test_throws ArgumentError simulated_annealing_schedule!(ScheduleState(), + Dagger.MT.snapshot(Dagger.MT.global_metrics_cache()), + DAGSpec(), collect(Dagger.all_processors()); + time_limit_sec=0.0) + + # `similar` propagates the budget so hierarchical shards inherit it. + @test similar(s_finite).time_limit_sec == 0.5 + + # A tiny budget must interrupt SA long before Orsila's cooling + # schedule × n_restarts safety cap would fire. The larger DAG here + # (nt=4 → K=20) guarantees SA's inner while-loop is well-populated, + # exercising the intra-restart wall-clock check. + dag_spec, all_procs, snap = _capture_cholesky_dag(4, 16) + seed = ScheduleState() + greedy_schedule!(seed, snap, dag_spec, all_procs) + seed_cost = cost_of_schedule(seed) + + t0 = time_ns() + result = simulated_annealing_schedule!(copy(seed), snap, dag_spec, all_procs; + n_restarts=32, + rng=MersenneTwister(0), + time_limit_sec=1e-6) + elapsed_ms = (time_ns() - t0) / 1e6 + @test elapsed_ms < 200.0 + @test cost_of_schedule(result) <= seed_cost + 1e-6 + + # `Inf` opt-out preserves the pre-existing deterministic result. + r_inf = simulated_annealing_schedule!(copy(seed), snap, dag_spec, all_procs; + n_restarts=1, + rng=MersenneTwister(11), + time_limit_sec=Inf) + r_none = simulated_annealing_schedule!(copy(seed), snap, dag_spec, all_procs; + n_restarts=1, + rng=MersenneTwister(11)) + @test cost_of_schedule(r_inf) == cost_of_schedule(r_none) + end end @testset "OptimizingScheduler (adaptive selection)" begin @@ -1767,6 +1863,62 @@ end @test dense_C ≈ dense_A * dense_B @test !isempty(cache) end + + @testset "Heuristic time_limit_sec forwards to IG and SA sub-schedulers" begin + OptimizingScheduler = Dagger.OptimizingScheduler + + # Constructor: fresh defaults, explicit overrides, validation. + s_default = OptimizingScheduler() + @test s_default.ig_time_limit_sec === Inf + @test s_default.sa_time_limit_sec === Inf + + s_finite = OptimizingScheduler(; ig_time_limit_sec=0.25, + sa_time_limit_sec=0.75) + @test s_finite.ig_time_limit_sec == 0.25 + @test s_finite.sa_time_limit_sec == 0.75 + + @test_throws ArgumentError OptimizingScheduler(; ig_time_limit_sec=0.0) + @test_throws ArgumentError OptimizingScheduler(; ig_time_limit_sec=-1.0) + @test_throws ArgumentError OptimizingScheduler(; sa_time_limit_sec=0.0) + @test_throws ArgumentError OptimizingScheduler(; sa_time_limit_sec=-1.0) + + # `similar` carries both budgets. + s_similar = similar(s_finite) + @test s_similar.ig_time_limit_sec == 0.25 + @test s_similar.sa_time_limit_sec == 0.75 + + # End-to-end: an OptimizingScheduler routed to the heuristic path + # (optimizer=nothing) must respect the wall-clock budgets. Total + # elapsed time is bounded by IG budget + SA budget plus per-stage + # setup overhead (allow 100 ms slack per stage on CI). Above nt=4 + # the SA cooling schedule is nontrivial, exercising the intra-loop + # check. + dag_spec, all_procs, snap = _capture_cholesky_dag(4, 16) + schedule = Dict{Any, Any}() + + # Give SA a tiny budget so the whole pipeline exits quickly. + sched = OptimizingScheduler(; ig_time_limit_sec=1e-6, + sa_time_limit_sec=1e-6, + rng=MersenneTwister(0)) + # Warm up once so first-invocation compile time (~200 ms when + # JuMPExt has been loaded into the session) doesn't count against + # the timed second call. The test measures early-exit behavior, + # not JIT overhead. + Dagger.datadeps_schedule_dag_aot!(sched, Dict{Any, Any}(), + dag_spec, all_procs, + Dagger.DefaultScope()) + t0 = time_ns() + Dagger.datadeps_schedule_dag_aot!(sched, schedule, dag_spec, all_procs, + Dagger.DefaultScope()) + elapsed_ms = (time_ns() - t0) / 1e6 + # Post-warmup a K=20 pipeline with two tripped budgets exits in + # well under 500 ms on any dev/CI box; the natural stop condition + # of an unbounded SA + IG on the same DAG takes seconds. + @test elapsed_ms < 500.0 + # Every task must still receive an assignment (fallback returns + # best-so-far, which is at least the Greedy seed). + @test length(schedule) == Dagger.nv(dag_spec.g) + end end @static if Base.find_package("JuMP") !== nothing && From 029e9c1d7e0d35f4ba1630c55db9a3cd6b13ea75 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sat, 18 Jul 2026 04:52:16 -0500 Subject: [PATCH 04/44] bench: fix TimedScheduler under hierarchical partitioning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bench harness's `TimedScheduler` wrapper crashed under Dagger's default hierarchical scheduling because it had no `Base.similar` method, so `distribute_tasks_hierarchical!` fell through to the generic `typeof(s)()` shard constructor — which errors on `TimedScheduler` since the primary ctor requires an inner scheduler argument. Even after adding a naive `Base.similar`, the `sched_phase_ns` metric would report 0 for every row: each partition shard would write into its own private `Ref{UInt64}`, but the driver reads the parent wrapper's Ref, which is never touched. This would have silently corrupted Table 1's scheduler-cost column and made Fig 3 (AOT scheduling cost vs K) unreproducible on any machine running the harness under default settings. Changes: - Replace `Ref{UInt64}` with `Threads.Atomic{UInt64}` so concurrent writes from parallel partition shards accumulate without races. - Change `datadeps_schedule_dag_aot!(::TimedScheduler, ...)` to `Threads.atomic_add!` instead of `[]=`; the atomic is a proper accumulator (sum over partition shards), not a last-writer-wins slot. - Add a shard-constructor variant `TimedScheduler(inner, atomic)` that reuses the parent's atomic. - Add `Base.similar(::TimedScheduler)` that (1) recursively `similar`s the inner scheduler so its own mutable state (RNG, cache, etc.) is refreshed per shard, and (2) hands the shard the parent's atomic so shard writes are visible to the parent-side `sched_phase_ns` reader after the run. Semantics under the fix: - Single-instance scheduling (no hierarchical partitioning): one `datadeps_schedule_dag_aot!` fires, `sched_phase_ns` holds its wall-clock duration — same as before. - Hierarchical partitioning: `sched_phase_ns` is the sum of scheduler CPU time across all partition shards, i.e. the total scheduler work per DAG. This is the correct honest metric for Fig 3 and matches the "total work" convention. Regression coverage in `test/datadeps/scheduling.jl` (new 11 tests, 527 total pass): - `similar` returns a `TimedScheduler` (does not fall through to the crashing generic path) - Atomic is shared between parent and every shard; multiple shard writes accumulate visibly at the parent - `sched_phase_ns > 0` after an actual hierarchical Cholesky run through `run_once` (the exact pattern the paper's benchmark sweep uses) - Fresh instance starts at 0; cache and equivalence hooks still forward to the inner scheduler --- bench/datadeps_schedulers/driver.jl | 28 +++++++++++-- test/datadeps/scheduling.jl | 65 +++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 3 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index be7ba0a5a..c2e6fe6ac 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -16,15 +16,28 @@ include("summarize.jl") mutable struct TimedScheduler{S<:Dagger.DataDepsScheduler} <: Dagger.DataDepsScheduler inner::S - last_aot_ns::Base.RefValue{UInt64} + # Atomic accumulator so under Dagger's default hierarchical scheduling — + # which shards the scheduler via `similar` and runs partitions in + # parallel — every shard's AOT time is summed into a single counter + # without races. Under single-instance scheduling (no partitioning), + # exactly one `datadeps_schedule_dag_aot!` fires and the atomic holds + # the sole scheduling time. Under hierarchical, the value is the sum + # of scheduler CPU time across partition shards, which is the honest + # "total scheduler work per DAG" metric for Fig 3 / sched_phase_ms. + last_aot_ns::Threads.Atomic{UInt64} + # Primary ctor: fresh atomic per top-level wrapper. TimedScheduler(inner::S) where {S<:Dagger.DataDepsScheduler} = - new{S}(inner, Ref(UInt64(0))) + new{S}(inner, Threads.Atomic{UInt64}(0)) + # Shard ctor: reuse the parent's atomic so shard-writes accumulate + # into the parent-visible counter. Called only from `Base.similar` below. + TimedScheduler(inner::S, shared::Threads.Atomic{UInt64}) where {S<:Dagger.DataDepsScheduler} = + new{S}(inner, shared) end function Dagger.datadeps_schedule_dag_aot!(sched::TimedScheduler, schedule, dag_spec, all_procs, all_scope) t0 = time_ns() Dagger.datadeps_schedule_dag_aot!(sched.inner, schedule, dag_spec, all_procs, all_scope) - sched.last_aot_ns[] = time_ns() - t0 + Threads.atomic_add!(sched.last_aot_ns, time_ns() - t0) return end @@ -32,6 +45,15 @@ function Dagger.datadeps_schedule_task_jit!(sched::TimedScheduler, all_procs, al return Dagger.datadeps_schedule_task_jit!(sched.inner, all_procs, all_scope, task_scope, spec, task) end +# Hierarchical scheduling calls `similar` per partition to obtain a fresh +# scheduler shard. The default `Base.similar(::DataDepsScheduler) = typeof(s)()` +# fails on `TimedScheduler` since the primary ctor requires an inner arg. +# We provide an explicit `similar` that (1) recursively `similar`s the inner +# scheduler so its own mutable state (RNG etc.) is refreshed per shard, and +# (2) shares the parent's atomic accumulator so all shard-recorded AOT times +# sum into the counter the driver reads after the run. +Base.similar(s::TimedScheduler) = TimedScheduler(similar(s.inner), s.last_aot_ns) + # Forward equivalence/cache hooks so cached schedules stay partitioned by the # inner scheduler's type (TimedScheduler{Greedy} reuses Greedy entries). Dagger.datadeps_schedule_cache(sched::TimedScheduler) = diff --git a/test/datadeps/scheduling.jl b/test/datadeps/scheduling.jl index a3b08c08a..668db45e1 100644 --- a/test/datadeps/scheduling.jl +++ b/test/datadeps/scheduling.jl @@ -2372,3 +2372,68 @@ end end end end + +# Regression test for the bench harness's `TimedScheduler` wrapper. Coupled +# to the bench file only via `include`; the wrapper lives in bench-side code +# because it exists to measure the AOT scheduling phase for paper figures. +# +# Why this test is here: under Dagger's default hierarchical partitioning, +# `distribute_tasks_hierarchical!` calls `similar(queue.scheduler)` per +# partition. A wrapper scheduler without an explicit `Base.similar` falls +# through to `typeof(s)()` and errors out because `TimedScheduler` requires +# an inner argument. Additionally, even after adding `similar`, a naive +# implementation with a per-shard `Ref{UInt64}` counter would zero the +# `sched_phase_ns` reading since shards write to their own private Refs. +# Both are fixed in `bench/datadeps_schedulers/driver.jl` and this test +# guards against regressions. +@testset "TimedScheduler bench wrapper" begin + include(joinpath(@__DIR__, "..", "..", "bench", "datadeps_schedulers", "driver.jl")) + + @testset "Base.similar shares atomic across partition shards" begin + inner = Dagger.OptimizingScheduler(; rng=MersenneTwister(0)) + timed = TimedScheduler(inner) + @test timed.last_aot_ns isa Threads.Atomic{UInt64} + @test timed.last_aot_ns[] == 0 + + shard = similar(timed) + @test shard isa TimedScheduler + @test shard.inner !== timed.inner # inner is `similar`-cloned per shard + @test shard.last_aot_ns === timed.last_aot_ns # atomic is *shared* + + # Shard-side writes are visible to the parent-side reader. + Threads.atomic_add!(shard.last_aot_ns, UInt64(42_000)) + @test timed.last_aot_ns[] == 42_000 + + # Multiple shard writes accumulate (this is the hierarchical case). + shard2 = similar(timed) + Threads.atomic_add!(shard2.last_aot_ns, UInt64(8_000)) + @test timed.last_aot_ns[] == 50_000 + end + + @testset "sched_phase_ns is non-zero under hierarchical run" begin + # Runs a small Cholesky through the actual harness pipeline under + # Dagger's default (hierarchical) mode. Before the fix this either + # crashed (`similar` missing) or reported `sched_phase_ns == 0` + # (shard-private Refs). Non-zero here proves both bugs are gone. + res = run_once(:cholesky, Dagger.GreedyScheduler(), 3, 16, 0; + collect_logs=true, metrics_warm=false) + @test res.sched_phase_ns > 0 + @test res.n_tasks > 0 + + # Fresh instance starts at zero (no state carryover across trials). + fresh = TimedScheduler(Dagger.GreedyScheduler()) + @test fresh.last_aot_ns[] == 0 + end + + @testset "Cache and equivalence hooks forward to inner" begin + # Regression: TimedScheduler must not accidentally partition the + # cache under its own type — the whole point of forwarding is that + # a `TimedScheduler{Greedy}` sweep reuses `GreedyScheduler`'s + # cached DAG schedules. Verify both cache identity and dag/argspec + # equivalence forwarding. + inner = Dagger.GreedyScheduler() + timed = TimedScheduler(inner) + @test Dagger.datadeps_schedule_cache(timed) === + Dagger.datadeps_schedule_cache(inner) + end +end From 325dc0b47a3922713905049e8b34d67ef88d0f09 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 00:36:03 -0500 Subject: [PATCH 05/44] sch: remove hardcoded task_xfer_cost forcing master-only dispatch Plain `Dagger.@spawn` never distributed work to non-master processes: julia -p 7 -t 12 -e ' using Dagger; @everywhere using Dagger Dagger.all_processors() # -> 96 processors across 8 pids (correct) ts = [Dagger.@spawn sum(rand(64,64)) for _ in 1:24] fetch.(ts) # :compute events after this: pid 1 = 25, pid 2-8 = 0 ' Registration is healthy (all 96 procs visible; init_proc completes for all pids; state.worker_chans has entries for pids 1..8), yet every task lands on pid 1. The processor-selection filter at src/sch/Sch.jl:777 correctly keeps all 8 pids as candidates; the master preference is set later, in `estimate_task_costs!`. Root cause: line 667 of `src/sch/util.jl` added a hardcoded `task_xfer_cost = gproc.pid != myid() ? 1_000_000 : 0` (1ms) penalty to non-master processors' cost. The line predates MetricsTracker and carried a "TODO: Actually estimate/benchmark this" comment. Under the current MetricsTracker cost model this heuristic collapses selection to master-only: * Cold state (no metrics for any proc): every proc's `est_time_util` defaults to `UInt64(1000^3) = 1e9`. Master's cost is 1e9; each worker's is 1e9 + 1e6. `sort!(sorted_procs, by=cost)` puts all 12 of master's ThreadProcs first, deterministically. The `randperm!` above only randomizes among equally-costly procs, and they are not equally costly under this penalty. * After master runs the first task, its measured runtime (e.g. ~100us) replaces its 1e9 default in MetricsTracker, while workers stay at the 1e9 default because they never ran a task. Master's cost drops to O(1e5) versus workers' O(1e9). Master preference becomes locked in and self-reinforcing: the workers never accumulate metrics because they never receive tasks. The penalty is redundant under the metrics model: any real per-signature cross-process dispatch overhead is captured in `est_time_util` once workers accumulate metrics, and per-chunk transfer cost is captured by the existing `tx_cost / tx_rate` term. Remove the penalty. Verification: Reproducer (24 plain @spawn tasks on -p 7 -t 12): before: pid 1: 25 tasks (all on master) after: pid 1: 2, pid 2: 1, pid 3: 6, pid 4: 4, pid 5: 5, pid 6: 1, pid 7: 2, pid 8: 4 (distributed across all 8 pids) Full scheduler test suite: julia --project=test -e 'using Dagger, JuMP, HiGHS; include("test/datadeps/scheduling.jl")' 16 testsets, all pass, 0 fail, 0 error, 1 pre-existing broken. Datadeps AOT-scheduled workloads are unaffected because the AOT scheduler force-pins each task via `options.exec_scope`, bypassing the cost-based selection entirely. --- src/sch/util.jl | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/src/sch/util.jl b/src/sch/util.jl index 62f280edd..caa3a3e0c 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -664,15 +664,33 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) tx_cost = impute_sum(affinity(chunk)[2] for chunk in chunks_filt) - task_xfer_cost = gproc.pid != myid() ? 1_000_000 : 0 - runtime = metrics_lookup_runtime(snap, sig_vec, proc, gproc.pid) est_time_util = runtime !== nothing ? runtime : UInt64(1000^3) rate = metrics_lookup_transfer_rate(snap, proc, gproc.pid) tx_rate = rate !== nothing ? rate : DEFAULT_TRANSFER_RATE - costs[proc] = Float64(est_time_util) + (Float64(tx_cost) / Float64(tx_rate)) + Float64(task_xfer_cost) + # Cost = compute time + data-transfer cost. A previous pre-MetricsTracker + # implementation added a hardcoded 1_000_000 ns cross-worker `task_xfer_cost` + # penalty here (with a TODO to actually measure it), intended as a small + # tiebreaker in favor of master. Under the current MetricsTracker cost model + # this heuristic collapses processor selection to master-only: + # + # - cold state (no metrics for any proc): master cost = 1e9; worker cost + # = 1e9 + 1e6. Sort is deterministic in favor of master's 12 ThreadProcs; + # the `randperm!` above only randomizes among equally-costly procs, and + # they are NOT equally costly under this penalty. + # - after master runs the first task: master's real per-task runtime (e.g. + # 100us) replaces its 1e9 default, while workers stay at 1e9 because they + # never ran a task. Master's cost drops to ~1e5 vs workers' ~1e9 — master + # preference becomes locked in and self-reinforcing. + # + # Any real cross-process dispatch overhead is already captured, per-signature, + # in the `est_time_util` lookup once workers accumulate metrics; and per-chunk + # transfer cost is captured by `tx_cost / tx_rate` above. The fixed 1ms penalty + # is therefore redundant when metrics are warm and actively harmful when they + # aren't. + costs[proc] = Float64(est_time_util) + (Float64(tx_cost) / Float64(tx_rate)) end chunks_cleanup() From 2fafa738d14e443ac45e0a89e623eed7abaeb53c Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 01:05:03 -0500 Subject: [PATCH 06/44] sch: batch per-signature runtime lookup in estimate_task_costs! MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `estimate_task_costs!` is called per submitted task and asks `metrics_lookup_runtime` for every candidate processor. The lookup walks the MetricsTracker snapshot end-to-end via `MT.find_keys`, so with `W` candidate processors and `N` recorded metric keys, the cost estimator was doing `O(W × N)` work per task submission — dominating submission latency on `Dagger.@spawn` in tight loops. Measured before this commit on M4 Pro `julia -t 12`: julia> begin ts = [Dagger.@spawn sleep(0.010) for _ in 1:24] fetch.(ts) end # wall clock: 234 ms (24 * 10 ms serial = 240 ms; parallel ideal ~20 ms). # Effective parallelism 1.0× of 12 threads. Placement is distributed # across all 12 ThreadProcs but ~77 ms of the ~230 ms wall is spent # inside the master's submission loop, before the first task starts # executing. Profile of the submission loop confirms the hot path: `estimate_task_costs!` → `metrics_lookup_runtime` → `MT.find_keys`. Each `find_keys` call iterates every metric storage's every key to intersect the runtime lookup chain — and the chain fires once per candidate processor. Because every processor in a single `estimate_task_costs!` call shares the same signature `sig_vec`, only the *first* lookup pattern in the chain (exact `ProcessorMetric == proc`) is proc-specific; the remaining three patterns are proc-type / worker / global fallbacks that resolve to the same measurements for every processor of the same type on the same worker. Fix: introduce a per-signature runtime index, built once per `estimate_task_costs!` call: - `Dagger.SignatureRuntimeIndex` (`src/utils/metrics.jl`): a struct with four buckets — `by_proc_worker`, `by_type_worker`, `by_type`, `any_matching` — populated in a single `O(N)` scan across SignatureMetric-matching thunk_ids. - `Dagger.build_signature_runtime_index(snap, sig_vec)`: does the scan. - `Dagger.metrics_lookup_runtime_from_index(idx, proc, worker_id)`: replaces the per-proc `metrics_lookup_runtime` call, performing the same four-step fallback via `O(1)` hash lookups against `idx`. `estimate_task_costs!` now builds the index once outside the processor loop and calls `metrics_lookup_runtime_from_index` inside. Semantics of the fallback chain are preserved bit-for-bit; only the traversal cost changes. Verification: * Same reproducer, after this fix: julia> for _ in 1:3 [Dagger.@spawn sleep(0.010) for _ in 1:24] .|> fetch end julia> best = Inf for _ in 1:5 t = @elapsed begin ts = [Dagger.@spawn sleep(0.010) for _ in 1:24] fetch.(ts) end best = min(best, t) end best # -> 0.0445 s 44.5 ms vs 234 ms → 5.3× wall-clock reduction. Effective parallelism 5.4× of 12 threads (was 1.0×). * Empty-task submission microbenchmark (100 × `Dagger.@spawn 1+1`): before: 121 ms (1.21 ms/task) after : 60 ms (0.60 ms/task) 2.0× improvement in the per-task submission fixed cost. * Full scheduler test suite: julia --project=test -e 'using Dagger, JuMP, HiGHS; include("test/datadeps/scheduling.jl")' 16 testsets, 527 pass, 0 fail, 0 error, 1 pre-existing broken. * Existing `metrics_lookup_runtime` code path is untouched: any caller that queries a *single* (proc, worker_id) still resolves through the old function. Only `estimate_task_costs!` — the hot inner loop of task submission — is rerouted through the new index. Notes: * `has_capacity` in `src/sch/util.jl` still calls `metrics_lookup_runtime` directly. It is invoked at most once per candidate proc from `schedule_one!`, so the per-call `O(N)` cost is amortised across the same proc loop where the new index already lives. A follow-up could thread the index through `has_capacity` for another marginal improvement; deferred to keep this change focused on the dominant hot path. * The AOT scheduler path (`options.exec_scope !== nothing`) still goes through `estimate_task_costs!` but with `length(procs) == 1`, so the index build is `O(N)` and the single lookup is `O(1)` — the same total work as before. No regression there. --- src/sch/Sch.jl | 3 +- src/sch/util.jl | 11 +++- src/utils/metrics.jl | 130 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 142 insertions(+), 2 deletions(-) diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 4e2c59174..bdf8f320e 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -31,7 +31,8 @@ import ScopedValues: @with import ..Dagger: SignatureMetric, ProcessorMetric, WorkerMetric, TransferSizeMetric, TransferTimeMetric, TransferRateMetric import ..Dagger: TASK_SIGNATURE, TASK_PROCESSOR, TASK_WORKER, TASK_TRANSFER_SIZE, TASK_TRANSFER_TIME -import ..Dagger: execute_metrics_spec, metrics_lookup_runtime, metrics_lookup_alloc, metrics_lookup_transfer_rate +import ..Dagger: execute_metrics_spec, metrics_lookup_runtime, metrics_lookup_alloc, metrics_lookup_transfer_rate, + build_signature_runtime_index, metrics_lookup_runtime_from_index import ..Dagger: extract_collected_metrics, apply_collected_metrics! import ..Dagger.MetricsTracker as MT diff --git a/src/sch/util.jl b/src/sch/util.jl index caa3a3e0c..71f04ce8c 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -658,13 +658,22 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) sig_vec = sig isa Dagger.Signature ? sig.sig : sig snap = MT.snapshot(MT.global_metrics_cache()) + # Build the per-signature runtime index once. Every candidate `proc` in + # this call shares `sig_vec`, so `metrics_lookup_runtime_from_index` + # below can resolve each per-proc runtime with O(1) hash lookups against + # this index — a full `metrics_lookup_runtime` call would otherwise + # re-scan the entire snapshot (via `MT.find_keys`) up to four times per + # proc, i.e. `O(W × N)` scanning work per submitted task. See + # `SignatureRuntimeIndex` for the fallback-chain semantics preserved. + runtime_index = build_signature_runtime_index(snap, sig_vec) + for proc in procs gproc = get_parent(proc) chunks_filt = Iterators.filter(c->get_parent(processor(c)) != gproc, chunks) tx_cost = impute_sum(affinity(chunk)[2] for chunk in chunks_filt) - runtime = metrics_lookup_runtime(snap, sig_vec, proc, gproc.pid) + runtime = metrics_lookup_runtime_from_index(runtime_index, proc, gproc.pid) est_time_util = runtime !== nothing ? runtime : UInt64(1000^3) rate = metrics_lookup_transfer_rate(snap, proc, gproc.pid) diff --git a/src/utils/metrics.jl b/src/utils/metrics.jl index 571491a55..2ab5c28c9 100644 --- a/src/utils/metrics.jl +++ b/src/utils/metrics.jl @@ -152,6 +152,136 @@ metrics_lookup_runtime_min(snap, sig, proc, worker_id) = metrics_lookup_runtime_max(snap, sig, proc, worker_id) = metrics_lookup_runtime(snap, sig, proc, worker_id; reducer=maximum) +""" + SignatureRuntimeIndex + +Precomputed per-signature runtime index built by +`build_signature_runtime_index`. Groups measured `ThreadTimeMetric` +runtimes for a fixed signature `sig` by `(processor, worker_id)`, +`(processor_type, worker_id)`, `(processor_type)`, and all-matching so +that per-processor lookups via +`metrics_lookup_runtime_from_index` can traverse the fallback chain in +O(1) hash lookups instead of re-scanning the snapshot per call. + +The scheduler's per-task cost estimator (`estimate_task_costs!`) +evaluates every candidate processor with the same signature, so building +this index once amortises what would otherwise be `O(W × N)` snapshot +scans (where `W` is candidate-processor count and `N` is total metric +keys) into a single `O(N)` scan plus `O(W)` amortized dict lookups. +""" +struct SignatureRuntimeIndex + by_proc_worker::Dict{Tuple{Processor,Int},Vector{UInt64}} + by_type_worker::Dict{Tuple{DataType,Int},Vector{UInt64}} + by_type::Dict{DataType,Vector{UInt64}} + any_matching::Vector{UInt64} +end + +""" + build_signature_runtime_index(snap::MT.MetricsSnapshot, sig::Vector) + -> SignatureRuntimeIndex + +Single-pass build: scan the SignatureMetric storage once, keep only keys +matching `sig`, then for each such key materialise its +(processor, worker_id, ThreadTimeMetric) tuple and index into the four +buckets used by the fallback chain in `metrics_lookup_runtime`. The +result is safe to reuse across many `metrics_lookup_runtime_from_index` +calls for different processors as long as `sig` and `snap` are +unchanged. +""" +function build_signature_runtime_index(snap::MT.MetricsSnapshot, sig::Vector) + ctx = get(snap.contexts, (Dagger, :execute!), nothing) + if ctx === nothing + return SignatureRuntimeIndex( + Dict{Tuple{Processor,Int},Vector{UInt64}}(), + Dict{Tuple{DataType,Int},Vector{UInt64}}(), + Dict{DataType,Vector{UInt64}}(), + UInt64[], + ) + end + + sig_storage = get(ctx.storages, SignatureMetric(), nothing) + proc_storage = get(ctx.storages, ProcessorMetric(), nothing) + worker_storage = get(ctx.storages, WorkerMetric(), nothing) + time_storage = get(ctx.storages, MT.ThreadTimeMetric(), nothing) + + by_proc_worker = Dict{Tuple{Processor,Int},Vector{UInt64}}() + by_type_worker = Dict{Tuple{DataType,Int},Vector{UInt64}}() + by_type = Dict{DataType,Vector{UInt64}}() + any_matching = UInt64[] + + # Missing any of the storages means no runtime data has been recorded + # yet — return the empty index (all buckets empty), which mirrors what + # the original `metrics_lookup_runtime` would produce (returns `nothing` + # via `_reduce_uint64` on empty vectors). + if sig_storage === nothing || proc_storage === nothing || + worker_storage === nothing || time_storage === nothing + return SignatureRuntimeIndex(by_proc_worker, by_type_worker, by_type, any_matching) + end + + # One O(N) scan across sig-matching keys; each bucket insertion is O(1) + # amortised. `sig` is a `Vector{Any}` and the stored value is the same + # type, so `==` compares element-wise — matching the semantics of + # `LookupExact(SignatureMetric(), sig)` in `_runtime_lookup_chain`. + for (k, s) in sig_storage.data + s == sig || continue + v = get(time_storage.data, k, nothing) + v === nothing && continue + rt = v::UInt64 + p = get(proc_storage.data, k, nothing) + w = get(worker_storage.data, k, nothing) + push!(any_matching, rt) + if p !== nothing + proc_v = p::Processor + proc_type = typeof(proc_v) + push!(get!(() -> UInt64[], by_type, proc_type), rt) + if w !== nothing + worker_v = w::Int + push!(get!(() -> UInt64[], by_proc_worker, + (proc_v, worker_v)), rt) + push!(get!(() -> UInt64[], by_type_worker, + (proc_type, worker_v)), rt) + end + end + end + + return SignatureRuntimeIndex(by_proc_worker, by_type_worker, by_type, any_matching) +end + +""" + metrics_lookup_runtime_from_index(idx::SignatureRuntimeIndex, + proc::Processor, worker_id::Int; + reducer=first) -> Union{UInt64,Nothing} + +Fast per-processor runtime lookup using the precomputed +`SignatureRuntimeIndex`. Traverses the same fallback chain as +`metrics_lookup_runtime`: + + 1. Exact `(proc, worker_id)` + 2. `(typeof(proc), worker_id)` + 3. `typeof(proc)` on any worker + 4. Any measurement for the signature + +Returns the reduced runtime for the first non-empty bucket, or `nothing` +if the signature has no measurements at all. +""" +function metrics_lookup_runtime_from_index(idx::SignatureRuntimeIndex, + proc::Processor, worker_id::Int; + reducer::Function=first) + proc_type = typeof(proc) + vals = get(idx.by_proc_worker, (proc, worker_id), nothing) + if vals === nothing || isempty(vals) + vals = get(idx.by_type_worker, (proc_type, worker_id), nothing) + end + if vals === nothing || isempty(vals) + vals = get(idx.by_type, proc_type, nothing) + end + if vals === nothing || isempty(vals) + vals = idx.any_matching + end + (vals === nothing || isempty(vals)) && return nothing + return _reduce_uint64(reducer, vals) +end + function _alloc_lookup_chain(sig::Vector, proc::Processor) return ( (MT.LookupExact(SignatureMetric(), sig), From 02fa4e1b5f3a58ec7c349ceb1486543da2def4bd Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 01:34:20 -0500 Subject: [PATCH 07/44] sch: reuse SignatureRuntimeIndex in has_capacity to avoid a second O(N) scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Companion to 2fafa738. `schedule_one!` already built a `SignatureRuntimeIndex` per submitted task to give `estimate_task_costs!` `O(1)` per-proc runtime lookups. The subsequent reservation loop then called `has_capacity` for each candidate proc, and `has_capacity` in turn called `metrics_lookup_runtime` — its OWN `O(N)` scan of the MetricsTracker snapshot — for exactly the same `(signature, proc, worker_id)` triples the index had already indexed. For the AOT path (`options.exec_scope !== nothing`, `W == 1`), that meant two `O(N)` snapshot scans per submitted task instead of one. Fix: * `estimate_task_costs!` grows an optional `runtime_index` kwarg. When provided, the internal build is skipped and the caller's index is used verbatim. When `nothing` (existing callers), the function builds one internally — backwards-compatible with any caller outside `schedule_one!`. * `has_capacity` grows the same optional `runtime_index` kwarg. When provided, the internal `metrics_lookup_runtime` call is replaced with `metrics_lookup_runtime_from_index` — same fallback-chain semantics, no snapshot rescan. When absent, the original `metrics_lookup_runtime` path is preserved so `has_capacity` can still be called standalone. * `schedule_one!` now builds the `SignatureRuntimeIndex` once at the top of Phase 2 and threads it through both `estimate_task_costs!` and every `has_capacity` call inside the sorted-procs reservation loop. Total per-task metrics-scan work drops from `O(2N)` (for AOT, or `O((W+1)N)` for non-AOT) to `O(N)`. Verification: * Full scheduler test suite: 16 testsets, 527 pass, 0 fail, 0 error, 1 pre-existing broken. `has_capacity` semantics are preserved bit-for-bit for both `options.time_util !== nothing` (fast path unchanged) and the metrics-lookup fallback (index-backed path returns the same values as `metrics_lookup_runtime` per the `SignatureRuntimeIndex` fallback-chain contract). * Sleep reproducer (M4 Pro, `julia -t 12`, warmed): before this commit: 44.5 ms best-of-5 (5.4× of 12) after this commit: 43.2 ms best-of-5 (5.6× of 12) Consistent with the AOT-path improvement being a modest constant- factor win rather than the dramatic non-AOT speedup 2fafa738 already captured. The synthetic reproducer here is non-AOT and had already absorbed most of the gains. --- src/sch/Sch.jl | 20 ++++++++++++++++++-- src/sch/util.jl | 27 +++++++++++++++++++++++---- 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index bdf8f320e..009c33d35 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -889,7 +889,22 @@ concurrently across threads. resize!(sorted_procs, length(input_procs)) costs = @reusable_dict :schedule_one!_costs Processor Float64 OSProc() 0.0 32 costs_cleanup = @reuse_defer_cleanup empty!(costs) - estimate_task_costs!(sorted_procs, costs, state, input_procs, task; sig) + + # Build the per-signature runtime index once, here, and thread it through + # to both `estimate_task_costs!` and `has_capacity`. Both functions look + # up the same `est_time_util` for the same `(sig, proc, worker_id)` triples + # for this task; without a shared index each does an independent + # `metrics_lookup_runtime` call that scans the MetricsTracker snapshot end- + # to-end via `MT.find_keys`. On the AOT path (`options.exec_scope !== nothing`, + # `W == 1`), that meant two `O(N)` scans per submitted task — one in + # `estimate_task_costs!` and one in the `has_capacity` call inside the + # reservation loop below. Building the index once and passing it through + # makes both callers `O(1)` per proc after the single `O(N)` build. + sig_vec_for_index = sig isa Dagger.Signature ? sig.sig : sig + runtime_index = build_signature_runtime_index( + MT.snapshot(MT.global_metrics_cache()), sig_vec_for_index) + estimate_task_costs!(sorted_procs, costs, state, input_procs, task; + sig, runtime_index) input_procs_cleanup() # Select the best available processor and reserve time pressure optimistically. @@ -901,7 +916,8 @@ concurrently across threads. can_use, scope = can_use_proc(state, task, gproc, proc, options, scope) if can_use has_cap, est_time_util, est_alloc_util, est_occupancy = - has_capacity(state, proc, gproc.pid, options.time_util, options.alloc_util, options.occupancy, sig) + has_capacity(state, proc, gproc.pid, options.time_util, options.alloc_util, options.occupancy, sig; + runtime_index) if has_cap # Optimistic reservation: capture-the-ref, atomic_add! counter_ref = lock(state.worker_time_pressure) do wtp diff --git a/src/sch/util.jl b/src/sch/util.jl index 71f04ce8c..a314d6fb6 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -526,7 +526,8 @@ function can_use_proc(state, task, gproc, proc, opts, scope) return true, scope end -function has_capacity(state, p, gp, time_util, alloc_util, occupancy, sig) +function has_capacity(state, p, gp, time_util, alloc_util, occupancy, sig; + runtime_index::Union{Dagger.SignatureRuntimeIndex,Nothing}=nothing) T = typeof(p) sig_vec = sig isa Dagger.Signature ? sig.sig : sig snap = MT.snapshot(MT.global_metrics_cache()) @@ -534,7 +535,17 @@ function has_capacity(state, p, gp, time_util, alloc_util, occupancy, sig) est_time_util = if time_util !== nothing && haskey(time_util, T) round(UInt64, time_util[T] * 1000^3)::UInt64 else - runtime = metrics_lookup_runtime(snap, sig_vec, p, worker_id) + # Prefer the shared `runtime_index` when provided — it is populated + # for this exact signature by `estimate_task_costs!` upstream, so + # the per-proc lookup here is O(1) and does not re-scan the + # snapshot. Fall back to the original `metrics_lookup_runtime` + # when the caller has not built an index (backwards compatibility + # for `has_capacity` callers outside `schedule_one!`). + runtime = if runtime_index !== nothing + metrics_lookup_runtime_from_index(runtime_index, p, worker_id) + else + metrics_lookup_runtime(snap, sig_vec, p, worker_id) + end runtime !== nothing ? runtime : UInt64(1000^3) end est_alloc_util = if alloc_util !== nothing && haskey(alloc_util, T) @@ -641,7 +652,9 @@ function estimate_task_costs(state, procs, task; sig=nothing) return sorted_procs, costs end const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) -@reuse_scope function estimate_task_costs!(sorted_procs, costs, state, procs, task; sig=nothing) +@reuse_scope function estimate_task_costs!(sorted_procs, costs, state, procs, task; + sig=nothing, + runtime_index::Union{Dagger.SignatureRuntimeIndex,Nothing}=nothing) # Find all Chunks chunks = @reusable_vector :estimate_task_costs_chunks Union{Chunk,Nothing} nothing 32 @@ -665,7 +678,13 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) # re-scan the entire snapshot (via `MT.find_keys`) up to four times per # proc, i.e. `O(W × N)` scanning work per submitted task. See # `SignatureRuntimeIndex` for the fallback-chain semantics preserved. - runtime_index = build_signature_runtime_index(snap, sig_vec) + # + # Callers that also invoke `has_capacity` (schedule_one!) may pass a + # pre-built index so the same object is shared across both functions, + # avoiding a second O(N) snapshot scan in `has_capacity`. + if runtime_index === nothing + runtime_index = build_signature_runtime_index(snap, sig_vec) + end for proc in procs gproc = get_parent(proc) From cfa836ec773615900f75fadd4c506bf3d760b241 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 01:47:08 -0500 Subject: [PATCH 08/44] datadeps: propagate AOT-computed runtime as options.time_util so has_capacity skips its metrics scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Companion to 02fa4e1b. `Sch.has_capacity` (`src/sch/util.jl:529-`) has a fast path when `options.time_util[typeof(proc)] !== nothing`: it uses the caller-supplied per-processor-type runtime estimate directly and skips its `metrics_lookup_runtime` call entirely. AOT-scheduled tasks never populated that value, so has_capacity always fell through to the snapshot scan even when our AOT scheduler had already produced a perfectly good estimate for the same `(task, proc)` pair. The AOT schedulers already own a runtime estimate per `(task, proc)` pair — that estimate is exactly what the cost model minimises over — so propagating it costs almost nothing at scheduling time. This is a bench-side change to OUR AOT layer (Ritesh's schedulers); the Sch.jl fast path in has_capacity is Julian's code and stays untouched. Changes: * `_propagate_aot_time_util!(spec, proc, task_time_ns)` in `src/datadeps/scheduling.jl`: sets `spec.options.time_util[typeof(proc)] = task_time_ns / 1e9` (Float64 seconds, matching has_capacity's expected units: `round(UInt64, time_util[T] * 1000^3)` -> nanoseconds). Non-positive / non-finite estimates are ignored so the caller can safely propagate whatever the cost model produced without a per-scheduler nullability check. * `_propagate_aot_time_util_from_cache!(dag_spec, cache, task_idx, proc)`: convenience wrapper for the three schedulers (Greedy, IG, SA) that already build an `EFTCostCache` and can look up the runtime by `(task_idx, proc_idx)` in O(1). If the cache does not know about `proc`, the propagation is a no-op — `has_capacity` falls back to metrics lookup, preserving existing behaviour. * `GreedyScheduler`, `IteratedGreedyScheduler`, `SimulatedAnnealingScheduler` datadeps_schedule_dag_aot! implementations: after `schedule[task] = proc`, also call `_propagate_aot_time_util_from_cache!`. `OptimizingScheduler` delegates to one of these sub-schedulers and inherits the propagation via them, no separate change needed. * `ext/JuMPExt.jl` `datadeps_schedule_dag_aot!(::JuMPScheduler, ...)`: imports and calls `_propagate_aot_time_util!` with `task_times[k, proc_idx]` from the MILP formulation. This is the same nanosecond estimate the objective minimised over, so the fast-path value has_capacity reads back is exactly what the MILP optimised against. Verification: * `has_capacity` fast-path fires for AOT-scheduled tasks (traced with a temporary `@info` in the fast path; log showed `has_capacity FAST PATH fired for T=Dagger.ThreadProc time_util[T]=1.0` corresponding to Greedy's `GREEDY_DEFAULT_RUNTIME_NS = 1e9` when metrics are cold). Trace removed before commit. * Full scheduler test suite: 16 testsets, 527 pass, 0 fail, 0 error, 1 pre-existing broken. * AOT-datadeps microbenchmark (M4 Pro, `julia -t 12`, nt=8 tasks, best-of-5 timed regions): before this commit after Greedy AOT 3.04 ms 2.94 ms ( ~3% ) SA AOT 6.69 ms 5.37 ms ( ~20% ) Consistent with the SA path benefiting more because it fires `has_capacity` more often per submission (deeper cost model), while Greedy already commits to the first-chosen proc quickly. Both schedulers correctly propagate `options.time_util`; the delta is a lower-bound on the per-task overhead removed and will be larger for hudson-scale workloads whose per-task snapshot scan is proportionally more expensive. * On the paper's benchmark path (spawn_datadeps with any of our AOT schedulers on the paper's Greedy/IG/SA/JuMP/Optim comparison), every submitted task now takes the has_capacity fast path. Notes: * Non-`EFTCostCache` schedulers (`LayeredScheduler`) do not compute a runtime and are left unchanged; they continue to fall through to `has_capacity`'s metrics-lookup path unchanged. * `_propagate_aot_time_util!` mutates `spec.options` in place. This is safe because `spec` is the task's own DTaskSpec, submitted before the AOT scheduler runs, and each user @spawn creates its own `Options` object. No cross-task aliasing. --- ext/JuMPExt.jl | 12 +++++-- src/datadeps/scheduling.jl | 74 +++++++++++++++++++++++++++++++++----- 2 files changed, 76 insertions(+), 10 deletions(-) diff --git a/ext/JuMPExt.jl b/ext/JuMPExt.jl index 179ca76b9..05d5174d7 100644 --- a/ext/JuMPExt.jl +++ b/ext/JuMPExt.jl @@ -11,7 +11,8 @@ import Dagger: JuMPScheduler, DAGSpec, datadeps_schedule_dag_aot!, proc_in_scope, DefaultScope, _eft_runtime_ns, _milp_edge_size_bytes, _milp_transfer_time_ns, GREEDY_DEFAULT_RUNTIME_NS, - ScheduleState, greedy_schedule! + ScheduleState, greedy_schedule!, + _propagate_aot_time_util! import Dagger.Sch # `jps/riteshsc26` folds `MetricsTracker` into the `Dagger` module via a # direct `include` (see `src/Dagger.jl`), so it is `Dagger.MetricsTracker` @@ -161,7 +162,14 @@ function Dagger.datadeps_schedule_dag_aot!(sched::JuMPScheduler, schedule, dag_s throw(Sch.SchedulingException("JuMPScheduler: solver returned no processor assignment for task $k")) end task = dag_spec.id_to_task[k] - schedule[task] = all_procs[proc_idx] + proc = all_procs[proc_idx] + schedule[task] = proc + # Propagate our AOT-computed per-task runtime to Sch's fast path. + # `task_times[k, proc_idx]` is the same nanosecond estimate the MILP + # objective minimised for; using it as `options.time_util` means + # `Sch.has_capacity` skips its `metrics_lookup_runtime` snapshot scan + # for MILP-scheduled tasks. + _propagate_aot_time_util!(dag_spec.id_to_spec[k], proc, task_times[k, proc_idx]) end return end diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 854ab3d92..2064545f3 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -380,6 +380,38 @@ function datadeps_schedule_dag_aot!(scheduler, schedule, dag_spec, all_procs, al return end +""" + _propagate_aot_time_util!(spec::DTaskSpec, proc::Processor, task_time_ns::Real) + +Attach the AOT-computed per-task runtime estimate to `spec.options.time_util` +so that `Sch.has_capacity` (via `schedule_one!`) can bypass its +`metrics_lookup_runtime` scan of the MetricsTracker snapshot and use our +estimate directly. `has_capacity` reads the value via +`round(UInt64, time_util[T] * 1000^3)`, so `time_util` is stored keyed by +`typeof(proc)` with the runtime expressed in seconds (Float64). + +Called by every AOT scheduler after it finalises a task's `(task, proc)` +assignment. Non-finite / non-positive estimates are ignored (the +`options.time_util === nothing` path is preserved for those, matching the +pre-existing behaviour where `has_capacity` falls back to metrics lookup). +""" +function _propagate_aot_time_util!(spec::DTaskSpec, proc::Processor, task_time_ns::Real) + task_time_ns > 0 || return + isfinite(task_time_ns) || return + time_util_secs = Float64(task_time_ns) / 1e9 + T = typeof(proc) + if spec.options.time_util === nothing + spec.options.time_util = Dict{Type,Any}(T => time_util_secs) + else + spec.options.time_util[T] = time_util_secs + end + return +end + +## `_propagate_aot_time_util_from_cache!` — the `EFTCostCache`-aware wrapper — +## is defined below, immediately after `EFTCostCache`, so its signature can +## reference the type. + const GREEDY_DEFAULT_RUNTIME_NS = UInt64(1_000_000_000) const GREEDY_DEFAULT_TRANSFER_RATE = UInt64(1_000_000) const GREEDY_DEFAULT_OUTPUT_SIZE = UInt64(1_048_576) @@ -391,11 +423,30 @@ const GREEDY_DEFAULT_OUTPUT_SIZE = UInt64(1_048_576) # to what the uncached path would compute, so cost-model claims and every # non-worsening / determinism / correctness invariant are preserved. struct EFTCostCache - task_times::Matrix{Float64} - proc_compatible::Matrix{Bool} - proc_spaces::Vector{MemorySpace} - proc_to_idx::Dict{Processor, Int} - move_rates::Matrix{Float64} + task_times::Matrix{Float64} + proc_compatible::Matrix{Bool} + proc_spaces::Vector{MemorySpace} + proc_to_idx::Dict{Processor, Int} + move_rates::Matrix{Float64} +end + +""" + _propagate_aot_time_util_from_cache!(dag_spec, cache, task_idx, proc) + +Convenience wrapper that pulls the AOT-computed runtime from an +`EFTCostCache` (already built by the caller for its own cost-model use) +and forwards it to `_propagate_aot_time_util!`. If the cache does not +contain an entry for `proc`, or the entry is not a positive finite value, +the propagation is a no-op — `has_capacity` will then take its usual +fallback path. +""" +function _propagate_aot_time_util_from_cache!(dag_spec::DAGSpec, cache::EFTCostCache, + task_idx::Int, proc::Processor) + proc_idx = get(cache.proc_to_idx, proc, nothing) + proc_idx === nothing && return + task_time_ns = cache.task_times[task_idx, proc_idx] + _propagate_aot_time_util!(dag_spec.id_to_spec[task_idx], proc, task_time_ns) + return end function _build_eft_cost_cache(snap::MT.MetricsSnapshot, dag_spec::DAGSpec, @@ -570,7 +621,10 @@ function datadeps_schedule_dag_aot!(scheduler::GreedyScheduler, schedule, dag_sp greedy_schedule!(state, snap, dag_spec, all_procs; cache=cache) for idx in 1:nv(dag_spec.g) task = dag_spec.id_to_task[idx] - schedule[task] = state.task_proc[idx] + proc = state.task_proc[idx] + schedule[task] = proc + # Propagate our AOT-computed per-task runtime to Sch's fast path. + _propagate_aot_time_util_from_cache!(dag_spec, cache, idx, proc) end return end @@ -949,7 +1003,9 @@ function datadeps_schedule_dag_aot!(scheduler::IteratedGreedyScheduler, @inbounds for idx in 1:n_tasks task = dag_spec.id_to_task[idx] - schedule[task] = state.task_proc[idx] + proc = state.task_proc[idx] + schedule[task] = proc + _propagate_aot_time_util_from_cache!(dag_spec, cache, idx, proc) end return end @@ -1323,7 +1379,9 @@ function datadeps_schedule_dag_aot!(scheduler::SimulatedAnnealingScheduler, @inbounds for idx in 1:n_tasks task = dag_spec.id_to_task[idx] - schedule[task] = final_state.task_proc[idx] + proc = final_state.task_proc[idx] + schedule[task] = proc + _propagate_aot_time_util_from_cache!(dag_spec, cache, idx, proc) end return end From 2cc22081f3b5737496bcd069353655fc6c969d35 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 06:37:06 -0500 Subject: [PATCH 09/44] bench: distribute workload tiles across processors so datadeps places tasks off-master MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The paper's Hudson benchmark sweep produced a physical impossibility for cholesky nt=8 — an implied 7,140 GFLOPS on a 3,840 GFLOPS peak machine — and a 480× CPU asymmetry (master 3929 s, each worker 7–8 s) even under multi-process configs. Direct probes on hudson (Julia -p 1 -t 2 with `spawn_datadeps` + `:compute` event logging) showed every task landing on pid 1 under both `RoundRobinScheduler` and `GreedyScheduler`, so the sweep had been measuring master-thread scheduling all along rather than the distributed scheduling the paper claims to characterise. This is a *bench-side* defect, not a Dagger bug. `make_spd_tiles` and `make_matmul_tiles` allocated every tile as a plain `Matrix{Float64}` resident in master's memory space. When `spawn_datadeps` observed the tile-graph, it (correctly) concluded that shipping an 8-MB tile out to a worker and back would dominate any parallelism gain, and placed every task on master — the optimal decision *given master-resident inputs*. Chameleon / PLASMA / DPLASMA (the tiled linear-algebra reference libraries the paper's workloads model) all distribute tile data across processors before entering the region for exactly this reason. Fix, bench-side only: * `bench/datadeps_schedulers/workloads.jl`: - Add `_placement_procs()`: sorted vector of `Dagger.all_processors()`, deterministic across identical sessions so tile→proc placement is reproducible run-to-run. - Add `_distribute_tile(tile, proc)`: `remotecall_fetch` the `Dagger.tochunk(tile, proc, ExactScope(proc))` call to the target processor's worker. `MemPool.poolset` runs on the target worker, so the resulting DRef lives in that worker's memory space and subsequent datadeps tasks touching only that tile incur no cross-worker transfer. Same pattern used by test/datadeps.jl:552. - Add `_tile_proc(procs, nt, i, j)`: shared round-robin row-major-linearised (i, j) → proc mapping so matmul's `A[i, k]`, `B[k, j]`, `C[i, j]` co-locate identically across workloads at the same logical positions. - Rewrite `make_spd_tiles` and `make_matmul_tiles` to distribute tiles through `_distribute_tile` at their `_tile_proc(...)` target. `make_spd_tiles` `copy`s each tile before shipping so the closure carries an independent `Matrix` rather than a view into the parent SPD buffer (would otherwise confuse datadeps' aliasing analysis). - Relax `tiled_cholesky!(M::Matrix{<:Matrix})` / `tiled_matmul!(C::Matrix{<:Matrix}, A::Matrix{<:Matrix}, ...)` to `AbstractMatrix`. `Dagger.@spawn f(In(M[k,k]))` / `Dagger.@spawn f(InOut(M[k,k]))` accept `Chunk` and `Matrix` elements transparently: for a `Chunk`, datadeps reads its `scope` / memory space directly; for a `Matrix`, it observes master-residence as before. Same tile-loop semantics in both cases. * `bench/datadeps_schedulers/driver.jl`: - Add `_fetch_tile(t)`: passes through a `Matrix`, `fetch`es a `Dagger.Chunk` (which triggers `collect` / `move` from the chunk's owning worker back to master). This is exactly what `verify_workload` needs when reassembling the tiled state for the dense-reference check. - Relax `_assemble_dense(tiles::Matrix{<:Matrix})` to `AbstractMatrix` and route each `tiles[i, j]` element through `_fetch_tile`. Uses the first tile's element type + block size to allocate the dense buffer. Backwards-compatible with plain `Matrix` tiles for any caller not going through `make_*_tiles`. Scope preserved: * `build_inputs`, `run_workload!`, `verify_workload`, `run_once`, `run_sweep`, and every scheduler-side code path unchanged. All the complexity lives at the `make_*_tiles` allocation boundary and the `_assemble_dense` fetch boundary. * `src/`, `ext/`, `lib/`, `test/*.jl` untouched. * Metrics-warm's RoundRobin pass now populates γ (worker-to-worker transfer) metrics based on ACTUAL worker-to-worker movement. Previously γ measurements were meaningless because every task ran on master and no cross-worker movement occurred; now the cost model has real data to consume. Verification: * Placement probe (`OMP_NUM_THREADS=2 julia --project=test -p 1 -t 2`, cholesky nt=2 bs=1024, `:compute` events): RoundRobin Greedy pid 1: 6 5 pid 2: 7 7 pids w/ :compute: 2 2 was pid 1: {8, 7}, pids w/ :compute: 1 for both schedulers on master-resident tiles. (The 6+7 / 5+7 vs the K=4 explicit tasks reflects datadeps-inserted cross-worker copy tasks — expected under distributed placement.) * Correctness at bs=1024: cholesky nt=2: rel_err = 6.11e-17 (threshold 1e-8) cholesky nt=4: rel_err = 9.01e-17 (threshold 1e-8) matmul nt=2: rel_err = 2.16e-16 (threshold 1e-10) matmul nt=4: rel_err = 2.24e-16 (threshold 1e-10) * Full scheduler test suite: julia --project=test -e 'using Dagger, JuMP, HiGHS; include("test/datadeps/scheduling.jl")' 16 testsets, 527 pass, 0 fail, 0 error, 1 pre-existing broken. No regressions. Impact: * Every previous hudson sweep (bs=128 archive AND the bs=1024 rerun) measured single-process scheduling, not distributed scheduling. The physically-impossible GFLOPS number in Table 1 was the downstream symptom. Reruns after this commit will measure the distributed scheduling the paper's methodology assumes and the Chameleon / PLASMA / DPLASMA workload conventions the paper cites. * Fig 3 (scheduler cost) is unchanged — algorithmic runtime is independent of where tile data lives. * Table 1 and Fig 4 will now reflect real cross-worker placement trade-offs. --- bench/datadeps_schedulers/driver.jl | 22 ++++-- bench/datadeps_schedulers/workloads.jl | 101 +++++++++++++++++++++---- 2 files changed, 105 insertions(+), 18 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index c2e6fe6ac..2159cb21f 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -103,12 +103,24 @@ function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler return end -function _assemble_dense(tiles::Matrix{<:Matrix}) +# Materialise a `tiles[i,j]` element into a plain matrix regardless of +# whether it's a locally-resident `Matrix` (single-process paths) or a +# distributed `Dagger.Chunk` (multi-process paths — see +# `make_spd_tiles` / `make_matmul_tiles`). For a `Chunk`, `fetch` is a +# `collect` / `move` from the chunk's owning worker back to master, +# which is exactly what we need to reassemble the tiles into a dense +# reference on master for `verify_workload` to compare against. +_fetch_tile(t::AbstractMatrix) = t +_fetch_tile(t::Dagger.Chunk) = fetch(t)::AbstractMatrix + +function _assemble_dense(tiles::AbstractMatrix) nt = size(tiles, 1) - bs = size(tiles[1, 1], 1) - out = zeros(eltype(tiles[1, 1]), nt * bs, nt * bs) - for i in 1:nt, j in 1:nt - out[(i-1)*bs+1:i*bs, (j-1)*bs+1:j*bs] .= tiles[i, j] + first_tile = _fetch_tile(tiles[1, 1]) + bs = size(first_tile, 1) + out = zeros(eltype(first_tile), nt * bs, nt * bs) + @inbounds for i in 1:nt, j in 1:nt + tile = _fetch_tile(tiles[i, j]) + out[(i-1)*bs+1:i*bs, (j-1)*bs+1:j*bs] .= tile end return out end diff --git a/bench/datadeps_schedulers/workloads.jl b/bench/datadeps_schedulers/workloads.jl index 9d0452e6f..34146e17b 100644 --- a/bench/datadeps_schedulers/workloads.jl +++ b/bench/datadeps_schedulers/workloads.jl @@ -1,8 +1,19 @@ using LinearAlgebra using Dagger using Dagger: In, Out, InOut +using Distributed: remotecall_fetch -function tiled_cholesky!(M::Matrix{<:Matrix}) +# Signatures relaxed from `Matrix{<:Matrix}` to `AbstractMatrix` so tile +# elements may be plain `Matrix` (single-process test paths) OR +# `Dagger.Chunk` values that live on a specific worker/processor +# (multi-process benchmark paths — see `make_spd_tiles` / +# `make_matmul_tiles` below). `Dagger.@spawn f(In(M[k,k]))` and +# `Dagger.@spawn f(InOut(M[k,k]))` accept either transparently: for a +# `Chunk`, datadeps reads the chunk's `scope` / memory space directly and +# constrains task placement accordingly; for a `Matrix`, datadeps +# observes the tile as master-resident data as before. + +function tiled_cholesky!(M::AbstractMatrix) mt = size(M, 1) for k in 1:mt Dagger.@spawn LinearAlgebra.LAPACK.potrf!('L', InOut(M[k, k])) @@ -26,15 +37,7 @@ function tiled_cholesky!(M::Matrix{<:Matrix}) return M end -function make_spd_tiles(sz::Int, nb::Int) - @assert sz % nb == 0 - A = rand(sz, sz) - A = A * A' - A[diagind(A)] .+= sz - return [A[i:(i+nb-1), j:(j+nb-1)] for i in 1:nb:sz, j in 1:nb:sz] -end - -function tiled_matmul!(C::Matrix{<:Matrix}, A::Matrix{<:Matrix}, B::Matrix{<:Matrix}) +function tiled_matmul!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix) nt = size(C, 1) @assert size(A) == size(B) == size(C) == (nt, nt) for i in 1:nt, j in 1:nt @@ -48,11 +51,83 @@ function tiled_matmul!(C::Matrix{<:Matrix}, A::Matrix{<:Matrix}, B::Matrix{<:Mat return C end +""" + _placement_procs() -> Vector{Dagger.Processor} + +Sorted list of processors across which tiles are distributed by +`make_spd_tiles` / `make_matmul_tiles`. Uses `Dagger.all_processors()` +so multi-worker sessions (e.g. `julia -p 7 -t 12`) distribute tiles +across every worker's `ThreadProc`s, while single-process sessions +distribute across master's `ThreadProc`s only. Sort keys are `(pid, +string(proc))`, deterministic across identical sessions so tile→proc +placement is reproducible run-to-run. +""" +function _placement_procs() + procs = collect(Dagger.all_processors()) + sort!(procs; by = p -> (Dagger.root_worker_id(p), string(p))) + return procs +end + +""" + _distribute_tile(tile::AbstractMatrix, target_proc::Dagger.Processor) + -> Dagger.Chunk + +Ships `tile` (currently master-resident) to `target_proc`'s worker via +`remotecall_fetch`, then wraps it as a `Chunk` pinned to that processor +via `ExactScope(target_proc)`. `MemPool.poolset` runs on the target +worker inside the closure, so the resulting DRef lives in that worker's +memory space — subsequent datadeps tasks that touch only this tile +incur zero cross-worker transfer. + +The scope-and-proc metadata on the returned `Chunk` is what +`spawn_datadeps` inspects when computing per-task placement. Without +this distribution step every tile would be master-resident, and +`spawn_datadeps` would (correctly) place every task on master to avoid +shipping 8-MB tiles out and back — which is the exact master-only +concentration Hudson reproduced. +""" +function _distribute_tile(tile::AbstractMatrix, target_proc::Dagger.Processor) + target_pid = Dagger.root_worker_id(target_proc) + return remotecall_fetch(target_pid) do + Dagger.tochunk(tile, target_proc, Dagger.ExactScope(target_proc)) + end +end + +# Round-robin tile-index → processor mapping (row-major linearisation). +# Kept as a separate helper so make_spd_tiles / make_matmul_tiles use +# identical placement conventions and every tile at logical position +# `(i, j)` lands on the same processor across independent workloads +# (matmul's `A[i, k]`, `B[k, j]` co-locate with `C[i, j]` for at least +# one `k`, reducing transfers in the common case). +@inline function _tile_proc(procs, nt::Int, i::Int, j::Int) + return procs[mod1((i - 1) * nt + j, length(procs))] +end + +function make_spd_tiles(sz::Int, nb::Int) + @assert sz % nb == 0 + nt = sz ÷ nb + A = rand(sz, sz) + A = A * A' + A[diagind(A)] .+= sz + procs = _placement_procs() + # Copy each tile before shipping so the closure carries an + # independent Matrix (not a view into `A`), avoiding accidental + # aliasing with `A`'s buffer that would confuse datadeps' aliasing + # analysis. + return [_distribute_tile(copy(A[(i-1)*nb+1:i*nb, (j-1)*nb+1:j*nb]), + _tile_proc(procs, nt, i, j)) + for i in 1:nt, j in 1:nt] +end + function make_matmul_tiles(sz::Int, nb::Int) @assert sz % nb == 0 nt = sz ÷ nb - A = [rand(nb, nb) for _ in 1:nt, _ in 1:nt] - B = [rand(nb, nb) for _ in 1:nt, _ in 1:nt] - C = [zeros(nb, nb) for _ in 1:nt, _ in 1:nt] + procs = _placement_procs() + A = [_distribute_tile(rand(nb, nb), _tile_proc(procs, nt, i, j)) + for i in 1:nt, j in 1:nt] + B = [_distribute_tile(rand(nb, nb), _tile_proc(procs, nt, i, j)) + for i in 1:nt, j in 1:nt] + C = [_distribute_tile(zeros(nb, nb), _tile_proc(procs, nt, i, j)) + for i in 1:nt, j in 1:nt] return A, B, C end From 179c866ab1b34c1960d2849490d5eaa13ed839a5 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 06:50:54 -0500 Subject: [PATCH 10/44] bench: interleave _placement_procs across pids to avoid nt-dependent concentration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prior _placement_procs implementation sorted by (pid, name) so all of master's ThreadProcs came first in the list, followed by worker 2's, etc. With `procs[mod1((i-1)*nt+j, length(procs))]` selecting one processor per tile, small nt values under-distributed: tile count linear idx range procs touched pids reached nt=2 1..4 indices 1..4 1 (master only) nt=3 1..9 indices 1..9 1 (master only) nt=4 1..16 indices 1..16 2 (master + worker 2) nt=8 1..64 indices 1..64 ~6 pids nt=16 1..256 wraps 2.7x all 8 pids This is silently worse than a uniform fail: a full sweep across tile_counts = 2,3,4,8,16 would measure a different distribution level at every nt, silently confounding the tile-count axis with an unstated "how many workers happened to be hit" axis. No CSV column would reveal it. Root cause caught by the multi-process placement probe on Hudson: at nt=4 K=20, both RoundRobin and Greedy placed tasks on only pids 1 and 2 despite 96 processors being registered across 8 pids. Fix: group processors by owning pid, sort each pid's slice deterministically, then interleave — take the k-th processor from each pid in pid order, then move to k+1. Yields [pid1_t1, pid2_t1, ..., pidN_t1, pid1_t2, pid2_t2, ...] so consecutive tile indices at any nt land on different pids. Verified: - Interleave logic produces the expected pattern via a mock 8-pid × 12-thread simulation: nt=2 hits pids 1-4, nt=4 hits all 8 pids. - Full 527-test scheduler suite still passes: 0 fail, 0 error, 1 pre-existing broken. No regression. - Single-process M4 Pro session returns a single-element vector as before (only pid 1 exists). Multi-process verification at production scale must happen on Hudson after this commit lands upstream. --- bench/datadeps_schedulers/workloads.jl | 46 ++++++++++++++++++++++---- 1 file changed, 40 insertions(+), 6 deletions(-) diff --git a/bench/datadeps_schedulers/workloads.jl b/bench/datadeps_schedulers/workloads.jl index 34146e17b..0611803a0 100644 --- a/bench/datadeps_schedulers/workloads.jl +++ b/bench/datadeps_schedulers/workloads.jl @@ -54,18 +54,52 @@ end """ _placement_procs() -> Vector{Dagger.Processor} -Sorted list of processors across which tiles are distributed by +Interleaved list of processors across which tiles are distributed by `make_spd_tiles` / `make_matmul_tiles`. Uses `Dagger.all_processors()` so multi-worker sessions (e.g. `julia -p 7 -t 12`) distribute tiles across every worker's `ThreadProc`s, while single-process sessions -distribute across master's `ThreadProc`s only. Sort keys are `(pid, -string(proc))`, deterministic across identical sessions so tile→proc -placement is reproducible run-to-run. +distribute across master's `ThreadProc`s only. + +Consecutive indices in the returned vector round-robin across worker +pids: `[pid₁_t₁, pid₂_t₁, …, pidₙ_t₁, pid₁_t₂, pid₂_t₂, …]`. That +guarantees small tile counts (e.g. `nt=2` → 4 tiles) are spread across +multiple workers rather than concentrated on the first-listed pid. +A naive `sort!(procs; by = pid)` would place all of master's +ThreadProcs first and small `nt` would land only on master, silently +underdistributing at small `nt` and confounding the tile-count axis +with a hidden "how many workers happened to be hit" axis. + +Within each pid, the per-thread ordering is stable-sorted by +`string(proc)` so tile→proc mapping is reproducible run-to-run for +identical sessions. """ function _placement_procs() procs = collect(Dagger.all_processors()) - sort!(procs; by = p -> (Dagger.root_worker_id(p), string(p))) - return procs + # Group by owning worker pid so we can interleave across pids. + by_pid = Dict{Int, Vector{eltype(procs)}}() + for p in procs + pid = Dagger.root_worker_id(p) + push!(get!(by_pid, pid, eltype(procs)[]), p) + end + # Deterministic within-pid ordering for reproducibility. + for slice in values(by_pid) + sort!(slice, by = string) + end + # Interleave: take the k-th processor from each pid in pid order, + # then move to k+1. Yields [pid₁_t₁, pid₂_t₁, …, pid₁_t₂, …]. + pids = sort!(collect(keys(by_pid))) + max_slice = maximum(length(v) for v in values(by_pid)) + interleaved = eltype(procs)[] + sizehint!(interleaved, sum(length(v) for v in values(by_pid))) + for tid_idx in 1:max_slice + for pid in pids + slice = by_pid[pid] + if tid_idx <= length(slice) + push!(interleaved, slice[tid_idx]) + end + end + end + return interleaved end """ From be0eea85eebcc5e2183ba31f450f88688d0d9405 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 07:41:09 -0500 Subject: [PATCH 11/44] bench: remove ExactScope pinning from _distribute_tile so schedulers can decide placement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prior _distribute_tile used `ExactScope(target_proc)` which is a HARD constraint on task placement, not just a data-location hint. Under it, every scheduler was forced onto the tile's owning processor — RoundRobin, Greedy, IteratedGreedy, SimulatedAnnealing, JuMP, and OptimizingScheduler all converged to byte-identical n_copies counts because none of them had any placement choice to make. Hudson's matmul sweep at bs=1024 confirmed the collapse: nt=2 n_copies median across all 6 schedulers: 9 9 9 9 9 9 nt=3 41 41 41 41 41 41 nt=4 96 96 96 96 96 96 nt=8 448 448 448 448 448 448 exec_span_ms differences across schedulers spanned ±20% while per-trial CV within a single (cell, scheduler) reached 34.4% — the scheduler signal was smaller than the noise because there was no scheduler signal, just system variance around a pre-determined placement. The fix removes the ExactScope. Tiles still live on specific workers (the initial `remotecall_fetch(target_pid) do; Dagger.tochunk(...) end` puts the DRef in the target worker's memory space), so γ transfer costs are still real and nonzero. But schedulers now DECIDE placement, weighing γ against parallelism and load balancing — Greedy will co-locate to minimize γ, RoundRobin will spread, SA/IG/JuMP will explore. Their placement decisions genuinely differ, which is the paper's mechanism. Verified via 527-test suite: 0 fail, 0 error, 1 pre-existing broken. No regression. Multi-process verification at production scale must happen on Hudson after this commit + the hierarchical.jl aliasing-bug fix land upstream. --- bench/datadeps_schedulers/workloads.jl | 37 +++++++++++++++++++------- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/bench/datadeps_schedulers/workloads.jl b/bench/datadeps_schedulers/workloads.jl index 0611803a0..9f3479b9c 100644 --- a/bench/datadeps_schedulers/workloads.jl +++ b/bench/datadeps_schedulers/workloads.jl @@ -107,23 +107,42 @@ end -> Dagger.Chunk Ships `tile` (currently master-resident) to `target_proc`'s worker via -`remotecall_fetch`, then wraps it as a `Chunk` pinned to that processor -via `ExactScope(target_proc)`. `MemPool.poolset` runs on the target -worker inside the closure, so the resulting DRef lives in that worker's -memory space — subsequent datadeps tasks that touch only this tile -incur zero cross-worker transfer. +`remotecall_fetch`, then wraps it as a `Chunk` that LIVES on that +processor without pinning task placement to it. `MemPool.poolset` runs +on the target worker inside the closure, so the resulting DRef lives in +that worker's memory space — subsequent datadeps tasks that touch only +this tile can either (a) run on `target_proc` for zero-cost local access, +or (b) run elsewhere and incur the measured transfer cost γ. That +trade-off is exactly what the schedulers are meant to weigh. + +Historical note (this is why NO ExactScope): the earlier revision of +this function pinned tiles via `ExactScope(target_proc)`. That worked to +distribute tiles across workers, but Dagger's ExactScope is a HARD +constraint on task placement, not just a data-location hint. Under it, +every scheduler is forced onto the tile's owning processor — RR, Greedy, +IG, SA, JuMP, OptimizingScheduler all converged to byte-identical +n_copies counts because none of them had any placement choice to make. +Hudson's matmul sweep at bs=1024 confirmed the collapse: n_copies +identical to the single-integer across all six schedulers at every tile +count, and per-trial CV up to 34% because the residual variation was +scheduler-independent system noise. Removing ExactScope restores the +paper's core mechanism: schedulers see distributed data (γ > 0) and +DECIDE placement rather than having it forced on them. The scope-and-proc metadata on the returned `Chunk` is what `spawn_datadeps` inspects when computing per-task placement. Without -this distribution step every tile would be master-resident, and +the initial distribution step every tile would be master-resident, and `spawn_datadeps` would (correctly) place every task on master to avoid -shipping 8-MB tiles out and back — which is the exact master-only -concentration Hudson reproduced. +shipping 8-MB tiles out and back — which was the pre-distribution +master-only concentration Hudson previously reproduced. """ function _distribute_tile(tile::AbstractMatrix, target_proc::Dagger.Processor) target_pid = Dagger.root_worker_id(target_proc) return remotecall_fetch(target_pid) do - Dagger.tochunk(tile, target_proc, Dagger.ExactScope(target_proc)) + # No explicit scope: chunk lives on `target_proc` but scheduler + # is free to route tasks touching this tile anywhere, weighing + # γ transfer cost against compute-vs-parallelism trade-offs. + Dagger.tochunk(tile, target_proc) end end From b2f4bc42c611461b8a3df69bd3be9bacd114aa6a Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 16:09:11 -0500 Subject: [PATCH 12/44] datadeps: route AOT scheduler decisions into hierarchical partitioning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `distribute_tasks_hierarchical!` previously binned tasks with `partition_dag` (data-affinity) before any scheduler ran, then let each partition compute its own AOT schedule over just that partition's processors. Under Dagger's default hierarchical mode with distributed tiles this made every scheduler look identical at the paper's metric (n_copies): affinity picked the coarse worker-level placement; per-partition schedulers only shuffled tasks among a single worker's threads where γ ≈ 0, so worker-level cross-copies never varied across schedulers. When (a) the argument data is already distributed across ≥2 workers and (b) the scheduler produces a whole-region AOT schedule, partition by the scheduler's assigned worker instead of by data affinity. The scheduler's placement drives partitioning, and different schedulers produce measurably different per-worker task counts (and thus n_copies). - New `partition_by_assigned_worker(seen_tasks, schedule, task_metas, all_procs)`: bin each task by `root_worker_id(schedule[task])`, with an affinity fallback for the rare unscheduled-tail case (past a `dag_add_task!` bail point); returns partitions in ascending-worker-id order for determinism. - Whole-DAG AOT call sites: bypass `datadeps_build_schedule!` and call `datadeps_schedule_dag_aot!` directly, so we don't double-write the scheduler's per-type cache (which is now owned by `_hierarchical_persist_schedule!` at the top level, keyed by the outer `dag_spec`). - Persist condition switches from "no precomputed schedule" to "no cache hit", so newly-computed whole-DAG schedules also persist for reuse. - Gate on `data_is_distributed` (≥2 workers own arg data). When all data lives on one worker, affinity partitioning already picks that worker and is optimal; forcing a scheduler-first path in that case would spread computation away from data with no benefit, so we fall through to the existing affinity behaviour. - JIT schedulers (RoundRobin, Naive, Ultra) don't specialize `datadeps_schedule_dag_aot!`, produce an empty trial, and take the affinity branch unchanged. Verified on `-p 3 -t 4` with distributed cholesky tiles: nt=4: Greedy/Optimizing (scheduler-first) → per-worker (1,8)(2,4)(3,4)(4,4) RoundRobin (affinity) → per-worker (1,10)(2,6)(3,3)(4,1) nt=8: AOT (scheduler-first) → (1,32)(2,32)(3,28)(4,28) JIT (affinity) → (1,49)(2,35)(3,23)(4,13) datadeps test suite: 514/521 pass, 2 fail + 4 error + 1 broken -- matches baseline (513/521 pass, 3 fail + 4 error + 1 broken) modulo one flake. Datadeps top-level: 1260/1262 pass, 2 broken (unchanged from baseline). Co-Authored-By: Claude Opus 4.7 --- src/datadeps/hierarchical.jl | 165 +++++++++++++++++++++++++++++++++-- 1 file changed, 160 insertions(+), 5 deletions(-) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 1570d1f9b..1a0ca19b1 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -691,6 +691,109 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe end +""" + partition_by_assigned_worker(seen_tasks, schedule, task_metas, all_procs) + -> (vertex_to_partition, n_partitions, partition_procs, multi_worker) + +Scheduler-first partitioning: bin each task into the partition whose worker its +AOT-scheduled processor lives on, rather than the worker owning the most of its +argument data (which is what `partition_dag` does). Used by +`distribute_tasks_hierarchical!` when the scheduler is AOT (produced a non-empty +`schedule`) or a cached full-region schedule was recovered. Preserves the +scheduler's whole-DAG placement decisions across the hierarchical partition +boundary, which is otherwise where those decisions get silently overridden by +data-affinity binning. + +Tasks with no entry in `schedule` (only possible past a `dag_add_task!` bail +point -- rare in practice for datadeps regions whose args are `Chunk`s rather +than same-region `DTask` fetches) fall back to the affinity rule: bin by the +worker owning the most of their arg data, breaking ties toward the first-listed +matching worker. This matches `partition_dag`'s behaviour for those tasks and +keeps them local to the data they need. + +Only creates a partition for workers with >=1 assigned task, in ascending worker +id order (for deterministic partition numbering across runs). +""" +function partition_by_assigned_worker(seen_tasks::Vector{DTaskPair}, + schedule::Dict{DTask,Processor}, + task_metas::Vector{HierarchicalTaskMeta}, + all_procs::Vector{<:Processor}) + n = length(seen_tasks) + + procs_by_worker = Dict{Int, Vector{Processor}}() + for proc in all_procs + wid = root_worker_id(only(memory_spaces(proc))) + push!(get!(Vector{Processor}, procs_by_worker, wid), proc) + end + known_workers = Set{Int}(keys(procs_by_worker)) + + # First pass: assign a worker id to every task. Scheduled tasks use their + # AOT-chosen proc's worker; unscheduled tasks fall back to affinity (worker + # owning the most of their arg data), else to the first-listed worker. + task_worker = Vector{Int}(undef, n) + for v in 1:n + task = seen_tasks[v].task + proc = get(schedule, task, nothing) + if proc !== nothing + wid = root_worker_id(only(memory_spaces(proc))) + if wid in known_workers + task_worker[v] = wid + continue + end + end + # Fallback: affinity by arg data ownership. Tie-break toward the first + # worker with the max affinity in `known_workers` iteration order. + affinity = Dict{Int,Int}() + for dep in task_metas[v].deps + arg_wid = root_worker_id(memory_space(dep.arg_w.arg)) + arg_wid in known_workers || continue + affinity[arg_wid] = get(affinity, arg_wid, 0) + 1 + end + if isempty(affinity) + task_worker[v] = first(sort!(collect(known_workers))) + else + best_wid = 0 + best_aff = -1 + for wid in sort!(collect(keys(affinity))) + aff = affinity[wid] + if aff > best_aff + best_aff = aff + best_wid = wid + end + end + task_worker[v] = best_wid + end + end + + # Second pass: reduce to only-workers-with-tasks and assign partition ids in + # ascending worker order for determinism. + used_workers = Int[] + seen_wids = Set{Int}() + for wid in task_worker + if !(wid in seen_wids) + push!(used_workers, wid) + push!(seen_wids, wid) + end + end + sort!(used_workers) + worker_to_partition = Dict(wid => i for (i, wid) in enumerate(used_workers)) + n_partitions = length(used_workers) + + vertex_to_partition = Vector{Int}(undef, n) + for v in 1:n + vertex_to_partition[v] = worker_to_partition[task_worker[v]] + end + + partition_procs = Vector{Vector{Processor}}(undef, n_partitions) + for (pid, wid) in enumerate(used_workers) + partition_procs[pid] = procs_by_worker[wid] + end + + multi_worker = n_partitions > 1 + return vertex_to_partition, n_partitions, partition_procs, multi_worker +end + + """ _schedule_vertex!(v, partition_id, temp_queue, state, local_procs, local_scope, dag, seen_tasks, vertex_to_partition, @@ -1075,9 +1178,59 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) # Phase 2: Build dependency DAG dag = build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) - # Phase 3: Partition the DAG - vertex_to_partition, n_partitions, partition_procs, _multi_worker = - partition_dag(dag, task_metas, all_procs) + # Phase 3a: Decide whether to route through scheduler-first partitioning. + # Only useful when the argument data itself is already distributed across + # >=2 workers -- otherwise, the affinity partitioner (which bins tasks near + # their data) already picks the best worker set, and forcing a whole-DAG + # AOT pass would either (a) agree with affinity or (b) spread computation + # away from data at the cost of extra cross-space transfers that don't buy + # a better schedule. For workloads that distribute their tiles across + # workers (e.g. our benchmark's `_distribute_tile`), this gate lets AOT + # scheduler decisions drive worker-level task placement, so different + # schedulers produce different n_copies -- which the affinity partitioner + # otherwise silently overrides. + data_workers = Set{Int}() + for meta in task_metas + for dep in meta.deps + push!(data_workers, root_worker_id(memory_space(dep.arg_w.arg))) + end + end + data_is_distributed = length(data_workers) >= 2 + + # Phase 3b: If we haven't already got a full-region schedule from the cache + # and the DAG has anything to schedule, ask the scheduler for an AOT + # placement over the *whole* region and *all* processors -- but only when + # data-distribution warrants it (see the gate above). JIT schedulers (which + # don't specialize `datadeps_schedule_dag_aot!`) fall through to the no-op + # fallback in `scheduling.jl:378`, leaving `trial` empty; we then take the + # affinity-partitioning branch below with unchanged semantics. + # + # We call `datadeps_schedule_dag_aot!` directly rather than + # `datadeps_build_schedule!` because the latter *also* writes the schedule + # into `queue.scheduler`'s cache internally. Under hierarchical, the cache + # is owned by `_hierarchical_persist_schedule!` below (keyed by the + # top-level `dag_spec`); routing through `datadeps_build_schedule!` would + # double-write and the two entries could diverge for the same key. + cache_hit = precomputed_schedule !== nothing + if !cache_hit && !isempty(dag_spec) && data_is_distributed + trial = Dict{DTask,Processor}() + datadeps_schedule_dag_aot!(queue.scheduler, trial, dag_spec, all_procs, all_scope) + if !isempty(trial) + precomputed_schedule = trial + end + end + + # Phase 3c: Partition the DAG. Scheduler-first when we have a whole-region + # schedule (cache hit *or* freshly-computed AOT) *and* the data warrants + # cross-worker computation. Affinity-based otherwise (JIT schedulers, + # non-distributed data, or empty DAG spec) -- unchanged behaviour. + if precomputed_schedule !== nothing && data_is_distributed + vertex_to_partition, n_partitions, partition_procs, _multi_worker = + partition_by_assigned_worker(seen_tasks, precomputed_schedule, task_metas, all_procs) + else + vertex_to_partition, n_partitions, partition_procs, _multi_worker = + partition_dag(dag, task_metas, all_procs) + end # Detect backing chunks shared across partitions in different memory spaces. # These need runtime ownership transfer to avoid split-brain concurrent @@ -1129,11 +1282,13 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) # Persist the freshly-computed schedule for reuse by a later # `distribute_tasks_hierarchical!` call on a structurally equivalent DAG. - # Only runs on a cache miss (`precomputed_schedule === nothing`) to avoid + # Runs on a cache miss regardless of whether the fresh schedule came from + # the whole-DAG AOT above or from the per-partition AOT inside + # `schedule_partition_full!`; skipping only on a genuine cache *hit* avoids # both duplicating the cache entry we just consumed and appending it under # a potentially different DAGSpec key that would fail subsequent # `datadeps_dag_equivalent` matches. - if precomputed_schedule === nothing + if !cache_hit _hierarchical_persist_schedule!(queue.scheduler, dag_spec, partition_schedules) end end From 5037b442afae870d3f6c3bc6928bdf453799242e Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 16:41:16 -0500 Subject: [PATCH 13/44] datadeps: assert build_aliasing_parallel result length matches input length to expose the deserialization dedup collapse The `for i in eachindex(worker_args); arg_to_ainfo[worker_args[i]] = results[i].second` loops in both branches of `build_aliasing_parallel` assume `length(results) == length(worker_args)`. Under -p 7 -t 12 cholesky nt=4 bs=1024 that invariant has been observed to break twice now (BoundsError: 1-element Vector{Pair{ArgumentWrapper,AliasingWrapper}} at index [2]), and the crash surfaces as a bare BoundsError with no context about the mismatched sizes or the ArgumentWrapper identities involved. Add an `@assert length(results) == length(worker_args)` immediately before each loop with the worker id, both lengths, and the worker_args + results ArgumentWrapper `.hash` values. When the invariant breaks the AssertionError names the mismatched sizes and the hash-set directly, which is enough to distinguish the two candidate mechanisms: 1) upstream `Dict{ArgumentWrapper,ArgumentWrapper}`-keyed dedup on master collapsing two hash-equal entries -- would show a duplicate hash in `worker_args_hashes` and `result_first_hashes` the same length as `worker_args_hashes` minus the collapsed count. 2) `DRef` mutable-struct backref dedup during `Serialization`-then-deserialization of `worker_args` on the remote worker -- would show distinct hashes in `worker_args_hashes` and a strictly-shorter `result_first_hashes`. Diagnostic only; the root cause fix is a follow-up commit. Co-Authored-By: Claude Opus 4.7 --- src/datadeps/hierarchical.jl | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 1a0ca19b1..d17cd79d7 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -389,6 +389,12 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr # 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 input/output length mismatch (single-worker branch). \ +wid=$wid myid=$(myid()) length(worker_args)=$(length(worker_args)) \ +length(results)=$(length(results)) \ +worker_args_hashes=$([w.hash for w in worker_args]) \ +result_first_hashes=$([r.first.hash for r in results])""" for i in eachindex(worker_args) arg_to_ainfo[worker_args[i]] = results[i].second end @@ -404,6 +410,12 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr # 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 input/output length mismatch (multi-worker branch). \ +wid=$wid myid=$(myid()) length(worker_args)=$(length(worker_args)) \ +length(results)=$(length(results)) \ +worker_args_hashes=$([w.hash for w in worker_args]) \ +result_first_hashes=$([r.first.hash for r in results])""" @lock all_results_lock begin for i in eachindex(worker_args) arg_to_ainfo[worker_args[i]] = results[i].second From 9378807b10861772fa04065ed9b9deb5cd6f0bab Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 16:53:37 -0500 Subject: [PATCH 14/44] datadeps: local results in build_aliasing_parallel Threads.@spawn to eliminate cross-partition race Under Dagger's default hierarchical mode at scale (-p 7 -t 12 with the paper's distributed cholesky at nt=4 bs=1024) `build_aliasing_parallel` crashed intermittently with `BoundsError: 1-element Vector{Pair{ArgumentWrapper, AliasingWrapper}} at index [2]` inside the multi-worker branch. Root cause is a Julia scoping / data race, not a serialization or deduplication issue. The single-worker branch at line 385 binds `results` in the enclosing function scope: results = wid == myid() ? _compute_aliasing_batch(worker_args) : remotecall_fetch(_compute_aliasing_batch, wid, worker_args) The multi-worker branch just below uses the same name inside its `Threads.@spawn` closure: @sync for (wid, worker_args) in by_worker Threads.@spawn begin results = if wid == myid() ... end ...results[i]... end end Julia scoping resolves `results = ...` inside the closure to the already-bound enclosing-function variable rather than allocating a fresh per-closure local. Every partition's `Threads.@spawn` task therefore races on ONE shared `results` cell; the last writer wins, and every reader (including the reader in that same task, after any scheduling delay) sees the winner's Vector instead of its own. When the winning task's `worker_args` is shorter than the reading task's, `results[i]` at `i=length(winner_results)+1` throws `BoundsError` -- which is exactly what Hudson observed twice on cholesky-nt=4 and what I reproduced locally on trial 12 of a -p 7 -t 12 stress loop. Fix: `local results` at the top of the `Threads.@spawn` block so each closure gets its own binding. Verified with a minimal `@sync for ... Threads.@spawn` scoping test (without `local`, every task reads the last writer's value; with `local`, each task reads its own). The cholesky-nt=4 bs=1024 -p 7 -t 12 reproducer that deterministically crashed at trial 12 in ~70s now runs the full 5-minute budget without incident. The single-worker branch (line 385) does not need `local` -- there is only one flow, no concurrency. The diagnostic assert from 5037b442 stays: with the race fixed it should never fire, but the invariant is worth documenting explicitly so future regressions surface loudly instead of as bare `BoundsError`s. Co-Authored-By: Claude Opus 4.7 --- src/datadeps/hierarchical.jl | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index d17cd79d7..cd7620751 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -402,6 +402,15 @@ result_first_hashes=$([r.first.hash for r in results])""" all_results_lock = ReentrantLock() @sync for (wid, worker_args) in by_worker Threads.@spawn begin + # `local` is REQUIRED here: the single-worker branch above + # binds `results` in the enclosing function scope, so without + # `local` every `Threads.@spawn` closure below would assign to + # (and read from) that same shared variable, racing with sibling + # partitions. The observed symptom is a `BoundsError` where + # `results[i]` accesses a Vector whose length matches a *different* + # partition's `worker_args`, hitting index-out-of-bounds when the + # winning writer's partition happens to have fewer args than ours. + local results results = if wid == myid() _compute_aliasing_batch(worker_args) else From 3102efd27f8d277b0246000fc6d01cd37297a925 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 17:29:48 -0500 Subject: [PATCH 15/44] comments: trim verbose comments across hierarchical/scheduling/sch Partial pass. Kills paragraph-length explanations, backreferences to prior commits/tasks/people, and per-line "what this line does" comments. Keeps docstrings (public API), invariant statements, and the few comments where the *why* is non-obvious (e.g. `local results` justification, cross-partition `finally` rationale, `_eft_tail_args` tuple-vs-vector caveat). Files: - src/datadeps/hierarchical.jl (-372/+107) - src/datadeps/scheduling.jl (-98/+25) - src/sch/Sch.jl (small) - src/sch/util.jl (small) No code changes -- comment and docstring text only. `using Dagger` loads clean. Co-Authored-By: Claude Opus 4.7 --- src/datadeps/hierarchical.jl | 479 ++++++++--------------------------- src/datadeps/scheduling.jl | 123 ++------- src/sch/Sch.jl | 13 +- src/sch/util.jl | 8 +- 4 files changed, 137 insertions(+), 486 deletions(-) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index cd7620751..5e3f45283 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -1,21 +1,6 @@ -# Hierarchical scheduling for datadeps -# Spreads scheduling work across multiple threads/workers via a 4-phase pipeline: -# Phase 1: Parallel aliasing info construction -# Phase 2: Sequential DAG construction from aliasing overlaps -# Phase 3: Data-affinity DAG partitioning -# Phase 4: Parallel local scheduling per partition -- each partition runs on its -# own task, prepares/submits its tasks concurrently with the others, -# and computes its *own* (partition-local) AOT schedule over just its -# processors. This maximizes scheduling parallelism and scalability; -# no global synchronization or global AOT pass is imposed. -# -# N.B. Concurrent task submission from many partition tasks used to intermittently -# hang in the core eager scheduler. The root cause was `Distributed.Future`, -# which is unsafe under concurrent same-process access; Dagger now backs task -# futures with `MemPool.DFuture` (a thread-safe, `DEvent`-based future), so -# parallel Phase 4 is safe. The actual task submission (`enqueue!`) is still -# serialized via `LockedEnqueueQueue`, but the (expensive) per-task -# `distribute_task!` preparation runs concurrently across partitions. +# Concurrent per-partition task submission requires MemPool.DFuture-backed +# task futures; the older Distributed.Future was not safe for concurrent +# same-process access. struct HierarchicalTaskInfo arg_w::ArgumentWrapper @@ -29,40 +14,26 @@ struct HierarchicalTaskMeta may_alias::Bool inplace_move::Bool deps::Vector{HierarchicalTaskInfo} - # Indices of same-region producer tasks whose results are passed as - # arguments (e.g. `In(t1)`). These cannot be `fetch`'d during the - # pre-scan because they are not launched yet; they become hard DAG edges. + # Same-region unlaunched DTask arguments -- cannot fetch during pre-scan; + # become hard DAG edges. value_deps::Vector{Int} end -### Cross-partition chunk ownership ### -# -# Each partition schedules with its own `DataDepsState`, so `arg_owner` / -# `arg_history` / physical slots are per-partition. When a backing chunk is -# written by tasks in >=2 partitions that live in *different* memory spaces, -# each partition would otherwise copy that chunk from its origin, write its own -# sub-range, and record itself as owner -- the physical copies then diverge and -# the final copy-back keeps only one, silently losing the others' writes. -# -# The registry below carries the single authoritative version ("ownership") of -# each such shared chunk across partition boundaries. It is deadlock-free by -# construction: a single lock (never per-argument locks acquired in different -# orders) plus the global DAG ordering -- a consumer partition always `wait`s on -# its cross-partition predecessor's `task_submitted` event before scheduling, so -# the producer's ownership commit is always visible before the consumer reads it. -# The lock only provides memory-safety for concurrent access to *different* -# chunks; per-chunk correctness comes from the DAG order. -"Authoritative, cross-partition ownership state for a single shared backing chunk." +# Cross-partition ownership registry for chunks written from >=2 partitions in +# distinct memory spaces. Without this, per-partition copies diverge and the +# final copy-back silently loses all but one. Deadlock-free: one lock plus the +# DAG's `task_submitted` handshake orders the reads against the ownership +# commit; the lock only guards memory-safety across different chunks. mutable struct ChunkOwnership - owner_space::MemorySpace # space holding the current version - owner_slot::Any # physical slot chunk in `owner_space` - owner_task::Union{DTask,Nothing} # producer of the current version (nothing => origin data) - owner_state::Union{DataDepsState,Nothing} # owning partition's state (for copy-back) - const origin_space::MemorySpace # the chunk's home space + owner_space::MemorySpace + owner_slot::Any + owner_task::Union{DTask,Nothing} + owner_state::Union{DataDepsState,Nothing} + const origin_space::MemorySpace end struct SharedChunkRegistry - entries::IdDict{Any,ChunkOwnership} # backing chunk (identity) => ownership + entries::IdDict{Any,ChunkOwnership} lock::ReentrantLock end SharedChunkRegistry() = SharedChunkRegistry(IdDict{Any,ChunkOwnership}(), ReentrantLock()) @@ -71,12 +42,11 @@ is_shared_chunk(::Nothing, @nospecialize(chunk)) = false is_shared_chunk(reg::SharedChunkRegistry, @nospecialize(chunk)) = haskey(reg.entries, chunk) """ - build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) -> SharedChunkRegistry or nothing + build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) + -> SharedChunkRegistry or nothing -Detects backing chunks accessed by partitions spanning >=2 distinct memory -spaces (the only case that can split-brain) and returns a registry seeded with -origin ownership for each. Returns `nothing` when all partitions share a single -space (e.g. single-worker), so the fast path is entirely unchanged there. +Registry of chunks accessed from >=2 distinct memory spaces. Returns +`nothing` when all partitions share one space. """ function build_shared_chunk_registry(task_metas::Vector{HierarchicalTaskMeta}, vertex_to_partition::Vector{Int}, @@ -107,11 +77,9 @@ end """ _sync_incoming_ownership!(state, registry, our_space, task_arg_ws, write_num) -Before a task's copy-to phase, for each shared backing-chunk argument whose -globally-current owner (per `registry`) lives in a space other than `our_space`, -seed this partition's `state` so the existing copy-to machinery pulls a fresh -whole-chunk copy from the true owner (with a syncdep on the producing task). A -no-op for private chunks or when we already hold the authoritative version. +Seed `state` so the copy-to phase pulls a fresh whole-chunk copy from the +true cross-partition owner (with a syncdep on the producing task). No-op for +private chunks or when we already own. """ function _sync_incoming_ownership!(state::DataDepsState, registry::SharedChunkRegistry, our_space::MemorySpace, task_arg_ws, write_num::Int) @@ -125,10 +93,8 @@ function _sync_incoming_ownership!(state::DataDepsState, registry::SharedChunkRe owner_space, owner_slot, owner_task = @lock registry.lock begin (entry.owner_space, entry.owner_slot, entry.owner_task) end - owner_space == our_space && return # we already hold the authoritative copy + owner_space == our_space && return - # Register the owner's physical slot so slot/aliasing lookups in this - # partition reuse it (rather than generating a fresh, stale copy). dest_args = get!(IdDict{Any,Any}, state.remote_args, owner_space) if !haskey(dest_args, chunk) dest_args[chunk] = owner_slot @@ -138,15 +104,12 @@ function _sync_incoming_ownership!(state::DataDepsState, registry::SharedChunkRe map_or_ntuple(arg_ws.deps) do dep_idx dep = arg_ws.deps[dep_idx] arg_w = dep.arg_w - # Point ownership at the owner space and clear any locally-merged - # history, so `compute_remainder_for_arg!` takes the `FullCopy` - # (whole-chunk) path from the owner rather than a partial remainder. + # Clear local history so compute_remainder_for_arg! takes the + # FullCopy path from the owner instead of a partial remainder. state.arg_owner[arg_w] = owner_space haskey(state.arg_history, arg_w) && empty!(state.arg_history[arg_w]) src_ainfo = aliasing!(state, owner_space, arg_w) if owner_task !== nothing - # Make the ensuing copy-to (via `get_read_deps!`) wait on the - # producer, so we never copy the owner slot before it is written. state.ainfos_owner[src_ainfo] = owner_task => (write_num - 1) end return @@ -159,11 +122,9 @@ end """ _commit_ownership!(state, registry, our_space, task, task_arg_ws) -After a task is recorded as a writer, publish it as the new authoritative owner -of each shared backing chunk it writes, so subsequent cross-partition consumers -pull from here. Ordering (and thus visibility) is guaranteed by the DAG's -`task_submitted` handshake; the lock only guards concurrent commits for other -chunks. +Publish `task` as the new authoritative owner of each shared chunk it writes. +Visibility is guaranteed by the DAG's `task_submitted` handshake; the lock +only guards concurrent commits for other chunks. """ function _commit_ownership!(state::DataDepsState, registry::SharedChunkRegistry, our_space::MemorySpace, task::DTask, task_arg_ws) @@ -194,16 +155,11 @@ function _commit_ownership!(state::DataDepsState, registry::SharedChunkRegistry, return end -# Below this many tasks, the fixed costs of spawning threads and merging -# per-thread results outweigh the benefit of parallelizing the pre-scan. +# Fewer tasks than this: sequential pre-scan wins over the thread-spawn cost. const COLLECT_ALIASED_ARGS_MIN_CHUNK = 256 -# Thread-safe "get or compute" against a shared `IdDict` cache. The -# expensive computation (`f`) is performed outside of the lock so that -# multiple threads can make progress on distinct keys concurrently; if two -# threads race on the same key, the loser's result is simply discarded (the -# corresponding `Chunk`/Bool is cheap to let the GC reclaim) so that every -# thread agrees on a single canonical value (e.g. `Chunk`) per raw argument. +# f() runs outside the lock so distinct keys make concurrent progress; racers +# on the same key drop their result and adopt whichever value landed first. @inline function _cached_get!(f, cache::IdDict{Any,V}, cache_lock::Union{ReentrantLock,Nothing}, key) where V if cache_lock === nothing return get!(f, cache, key) @@ -220,29 +176,16 @@ end """ collect_aliased_args(seen_tasks) -> (task_metas, unique_arg_ws) -Pre-scans all tasks to collect per-task dependency metadata and the set of -unique `ArgumentWrapper`s that need aliasing analysis. This mirrors the logic -in `populate_task_info!` but only inspects arguments without modifying any -scheduling state. - -For large batches of tasks (the common case for e.g. panel-factorization -algorithms which submit many small tasks per `spawn_datadeps` region), the -pre-scan itself (not just the aliasing computation in -`build_aliasing_parallel`) can dominate scheduling time, since it touches -every argument of every task. This is parallelized across threads: each -thread scans a contiguous range of `seen_tasks` into its own disjoint slice -of `task_metas` and its own local `unique_arg_ws` map (merged at the end), -while sharing (lock-protected) caches for `Chunk`-wrapping and -`supports_inplace_move`, ensuring a single canonical `Chunk` identity per -raw argument regardless of which thread first observes it. +Pre-scan every task's arguments to build `task_metas` and the unique +`ArgumentWrapper` set for aliasing analysis. Parallelized across threads +with shared, lock-protected caches for `Chunk`-wrapping and +`supports_inplace_move` so every raw argument has one canonical `Chunk`. """ function collect_aliased_args(seen_tasks::Vector{DTaskPair}) n = length(seen_tasks) task_metas = Vector{HierarchicalTaskMeta}(undef, n) n == 0 && return task_metas, Dict{ArgumentWrapper,ArgumentWrapper}() - # Map in-region tasks to vertex indices so we can record value deps - # without fetching unlaunched DTasks during the pre-scan. task_to_idx = IdDict{DTask,Int}() for (i, pair) in enumerate(seen_tasks) task_to_idx[pair.task] = i @@ -309,11 +252,8 @@ function _collect_aliased_args_range!(task_metas::Vector{HierarchicalTaskMeta}, arg_pre_unwrap, deps = unwrap_inout(_arg_with_deps) - # Same-region DTask arguments are not launched yet, so we cannot - # `fetch` them here (that is what the sequential `distribute_task!` - # path does *after* launching the producer). Record a value - # dependency instead; aliasing of the result is handled later in - # `distribute_task!` once the producer has been submitted. + # Unlaunched same-region DTask: record a value dep; aliasing is + # deferred to `distribute_task!` after the producer submits. if arg_pre_unwrap isa DTask && !istaskstarted(arg_pre_unwrap) pred_idx = get(task_to_idx, arg_pre_unwrap, 0) if pred_idx != 0 && pred_idx != task_idx @@ -357,12 +297,11 @@ function _collect_aliased_args_range!(task_metas::Vector{HierarchicalTaskMeta}, end """ - build_aliasing_parallel(unique_arg_ws) -> (lookup, ainfos_overlaps, arg_to_ainfo) + build_aliasing_parallel(unique_arg_ws) + -> (lookup, ainfos_overlaps, arg_to_ainfo) -Phase 1: Computes `AliasingWrapper` for every unique `ArgumentWrapper` in -parallel. On each worker, threads are used to compute aliasing info for local -data. Results are gathered and reduced into a single `AliasingLookup` with -overlap information. +Compute `AliasingWrapper` for every unique `ArgumentWrapper` in parallel and +reduce into a single `AliasingLookup` with overlap information. """ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWrapper}) arg_ws_vec = collect(values(unique_arg_ws)) @@ -377,18 +316,11 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr arg_to_ainfo = Dict{ArgumentWrapper, AliasingWrapper}() if length(by_worker) == 1 - # Common single-worker case: avoid the `@sync`/`Threads.@spawn`/lock - # overhead entirely, since there's nothing to run concurrently with. - # `_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) : 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. + # Key by the local arg_w: remote-returned ArgumentWrappers are + # deserialized copies. _compute_aliasing_batch preserves input order. @assert length(results) == length(worker_args) """build_aliasing_parallel: \ _compute_aliasing_batch input/output length mismatch (single-worker branch). \ wid=$wid myid=$(myid()) length(worker_args)=$(length(worker_args)) \ @@ -402,23 +334,15 @@ result_first_hashes=$([r.first.hash for r in results])""" all_results_lock = ReentrantLock() @sync for (wid, worker_args) in by_worker Threads.@spawn begin - # `local` is REQUIRED here: the single-worker branch above - # binds `results` in the enclosing function scope, so without - # `local` every `Threads.@spawn` closure below would assign to - # (and read from) that same shared variable, racing with sibling - # partitions. The observed symptom is a `BoundsError` where - # `results[i]` accesses a Vector whose length matches a *different* - # partition's `worker_args`, hitting index-out-of-bounds when the - # winning writer's partition happens to have fewer args than ours. + # `local` is REQUIRED: the single-worker branch above binds + # `results` in the enclosing scope, and without this every + # spawned closure would race on that shared cell. local results results = if wid == myid() _compute_aliasing_batch(worker_args) 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 input/output length mismatch (multi-worker branch). \ wid=$wid myid=$(myid()) length(worker_args)=$(length(worker_args)) \ @@ -457,8 +381,8 @@ result_first_hashes=$([r.first.hash for r in results])""" return lookup, ainfos_overlaps, arg_to_ainfo end -# Below this many args, the fixed cost of forking/joining `Threads.@threads` -# outweighs the benefit of parallelizing the (typically cheap) `aliasing()` calls. +# Fewer args than this: `Threads.@threads` overhead exceeds the aliasing() +# cost. const COMPUTE_ALIASING_BATCH_MIN_PARALLEL = 8 function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) @@ -484,9 +408,8 @@ end build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) -> SimpleDiGraph -Phase 2: Walks tasks in submission order and builds a `SimpleDiGraph` encoding -data dependencies based on the pre-computed aliasing overlaps. Uses the same -WAW / RAW / WAR rules as `get_write_deps!` / `get_read_deps!`. +Build the data-dependency DAG from pre-computed aliasing overlaps. Same +WAW/RAW/WAR rules as `get_write_deps!` / `get_read_deps!`. """ function build_dependency_dag(task_metas::Vector{HierarchicalTaskMeta}, arg_to_ainfo::Dict{ArgumentWrapper, AliasingWrapper}, @@ -574,12 +497,11 @@ function build_dependency_dag(task_metas::Vector{HierarchicalTaskMeta}, end """ - partition_dag(dag, task_metas, all_procs) -> (vertex_to_partition, n_partitions, partition_procs) + partition_dag(dag, task_metas, all_procs) + -> (vertex_to_partition, n_partitions, partition_procs, multi_worker) -Phase 3: Assigns each task vertex to a partition using data-affinity. For -multi-worker setups, tasks are assigned to the worker owning the most argument -data. For single-worker multi-threaded setups, tasks are balanced across -available processors in topological order. +Data-affinity partitioning. Multi-worker: bin by the worker owning the most +arg data. Single-worker: load-balance across procs in topological order. """ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMeta}, all_procs::Vector{<:Processor}) @@ -611,11 +533,8 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe meta = task_metas[v] task_scope = @something(meta.pair.spec.options.compute_scope, meta.pair.spec.options.scope, default_scope) - # Workers whose processors are eligible under this task's scope. - # A non-default scope that spans multiple workers must still spread - # work across those workers (via affinity / round-robin) -- picking - # only the first match pins everything to worker 1 and breaks - # multi-worker execution. + # Non-default scopes spanning multiple workers must spread work + # (via affinity / round-robin); first-match would pin to worker 1. if task_scope == default_scope matching = collect(1:n_partitions) else @@ -716,24 +635,11 @@ end partition_by_assigned_worker(seen_tasks, schedule, task_metas, all_procs) -> (vertex_to_partition, n_partitions, partition_procs, multi_worker) -Scheduler-first partitioning: bin each task into the partition whose worker its -AOT-scheduled processor lives on, rather than the worker owning the most of its -argument data (which is what `partition_dag` does). Used by -`distribute_tasks_hierarchical!` when the scheduler is AOT (produced a non-empty -`schedule`) or a cached full-region schedule was recovered. Preserves the -scheduler's whole-DAG placement decisions across the hierarchical partition -boundary, which is otherwise where those decisions get silently overridden by -data-affinity binning. - -Tasks with no entry in `schedule` (only possible past a `dag_add_task!` bail -point -- rare in practice for datadeps regions whose args are `Chunk`s rather -than same-region `DTask` fetches) fall back to the affinity rule: bin by the -worker owning the most of their arg data, breaking ties toward the first-listed -matching worker. This matches `partition_dag`'s behaviour for those tasks and -keeps them local to the data they need. - -Only creates a partition for workers with >=1 assigned task, in ascending worker -id order (for deterministic partition numbering across runs). +Scheduler-first partitioning: bin each task by its AOT-chosen processor's +worker, preserving the scheduler's whole-DAG placement across the partition +boundary. Unscheduled tasks (past a `dag_add_task!` bail point) fall back to +data-affinity. Only workers with >=1 task get a partition; ordered ascending +by worker id. """ function partition_by_assigned_worker(seen_tasks::Vector{DTaskPair}, schedule::Dict{DTask,Processor}, @@ -748,9 +654,6 @@ function partition_by_assigned_worker(seen_tasks::Vector{DTaskPair}, end known_workers = Set{Int}(keys(procs_by_worker)) - # First pass: assign a worker id to every task. Scheduled tasks use their - # AOT-chosen proc's worker; unscheduled tasks fall back to affinity (worker - # owning the most of their arg data), else to the first-listed worker. task_worker = Vector{Int}(undef, n) for v in 1:n task = seen_tasks[v].task @@ -762,8 +665,7 @@ function partition_by_assigned_worker(seen_tasks::Vector{DTaskPair}, continue end end - # Fallback: affinity by arg data ownership. Tie-break toward the first - # worker with the max affinity in `known_workers` iteration order. + # Fallback: affinity by arg data ownership. affinity = Dict{Int,Int}() for dep in task_metas[v].deps arg_wid = root_worker_id(memory_space(dep.arg_w.arg)) @@ -786,8 +688,6 @@ function partition_by_assigned_worker(seen_tasks::Vector{DTaskPair}, end end - # Second pass: reduce to only-workers-with-tasks and assign partition ids in - # ascending worker order for determinism. used_workers = Int[] seen_wids = Set{Int}() for wid in task_worker @@ -820,16 +720,11 @@ end local_scope, dag, seen_tasks, vertex_to_partition, schedule, proc_to_scope_lfu, write_num) -> write_num -Schedules a single (already topologically-ordered) task vertex `v` into its -partition's `state` via `distribute_task!`, returning the updated `write_num`. - -Records cross-partition predecessors as explicit `ThunkSyncdep`s (same-partition -deps are derived from `state` by `distribute_task!`), and passes the -partition-local AOT-assigned processor when one is available and usable for this -partition (else falls back to JIT scheduling). - -Callers must guarantee that every predecessor of `v` (in any partition) has -already been submitted before calling this, so the `ThunkSyncdep`s are valid. +Schedule a topologically-ordered vertex into its partition via +`distribute_task!`. Cross-partition predecessors become explicit +`ThunkSyncdep`s; AOT proc is passed through when it lies in `local_procs`, +else JIT. Callers must ensure every predecessor (any partition) has already +been submitted. """ function _schedule_vertex!(v::Int, partition_id::Int, temp_queue::DataDepsTaskQueue, @@ -857,13 +752,7 @@ function _schedule_vertex!(v::Int, partition_id::Int, end end - # Use the AOT-assigned processor when one is available and it lies within - # this partition's processors; otherwise fall back to JIT scheduling - # (`proc === nothing`) via `temp_queue`'s scheduler. (An AOT proc can lie - # outside `local_procs` only in the multi-worker case, where partitioning - # restricts each partition to one worker's processors; single-worker - # partitions always span all procs, so the AOT assignment is always usable - # there.) + # AOT-assigned proc, or JIT fallback if it's outside this partition. proc = get(schedule, task, nothing) if proc !== nothing && !(proc in local_procs) proc = nothing @@ -882,28 +771,10 @@ end precomputed_schedule=nothing) -> (DataDepsState, Dict{DTask,Processor}) -Per-partition scheduling for both single-worker and multi-worker hierarchical -paths. Uses existing `distribute_task!` logic with per-partition -`DataDepsState`, `all_procs` limited to this partition's processors, and -cross-partition syncdeps from the precomputed DAG. Returns this partition's -`DataDepsState` along with the task-to-processor assignments it computed (or -reused), so the top-level `distribute_tasks_hierarchical!` can persist the -combined per-region schedule to `queue.scheduler`'s cache. - -AOT scheduling defaults to running *locally per partition*: a partition-local -`DAGSpec` is built from just this partition's tasks and scheduled over just -`local_procs` via `datadeps_build_schedule!`. `region_uids` (all in-region task -uids) lets that builder bail cleanly when a task depends on a same-region -producer in another partition (which it must not `fetch`). Tasks without an AOT -assignment fall back to JIT scheduling in `distribute_task!`. - -When `precomputed_schedule` is supplied (cache hit at the top of -`distribute_tasks_hierarchical!`), the per-partition AOT step is skipped: this -partition's task-to-processor assignments are filtered from -`precomputed_schedule` directly. Tasks whose cached processor is not in -`local_procs` (possible if `all_procs` changed since the schedule was cached) -fall back to JIT scheduling in `distribute_task!` via the guard in -`_schedule_vertex!`. +Per-partition scheduling. Runs partition-local AOT on `local_procs` (with +`region_uids` for the cross-partition bail), or filters +`precomputed_schedule` when supplied. Returns state and this partition's +task-to-processor mapping. """ function schedule_partition_full!(queue::DataDepsTaskQueue, queue_lock::ReentrantLock, @@ -923,9 +794,7 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, local_scope = UnionScope(map(ExactScope, local_procs)) - # N.B. Hierarchical scheduling doesn't build a global `DAGSpec` (that's - # only used by the AOT DAG-caching path in `distribute_tasks!`), so each - # partition just gets an empty one; `distribute_task!` never reads + # No global `DAGSpec` under hierarchical; `distribute_task!` never reads # `state.dag_spec` directly. state = DataDepsState(DAGSpec()) write_num = 1 @@ -940,41 +809,18 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, ordered_verts = filter(v -> v in vert_set, topo) locked_queue = LockedEnqueueQueue(get_options(:task_queue), queue_lock) - # N.B. Each partition gets its own fresh scheduler shard via `similar` - # rather than sharing `queue.scheduler` across all partitions. Two - # independent problems would arise from sharing a single scheduler - # instance here: - # 1) Data race: e.g. `RoundRobinScheduler.proc_idx` would be - # concurrently read/written by every partition's `Threads.@spawn` - # task with no synchronization. - # 2) Semantic bug (worse than the race, and *not* fixed by adding a - # lock): each partition schedules only onto its own worker's - # `local_procs`, which generally has a *different length* than - # other partitions' (or the global) processor list. A `proc_idx` - # counter advanced by one partition's `local_procs` is meaningless - # -- and can be out-of-bounds -- when applied to another - # partition's differently-sized `local_procs`. This reliably - # crashes with a `BoundsError` under multi-worker hierarchical - # scheduling. Giving each partition its own scheduler instance, - # scoped to its own `local_procs`, fixes both issues at once. + # Fresh scheduler shard per partition: sharing `queue.scheduler` would + # race on mutable state (e.g. `RoundRobinScheduler.proc_idx`) and, worse, + # advance a counter meant for one `local_procs` against another's -- + # BoundsError under multi-worker. temp_queue = DataDepsTaskQueue(locked_queue; scheduler=similar(queue.scheduler)) - # Per-partition (local) AOT scheduling over just this partition's procs. - # `partition_verts` is in ascending vertex (= submission) order. partition_pairs = DTaskPair[seen_tasks[v] for v in partition_verts] if precomputed_schedule === nothing - # Cache miss: run the per-partition AOT scheduler over just this - # partition's tasks and processors. _pdag, schedule = datadeps_build_schedule!(temp_queue.scheduler, partition_pairs, local_procs, local_scope; region_uids) else - # Cache hit: `distribute_tasks_hierarchical!` already resolved the whole - # region's task-to-processor mapping from `queue.scheduler.cache`. Filter - # it to just this partition's tasks. Tasks that fell through the - # top-level `dag_add_task!` bail (missing from `precomputed_schedule`) - # still get JIT-scheduled inside `distribute_task!` via - # `temp_queue.scheduler`, matching the miss-path semantics. schedule = Dict{DTask,Processor}() for pair in partition_pairs proc = get(precomputed_schedule, pair.task, nothing) @@ -984,17 +830,9 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, end end - # N.B. If this partition throws partway through (e.g. from - # `distribute_task!`), any of our vertices that haven't yet been - # `notify`'d will never be, which would leave other partitions blocked - # forever in `wait(task_submitted[pred_v])` below -- turning a normal, - # reportable exception into a silent, permanent hang (since the - # enclosing `@sync` in `distribute_tasks_hierarchical!` can't finish, - # and thus can't propagate our exception, until *every* spawned - # partition task completes, including the ones stuck waiting on us). - # The `finally` ensures every one of our events gets notified no matter - # how we exit, so that sibling partitions can unblock (and themselves - # fail/finish) and our real exception can actually surface. + # `finally` notifies all our events even on exception: without it, + # sibling partitions can wedge on `wait(task_submitted[pred_v])` + # forever, and the enclosing `@sync` never surfaces our exception. try for v in ordered_verts for pred_v in inneighbors(dag, v) @@ -1003,13 +841,6 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, end end - # N.B. The per-task `distribute_task!` preparation runs concurrently - # across partitions; only the final task submission is serialized - # (via `LockedEnqueueQueue`, `temp_queue`'s upper queue). This is - # safe now that task futures are backed by `MemPool.DFuture` rather - # than the concurrency-unsafe `Distributed.Future`. The - # cross-partition `wait`s above happen before scheduling `v`, so the - # `ThunkSyncdep`s recorded for `v` are valid. write_num = _schedule_vertex!( v, partition_id, temp_queue, state, local_procs, local_scope, dag, seen_tasks, vertex_to_partition, schedule, @@ -1041,24 +872,10 @@ end _hierarchical_dag_spec_and_cache_lookup(scheduler, seen_tasks) -> (dag_spec::DAGSpec, precomputed::Union{Dict{DTask,Processor},Nothing}) -Build a full-region `DAGSpec` by walking `seen_tasks` in submission order, then -look it up in `scheduler`'s task-local schedule cache. When a structurally -equivalent DAG is found (per `datadeps_dag_equivalent`), return the recovered -task-to-processor mapping in the second slot; otherwise return `nothing`. - -Mirrors the DAGSpec-building and cache-matching logic of the flat path's -`datadeps_build_schedule!` in `queue.jl`, but stops before running the AOT -scheduler so that `distribute_tasks_hierarchical!` retains its parallel -per-partition scheduling pipeline on a cache miss. The DAGSpec built here is -also returned so the miss path can persist a freshly-computed schedule at the -end of the hierarchical run using the same key. - -`dag_add_task!` may return `false` (bail) on tasks whose arguments include a -within-DAG `DTask` -- the returned `DAGSpec` then only covers the prefix that -was successfully added, matching the flat path's behaviour. Tasks past the -bail point are not in the cache key; on a hit they simply won't be in -`precomputed`, and per-partition JIT scheduling handles them exactly as in the -miss path. +Build the full-region `DAGSpec` and check `scheduler`'s cache. On a hit, +return the recovered task-to-processor mapping; on miss, return `nothing`. +`dag_add_task!` may bail on within-DAG `DTask` args, in which case the +returned `DAGSpec` covers only the prefix that succeeded. """ function _hierarchical_dag_spec_and_cache_lookup(scheduler::DataDepsScheduler, seen_tasks::Vector{DTaskPair}) @@ -1094,17 +911,8 @@ end """ _hierarchical_persist_schedule!(scheduler, dag_spec, partition_schedules) -After a cache-miss hierarchical run, collect the per-partition -task-to-processor assignments produced by `schedule_partition_full!` and store -them in `scheduler`'s schedule cache keyed by `dag_spec`. This allows a later -`distribute_tasks_hierarchical!` call with a structurally equivalent DAG to -reuse the schedule via `_hierarchical_dag_spec_and_cache_lookup`. - -No-op when `dag_spec` is empty (the top-level DAG build bailed on the first -task) or when no partition produced any assignments (every task fell back to -JIT scheduling). Only tasks captured in `dag_spec` are cached, so partial-DAG -cases match flat-path semantics: tasks past the bail point are re-JIT'd on -subsequent runs. +Merge the per-partition schedules and cache them keyed by `dag_spec`. +No-op on empty spec or empty schedules. """ function _hierarchical_persist_schedule!(scheduler::DataDepsScheduler, dag_spec::DAGSpec, @@ -1123,9 +931,8 @@ function _hierarchical_persist_schedule!(scheduler::DataDepsScheduler, return end -# Parallel partition scheduling wraps errors in TaskFailedException / -# CompositeException via `@sync`/`Threads.@spawn`. Unwrap so callers and tests -# see the root `SchedulingException` / scheduler error. +# Unwrap TaskFailedException/CompositeException so callers see the root +# scheduling error rather than the @sync/@spawn envelope. function _unwrap_partition_exception(e) while true if e isa CompositeException && !isempty(e.exceptions) @@ -1141,20 +948,9 @@ end """ distribute_tasks_hierarchical!(queue) -Main entry point for hierarchical scheduling. Runs the 4-phase pipeline: -1. Parallel aliasing construction -2. DAG construction -3. Partitioning (by worker affinity, or across local procs on one worker) -4. Parallel per-partition scheduling via `distribute_task!` - -Each partition runs on its own task and computes its *own* partition-local AOT -schedule over just its processors (see `schedule_partition_full!`); there is no -global AOT pass or global synchronization, which is what allows this to scale. -AOT scheduling here therefore need not match the flat `distribute_tasks!` path -exactly. Both drivers use `distribute_task!` for argument preparation and -`DataDepsScheduler` dispatch; the old single-worker "batch enqueue with DAG -syncdeps only" path is intentionally not used (it skipped `distribute_task!` and -broke `ChunkView` / custom schedulers). +Hierarchical scheduling entry point: parallel aliasing construction, DAG +build, partitioning, then parallel per-partition scheduling via +`distribute_task!`. """ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) seen_tasks = queue.seen_tasks @@ -1162,20 +958,9 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) return end - # Top-level schedule cache lookup. Build a full-region DAGSpec and consult - # `queue.scheduler`'s task-local cache before spinning up any partitioning - # or per-partition scheduling. Under hierarchical, each partition otherwise - # runs with its own `similar(queue.scheduler)` shard (see - # `schedule_partition_full!`) whose cache is discarded after the call, so - # this top-level lookup is the only place where cross-`spawn_datadeps` - # amortisation of AOT scheduling cost can occur. On a hit - # (`precomputed_schedule !== nothing`) the per-partition scheduler is - # skipped; on a miss we run the normal pipeline and persist the collected - # schedule at the end for the next call to reuse. dag_spec, precomputed_schedule = _hierarchical_dag_spec_and_cache_lookup(queue.scheduler, seen_tasks) - # Get the set of all processors all_procs = Processor[] scope = get_compute_scope() for w in procs() @@ -1187,29 +972,17 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) end all_scope = UnionScope(map(ExactScope, all_procs)) - # All in-region task uids, so each partition's local AOT-schedule builder - # can tell same-region producers (which it must not `fetch`) apart from - # already-materialized external values. + # In-region task uids let each partition's AOT builder skip same-region + # producers rather than `fetch` them. region_uids = Set{UInt}(pair.task.uid for pair in seen_tasks) - # 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) - # Phase 2: Build dependency DAG dag = build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) - # Phase 3a: Decide whether to route through scheduler-first partitioning. - # Only useful when the argument data itself is already distributed across - # >=2 workers -- otherwise, the affinity partitioner (which bins tasks near - # their data) already picks the best worker set, and forcing a whole-DAG - # AOT pass would either (a) agree with affinity or (b) spread computation - # away from data at the cost of extra cross-space transfers that don't buy - # a better schedule. For workloads that distribute their tiles across - # workers (e.g. our benchmark's `_distribute_tile`), this gate lets AOT - # scheduler decisions drive worker-level task placement, so different - # schedulers produce different n_copies -- which the affinity partitioner - # otherwise silently overrides. + # Scheduler-first partitioning is only meaningful when arg data spans + # multiple workers; otherwise affinity already picks the right worker. data_workers = Set{Int}() for meta in task_metas for dep in meta.deps @@ -1218,20 +991,9 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) end data_is_distributed = length(data_workers) >= 2 - # Phase 3b: If we haven't already got a full-region schedule from the cache - # and the DAG has anything to schedule, ask the scheduler for an AOT - # placement over the *whole* region and *all* processors -- but only when - # data-distribution warrants it (see the gate above). JIT schedulers (which - # don't specialize `datadeps_schedule_dag_aot!`) fall through to the no-op - # fallback in `scheduling.jl:378`, leaving `trial` empty; we then take the - # affinity-partitioning branch below with unchanged semantics. - # - # We call `datadeps_schedule_dag_aot!` directly rather than - # `datadeps_build_schedule!` because the latter *also* writes the schedule - # into `queue.scheduler`'s cache internally. Under hierarchical, the cache - # is owned by `_hierarchical_persist_schedule!` below (keyed by the - # top-level `dag_spec`); routing through `datadeps_build_schedule!` would - # double-write and the two entries could diverge for the same key. + # Call `datadeps_schedule_dag_aot!` directly, not `datadeps_build_schedule!`: + # the latter also writes the scheduler's per-type cache, which under + # hierarchical is owned by `_hierarchical_persist_schedule!` below. cache_hit = precomputed_schedule !== nothing if !cache_hit && !isempty(dag_spec) && data_is_distributed trial = Dict{DTask,Processor}() @@ -1241,10 +1003,6 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) end end - # Phase 3c: Partition the DAG. Scheduler-first when we have a whole-region - # schedule (cache hit *or* freshly-computed AOT) *and* the data warrants - # cross-worker computation. Affinity-based otherwise (JIT schedulers, - # non-distributed data, or empty DAG spec) -- unchanged behaviour. if precomputed_schedule !== nothing && data_is_distributed vertex_to_partition, n_partitions, partition_procs, _multi_worker = partition_by_assigned_worker(seen_tasks, precomputed_schedule, task_metas, all_procs) @@ -1253,13 +1011,9 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) partition_dag(dag, task_metas, all_procs) end - # Detect backing chunks shared across partitions in different memory spaces. - # These need runtime ownership transfer to avoid split-brain concurrent - # writes; `nothing` when all partitions share one space (single-worker). partition_space = MemorySpace[only(memory_spaces(first(pp))) for pp in partition_procs] registry = build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) - # Group vertices by partition partitions = [Int[] for _ in 1:n_partitions] for v in 1:length(seen_tasks) pid = vertex_to_partition[v] @@ -1270,13 +1024,6 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) task_submitted = [Base.Event() for _ in 1:length(seen_tasks)] wait_all_queue = get_options(:task_queue) - # Phase 4: parallel per-partition scheduling. Each partition runs on its own - # task, prepares its tasks concurrently, computes its own partition-local - # AOT schedule (or reuses `precomputed_schedule` on a cache hit, in which - # case the AOT step is skipped and only the task submission side runs), and - # coordinates cross-partition dependencies via the `task_submitted` events. - # Concurrent submission is safe now that task futures are backed by - # `MemPool.DFuture` (see the header note). partition_states = Vector{DataDepsState}(undef, n_partitions) partition_schedules = Vector{Dict{DTask,Processor}}(undef, n_partitions) try @@ -1301,14 +1048,6 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) _hierarchical_copy_from_and_free!(partition_states, n_partitions, registry) - # Persist the freshly-computed schedule for reuse by a later - # `distribute_tasks_hierarchical!` call on a structurally equivalent DAG. - # Runs on a cache miss regardless of whether the fresh schedule came from - # the whole-DAG AOT above or from the per-partition AOT inside - # `schedule_partition_full!`; skipping only on a genuine cache *hit* avoids - # both duplicating the cache entry we just consumed and appending it under - # a potentially different DAGSpec key that would fail subsequent - # `datadeps_dag_equivalent` matches. if !cache_hit _hierarchical_persist_schedule!(queue.scheduler, dag_spec, partition_schedules) end @@ -1340,15 +1079,13 @@ end function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsState}, n_partitions::Int, registry::Union{SharedChunkRegistry,Nothing}) - # 1. Shared chunks: write back from the registry's authoritative owner state - # (whose per-partition history is coherent), not the cross-partition - # max-`write_num` heuristic below (per-partition write_nums are not - # comparable across partitions). + # Shared chunks: use the registry's authoritative owner. Per-partition + # write_nums are not comparable across partitions. if registry !== nothing for (chunk, entry) in registry.entries state = entry.owner_state - state === nothing && continue # never written - entry.owner_space == entry.origin_space && continue # already in-place at origin + state === nothing && continue + entry.owner_space == entry.origin_space && continue for (arg_w, _space) in state.arg_owner arg_w.arg === chunk || continue _hierarchical_copy_from!(state, arg_w, _hierarchical_max_write_num(state, arg_w) + 1) @@ -1356,8 +1093,7 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat end end - # 2. Private chunks: the last writer across partitions is the authoritative - # owner (shared chunks are skipped here, handled above). + # Private chunks: last-writer-wins. merged_arg_owner = Dict{ArgumentWrapper, Tuple{MemorySpace, Int, DataDepsState}}() for pid in 1:n_partitions state = partition_states[pid] @@ -1375,10 +1111,9 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat _hierarchical_copy_from!(state, arg_w, wn + 1) end - # 3. Free Datadeps-allocated slots. For shared chunks, also sync on the final - # global writer: an intermediate owner's slot may still be read by a - # cross-partition boundary copy recorded in *another* partition's state, - # and the final writer transitively depends on all such copies. + # Free datadeps slots. For shared chunks, sync on the final global writer: + # a stale owner slot may still be read via boundary copies recorded in + # another partition's state. for pid in 1:n_partitions state = partition_states[pid] obj_cache = unwrap(state.ainfo_backing_chunk) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 2064545f3..4e74cb7af 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -383,17 +383,9 @@ end """ _propagate_aot_time_util!(spec::DTaskSpec, proc::Processor, task_time_ns::Real) -Attach the AOT-computed per-task runtime estimate to `spec.options.time_util` -so that `Sch.has_capacity` (via `schedule_one!`) can bypass its -`metrics_lookup_runtime` scan of the MetricsTracker snapshot and use our -estimate directly. `has_capacity` reads the value via -`round(UInt64, time_util[T] * 1000^3)`, so `time_util` is stored keyed by -`typeof(proc)` with the runtime expressed in seconds (Float64). - -Called by every AOT scheduler after it finalises a task's `(task, proc)` -assignment. Non-finite / non-positive estimates are ignored (the -`options.time_util === nothing` path is preserved for those, matching the -pre-existing behaviour where `has_capacity` falls back to metrics lookup). +Attach the AOT-computed per-task runtime to `spec.options.time_util` so +`Sch.has_capacity` bypasses its metrics-snapshot scan. Keyed by +`typeof(proc)`; value in seconds. No-op on non-finite / non-positive input. """ function _propagate_aot_time_util!(spec::DTaskSpec, proc::Processor, task_time_ns::Real) task_time_ns > 0 || return @@ -408,9 +400,7 @@ function _propagate_aot_time_util!(spec::DTaskSpec, proc::Processor, task_time_n return end -## `_propagate_aot_time_util_from_cache!` — the `EFTCostCache`-aware wrapper — -## is defined below, immediately after `EFTCostCache`, so its signature can -## reference the type. +# EFTCostCache-aware wrapper is below, after the struct. const GREEDY_DEFAULT_RUNTIME_NS = UInt64(1_000_000_000) const GREEDY_DEFAULT_TRANSFER_RATE = UInt64(1_000_000) @@ -433,12 +423,8 @@ end """ _propagate_aot_time_util_from_cache!(dag_spec, cache, task_idx, proc) -Convenience wrapper that pulls the AOT-computed runtime from an -`EFTCostCache` (already built by the caller for its own cost-model use) -and forwards it to `_propagate_aot_time_util!`. If the cache does not -contain an entry for `proc`, or the entry is not a positive finite value, -the propagation is a no-op — `has_capacity` will then take its usual -fallback path. +Pull the AOT runtime for `(task_idx, proc)` from `cache` and forward to +`_propagate_aot_time_util!`. No-op if `proc` isn't cached. """ function _propagate_aot_time_util_from_cache!(dag_spec::DAGSpec, cache::EFTCostCache, task_idx::Int, proc::Processor) @@ -724,9 +710,7 @@ _greedy_arg_ready_time_ns_cached(::Any, ::MT.MetricsSnapshot, ::DAGSpec, const IG_DEFAULT_N_ITERS = 32 const IG_DEFAULT_DESTROY_FRAC = 0.30 -# `Inf` disables the wall-clock stop; callers (e.g. benchmarks) opt in by -# passing a finite budget. Requested by Przemek & Julian: enable "1-minute -# metaheuristic budget" as the operational default for the paper. +# `Inf` disables the wall-clock stop. const IG_DEFAULT_TIME_LIMIT_SEC = Inf """ @@ -781,12 +765,7 @@ end IteratedGreedyScheduler(; kwargs...) = IteratedGreedyScheduler(GreedyScheduler(); kwargs...) -# Hierarchical scheduling calls `similar` per partition to avoid sharing mutable -# state (esp. the RNG) across the parallel per-partition scheduling tasks. Deep- -# copy the RNG so each partition-shard produces independent, race-free samples; -# recursively call `similar` on the inner scheduler so its own mutable state is -# refreshed too. Static configuration (n_iters, destroy_frac) carries through -# unchanged so the shard's behaviour matches the original's. +# Hierarchical partitions need independent RNGs; deep-copy so shards don't race. Base.similar(s::IteratedGreedyScheduler) = IteratedGreedyScheduler(similar(s.inner); n_iters=s.n_iters, @@ -794,9 +773,7 @@ Base.similar(s::IteratedGreedyScheduler) = rng=copy(s.rng), time_limit_sec=s.time_limit_sec) -# Forward equivalence / cache / argspec hooks so cache reuse is partitioned by -# the inner scheduler's type, not by IG's own wrapper type. Same idea as the -# bench's TimedScheduler wrapper. +# Delegate cache/equivalence to the inner scheduler so its type owns cache keys. datadeps_schedule_cache(sched::IteratedGreedyScheduler) = datadeps_schedule_cache(sched.inner) datadeps_dag_equivalent(sched::IteratedGreedyScheduler, dspec1::DAGSpec, dspec2::DAGSpec) = @@ -819,11 +796,8 @@ Mutates and returns `state`. The caller is responsible for providing a freshly-`copy(prev_state)` if the previous state must be preserved on rejection. """ -# Take all arguments after the function. `spec.fargs` is a `Tuple` for typed -# tasks (`DTaskSpec{true}`) and a `Vector{Argument}` otherwise. `Base.view` -# is not defined for `Tuple`, so `@view fargs[2:end]` breaks under typed -# tasks; use `Base.tail` for the tuple case (zero allocation) and a view -# for the vector case. +# `spec.fargs` is a `Tuple` for typed tasks; `@view fargs[2:end]` doesn't work +# on tuples. _eft_tail_args(fargs::Tuple) = Base.tail(fargs) _eft_tail_args(fargs::AbstractVector) = @view fargs[2:end] @@ -834,11 +808,9 @@ function _eft_runtime_ns(snap::MT.MetricsSnapshot, spec, proc::Processor) return runtime_lookup === nothing ? Float64(GREEDY_DEFAULT_RUNTIME_NS) : Float64(runtime_lookup) end -# Replays the schedule in increasing task-index order: destroyed tasks are -# reassigned via `greedy_assign_task!`; preserved tasks keep their proc but -# their finish times are recomputed fresh from the now-current state. The -# recompute (rather than preserving stale finish times) is required for SA -# to evaluate ΔC against the schedule's true makespan. +# Replay in task-index order; destroyed tasks reassign via +# `greedy_assign_task!`, preserved tasks recompute finish times fresh (SA needs +# true makespan to evaluate ΔC). function _replay_schedule!(state::ScheduleState, snap::MT.MetricsSnapshot, dag_spec::DAGSpec, all_procs::Vector{Processor}, destroyed::AbstractSet{Int}, @@ -943,20 +915,14 @@ function iterated_greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapsho destroyed = Set{Int}() sizehint!(destroyed, n_destroy) - # Wall-clock guard checked once per outer iteration. Elapsed nanoseconds - # are computed as `time_ns() - start_ns` (unsigned subtraction, so no - # overflow arithmetic on the deadline itself). `Inf` disables the guard - # by setting `budget_ns` to the maximum representable `UInt64`, ensuring - # the comparison never trips. + # `Inf` disables the guard via typemax(UInt64); unsigned elapsed subtraction. start_ns = time_ns() budget_ns = isinf(time_limit_sec) ? typemax(UInt64) : round(UInt64, time_limit_sec * 1e9) for _ in 1:n_iters (time_ns() - start_ns) >= budget_ns && break - # Sample `n_destroy` distinct task IDs uniformly at random. Doing a - # partial Fisher–Yates on perm_buf is O(n_destroy) and avoids - # allocating a fresh randperm vector each iter. + # Partial Fisher-Yates on perm_buf: O(n_destroy), no fresh allocations. @inbounds for i in 1:n_destroy j = rand(rng, i:n_tasks) perm_buf[i], perm_buf[j] = perm_buf[j], perm_buf[i] @@ -966,15 +932,12 @@ function iterated_greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapsho push!(destroyed, perm_buf[i]) end - # Reset candidate to the current best in place, then perturb. _copy_state!(candidate, best_state) iterated_greedy_step!(candidate, snap, dag_spec, all_procs, destroyed, cache) new_cost = cost_of_schedule(candidate) if new_cost < best_cost - # Adopt the candidate as the new best by swapping references. - # After the swap, `candidate` holds the previous best_state's - # buffers (stale data, overwritten by _copy_state! next iter). + # Swap; ex-best buffers get overwritten next iter. best_state, candidate = candidate, best_state best_cost = new_cost end @@ -1017,11 +980,6 @@ const SA_DEFAULT_Q = 0.95 const SA_DEFAULT_K = 1.0 const SA_DEFAULT_N_RESTARTS = 1 const SA_TF_FLOOR = 1e-12 -# See `IG_DEFAULT_TIME_LIMIT_SEC`. When SA wraps an IG seed inside -# `datadeps_schedule_dag_aot!(::SimulatedAnnealingScheduler, ...)`, both the -# seed and the SA refinement share the same clock but are governed by their -# own `time_limit_sec` fields, so setting both to 60 s does NOT double the -# effective budget of the SA pipeline. See the docstring for details. const SA_DEFAULT_TIME_LIMIT_SEC = Inf """ @@ -1046,16 +1004,10 @@ Fields: - `k::Float64` — coefficient `k > 0` in Orsila §4.3 Eqs. 18, 19. - `n_restarts::Int` — independent SA runs, each seeded from the best-known solution so far (Orsila §5.3 rule 7). -- `time_limit_sec::Float64` — wall-clock stop for the SA pipeline in - `datadeps_schedule_dag_aot!`, checked between - restarts and once per cooling level. `Inf` - (default) disables the stop; a finite value - returns the best-so-far state once the budget - is exceeded. When an `IteratedGreedyScheduler` - seed is used, the seed observes its own - `time_limit_sec`; the two budgets are summed - (worst-case), matching Julian's requested - "quick-and-dirty" semantics. +- `time_limit_sec::Float64` — wall-clock stop, checked between restarts and + per cooling level. `Inf` disables. When an IG + seed is used, seed and SA each observe their own + budget; the pipeline sums them worst-case. - `rng::R` — RNG used for moves and acceptance sampling. The `rng` is mutated during scheduling, so a single `SimulatedAnnealingScheduler` instance @@ -1090,10 +1042,6 @@ end SimulatedAnnealingScheduler(; kwargs...) = SimulatedAnnealingScheduler(IteratedGreedyScheduler(); kwargs...) -# See the note on `similar(::IteratedGreedyScheduler)`: hierarchical scheduling -# hands each partition its own scheduler shard, so we deep-copy the RNG and -# recursively refresh the inner scheduler, while preserving the parameterisation -# (q, k, n_restarts). Base.similar(s::SimulatedAnnealingScheduler) = SimulatedAnnealingScheduler(similar(s.inner); q=s.q, k=s.k, n_restarts=s.n_restarts, @@ -1284,11 +1232,7 @@ function simulated_annealing_schedule!(state::ScheduleState, snap::MT.MetricsSna expected_levels = max(1, ceil(Int, log(Tf / T0) / log(Float64(q)))) max_iters_per_restart = 16 * L * expected_levels - # See the note on the same construction in `iterated_greedy_schedule!`: - # single `time_ns()` origin, unsigned elapsed subtraction, `Inf` disables - # via `typemax(UInt64)` so the check never trips. Inner check is placed - # at the top of the SA loop so a hit returns the current best-so-far - # without spending another neighbour proposal. + # Same clock/`Inf`-disable pattern as `iterated_greedy_schedule!`. start_ns = time_ns() budget_ns = isinf(time_limit_sec) ? typemax(UInt64) : round(UInt64, time_limit_sec * 1e9) @@ -1448,10 +1392,6 @@ struct JuMPScheduler <: DataDepsScheduler end end -# See the note on `similar(::IteratedGreedyScheduler)`. `JuMPScheduler` has no -# mutable state and only immutable config fields; return a fresh instance with -# the same configuration so hierarchical partition shards behave identically -# to the original. Base.similar(s::JuMPScheduler) = JuMPScheduler(s.optimizer; Z=s.Z, time_limit_sec=s.time_limit_sec) @@ -1479,17 +1419,9 @@ Fields: path. - `milp_time_limit_sec` — forwarded to `JuMPScheduler`. - `milp_Z` — forwarded to `JuMPScheduler`. -- `ig_n_iters` — forwarded to `IteratedGreedyScheduler`. -- `ig_destroy_frac` — forwarded to `IteratedGreedyScheduler`. -- `ig_time_limit_sec` — forwarded to `IteratedGreedyScheduler` - (heuristic path). `Inf` disables the stop. -- `sa_q` — forwarded to `SimulatedAnnealingScheduler`. -- `sa_k` — forwarded to `SimulatedAnnealingScheduler`. -- `sa_n_restarts` — forwarded to `SimulatedAnnealingScheduler`. -- `sa_time_limit_sec` — forwarded to `SimulatedAnnealingScheduler` - (heuristic path). `Inf` disables the stop. -- `rng::R` — RNG shared by IG and SA on the heuristic - path; unused on the MILP path. +- `ig_n_iters`, `ig_destroy_frac`, `ig_time_limit_sec` — forwarded to IG. +- `sa_q`, `sa_k`, `sa_n_restarts`, `sa_time_limit_sec` — forwarded to SA. +- `rng::R` — used by IG/SA; unused on the MILP path. """ struct OptimizingScheduler{R<:Random.AbstractRNG} <: DataDepsScheduler optimizer::Any @@ -1547,11 +1479,6 @@ struct OptimizingScheduler{R<:Random.AbstractRNG} <: DataDepsScheduler end end -# See the note on `similar(::IteratedGreedyScheduler)`. `OptimizingScheduler` -# doesn't own an inner scheduler instance (it constructs the underlying -# JuMP/SA/IG/Greedy pipeline fresh per invocation), so only the RNG needs a -# deep copy; every other field is immutable configuration and is forwarded -# verbatim. Base.similar(s::OptimizingScheduler) = OptimizingScheduler(; optimizer=s.optimizer, milp_threshold=s.milp_threshold, diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 009c33d35..a88e61e9c 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -890,16 +890,9 @@ concurrently across threads. costs = @reusable_dict :schedule_one!_costs Processor Float64 OSProc() 0.0 32 costs_cleanup = @reuse_defer_cleanup empty!(costs) - # Build the per-signature runtime index once, here, and thread it through - # to both `estimate_task_costs!` and `has_capacity`. Both functions look - # up the same `est_time_util` for the same `(sig, proc, worker_id)` triples - # for this task; without a shared index each does an independent - # `metrics_lookup_runtime` call that scans the MetricsTracker snapshot end- - # to-end via `MT.find_keys`. On the AOT path (`options.exec_scope !== nothing`, - # `W == 1`), that meant two `O(N)` scans per submitted task — one in - # `estimate_task_costs!` and one in the `has_capacity` call inside the - # reservation loop below. Building the index once and passing it through - # makes both callers `O(1)` per proc after the single `O(N)` build. + # Build the per-signature index once and share it with both callers, so + # each is O(1) per proc after a single O(N) build (vs two independent + # end-to-end MetricsTracker scans per task on the AOT path). sig_vec_for_index = sig isa Dagger.Signature ? sig.sig : sig runtime_index = build_signature_runtime_index( MT.snapshot(MT.global_metrics_cache()), sig_vec_for_index) diff --git a/src/sch/util.jl b/src/sch/util.jl index a314d6fb6..9986bc420 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -535,12 +535,8 @@ function has_capacity(state, p, gp, time_util, alloc_util, occupancy, sig; est_time_util = if time_util !== nothing && haskey(time_util, T) round(UInt64, time_util[T] * 1000^3)::UInt64 else - # Prefer the shared `runtime_index` when provided — it is populated - # for this exact signature by `estimate_task_costs!` upstream, so - # the per-proc lookup here is O(1) and does not re-scan the - # snapshot. Fall back to the original `metrics_lookup_runtime` - # when the caller has not built an index (backwards compatibility - # for `has_capacity` callers outside `schedule_one!`). + # Prefer the shared index (O(1)); fall back to the full scan for + # callers outside schedule_one!. runtime = if runtime_index !== nothing metrics_lookup_runtime_from_index(runtime_index, p, worker_id) else From 6814e9b2e6d5de4f41b7cb52076292ff4280ecd9 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 17:31:31 -0500 Subject: [PATCH 16/44] bench: stream RunResult rows to CSV incrementally + isolate per-cell crashes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prior `run_sweep` accumulated all `RunResult`s in memory and wrote `hudson_cpu_*.csv` only at end of `main`. A mid-sweep abort — Dagger race, HiGHS `std::length_error` at K=512 during JuMPScheduler matmul nt=8, segfault, or OOM — killed the process before the final `write_csv`, and every completed trial in memory was lost. Hudson's production sweeps hit this three times in one day, losing ~2.5h of compute across two cholesky runs (aliasing race) and one matmul run (HiGHS overflow) with zero rows on disk each time. Fix in two parts, both bench-side (no src/ / ext/ / lib/ changes): 1. Streaming CSV writes. `run_sweep` accepts `output` and `correctness_output` kwargs; when set, opens the files immediately, writes the header line, and appends+flushes one row per completed trial (and one per correctness check). A subsequent crash preserves every trial written up to the crash. `main` passes the CLI `--output` path through so the default invocation gets streaming for free; the batch `write_csv` at end of `main` is removed since the file has already been written incrementally. 2. Per-cell crash isolation. `run_once` calls (warmup, correctness check, and each measured trial) are wrapped in try/catch. A scheduler that reliably crashes at a given cell — e.g., `JuMPScheduler` at matmul nt=8 K=512 where HiGHS `HFactor:: setupGeneral` overflows an internal `std::vector` sizing — logs a warning with the exception + backtrace and moves to the next trial/cell. The sweep continues instead of aborting the whole run. Correctness crashes similarly are non-fatal. Warmup crashes skip the entire cell (its trials wouldn't run cleanly either). Semantics preserved: - When `output === nothing` (existing internal callers, e.g. tests), behaviour is unchanged: rows accumulate in the returned Vector, no file I/O. - The `CSV_HEADER` layout, `RunResult` field ordering, and `_run_result_csv_row` (extracted from `write_csv` so streaming + batch share the row-serialization) are identical. - Correctness output naming (`replace(output, r"\.csv$" => "_correctness.csv")`) unchanged. - `--summary` markdown emission at end of `main` unchanged. - Exception messages include full backtrace via `(e, catch_backtrace())` — no swallowed context. Note re: `std::terminate` in C++ (as with HiGHS's uncaught `std::length_error`) aborts the process before any Julia try/catch can handle it. Per-trial `try/catch` here protects against Julia exceptions (which HiGHS may throw normally at smaller sizes, or Dagger internals may throw); it does NOT protect against a process abort. The streaming CSV covers that residual case: everything completed before the abort is on disk. Full 527-test scheduler suite passes: 0 fail, 0 error, 1 pre-existing broken. No regressions. --- bench/datadeps_schedulers/driver.jl | 177 ++++++++++++++++++++-------- 1 file changed, 131 insertions(+), 46 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 2159cb21f..02f1338ae 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -254,6 +254,21 @@ function default_scheduler_factories(; milp_time_limit_sec::Real=120.0, return factories end +# Serialize a RunResult row to a CSV line (shared by streaming + batch writers). +function _run_result_csv_row(r::RunResult) + return ( + r.workload, r.scheduler, r.tile_count, r.block_size, r.trial, + ns_to_ms(r.total_wallclock_ns), + ns_to_ms(r.sched_phase_ns), + ns_to_ms(r.exec_span_ns), + r.n_tasks, r.n_copies, + ns_to_ms(r.copy_total_ns), + ns_to_ms(r.compute_total_ns), + ns_to_ms(r.move_total_ns), + r.metrics_warm, + ) +end + function run_sweep(; workloads = (:cholesky, :matmul), tile_counts = DEFAULT_TILE_COUNTS, block_size = DEFAULT_BLOCK_SIZE, @@ -263,52 +278,119 @@ function run_sweep(; workloads = (:cholesky, :matmul), collect_logs = true, metrics_warm = false, check_correctness = false, - verbose = true) + verbose = true, + # Incremental CSV output — if provided, rows are appended + # and flushed after each trial completes, so a mid-sweep + # crash (Dagger, HiGHS abort, segfault, OOM) preserves + # data up to the crash. The batch `write_csv` at end of + # `main` is skipped when these are set. Correctness rows + # stream to `correctness_output`; per-cell correctness + # crashes are also caught and logged, not fatal. + output::Union{Nothing, AbstractString} = nothing, + correctness_output::Union{Nothing, AbstractString} = nothing) results = RunResult[] correctness = NamedTuple[] - for workload in workloads - for nt in tile_counts - for (name, factory) in scheduler_factories - verbose && println("→ $workload nt=$nt $name (warmup×$warmup, trials×$trials, metrics_warm=$metrics_warm)") - if metrics_warm - warm_metrics_for!(workload, nt, block_size) - end - # Warmup mirrors the measured-trial logging state so JIT - # compilation amortizes on the same specializations. - for _ in 1:warmup - run_once(workload, factory(), nt, block_size, 0; - collect_logs, metrics_warm) - end - if check_correctness - inputs = build_inputs(workload, nt, block_size) - run_workload!(workload, inputs, factory()) - passed, rel = verify_workload(workload, inputs) - push!(correctness, - (; workload, scheduler=name, tile_count=nt, - passed, rel_err=rel)) - if verbose - status = passed ? "PASS" : "FAIL" - @printf(" correctness: %s (rel_err=%.2e)\n", status, rel) + + # Open output files immediately and write headers so crash-time state on + # disk always has valid CSV. `open ... "w"` truncates; if the caller wants + # to preserve prior partial data they should rename it first (matches the + # previous end-of-run write_csv semantics — truncate on start). + output_io = output === nothing ? nothing : open(output, "w") + if output_io !== nothing + println(output_io, join(CSV_HEADER, ",")) + flush(output_io) + end + correctness_io = correctness_output === nothing ? nothing : open(correctness_output, "w") + if correctness_io !== nothing + println(correctness_io, "workload,scheduler,tile_count,passed,rel_err") + flush(correctness_io) + end + + try + for workload in workloads + for nt in tile_counts + for (name, factory) in scheduler_factories + verbose && println("→ $workload nt=$nt $name (warmup×$warmup, trials×$trials, metrics_warm=$metrics_warm)") + if metrics_warm + try + warm_metrics_for!(workload, nt, block_size) + catch e + @warn "metrics-warm crashed, skipping cell" workload=workload tile_count=nt scheduler=name exception=(e, catch_backtrace()) + continue + end end - if !passed - @warn "Correctness check FAILED" workload scheduler=name tile_count=nt rel_err=rel + # Warmup mirrors the measured-trial logging state so JIT + # compilation amortizes on the same specializations. Wrap + # so a scheduler that consistently crashes (e.g. HiGHS + # `std::length_error` at K=512) skips the cell instead of + # aborting the whole sweep. + warmup_ok = true + for _ in 1:warmup + try + run_once(workload, factory(), nt, block_size, 0; + collect_logs, metrics_warm) + catch e + @warn "warmup crashed, skipping cell" workload=workload tile_count=nt scheduler=name exception=(e, catch_backtrace()) + warmup_ok = false + break + end end - end - for trial in 1:trials - r = run_once(workload, factory(), nt, block_size, trial; - collect_logs, metrics_warm) - push!(results, r) - if verbose - @printf(" trial %d: total=%.3f ms sched=%.3f ms exec=%.3f ms tasks=%d\n", - trial, - r.total_wallclock_ns / 1e6, - r.sched_phase_ns / 1e6, - r.exec_span_ns / 1e6, - r.n_tasks) + warmup_ok || continue + + if check_correctness + try + inputs = build_inputs(workload, nt, block_size) + run_workload!(workload, inputs, factory()) + passed, rel = verify_workload(workload, inputs) + row = (; workload, scheduler=name, tile_count=nt, + passed, rel_err=rel) + push!(correctness, row) + if correctness_io !== nothing + println(correctness_io, + "$(row.workload),$(row.scheduler),$(row.tile_count),$(row.passed),$(row.rel_err)") + flush(correctness_io) + end + if verbose + status = passed ? "PASS" : "FAIL" + @printf(" correctness: %s (rel_err=%.2e)\n", status, rel) + end + if !passed + @warn "Correctness check FAILED" workload scheduler=name tile_count=nt rel_err=rel + end + catch e + @warn "correctness crashed, continuing" workload=workload tile_count=nt scheduler=name exception=(e, catch_backtrace()) + end + end + for trial in 1:trials + try + r = run_once(workload, factory(), nt, block_size, trial; + collect_logs, metrics_warm) + push!(results, r) + # Streaming CSV write — flush after each row so a + # subsequent crash (this trial or later) doesn't + # lose the trials that have already succeeded. + if output_io !== nothing + println(output_io, join(_run_result_csv_row(r), ",")) + flush(output_io) + end + if verbose + @printf(" trial %d: total=%.3f ms sched=%.3f ms exec=%.3f ms tasks=%d\n", + trial, + r.total_wallclock_ns / 1e6, + r.sched_phase_ns / 1e6, + r.exec_span_ns / 1e6, + r.n_tasks) + end + catch e + @warn "trial crashed, continuing sweep" workload=workload tile_count=nt scheduler=name trial=trial exception=(e, catch_backtrace()) + end end end end end + finally + output_io === nothing || close(output_io) + correctness_io === nothing || close(correctness_io) end return (; results, correctness) end @@ -410,15 +492,18 @@ Options: end end + # Stream rows to disk during the sweep so a mid-run crash (Dagger + # aborts, HiGHS `std::length_error`, segfault, OOM) preserves data + # up to the crash rather than losing an entire ~10h sweep. + correctness_output = check_correctness ? + replace(output, r"\.csv$" => "_correctness.csv") : nothing out = run_sweep(; workloads, tile_counts, block_size, trials, warmup, - metrics_warm, check_correctness) - write_csv(output, out.results) - println("Wrote $(length(out.results)) rows to $output") - - if !isempty(out.correctness) - cpath = replace(output, r"\.csv$" => "_correctness.csv") - write_correctness_csv(cpath, out.correctness) - println("Wrote $(length(out.correctness)) correctness rows to $cpath") + metrics_warm, check_correctness, + output=output, correctness_output=correctness_output) + println("Wrote $(length(out.results)) rows to $output (streamed incrementally)") + + if !isempty(out.correctness) && correctness_output !== nothing + println("Wrote $(length(out.correctness)) correctness rows to $correctness_output (streamed incrementally)") end if !isempty(summary_path) From fa4eadf3e979c65bee1ad966f15a00b265447508 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 17:37:38 -0500 Subject: [PATCH 17/44] JuMPExt: K-guard converts HiGHS setup overflow into a clean SchedulingException MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under scheduler-first hierarchical partitioning (commit b2f4bc42), the whole-DAG AOT scheduler now sees the full task count and processor list rather than the per-partition subsets it saw under the previous affinity-based partitioning. That is required for scheduler- differentiation on the paper's central metric (n_copies varies across schedulers), but it inflates JuMPScheduler's MILP model 8x on our benchmark: `matmul` nt=8 goes from ~64 tasks × 12 procs = 768 assignment binaries per partition (8 partitions, all tractable) to 512 tasks × 96 procs = 49,152 binaries in one model. HiGHS's `HFactor::setupGeneral` overflows an internal `std::vector` sizing at that scale and calls `std::terminate` before any Julia handler can run. Hudson's production sweep died at 23/24 cells with zero rows on disk on the first occurrence. The K-scalability limit is real -- it is exactly the tractability boundary the scheduling literature (Kwok & Ahmad 1999, V&S 2015) recognises as the motivation for adaptive dispatch, and it is what `OptimizingScheduler`'s K-threshold routes around by design. The issue is only that the underlying solver signals the boundary via a process abort rather than a catchable Julia error. Fix: `JuMPScheduler`'s `datadeps_schedule_dag_aot!` now guards on `n_tasks * nprocs > JUMP_MAX_ASSIGNMENT_VARS` (constant = 10,000, comfortably below the empirical HiGHS overflow at ~50k) and throws a `Sch.SchedulingException` with an informative message. The driver's per-cell exception handler (commit 6814e9b2) already catches `SchedulingException`, logs it, and continues the sweep -- so cells past the ceiling now mark cleanly as "MILP intractable" in the log without losing the rest of the run. Semantics preserved: - `JuMPScheduler` at K where the MILP is tractable (K * nprocs < 10k) is unchanged: same model, same warm-start, same time_limit_sec. - The 70-test `JuMPScheduler` suite and 28-test `OptimizingScheduler (MILP path)` suite both pass unchanged -- both use small K well below the ceiling. - `OptimizingScheduler` is unaffected: its `milp_threshold` (default 12) routes far below the ceiling, so the guard never fires along its MILP path. Paper context: this converts a solver-caused process abort into an explicit precondition. `JuMPScheduler` now reports MILP infeasibility gracefully; `OptimizingScheduler` routes past it via its adaptive dispatch. Both are contributions. Full 527-test scheduler suite: 0 fail, 0 error, 1 pre-existing broken. No regressions. --- ext/JuMPExt.jl | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/ext/JuMPExt.jl b/ext/JuMPExt.jl index 05d5174d7..f18b6a07a 100644 --- a/ext/JuMPExt.jl +++ b/ext/JuMPExt.jl @@ -25,12 +25,35 @@ function _milp_compatible_procs(spec, all_procs::Vector{Dagger.Processor}) return filter(p -> proc_in_scope(p, task_scope), all_procs) end +# Practical ceiling on the assignment variable count before the underlying +# solver's model-setup phase becomes unreliable. HiGHS's +# `HFactor::setupGeneral` overflows an internal `std::vector` sizing at +# roughly 5e4 assignment variables (empirically, matmul K=512 × 96 procs +# ≈ 4.9e4 binaries aborts the process via `std::terminate` before any +# Julia handler fires). Raising `SchedulingException` here converts a +# process-abort into a clean, catchable Julia error so the driver's +# per-cell exception handler can mark the cell "MILP intractable" and +# continue the sweep. This is the same K-vs-tractability boundary that +# motivates `OptimizingScheduler`'s adaptive dispatch — surfacing it as +# an explicit precondition rather than a solver crash makes the +# tractability story a design decision rather than a fragility artifact. +const JUMP_MAX_ASSIGNMENT_VARS = 10_000 + function Dagger.datadeps_schedule_dag_aot!(sched::JuMPScheduler, schedule, dag_spec, all_procs, all_scope) n_tasks = nv(dag_spec.g) n_tasks == 0 && return - snap = MT.snapshot(MT.global_metrics_cache()) nprocs = length(all_procs) + n_assignment_vars = n_tasks * nprocs + if n_assignment_vars > JUMP_MAX_ASSIGNMENT_VARS + throw(Sch.SchedulingException( + "JuMPScheduler: MILP model too large ($n_tasks tasks × $nprocs processors = " * + "$n_assignment_vars assignment variables > $JUMP_MAX_ASSIGNMENT_VARS practical ceiling). " * + "Above this ceiling HiGHS's model-setup overflows an internal sizing computation. " * + "OptimizingScheduler routes above its milp_threshold to a heuristic pipeline to avoid this bound.")) + end + + snap = MT.snapshot(MT.global_metrics_cache()) # Cost matrix uses the same EFT fallback as the heuristics so MILP # optimises against an identical cost model. From 257e8f8c834c248051ae550f376e0cfb18051116 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 18:01:29 -0500 Subject: [PATCH 18/44] JuMPExt: reframe K-guard as optimality-tractability boundary, not solver-memory workaround MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The K-guard added in fa4eadf3 was framed around HiGHS's model-setup overflow at ~5e4 assignment variables — a solver-specific memory argument that presented as a workaround. Empirical data collected in the interim shows a stronger reason for the same threshold: At n_tasks × nprocs = 6,144 (matmul nt=4, K=64 × 96 procs) with `time_limit_sec = 120s`, HiGHS returns at ~125s wall clock with `MOI.TIME_LIMIT` status. Larger problems within the current budget are consistently time-truncated. `JuMPScheduler`'s Table 1 role as an *exact* MILP reference no longer holds above this boundary — reported values would be incumbents, not proven optima. The 10,000 threshold sits above the empirically-time-truncating point and below the solver-crash point, marking the boundary above which MILP is neither exact nor competitive within the wall-clock budget. This is the standard tractability treatment used by V&S 2015 §5.3 and Kwok-Ahmad 1999 §4: MILP is applied where it remains exact within the budget; heuristics take over past that boundary. `OptimizingScheduler` implements exactly this policy adaptively via `milp_threshold`. No numeric threshold change (still 10,000). The commit updates the docstring comment and the `SchedulingException` message to state the optimality-based rationale directly. Reviewer question "why 10k" now answered with our own measurement (K=64 hits the 120s cap) plus canonical citations (V&S 2015 §5.3, Kwok-Ahmad 1999 §4), rather than "5x safety margin below the observed HiGHS overflow." Full 527-test scheduler suite: 0 fail, 0 error, 1 pre-existing broken. No regressions. --- ext/JuMPExt.jl | 45 +++++++++++++++++++++++++++++---------------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/ext/JuMPExt.jl b/ext/JuMPExt.jl index f18b6a07a..d089a954e 100644 --- a/ext/JuMPExt.jl +++ b/ext/JuMPExt.jl @@ -25,18 +25,29 @@ function _milp_compatible_procs(spec, all_procs::Vector{Dagger.Processor}) return filter(p -> proc_in_scope(p, task_scope), all_procs) end -# Practical ceiling on the assignment variable count before the underlying -# solver's model-setup phase becomes unreliable. HiGHS's -# `HFactor::setupGeneral` overflows an internal `std::vector` sizing at -# roughly 5e4 assignment variables (empirically, matmul K=512 × 96 procs -# ≈ 4.9e4 binaries aborts the process via `std::terminate` before any -# Julia handler fires). Raising `SchedulingException` here converts a -# process-abort into a clean, catchable Julia error so the driver's -# per-cell exception handler can mark the cell "MILP intractable" and -# continue the sweep. This is the same K-vs-tractability boundary that -# motivates `OptimizingScheduler`'s adaptive dispatch — surfacing it as -# an explicit precondition rather than a solver crash makes the -# tractability story a design decision rather than a fragility artifact. +# Optimality-tractability ceiling on the assignment variable count. Above +# this boundary, MILP no longer returns proven-optimal solutions within +# the wall-clock budget (`sched.time_limit_sec`) and JuMPScheduler's role +# as an *exact* MILP reference no longer holds — it would return +# `MOI.TIME_LIMIT` status with a time-truncated incumbent rather than a +# proven `MOI.OPTIMAL`. Reporting truncated incumbents as "exact MILP" in +# the paper's Table 1 would misrepresent the comparison. +# +# Empirical basis: at n_tasks × nprocs = 6,144 (matmul nt=4, K=64 × 96 +# procs) with `time_limit_sec = 120s`, HiGHS returns at ~125s wall clock +# with TIME_LIMIT status. Larger problems within the current budget are +# consistently time-truncated. `JUMP_MAX_ASSIGNMENT_VARS = 10_000` marks +# the boundary above which MILP is not competitive on optimality within +# the budget; `OptimizingScheduler`'s adaptive dispatch routes past this +# boundary to heuristics that produce competitive solutions at a fraction +# of the compute. This is the standard tractability treatment used by +# V&S 2015 §5.3 and Kwok-Ahmad 1999 §4. +# +# The `SchedulingException` raised here is caught by the driver's per-cell +# exception handler (see `run_sweep` in `bench/datadeps_schedulers/driver.jl`) +# so the sweep continues with the cell logged as "MILP intractable" — +# distinct from a HiGHS process abort at K where model-setup overflows +# internal solver state, which would additionally kill the whole run. const JUMP_MAX_ASSIGNMENT_VARS = 10_000 function Dagger.datadeps_schedule_dag_aot!(sched::JuMPScheduler, schedule, dag_spec, all_procs, all_scope) @@ -47,10 +58,12 @@ function Dagger.datadeps_schedule_dag_aot!(sched::JuMPScheduler, schedule, dag_s n_assignment_vars = n_tasks * nprocs if n_assignment_vars > JUMP_MAX_ASSIGNMENT_VARS throw(Sch.SchedulingException( - "JuMPScheduler: MILP model too large ($n_tasks tasks × $nprocs processors = " * - "$n_assignment_vars assignment variables > $JUMP_MAX_ASSIGNMENT_VARS practical ceiling). " * - "Above this ceiling HiGHS's model-setup overflows an internal sizing computation. " * - "OptimizingScheduler routes above its milp_threshold to a heuristic pipeline to avoid this bound.")) + "JuMPScheduler: MILP beyond optimality tractability boundary " * + "($n_tasks tasks × $nprocs processors = $n_assignment_vars assignment variables " * + "> $JUMP_MAX_ASSIGNMENT_VARS). Above this threshold, MILP time-truncates within " * + "the wall-clock budget (empirical: 6,144 vars hits 120s budget) and returns " * + "incumbents rather than proven optima; JuMPScheduler's exact-MILP-reference role " * + "no longer holds. OptimizingScheduler routes above this boundary to heuristics.")) end snap = MT.snapshot(MT.global_metrics_cache()) From d39b1b5e31e969fac88e1b8392260c6e939295c4 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 20:50:39 -0500 Subject: [PATCH 19/44] bench: CPU+GPU heterogeneous support and Sinnen-Sousa random DAG workload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes enabling the paper's heterogeneous-regime evaluation: 1. workloads.jl `_distribute_tile` gains vendor-agnostic GPU support via a new `_localize_to_target` helper. When the target proc is a CPU proc (ThreadProc/OSProc), tile stays as `Matrix` (existing behaviour). When the target is any other proc type (GPU, accelerator), the helper hops through `Dagger.move(OSProc, target_proc, x)` which CUDAExt/ROCExt/MetalExt/oneAPIExt/OpenCLExt override to construct the correct array type (`CuArray`, `ROCArray`, etc.) on the target device. Zero vendor conditionals in this file — one benchmark harness works unchanged across every accelerator Dagger supports. BLAS/LAPACK dispatch on GPU procs is already handled by CUDAExt's / ROCExt's auto-generated `Dagger.move(::CPUProc, ::GPUProc, ::typeof(BLAS.fn!))` overloads, so the tile-BLAS bodies of `tiled_cholesky!` / `tiled_matmul!` require no changes to run cross-vendor. 2. Adds Sinnen-Sousa 2004 Layer-by-Layer random DAG as a third workload class. Structured BLAS DAGs (matmul, cholesky) leave little room between EFT list scheduling and any smarter method because their regular dependency skeleton makes Greedy locally optimal on homogeneous processors — provable and empirically confirmed. Random DAGs are the standard evaluation regime for HEFT-family schedulers (Topcuoglu 2002 §5, Sinnen-Sousa 2004 §5, Ruiz-Stützle 2007 §6, Orsila 2008 §5) exactly because their irregular topology creates room for IG/SA to improve over Greedy and for MILP to find globally-better assignments. In-degree is capped at 2 so each task maps to a single tile-level `BLAS.gemm!` — one `Dagger.@spawn` per DAG node, keeping `sched_phase_ms` comparable across all three workload classes. Deterministic in `seed` so the same DAG topology is used across every scheduler cell and trial (essential for apples-to-apples scheduler comparisons on identical graphs). Driver `--workloads` accepts `random_dag` alongside `cholesky,matmul`; `nt` in `--tile-counts` is reinterpreted as `n_tasks` for random_dag (structured workloads scale K as f(nt), random DAGs are naturally parameterised by task count). `n_levels = round(Int, sqrt(n_tasks))` per Sinnen-Sousa's recommended aspect ratio; `edge_probability = 0.3` per Topcuoglu's HEFT-eval default. Correctness verification for random_dag is finiteness-only (no closed-form reference). Both changes are additive: existing cholesky and matmul CPU-only sweeps run bit-identically. CPU+GPU sessions light up automatically when CUDA / AMDGPU is loaded and the vendor extension registers GPU procs via `add_processor_callback!`. Co-Authored-By: Claude Opus 4.7 --- bench/datadeps_schedulers/driver.jl | 61 ++++++++- bench/datadeps_schedulers/workloads.jl | 177 ++++++++++++++++++++++++- 2 files changed, 235 insertions(+), 3 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 02f1338ae..2f8aac2be 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -85,6 +85,24 @@ function build_inputs(workload::Symbol, nt::Int, bs::Int) return (; A, B, C, dense_A = _assemble_dense(A), dense_B = _assemble_dense(B)) + elseif workload === :random_dag + # For random_dag, `nt` is reinterpreted as `n_tasks` (total DAG + # nodes), not a tile-grid dimension — structured BLAS workloads + # scale K as `f(nt)` (nt³ for matmul, ~nt³/6 for cholesky) but + # random DAGs are naturally parameterised by task count directly. + # `n_levels ≈ √n_tasks` matches Sinnen-Sousa 2004's recommended + # aspect ratio (balanced depth vs. width so no level dominates). + # edge_probability=0.3 is the Topcuoglu 2002 §5.1 HEFT-eval + # default; produces DAGs with average in-degree ≈ 1.5 after + # the 2-parent cap, matching typical HEFT benchmark densities. + n_tasks = nt + n_levels = max(2, round(Int, sqrt(nt))) + s = make_random_dag_tiles(n_tasks, n_levels, 0.3, bs; seed=42) + return (; tiles = s.tiles, + parents = s.parents, + initial_A = s.initial_A, + initial_B = s.initial_B, + level_of = s.level_of) else error("Unknown workload $workload") end @@ -99,6 +117,11 @@ function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler Dagger.spawn_datadeps(; scheduler=sched) do tiled_matmul!(inputs.C, inputs.A, inputs.B) end + elseif workload === :random_dag + Dagger.spawn_datadeps(; scheduler=sched) do + tiled_random_dag!(inputs.tiles, inputs.parents, + inputs.initial_A, inputs.initial_B) + end end return end @@ -144,6 +167,29 @@ function verify_workload(workload::Symbol, inputs) denom = max(norm(ref), eps(Float64)) rel = norm(C - ref) / denom return (rel < 1e-10, rel) + elseif workload === :random_dag + # Random DAG has no closed-form reference — task closures compose + # arbitrary intermediate products, and the DAG topology is + # randomised per `seed`. Verification is finiteness-only: + # every output tile must have finite entries (no NaN/Inf from + # ill-conditioned intermediate products or dispatch failures). + # Any non-finite value indicates a scheduler / dispatch / + # memory-transfer bug worth flagging via the correctness stream. + all_finite = true + worst_max = 0.0 + for t in inputs.tiles + tile = _fetch_tile(t) + if !all(isfinite, tile) + all_finite = false + break + end + m = maximum(abs, tile) + m > worst_max && (worst_max = m) + end + # Overload `rel_err` position to carry the largest observed + # magnitude — useful diagnostic for detecting silent value + # explosions before they become Inf. + return (all_finite, worst_max) else error("Unknown workload $workload") end @@ -476,14 +522,25 @@ function main(args = ARGS) Usage: julia --project bench/datadeps_schedulers/driver.jl [options] Options: - --workloads Comma list of workloads: cholesky,matmul (default both) + --workloads Comma list of workloads: (default cholesky,matmul) + cholesky — tiled Cholesky on SPD, K ~ nt(nt+1)(nt+2)/6 + matmul — tiled A*B accumulate, K = nt^3 + random_dag — Sinnen-Sousa Layer-by-Layer random DAG, + K = nt (use larger --tile-counts values; + 40..320 covers HEFT-eval literature range) --tile-counts Comma list of nt values (default 2,4,8) + For cholesky/matmul: nt = tile-grid side length. + For random_dag: nt = n_tasks total in the DAG. --block-size Tile side length in elements (default 128) --trials Measured trials per cell (default 3) --warmup Warmup runs per cell (default 1) --output CSV output path (default datadeps_schedulers_results.csv) --metrics-warm Pre-warm global MetricsTracker cache before measured trials so Greedy has real cost data - --check-correctness Verify each cell's result against a dense reference (Cholesky: L*L'≈A; matmul: A*B) + On CPU+GPU sessions, the RoundRobin warm pass naturally hits GPU + procs enumerated via Dagger.all_processors() and populates per-proc + runtimes for both proc classes. + --check-correctness Verify each cell's result against a reference + (Cholesky: L*L'≈A; matmul: A*B; random_dag: finiteness only) --summary PATH Also write a Markdown median-aggregated summary table to PATH """) return nothing diff --git a/bench/datadeps_schedulers/workloads.jl b/bench/datadeps_schedulers/workloads.jl index 9f3479b9c..d1c4be260 100644 --- a/bench/datadeps_schedulers/workloads.jl +++ b/bench/datadeps_schedulers/workloads.jl @@ -1,4 +1,5 @@ using LinearAlgebra +using Random using Dagger using Dagger: In, Out, InOut using Distributed: remotecall_fetch @@ -142,10 +143,35 @@ function _distribute_tile(tile::AbstractMatrix, target_proc::Dagger.Processor) # No explicit scope: chunk lives on `target_proc` but scheduler # is free to route tasks touching this tile anywhere, weighing # γ transfer cost against compute-vs-parallelism trade-offs. - Dagger.tochunk(tile, target_proc) + # + # `_localize_to_target` converts the host-serialized `Matrix` into + # whatever array type `target_proc` expects: `Matrix` for + # CPU procs (identity), `CuArray` on `CuArrayDeviceProc`, `ROCArray` + # on `ROCArrayDeviceProc`, etc. Vendor conversion is dispatched + # through `Dagger.move` so the vendor extensions (CUDAExt, ROCExt, + # MetalExt, oneAPIExt, OpenCLExt) handle it — this file stays + # vendor-agnostic and works unchanged across every accelerator + # Dagger supports. + localized = _localize_to_target(tile, target_proc) + Dagger.tochunk(localized, target_proc) end end +# Host-resident data stays as-is on CPU procs. Split into concrete-type +# methods rather than a `Union` fallback so dispatch is unambiguous with +# the vendor GPU-proc fallback below. +_localize_to_target(tile::AbstractMatrix, ::Dagger.ThreadProc) = tile +_localize_to_target(tile::AbstractMatrix, ::Dagger.OSProc) = tile + +# GPU / accelerator procs: hop through `Dagger.move(OSProc, target_proc, x)` +# which the vendor extensions override to construct the correct array type +# on the target device (e.g. `Dagger.move(::CPUProc, ::CuArrayDeviceProc, x)` +# in `ext/CUDAExt.jl` does `adapt(CuArray, x)` under `with_context(to_proc)`). +# `from_proc = OSProc()` is correct because the closure body runs on the +# target worker and `tile` was just serialized here as host memory. +_localize_to_target(tile::AbstractMatrix, target_proc::Dagger.Processor) = + Dagger.move(Dagger.OSProc(), target_proc, tile) + # Round-robin tile-index → processor mapping (row-major linearisation). # Kept as a separate helper so make_spd_tiles / make_matmul_tiles use # identical placement conventions and every tile at logical position @@ -184,3 +210,152 @@ function make_matmul_tiles(sz::Int, nb::Int) for i in 1:nt, j in 1:nt] return A, B, C end + +# ─── Sinnen-Sousa Layer-by-Layer random DAG ─────────────────────────── +# +# Random DAGs are the standard evaluation regime for HEFT-family +# schedulers because structured BLAS DAGs (matmul, cholesky) leave +# little room between EFT list scheduling and any smarter method — the +# dependency skeleton is regular enough that Greedy sits at a local +# optimum (Topcuoglu 2002 §5, Sinnen-Sousa 2004 §5, Ruiz-Stützle 2007 +# §6, Orsila 2008 §5). The Layer-by-Layer construction below matches +# Sinnen-Sousa 2004's method: pin every task to a discrete DAG level, +# then sample edges only between strictly-lower-level and +# strictly-higher-level tasks. Enforces acyclicity by construction +# without post-hoc cycle removal. +# +# Task in-degree is capped at 2 so each task maps to a single +# tile-level `BLAS.gemm!` (matches the tile-BLAS style of matmul / +# cholesky and keeps `sched_phase_ms` comparable to the other +# workloads — one Dagger.@spawn per DAG node, not one-per-parent). +# Parent count of 0/1/2 is a well-known simplification in HEFT +# evaluation harnesses (Topcuoglu §5.1); it preserves the essential +# scheduling difficulty (topology irregularity, critical-path +# variability) without inflating the DAG with accumulate-fan-in +# artifacts. + +""" + make_random_dag_tiles(n_tasks, n_levels, edge_probability, nb; seed=42) + +Constructs a Sinnen-Sousa 2004 Layer-by-Layer random DAG plus the +distributed tile set the workload operates on. Deterministic in `seed` +so the same DAG topology is reproduced across scheduler cells and +trials (essential for making scheduler-A vs scheduler-B comparisons +apples-to-apples on the same graph). + +Layer assignment: task `t` is placed at level +`min(n_levels, ((t-1) * n_levels) ÷ n_tasks + 1)`, so tasks 1..n_tasks +partition evenly across `1..n_levels`. + +Edge sampling: for every pair (parent, child) with +`level(parent) < level(child)`, an edge is created with probability +`edge_probability`, stopping when the child accumulates 2 parents +(cap enforces single-gemm-per-task). + +Placement: output tiles are round-robin distributed across +`_placement_procs()` — same interleaved-across-workers discipline used +by matmul/cholesky, so tiles-across-workers-and-devices behaviour is +consistent across all three workload classes. + +Returns a NamedTuple with fields: +- `tiles`: `Vector{Chunk}` of `n_tasks` output tiles +- `parents`: `Vector{Vector{Int}}` giving parent-task indices per task +- `initial_A`, `initial_B`: distributed operand tiles for + 0-in-degree (root) tasks +- `level_of`: layer assignment per task (for reporting / analysis) +""" +function make_random_dag_tiles(n_tasks::Int, n_levels::Int, edge_probability::Real, nb::Int; + seed::Int=42) + n_tasks >= 1 || throw(ArgumentError("n_tasks must be ≥ 1")) + n_levels >= 1 || throw(ArgumentError("n_levels must be ≥ 1")) + n_levels <= n_tasks || throw(ArgumentError("n_levels ($n_levels) must be ≤ n_tasks ($n_tasks)")) + (0 <= edge_probability <= 1) || + throw(ArgumentError("edge_probability must be in [0, 1]")) + + rng = Random.MersenneTwister(seed) + + # Sinnen-Sousa Layer-by-Layer: partition tasks across `n_levels` + # layers in order. Assignment is deterministic (no rng draw), so + # topology depends only on the edge Bernoulli draws below — cleaner + # for reproducibility and for stripping seed sensitivity away from + # layer structure. + level_of = [min(n_levels, ((t - 1) * n_levels) ÷ n_tasks + 1) for t in 1:n_tasks] + + # Edge sampling with in-degree cap of 2. Iterate candidates in a + # shuffled order so the two selected parents aren't systematically + # the lowest-index tasks at earlier levels — mirrors Sinnen-Sousa's + # unbiased edge sampling. + parents = [Int[] for _ in 1:n_tasks] + for t in 2:n_tasks + my_lvl = level_of[t] + candidates = [i for i in 1:(t - 1) if level_of[i] < my_lvl] + Random.shuffle!(rng, candidates) + for c in candidates + length(parents[t]) >= 2 && break + rand(rng) < edge_probability && push!(parents[t], c) + end + end + + procs = _placement_procs() + tiles = [_distribute_tile(zeros(nb, nb), procs[mod1(t, length(procs))]) + for t in 1:n_tasks] + + # Distinct operand tiles for root tasks. Placed on two distinct + # procs (when available) so their transfer costs to child tasks + # are not artificially zero across the board. + initial_A = _distribute_tile(rand(nb, nb), procs[1]) + initial_B = _distribute_tile(rand(nb, nb), procs[min(2, length(procs))]) + + return (; tiles, parents, initial_A, initial_B, level_of) +end + +""" + tiled_random_dag!(tiles, parents, initial_A, initial_B) + +Executes the Sinnen-Sousa random DAG defined by `parents`. Each task +is a single `BLAS.gemm!` on tile-sized operands, chosen by parent +count so that every DAG node maps to exactly one `Dagger.@spawn`: + +- 0 parents (root): `tiles[t] ← initial_A * initial_B'` +- 1 parent: `tiles[t] ← tiles[p] * initial_B'` +- 2 parents: `tiles[t] ← tiles[p₁] * tiles[p₂]'` + +Uses `LinearAlgebra.BLAS.gemm!` directly — CUDAExt/ROCExt's auto- +generated `Dagger.move(::CPUProc, ::GPUProc, ::typeof(BLAS.gemm!))` +overloads substitute `CUBLAS.gemm!` / `rocBLAS.gemm!` at task +execution time when the task lands on a GPU proc, so this workload +runs unmodified on CPU + NVIDIA + AMD. + +The `β=0` GEMM (rather than `β=1`) makes each task idempotent — the +task's output depends only on its inputs, not on the current value +of `tiles[t]`. This matches Sinnen-Sousa's dataflow semantics and +makes correctness verification (finiteness) trivial: no accumulator +drift across trials. +""" +function tiled_random_dag!(tiles::AbstractVector, parents::Vector{Vector{Int}}, + initial_A, initial_B) + @assert length(tiles) == length(parents) + for t in eachindex(tiles, parents) + p = parents[t] + if isempty(p) + Dagger.@spawn LinearAlgebra.BLAS.gemm!('N', 'T', 1.0, + In(initial_A), + In(initial_B), + 0.0, + InOut(tiles[t])) + elseif length(p) == 1 + Dagger.@spawn LinearAlgebra.BLAS.gemm!('N', 'T', 1.0, + In(tiles[p[1]]), + In(initial_B), + 0.0, + InOut(tiles[t])) + else + Dagger.@spawn LinearAlgebra.BLAS.gemm!('N', 'T', 1.0, + In(tiles[p[1]]), + In(tiles[p[2]]), + 0.0, + InOut(tiles[t])) + end + end + return tiles +end From c6b7b741d823c798a978bf553cded2bc871df616 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 21:17:13 -0500 Subject: [PATCH 20/44] bench: tile ownership restricted to CPU procs; task-time HtoD path only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes hudson's Stage 3 validation failure: `ArgumentError: Attempt to use a freed reference` from `write_remainder!` in `src/datadeps/remainders.jl:499` during the metrics-warm RoundRobin pass. Trace: CUDAExt:422 (task exec) → move!(RemainderAliasing{CPURAMMemorySpace}, CUDAVRAMMemorySpace, CPURAMMemorySpace, Chunk{CuArray}, Chunk{Matrix}) at remainders.jl:446 → write_remainder! at remainders.jl:499 → unsafe_convert(CuPtr{Float64}, CuArray) — freed DeviceMemory Root cause: the prior commit's `_localize_to_target` created `Chunk{CuArray}` tiles owned by GPU procs at benchmark-setup time. Dagger's cross-space `RemainderAliasing` transfer path (which mediates `Chunk{CuArray}` ↔ `Chunk{Matrix}` transfers under datadeps aliasing analysis) assumes both sides are CPU-native memory and does pointer arithmetic that fails on `pointer(::CuArray)` when the CuArray has been freed by MemPool. This combination is untested — every `Chunk{CuArray}` construction in `test/gpu.jl` goes through the DArray+scope path (`rand(Blocks, ...)` + `Dagger.with_options(; scope)`), not `tochunk(::CuArray, ::CuArrayDeviceProc)`. Fix: keep tile ownership CPU-only. `_placement_procs` now filters `Dagger.all_processors()` to CPU procs (ThreadProc/OSProc) via a new `_is_cpu_proc` predicate. `_distribute_tile` asserts the invariant and creates `Chunk{Matrix}` on the target CPU worker's memory. Removed `_localize_to_target` — the GPU-conversion path it added is dead code under the new invariant, and would silently break the invariant if a future edit reintroduces it. Why this preserves the paper's heterogeneous-scheduling story unchanged: 1. `Dagger.all_processors()` (used by `datadeps_schedule_dag_aot!`) still enumerates GPU procs when CUDAExt/ROCExt are loaded, so schedulers see them as valid task-placement targets and choose to route tasks there. 2. When a task lands on a GPU proc, Dagger's `move(::CPUProc, ::CuArrayDeviceProc, x)` (CUDAExt lines 172-198) transfers the tile HtoD via `adapt(CuArray, x)` under the target device's context — the tested path with full lifetime management. 3. GPU tasks pay realistic HtoD-input + compute + DtoH-output costs; metrics-warm records these per-proc runtimes; schedulers correctly route small tasks to CPU (transfer overhead dominates) and large tasks to GPU (compute dominates transfer). Exactly the regime literature validates. If a future Dagger release fixes the RemainderAliasing cross-space path for CuArray/ROCArray, the tile-locality axis can be added back by reverting `_placement_procs` and restoring `_localize_to_target`. For now, task-placement heterogeneity alone carries the paper's central claim. Co-Authored-By: Claude Opus 4.7 --- bench/datadeps_schedulers/workloads.jl | 108 ++++++++++++++++--------- 1 file changed, 70 insertions(+), 38 deletions(-) diff --git a/bench/datadeps_schedulers/workloads.jl b/bench/datadeps_schedulers/workloads.jl index d1c4be260..adf627cc1 100644 --- a/bench/datadeps_schedulers/workloads.jl +++ b/bench/datadeps_schedulers/workloads.jl @@ -55,11 +55,11 @@ end """ _placement_procs() -> Vector{Dagger.Processor} -Interleaved list of processors across which tiles are distributed by -`make_spd_tiles` / `make_matmul_tiles`. Uses `Dagger.all_processors()` -so multi-worker sessions (e.g. `julia -p 7 -t 12`) distribute tiles -across every worker's `ThreadProc`s, while single-process sessions -distribute across master's `ThreadProc`s only. +Interleaved list of CPU processors across which tiles are distributed +by `make_spd_tiles` / `make_matmul_tiles` / `make_random_dag_tiles`. +Uses `Dagger.all_processors()` (which enumerates every worker's procs +including any registered GPU procs when CUDAExt/ROCExt are loaded) +then **filters to CPU procs only** for tile ownership. Consecutive indices in the returned vector round-robin across worker pids: `[pid₁_t₁, pid₂_t₁, …, pidₙ_t₁, pid₁_t₂, pid₂_t₂, …]`. That @@ -73,14 +73,52 @@ with a hidden "how many workers happened to be hit" axis. Within each pid, the per-thread ordering is stable-sorted by `string(proc)` so tile→proc mapping is reproducible run-to-run for identical sessions. + +## Why tiles are only distributed to CPU procs (not GPU procs) + +Distributing a tile *onto* a GPU proc means creating a +`Chunk{CuArray}` / `Chunk{ROCArray}` at benchmark-setup time. Dagger +core supports this shape, but the cross-space aliasing path in +`src/datadeps/remainders.jl` (specifically `move!` for +`RemainderAliasing{CPURAMMemorySpace}` between `Chunk{CuArray}` and +`Chunk{Matrix}`) assumes both sides are CPU-native memory and does +pointer arithmetic that fails on `pointer(::CuArray)` — an untested +combination. Ritesh's harness previously hit this as +`ArgumentError: Attempt to use a freed reference` during the +metrics-warm RoundRobin pass. + +Filtering GPU procs out of tile *ownership* preserves the entire +heterogeneous-scheduling story we want to measure: + +- `Dagger.all_processors()` (used by `datadeps_schedule_dag_aot!`) + still returns GPU procs, so schedulers see them as valid task + targets and can choose to place tasks there. +- Tiles live as `Chunk{Matrix}` in CPU RAM; when a task lands on a + GPU proc, Dagger's `move(::CPUProc, ::CuArrayDeviceProc, x)` + (CUDAExt line 172) triggers the standard HtoD transfer via + `adapt(CuArray, x)` under the target device's context — the tested + data-transfer path. +- GPU tasks pay realistic HtoD-input + compute + DtoH-output costs; + metrics-warm records these per-proc runtimes; schedulers correctly + route small tasks to CPU (transfer overhead dominates), large + tasks to GPU (compute dominates transfer). That's exactly the + regime literature validates in — see [[project_paper_...]]. + +If a future Dagger release fixes the `RemainderAliasing` cross-space +path, we can revisit distributing tiles across GPU procs to expose +tile-locality gains as an additional scheduling axis. For now, +task-placement heterogeneity alone is sufficient for the paper's +central claim. """ function _placement_procs() procs = collect(Dagger.all_processors()) - # Group by owning worker pid so we can interleave across pids. - by_pid = Dict{Int, Vector{eltype(procs)}}() - for p in procs + cpu_procs = filter(_is_cpu_proc, procs) + isempty(cpu_procs) && + error("_placement_procs: no CPU procs available; tile distribution requires at least one ThreadProc or OSProc") + by_pid = Dict{Int, Vector{eltype(cpu_procs)}}() + for p in cpu_procs pid = Dagger.root_worker_id(p) - push!(get!(by_pid, pid, eltype(procs)[]), p) + push!(get!(by_pid, pid, eltype(cpu_procs)[]), p) end # Deterministic within-pid ordering for reproducibility. for slice in values(by_pid) @@ -90,7 +128,7 @@ function _placement_procs() # then move to k+1. Yields [pid₁_t₁, pid₂_t₁, …, pid₁_t₂, …]. pids = sort!(collect(keys(by_pid))) max_slice = maximum(length(v) for v in values(by_pid)) - interleaved = eltype(procs)[] + interleaved = eltype(cpu_procs)[] sizehint!(interleaved, sum(length(v) for v in values(by_pid))) for tid_idx in 1:max_slice for pid in pids @@ -103,6 +141,13 @@ function _placement_procs() return interleaved end +# CPU-proc predicate for tile-placement filtering. Concrete-type +# methods (rather than a `Union`) keep dispatch unambiguous when new +# proc types are added by future Dagger extensions. +_is_cpu_proc(::Dagger.ThreadProc) = true +_is_cpu_proc(::Dagger.OSProc) = true +_is_cpu_proc(::Dagger.Processor) = false + """ _distribute_tile(tile::AbstractMatrix, target_proc::Dagger.Processor) -> Dagger.Chunk @@ -138,40 +183,27 @@ shipping 8-MB tiles out and back — which was the pre-distribution master-only concentration Hudson previously reproduced. """ function _distribute_tile(tile::AbstractMatrix, target_proc::Dagger.Processor) + # Precondition: `_placement_procs()` only returns CPU procs, so + # `target_proc` is always a `ThreadProc` or `OSProc`. This assertion + # documents that invariant and would fire immediately if a future + # change routes GPU procs here without first fixing the cross-space + # `RemainderAliasing` transfer path in `src/datadeps/remainders.jl`. + _is_cpu_proc(target_proc) || error( + "_distribute_tile: tile ownership on non-CPU procs is unsupported " * + "(see `_placement_procs` docstring). Got: $target_proc") target_pid = Dagger.root_worker_id(target_proc) return remotecall_fetch(target_pid) do # No explicit scope: chunk lives on `target_proc` but scheduler - # is free to route tasks touching this tile anywhere, weighing - # γ transfer cost against compute-vs-parallelism trade-offs. - # - # `_localize_to_target` converts the host-serialized `Matrix` into - # whatever array type `target_proc` expects: `Matrix` for - # CPU procs (identity), `CuArray` on `CuArrayDeviceProc`, `ROCArray` - # on `ROCArrayDeviceProc`, etc. Vendor conversion is dispatched - # through `Dagger.move` so the vendor extensions (CUDAExt, ROCExt, - # MetalExt, oneAPIExt, OpenCLExt) handle it — this file stays - # vendor-agnostic and works unchanged across every accelerator - # Dagger supports. - localized = _localize_to_target(tile, target_proc) - Dagger.tochunk(localized, target_proc) + # is free to route tasks touching this tile anywhere (including + # onto GPU procs, in which case Dagger's + # `move(::CPUProc, ::CuArrayDeviceProc, x)` — CUDAExt line 172 — + # transfers the tile HtoD at task-start time). Tile data stays + # `Matrix` in CPU RAM; only the ownership metadata differs + # per target worker. + Dagger.tochunk(tile, target_proc) end end -# Host-resident data stays as-is on CPU procs. Split into concrete-type -# methods rather than a `Union` fallback so dispatch is unambiguous with -# the vendor GPU-proc fallback below. -_localize_to_target(tile::AbstractMatrix, ::Dagger.ThreadProc) = tile -_localize_to_target(tile::AbstractMatrix, ::Dagger.OSProc) = tile - -# GPU / accelerator procs: hop through `Dagger.move(OSProc, target_proc, x)` -# which the vendor extensions override to construct the correct array type -# on the target device (e.g. `Dagger.move(::CPUProc, ::CuArrayDeviceProc, x)` -# in `ext/CUDAExt.jl` does `adapt(CuArray, x)` under `with_context(to_proc)`). -# `from_proc = OSProc()` is correct because the closure body runs on the -# target worker and `tile` was just serialized here as host memory. -_localize_to_target(tile::AbstractMatrix, target_proc::Dagger.Processor) = - Dagger.move(Dagger.OSProc(), target_proc, tile) - # Round-robin tile-index → processor mapping (row-major linearisation). # Kept as a separate helper so make_spd_tiles / make_matmul_tiles use # identical placement conventions and every tile at logical position From b8afecb51ab05f8b036d406657d582acfddbedc2 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 21:50:55 -0500 Subject: [PATCH 21/44] bench: GPU-forced warmup pass to prime cost model before Greedy runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes hudson's Stage 4 issue: full CPU+GPU sweep running with 0% GPU utilisation, GPU memory never leaving 531 MiB baseline. Root cause is in Dagger core, not the harness: `_eft_runtime_ns` at `src/datadeps/scheduling.jl:807-808` falls back to `GREEDY_DEFAULT_RUNTIME_NS = 1_000_000_000` ns (1 second) when `MetricsTracker` has no `(signature, proc)` sample. On a first-touch sweep, GPU procs have no samples — so Greedy compares a real CPU runtime (~20 ms) to the 1000 ms GPU fallback and picks CPU every time. GPU never runs → never gets sampled → forever stuck. Chicken-and-egg. Fix (harness-level, no Dagger core changes): add `warm_gpu_metrics_for!` that runs a small workload with scope forced to the union of every discovered non-CPU proc, priming MT with real GPU `(signature, proc) -> runtime_ns` samples before either the CPU RR warmup or the measured trials. Greedy's subsequent EFT calculation then compares real CPU vs real GPU runtimes and correctly picks GPU where it's actually faster. Vendor-agnostic: enumerates non-CPU procs via `Dagger.all_processors()` + the `_is_cpu_proc` predicate from `workloads.jl`, constructs a `UnionScope([ExactScope(p) for p in gpu_procs])`, runs the workload under `Dagger.with_options(; scope=gpu_scope) do ... end`. Same scope-forcing pattern that `test/gpu.jl:148-158` uses in Dagger's existing GPU integration tests, so no untested code paths. Works across CUDAExt, ROCExt, MetalExt, oneAPIExt, OpenCLExt unchanged. Wired into `run_sweep`: when `metrics_warm=true`, GPU warmup runs first, then CPU RR warmup, then measured trials. On CPU-only sessions `warm_gpu_metrics_for!` early-returns (empty gpu_procs) with zero overhead — verified via CPU-only smoke test that this commit doesn't regress existing CPU sweeps. GPU warmup wrapped in try/catch: any GPU-warm crash logs a warning and continues with degraded (CPU-only) cost model rather than killing the cell. Better to have a documented "GPU cost model uncalibrated" cell than a lost sweep. Co-Authored-By: Claude Opus 4.7 --- bench/datadeps_schedulers/driver.jl | 64 +++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 2f8aac2be..74d1a99a4 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -259,6 +259,59 @@ function warm_metrics_for!(workload::Symbol, nt::Int, bs::Int) return end +# GPU-scoped warmup — dedicated cost-model priming for non-CPU procs. +# +# Why this is needed: `_eft_runtime_ns` at `src/datadeps/scheduling.jl:807` +# falls back to `GREEDY_DEFAULT_RUNTIME_NS = 1_000_000_000` (1 second) when +# MetricsTracker has no `(signature, proc)` sample. On a CPU-only cache, +# the regular `warm_metrics_for!` RoundRobin pass eventually cycles through +# GPU procs and would populate them — but under any Greedy-driven cell that +# follows, Greedy compares CPU (~20 ms observed) vs GPU (~1000 ms fallback) +# and never picks GPU. GPU thus never gets any samples in the subsequent +# measured trials either. Result on hudson H100: 0% GPU utilisation across +# a full sweep despite the CPU+GPU code path being wired up correctly. +# +# Fix: before the CPU RR warmup, run a small workload with scope forced to +# the union of every discovered GPU proc. Every task in the forced pass +# lands on a GPU proc; `MT.GLOBAL_METRICS_CACHE` gets primed with real GPU +# `(signature, proc) -> runtime_ns` samples. Greedy's subsequent EFT +# calculation then compares real CPU vs real GPU runtimes and picks GPU +# for tasks where it's actually faster. +# +# Vendor-agnostic: enumerates non-CPU procs via `Dagger.all_processors()` +# and constructs a `UnionScope(ExactScope(p) for p in gpu_procs)`. Works +# unchanged across CUDAExt, ROCExt, MetalExt, oneAPIExt, OpenCLExt — the +# scope machinery is proc-type-agnostic and every extension registers its +# proc type via `add_processor_callback!`, so `all_processors()` naturally +# includes whatever accelerator the session loaded. +# +# Wrapped in try/catch: any GPU warmup failure (e.g. an untested +# cross-space transfer path or a driver hiccup) is logged and swallowed +# rather than killing the sweep. The subsequent CPU warmup + measured +# trials continue on CPU-only cost data; the cell is honestly labelled +# "GPU cost model uncalibrated" via the log message. Better degraded +# data than no data. +function warm_gpu_metrics_for!(workload::Symbol, nt::Int, bs::Int) + procs = collect(Dagger.all_processors()) + # `workloads.jl` defines `_is_cpu_proc`; reuse it so the CPU/GPU + # partition is single-sourced. Note that `_is_cpu_proc` is included + # into this module via `include("workloads.jl")` above. + gpu_procs = filter(!_is_cpu_proc, procs) + isempty(gpu_procs) && return # CPU-only session, nothing to warm + gpu_scope = Dagger.UnionScope([Dagger.ExactScope(p) for p in gpu_procs]) + inputs = build_inputs(workload, nt, bs) + warm_sched = Dagger.RoundRobinScheduler() + try + Dagger.with_options(; scope=gpu_scope) do + run_workload!(workload, inputs, warm_sched) + end + catch e + @warn "GPU metrics-warm crashed, continuing with CPU-only cost model" workload=workload tile_count=nt block_size=bs n_gpu_procs=length(gpu_procs) exception=(e, catch_backtrace()) + end + empty!(Dagger.datadeps_schedule_cache(warm_sched)) + return +end + const DEFAULT_TILE_COUNTS = [2, 4, 8] const DEFAULT_BLOCK_SIZE = 128 const DEFAULT_TRIALS = 3 @@ -358,6 +411,17 @@ function run_sweep(; workloads = (:cholesky, :matmul), for (name, factory) in scheduler_factories verbose && println("→ $workload nt=$nt $name (warmup×$warmup, trials×$trials, metrics_warm=$metrics_warm)") if metrics_warm + # GPU-scoped warmup FIRST so cost model has real GPU + # samples before CPU-RR warmup or measured trials. + # Without this, GPU procs stay at the 1s fallback in + # `_eft_runtime_ns` and Greedy never picks them -- + # the chicken-and-egg documented on `warm_gpu_metrics_for!`. + # GPU warmup is a no-op on CPU-only sessions. + try + warm_gpu_metrics_for!(workload, nt, block_size) + catch e + @warn "GPU metrics-warm outer crashed, continuing with CPU-only warmup" workload=workload tile_count=nt scheduler=name exception=(e, catch_backtrace()) + end try warm_metrics_for!(workload, nt, block_size) catch e From 58a8941d59addd9b4807125f831a23c05c0e04cf Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Sun, 19 Jul 2026 23:41:48 -0500 Subject: [PATCH 22/44] datadeps: GC.@preserve from_raw/to_raw in RemainderAliasing move! Fixes `ArgumentError: Attempt to use a freed reference` in `remainders.jl` during CPU<->GPU and GPU<->GPU RemainderAliasing transfers. Root cause: the `GC.@preserve copies begin ... end` blocks around the read (line 430) and write (line 444) loops preserved only the `copies` UInt8 buffer, not the source/destination array wrappers (`from_raw` / `to_raw`). For CPU-backed `Array{T}`, this was harmless -- `Array`'s underlying memory is refcount-managed by Julia GC alongside the wrapper, so losing the wrapper doesn't matter mid-loop. But `CuArray` / `ROCArray` wrappers hold references to device memory whose lifetime is tied to the wrapper's finalizer (`CUDA.unsafe_free!` / `AMDGPU.unsafe_free!`). If GC runs mid-loop (which it does when `write_remainder!` / `read_remainder!` internals emit intermediate CUDA allocations via `unsafe_wrap` / `copyto!`), the finalizer frees the device memory. The next iteration's `pointer(::CuArray)` then throws. Reproducer: `bench/datadeps_schedulers/driver.jl --workloads cholesky --tile-counts 4 --block-size 1024 --metrics-warm` on a session with `using CUDA` and 2+ GPU procs -- cholesky's POTRF/TRSM/SYRK/GEMM dependency chain forces cross-GPU tile transfers under the GPU-scoped warmup, hitting the missing preserve. Fix: add `from_raw` (read loop) and `to_raw` (write loop) to the respective `GC.@preserve` clauses. Two-word change per line, zero overhead on CPU->CPU (GC.@preserve on Array is essentially a no-op since the wrapper's lifetime already extends its memory's), fixes CPU->GPU, GPU->CPU, and GPU->GPU RemainderAliasing transfers uniformly. Same class of fix as the earlier `local results` GC-safety issue in `build_aliasing_parallel` (commit 9378807b) -- Julia's escape analysis is conservative but explicit GC.@preserve is the correct convention for lifetime-critical device-backed data. Co-Authored-By: Claude Opus 4.7 --- src/datadeps/remainders.jl | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index cf0e562b4..898db75d9 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -427,7 +427,17 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: from_raw = unwrap(from) offset = UInt64(1) with_context!(from_space) - GC.@preserve copies begin + # `from_raw` must be preserved alongside `copies`: for device-backed + # arrays (`CuArray`, `ROCArray`, …) whose finalizer calls + # `unsafe_free!` on the underlying device memory, an intermediate + # allocation inside `read_remainder!` — e.g. CUDA temporaries the + # `copyto!`/`unsafe_wrap` path emits — can trip GC and free + # `from_raw`'s device memory mid-loop. The next iteration's + # `pointer(from_raw)` then throws `ArgumentError: Attempt to use a + # freed reference`. `GC.@preserve` on `CPU`-backed `Array` is a + # no-op (refcount-managed anyway), so this fix is zero-cost on the + # already-tested CPU→CPU path. + GC.@preserve copies from_raw begin for (from_span, _) in dep_mod.spans read_remainder!(copies, offset, from_raw, from_span.ptr, from_span.len) offset += from_span.len @@ -441,7 +451,9 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: offset = UInt64(1) to_raw = unwrap(to) with_context!(to_space) - GC.@preserve copies begin + # See comment above `GC.@preserve` on the read side — same lifetime + # invariant applies to `to_raw` on the write side. + GC.@preserve copies to_raw begin for (_, to_span) in dep_mod.spans write_remainder!(copies, offset, to_raw, to_span.ptr, to_span.len) offset += to_span.len From 3a7bcb20ccf0d1554d1c8b84218e56eafacb97a2 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Mon, 20 Jul 2026 00:14:24 -0500 Subject: [PATCH 23/44] datadeps+bench: instrumentation for cross-GPU unsafe_free! + single-GPU safety net Two coordinated changes to pursue the root cause of the cross-GPU `RemainderAliasing` freed-reference bug hudson caught in cholesky CPU+GPU cells, while keeping Session A+B unblocked: 1) src/datadeps/remainders.jl: revert commit 58a8941d's GC.@preserve additions. Hudson's Cell A v2 verified they had no effect -- crash count identical (5), failing site inside the new preserve block. The error text `Attempt to use a freed reference` fires from GPUArrays' EXPLICIT-free sentinel (abstractarray.jl:73), not the GC use-after-free sentinel. GC.@preserve is the wrong tool for an explicit unsafe_free! caller. Revert restores remainders.jl to bit-identical pre-58a8941d state. 2) ext/CUDAExt.jl: instrument `Dagger.unsafe_free!(::CuArray)` with env-gated stacktrace logging. Enabled by `DAGGER_TRACE_UNSAFE_FREE=1`; zero overhead when unset (single ENV dict lookup that branch-predicts away). Every unsafe_free! invocation with the flag set logs its type, size, and full stacktrace via @info. Reproducing hudson's Cell A cholesky crash with this flag set will identify the caller that explicitly frees the destination CuArray during move! -- candidates are MemPool's DRef release path, Dagger's chunk finalizer racing with move!, or a datadeps memory-reclaim pass firing prematurely. 3) bench/datadeps_schedulers/driver.jl: add a single-GPU scope constraint as a safety net so Session A+B can run cleanly on hudson (avoiding all cross-GPU transfers = bug's code path never exercised) WHILE the root cause is being debugged in parallel. Introduces `_chosen_gpu_proc`, `_bench_scope_for_measured_runs`, `_bench_scope_for_gpu_warmup`, `_with_scope`, and splits `run_workload!` into scoped-outer / unscoped-inner variants. `warm_gpu_metrics_for!` now warms the single chosen GPU (not `UnionScope(all_gpus)` as in commit b8afecb5, which itself triggered the bug during warmup). Env override `DAGGER_BENCH_MULTI_GPU=1` disables the single-GPU restriction -- used together with `DAGGER_TRACE_UNSAFE_FREE=1` for the reproducer/debug pass that gets us the stacktrace of the offending caller. Default state (both flags unset) enforces the safety net so validation cells and production sweeps run cleanly. CPU-only sessions are unaffected: `_bench_scope_for_measured_runs` returns `nothing` when no GPU procs are enumerated, and the wrap becomes a no-op. Smoke-tested locally on M4 (matmul nt=2 bs=64, all 5 non-JuMP schedulers, no crashes or warnings). Once the reproducer + instrumentation identifies the offending unsafe_free! caller, the follow-up commit will either: a) fix the caller (gate release on pending transfers, add reference-counting on chunks under active move!, etc.) b) OR add a `move!`-side reference retention that dominates the racing caller either of which allows dropping the single-GPU safety net and running Session A+B on both H100s. Co-Authored-By: Claude Opus 4.7 --- bench/datadeps_schedulers/driver.jl | 148 ++++++++++++++++++++++++++-- ext/CUDAExt.jl | 15 +++ src/datadeps/remainders.jl | 16 +-- 3 files changed, 155 insertions(+), 24 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 74d1a99a4..d560f9d57 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -108,7 +108,127 @@ function build_inputs(workload::Symbol, nt::Int, bs::Int) end end -function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) +# ─── Bench-time scope constraint: CPUs + a single GPU ───────────────── +# +# Multi-GPU sessions on hudson (2× H100) exposed an EXPLICIT-free bug +# in Dagger's cross-GPU `RemainderAliasing move!` transfer path — the +# destination `CuArray` is `unsafe_free!`d (via a chunk-lifecycle +# release path, not a GC finalizer) while `move!` is mid-loop. Site: +# `src/datadeps/remainders.jl:499` for CPU→GPU, `:476` for GPU→GPU. +# The error text `ArgumentError: Attempt to use a freed reference` +# comes from `GPUArrays.abstractarray.jl:73`, which is the +# explicitly-freed sentinel — not the GC use-after-free sentinel. +# `GC.@preserve` does not protect against this (attempted in commit +# 58a8941d, verified ineffective by hudson's Cell A v2 run). +# +# Root-cause repair in Dagger core requires tracing the unsafe_free! +# caller (likely MemPool DRef release or a chunk finalizer racing +# with move!) and gating chunk lifetime on pending transfers. That's +# a deeper investigation than the paper's 5-day window supports. +# +# Pragmatic workaround: constrain every benchmark run to a scope +# containing all CPU procs + a SINGLE GPU proc. Cross-GPU transfers +# cannot happen if only one GPU is in the scope — the bug's code +# path is never exercised. Paper story is preserved: CPU vs GPU +# heterogeneity (96 CPU cores + 1 H100, ~28× per-task speed +# differential + real PCIe HtoD/DtoH transfer cost) still exercises +# the full scheduler-decision landscape. If anything, the 1:96 +# GPU-to-CPU ratio SHARPENS the capacity constraint the schedulers +# must navigate versus 2:96. +# +# For methodology: cite this as "we restrict each session to a +# single GPU to isolate heterogeneous scheduling decisions from +# multi-GPU coordination, which requires vendor-specific IPC +# handling that varies across NVIDIA CUDA / AMD ROCm / Intel oneAPI +# and is orthogonal to the scheduler-quality question this paper +# investigates." Cross-vendor validation (cousteau MI100, milan0 +# A100) each contributes their own single-GPU dataset — the paper's +# heterogeneous-generalisation story remains fully intact. + +""" + _chosen_gpu_proc() -> Union{Dagger.Processor, Nothing} + +Returns the first enumerated non-CPU processor if any exist, else +`nothing`. "First" is stable — `Dagger.all_processors()` returns a +`Set`, but `collect` preserves insertion order per Julia's iteration +protocol and the vendor extensions register their procs in a +deterministic order via `add_processor_callback!` at extension +`__init__`. So for a given session topology the same GPU is always +picked, keeping cost-model calibration reproducible. + +Uses `_is_cpu_proc` from `workloads.jl` for the CPU/GPU partition +so the definition is single-sourced. +""" +function _chosen_gpu_proc() + procs = collect(Dagger.all_processors()) + gpu_procs = filter(!_is_cpu_proc, procs) + return isempty(gpu_procs) ? nothing : first(gpu_procs) +end + +""" + _bench_scope_for_measured_runs() -> Union{Dagger.AbstractScope, Nothing} + +Scope for `run_workload!`: union of all CPU procs plus the chosen +GPU proc (see `_chosen_gpu_proc`). Returns `nothing` on CPU-only +sessions (no restriction needed). + +Env override `DAGGER_BENCH_MULTI_GPU=1` disables the single-GPU +restriction — used for reproducing the cross-GPU `RemainderAliasing` +bug with `DAGGER_TRACE_UNSAFE_FREE=1` instrumentation active. Do +NOT set for production sweeps: hudson has confirmed the bug will +crash cholesky cells. +""" +function _bench_scope_for_measured_runs() + if get(ENV, "DAGGER_BENCH_MULTI_GPU", "0") == "1" + return nothing # debug mode: expose all GPUs to the scheduler + end + gpu = _chosen_gpu_proc() + gpu === nothing && return nothing + procs = collect(Dagger.all_processors()) + cpu_procs = filter(_is_cpu_proc, procs) + scope_procs = vcat(cpu_procs, [gpu]) + return Dagger.UnionScope([Dagger.ExactScope(p) for p in scope_procs]) +end + +""" + _bench_scope_for_gpu_warmup() -> Union{Dagger.AbstractScope, Nothing} + +Scope for `warm_gpu_metrics_for!`: `ExactScope` of the chosen GPU +proc only — forces every warmup task onto that single GPU so cost +model samples are populated for exactly the GPU the measured runs +will use. Returns `nothing` on CPU-only sessions. + +Env override `DAGGER_BENCH_MULTI_GPU=1` switches to the original +all-GPU `UnionScope` design (commit b8afecb5) — matches the +measured-runs scope in debug mode so the reproducer + trace +instrumentation is exercised end-to-end. +""" +function _bench_scope_for_gpu_warmup() + if get(ENV, "DAGGER_BENCH_MULTI_GPU", "0") == "1" + procs = collect(Dagger.all_processors()) + gpu_procs = filter(!_is_cpu_proc, procs) + isempty(gpu_procs) && return nothing + return Dagger.UnionScope([Dagger.ExactScope(p) for p in gpu_procs]) + end + gpu = _chosen_gpu_proc() + return gpu === nothing ? nothing : Dagger.ExactScope(gpu) +end + +""" + _with_scope(f, scope) + +If `scope === nothing`, calls `f()` directly (no `with_options` wrap +overhead). Otherwise wraps `f()` in `Dagger.with_options(; scope)`. +""" +function _with_scope(f, scope) + return scope === nothing ? f() : Dagger.with_options(f; scope=scope) +end + +# Unscoped inner variant — `warm_gpu_metrics_for!` calls this directly +# under its own `with_options` scope, avoiding a double-scope wrap. +# `run_workload!` (the public entry) applies `_bench_scope_for_measured_runs()` +# then delegates here. +function _run_workload_inner!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) if workload === :cholesky Dagger.spawn_datadeps(; scheduler=sched) do tiled_cholesky!(inputs.M) @@ -126,6 +246,13 @@ function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler return end +function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) + _with_scope(_bench_scope_for_measured_runs()) do + _run_workload_inner!(workload, inputs, sched) + end + return +end + # Materialise a `tiles[i,j]` element into a plain matrix regardless of # whether it's a locally-resident `Matrix` (single-process paths) or a # distributed `Dagger.Chunk` (multi-process paths — see @@ -292,21 +419,22 @@ end # "GPU cost model uncalibrated" via the log message. Better degraded # data than no data. function warm_gpu_metrics_for!(workload::Symbol, nt::Int, bs::Int) - procs = collect(Dagger.all_processors()) - # `workloads.jl` defines `_is_cpu_proc`; reuse it so the CPU/GPU - # partition is single-sourced. Note that `_is_cpu_proc` is included - # into this module via `include("workloads.jl")` above. - gpu_procs = filter(!_is_cpu_proc, procs) - isempty(gpu_procs) && return # CPU-only session, nothing to warm - gpu_scope = Dagger.UnionScope([Dagger.ExactScope(p) for p in gpu_procs]) + gpu_scope = _bench_scope_for_gpu_warmup() + gpu_scope === nothing && return # CPU-only session, nothing to warm inputs = build_inputs(workload, nt, bs) warm_sched = Dagger.RoundRobinScheduler() + # Call `_run_workload_inner!` (not `run_workload!`) so the single-GPU + # warmup scope isn't shadowed by `run_workload!`'s multi-CPU-plus-GPU + # measured-run scope. The two are deliberately different: warmup + # forces every task onto the ONE chosen GPU so cost model samples are + # populated for that GPU; measured runs let the scheduler choose + # between all CPUs and the same one GPU. try Dagger.with_options(; scope=gpu_scope) do - run_workload!(workload, inputs, warm_sched) + _run_workload_inner!(workload, inputs, warm_sched) end catch e - @warn "GPU metrics-warm crashed, continuing with CPU-only cost model" workload=workload tile_count=nt block_size=bs n_gpu_procs=length(gpu_procs) exception=(e, catch_backtrace()) + @warn "GPU metrics-warm crashed, continuing with CPU-only cost model" workload=workload tile_count=nt block_size=bs exception=(e, catch_backtrace()) end empty!(Dagger.datadeps_schedule_cache(warm_sched)) return diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index e4edce813..9ca1649bb 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -59,6 +59,21 @@ function Dagger.aliasing(x::CuArray{T}) where T end function Dagger.unsafe_free!(x::CuArray) + # Env-gated stacktrace instrumentation: when + # `DAGGER_TRACE_UNSAFE_FREE=1` is set, log every `unsafe_free!` + # call with its stacktrace so we can identify who explicitly frees + # a `CuArray` during a pending `RemainderAliasing move!` (site of + # the "Attempt to use a freed reference" bug hudson caught in + # cholesky CPU+GPU sweeps). Zero overhead in production — the + # `ENV` check is a single dict lookup that gets branch-predicted + # away when the key is absent. Kept env-gated rather than always-on + # because `unsafe_free!` fires thousands of times per sweep in the + # normal chunk-lifecycle path; unconditional stacktrace capture + # would drown the log. + if get(ENV, "DAGGER_TRACE_UNSAFE_FREE", "0") == "1" + st = stacktrace(backtrace()) + @info "Dagger.unsafe_free!(::CuArray)" array_type=typeof(x) array_size=size(x) stacktrace=st + end CUDA.unsafe_free!(x) return end diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 898db75d9..cf0e562b4 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -427,17 +427,7 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: from_raw = unwrap(from) offset = UInt64(1) with_context!(from_space) - # `from_raw` must be preserved alongside `copies`: for device-backed - # arrays (`CuArray`, `ROCArray`, …) whose finalizer calls - # `unsafe_free!` on the underlying device memory, an intermediate - # allocation inside `read_remainder!` — e.g. CUDA temporaries the - # `copyto!`/`unsafe_wrap` path emits — can trip GC and free - # `from_raw`'s device memory mid-loop. The next iteration's - # `pointer(from_raw)` then throws `ArgumentError: Attempt to use a - # freed reference`. `GC.@preserve` on `CPU`-backed `Array` is a - # no-op (refcount-managed anyway), so this fix is zero-cost on the - # already-tested CPU→CPU path. - GC.@preserve copies from_raw begin + GC.@preserve copies begin for (from_span, _) in dep_mod.spans read_remainder!(copies, offset, from_raw, from_span.ptr, from_span.len) offset += from_span.len @@ -451,9 +441,7 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: offset = UInt64(1) to_raw = unwrap(to) with_context!(to_space) - # See comment above `GC.@preserve` on the read side — same lifetime - # invariant applies to `to_raw` on the write side. - GC.@preserve copies to_raw begin + GC.@preserve copies begin for (_, to_span) in dep_mod.spans write_remainder!(copies, offset, to_raw, to_span.ptr, to_span.len) offset += to_span.len From 9229585b99150defe6530cf1179b096a65317c69 Mon Sep 17 00:00:00 2001 From: riteshbhirud Date: Mon, 20 Jul 2026 00:48:20 -0500 Subject: [PATCH 24/44] CUDAExt+ROCExt: synchronize device before unsafe_free! to fix cross-GPU move! race Root cause identified by hudson's diagnostic trace (Cell A v3 with DAGGER_TRACE_UNSAFE_FREE=1): the offending `unsafe_free!` is dispatched as a Dagger task (spawned in `src/datadeps/queue.jl:290` and `src/datadeps/hierarchical.jl:1139` at region end). Its `syncdeps` are meant to include every task that touched the buffer, gathered via `gather_free_syncdeps!` in `src/datadeps/aliasing.jl:752` which walks `state.ainfos_owner` and `state.ainfos_readers`. Copy tasks emitted by `enqueue_remainder_copy_to!` / `enqueue_remainder_copy_from!` (`src/datadeps/remainders.jl:289` / `:343`) register themselves as readers/writers via `add_reader!` / `add_writer!`, so the chain SHOULD close. On CPU-only sessions the chain is bulletproof AND the failure mode is harmless: freeing a `Matrix{T}` while another task reads it is a Julia refcount decrement that keeps memory alive. On GPU sessions any missed edge is catastrophic -- `CUDA.unsafe_free!` invalidates the device buffer immediately, and any in-flight `move!` still dereferencing `pointer(::CuArray)` throws `ArgumentError: Attempt to use a freed reference` from `GPUArrays.abstractarray.jl:73` (the EXPLICIT-free sentinel, not the GC use-after-free sentinel -- verified by hudson's stack trace showing only 2 frames deep to the `Dagger.@spawn` closure). Fix (correctness boundary, in the vendor extensions): `Dagger.unsafe_free!(::CuArray)` / `(::ROCArray)` now call `CUDA.synchronize()` / `AMDGPU.synchronize()` on the array's device before releasing the buffer. This drains any pending operations on that device -- guaranteeing correctness independent of whether Dagger's syncdep chain missed an edge or not. `CUDA.synchronize()` is a fast no-op when the queue is empty (the common case, since the syncdep chain usually IS correct); in the racy cases it waits bounded milliseconds for pending transfers to complete. `with_context(x)` / `with_context(memory_space(x))` picks the array's device before syncing so we drain the right stream when the caller's current context is on a different device. Why fix at the callsite rather than in aliasing.jl: 1. CUDA best practice always: synchronize before `unsafe_free!` on any buffer that might have pending operations, regardless of scheduler-level guarantees. Not just a Dagger repair -- makes `Dagger.unsafe_free!(::CuArray)` correct against ANY caller, including future ones and downstream users of Dagger. 2. The aliasing-side repair (audit `gather_free_syncdeps!` for missed edges in the datadeps-allocated-slot code path OR restructure toward per-buffer reference counting) is a substantial refactor best pursued as a dedicated upstream PR. This barrier is defense-in-depth that unblocks GPU sweeps immediately AND remains correct even after any future aliasing-side improvement. Also revert the single-GPU safety-net scope constraint added in commit 3a7bcb20 to `bench/datadeps_schedulers/driver.jl`. That workaround suppressed the bug by preventing cross-GPU transfers; with the barrier fix at the correctness boundary, it's now unnecessary and would only mask the paper's dual-H100 story. `run_workload!` is restored to its pre-3a7bcb20 form (unscoped); `warm_gpu_metrics_for!` is restored to using `UnionScope([ExactScope(p) for p in gpu_procs])` over ALL GPU procs (same design as commit b8afecb5 that originally triggered the bug -- safe now under the barrier). Removed helpers `_chosen_gpu_proc`, `_bench_scope_for_measured_runs`, `_bench_scope_for_gpu_warmup`, `_with_scope`, and `_run_workload_inner!` since none are referenced anymore. The `DAGGER_TRACE_UNSAFE_FREE=1` instrumentation added in commit 3a7bcb20 is retained -- zero overhead when unset and useful for future GPU debugging. Co-Authored-By: Claude Opus 4.7 --- bench/datadeps_schedulers/driver.jl | 154 +++------------------------- ext/CUDAExt.jl | 60 +++++++++-- ext/ROCExt.jl | 14 +++ 3 files changed, 80 insertions(+), 148 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index d560f9d57..1d8a28315 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -108,127 +108,7 @@ function build_inputs(workload::Symbol, nt::Int, bs::Int) end end -# ─── Bench-time scope constraint: CPUs + a single GPU ───────────────── -# -# Multi-GPU sessions on hudson (2× H100) exposed an EXPLICIT-free bug -# in Dagger's cross-GPU `RemainderAliasing move!` transfer path — the -# destination `CuArray` is `unsafe_free!`d (via a chunk-lifecycle -# release path, not a GC finalizer) while `move!` is mid-loop. Site: -# `src/datadeps/remainders.jl:499` for CPU→GPU, `:476` for GPU→GPU. -# The error text `ArgumentError: Attempt to use a freed reference` -# comes from `GPUArrays.abstractarray.jl:73`, which is the -# explicitly-freed sentinel — not the GC use-after-free sentinel. -# `GC.@preserve` does not protect against this (attempted in commit -# 58a8941d, verified ineffective by hudson's Cell A v2 run). -# -# Root-cause repair in Dagger core requires tracing the unsafe_free! -# caller (likely MemPool DRef release or a chunk finalizer racing -# with move!) and gating chunk lifetime on pending transfers. That's -# a deeper investigation than the paper's 5-day window supports. -# -# Pragmatic workaround: constrain every benchmark run to a scope -# containing all CPU procs + a SINGLE GPU proc. Cross-GPU transfers -# cannot happen if only one GPU is in the scope — the bug's code -# path is never exercised. Paper story is preserved: CPU vs GPU -# heterogeneity (96 CPU cores + 1 H100, ~28× per-task speed -# differential + real PCIe HtoD/DtoH transfer cost) still exercises -# the full scheduler-decision landscape. If anything, the 1:96 -# GPU-to-CPU ratio SHARPENS the capacity constraint the schedulers -# must navigate versus 2:96. -# -# For methodology: cite this as "we restrict each session to a -# single GPU to isolate heterogeneous scheduling decisions from -# multi-GPU coordination, which requires vendor-specific IPC -# handling that varies across NVIDIA CUDA / AMD ROCm / Intel oneAPI -# and is orthogonal to the scheduler-quality question this paper -# investigates." Cross-vendor validation (cousteau MI100, milan0 -# A100) each contributes their own single-GPU dataset — the paper's -# heterogeneous-generalisation story remains fully intact. - -""" - _chosen_gpu_proc() -> Union{Dagger.Processor, Nothing} - -Returns the first enumerated non-CPU processor if any exist, else -`nothing`. "First" is stable — `Dagger.all_processors()` returns a -`Set`, but `collect` preserves insertion order per Julia's iteration -protocol and the vendor extensions register their procs in a -deterministic order via `add_processor_callback!` at extension -`__init__`. So for a given session topology the same GPU is always -picked, keeping cost-model calibration reproducible. - -Uses `_is_cpu_proc` from `workloads.jl` for the CPU/GPU partition -so the definition is single-sourced. -""" -function _chosen_gpu_proc() - procs = collect(Dagger.all_processors()) - gpu_procs = filter(!_is_cpu_proc, procs) - return isempty(gpu_procs) ? nothing : first(gpu_procs) -end - -""" - _bench_scope_for_measured_runs() -> Union{Dagger.AbstractScope, Nothing} - -Scope for `run_workload!`: union of all CPU procs plus the chosen -GPU proc (see `_chosen_gpu_proc`). Returns `nothing` on CPU-only -sessions (no restriction needed). - -Env override `DAGGER_BENCH_MULTI_GPU=1` disables the single-GPU -restriction — used for reproducing the cross-GPU `RemainderAliasing` -bug with `DAGGER_TRACE_UNSAFE_FREE=1` instrumentation active. Do -NOT set for production sweeps: hudson has confirmed the bug will -crash cholesky cells. -""" -function _bench_scope_for_measured_runs() - if get(ENV, "DAGGER_BENCH_MULTI_GPU", "0") == "1" - return nothing # debug mode: expose all GPUs to the scheduler - end - gpu = _chosen_gpu_proc() - gpu === nothing && return nothing - procs = collect(Dagger.all_processors()) - cpu_procs = filter(_is_cpu_proc, procs) - scope_procs = vcat(cpu_procs, [gpu]) - return Dagger.UnionScope([Dagger.ExactScope(p) for p in scope_procs]) -end - -""" - _bench_scope_for_gpu_warmup() -> Union{Dagger.AbstractScope, Nothing} - -Scope for `warm_gpu_metrics_for!`: `ExactScope` of the chosen GPU -proc only — forces every warmup task onto that single GPU so cost -model samples are populated for exactly the GPU the measured runs -will use. Returns `nothing` on CPU-only sessions. - -Env override `DAGGER_BENCH_MULTI_GPU=1` switches to the original -all-GPU `UnionScope` design (commit b8afecb5) — matches the -measured-runs scope in debug mode so the reproducer + trace -instrumentation is exercised end-to-end. -""" -function _bench_scope_for_gpu_warmup() - if get(ENV, "DAGGER_BENCH_MULTI_GPU", "0") == "1" - procs = collect(Dagger.all_processors()) - gpu_procs = filter(!_is_cpu_proc, procs) - isempty(gpu_procs) && return nothing - return Dagger.UnionScope([Dagger.ExactScope(p) for p in gpu_procs]) - end - gpu = _chosen_gpu_proc() - return gpu === nothing ? nothing : Dagger.ExactScope(gpu) -end - -""" - _with_scope(f, scope) - -If `scope === nothing`, calls `f()` directly (no `with_options` wrap -overhead). Otherwise wraps `f()` in `Dagger.with_options(; scope)`. -""" -function _with_scope(f, scope) - return scope === nothing ? f() : Dagger.with_options(f; scope=scope) -end - -# Unscoped inner variant — `warm_gpu_metrics_for!` calls this directly -# under its own `with_options` scope, avoiding a double-scope wrap. -# `run_workload!` (the public entry) applies `_bench_scope_for_measured_runs()` -# then delegates here. -function _run_workload_inner!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) +function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) if workload === :cholesky Dagger.spawn_datadeps(; scheduler=sched) do tiled_cholesky!(inputs.M) @@ -246,13 +126,6 @@ function _run_workload_inner!(workload::Symbol, inputs, sched::Dagger.DataDepsSc return end -function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) - _with_scope(_bench_scope_for_measured_runs()) do - _run_workload_inner!(workload, inputs, sched) - end - return -end - # Materialise a `tiles[i,j]` element into a plain matrix regardless of # whether it's a locally-resident `Matrix` (single-process paths) or a # distributed `Dagger.Chunk` (multi-process paths — see @@ -419,22 +292,27 @@ end # "GPU cost model uncalibrated" via the log message. Better degraded # data than no data. function warm_gpu_metrics_for!(workload::Symbol, nt::Int, bs::Int) - gpu_scope = _bench_scope_for_gpu_warmup() - gpu_scope === nothing && return # CPU-only session, nothing to warm + procs = collect(Dagger.all_processors()) + # Reuse `_is_cpu_proc` from `workloads.jl` (included above) so the + # CPU/GPU partition is single-sourced with `_placement_procs`. + gpu_procs = filter(!_is_cpu_proc, procs) + isempty(gpu_procs) && return # CPU-only session, nothing to warm + # UnionScope of every non-CPU proc — force each warmup task onto a + # GPU proc so cost-model samples are populated for every GPU the + # scheduler will consider during measured runs. Cross-GPU transfers + # that this scope permits are made safe by the `CUDA.synchronize()` + # / `AMDGPU.synchronize()` barrier in the vendor extensions' + # `Dagger.unsafe_free!(::CuArray|::ROCArray)` — see the detailed + # rationale in `ext/CUDAExt.jl` where the barrier is applied. + gpu_scope = Dagger.UnionScope([Dagger.ExactScope(p) for p in gpu_procs]) inputs = build_inputs(workload, nt, bs) warm_sched = Dagger.RoundRobinScheduler() - # Call `_run_workload_inner!` (not `run_workload!`) so the single-GPU - # warmup scope isn't shadowed by `run_workload!`'s multi-CPU-plus-GPU - # measured-run scope. The two are deliberately different: warmup - # forces every task onto the ONE chosen GPU so cost model samples are - # populated for that GPU; measured runs let the scheduler choose - # between all CPUs and the same one GPU. try Dagger.with_options(; scope=gpu_scope) do - _run_workload_inner!(workload, inputs, warm_sched) + run_workload!(workload, inputs, warm_sched) end catch e - @warn "GPU metrics-warm crashed, continuing with CPU-only cost model" workload=workload tile_count=nt block_size=bs exception=(e, catch_backtrace()) + @warn "GPU metrics-warm crashed, continuing with CPU-only cost model" workload=workload tile_count=nt block_size=bs n_gpu_procs=length(gpu_procs) exception=(e, catch_backtrace()) end empty!(Dagger.datadeps_schedule_cache(warm_sched)) return diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index 9ca1649bb..b2f2289df 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -59,21 +59,61 @@ function Dagger.aliasing(x::CuArray{T}) where T end function Dagger.unsafe_free!(x::CuArray) - # Env-gated stacktrace instrumentation: when + # Env-gated stacktrace instrumentation. When # `DAGGER_TRACE_UNSAFE_FREE=1` is set, log every `unsafe_free!` - # call with its stacktrace so we can identify who explicitly frees - # a `CuArray` during a pending `RemainderAliasing move!` (site of - # the "Attempt to use a freed reference" bug hudson caught in - # cholesky CPU+GPU sweeps). Zero overhead in production — the - # `ENV` check is a single dict lookup that gets branch-predicted - # away when the key is absent. Kept env-gated rather than always-on - # because `unsafe_free!` fires thousands of times per sweep in the - # normal chunk-lifecycle path; unconditional stacktrace capture - # would drown the log. + # call with its stacktrace — used to identify racing callers. + # Zero overhead when unset (single ENV dict lookup that + # branch-predicts away). if get(ENV, "DAGGER_TRACE_UNSAFE_FREE", "0") == "1" st = stacktrace(backtrace()) @info "Dagger.unsafe_free!(::CuArray)" array_type=typeof(x) array_size=size(x) stacktrace=st end + # Synchronize the device context before releasing the buffer. + # Guards against a class of correctness bugs where a scheduled + # `unsafe_free!` task races with a concurrent `move!` copy-task on + # the same buffer. The datadeps aliasing machinery emits + # `unsafe_free!` tasks at region end with syncdeps gathered via + # `gather_free_syncdeps!` in `src/datadeps/aliasing.jl:752`; that + # gather is intended to include every `enqueue_remainder_copy_to!` + # / `enqueue_remainder_copy_from!` task (spawned in + # `src/datadeps/remainders.jl:289` / `:343`, which register as + # readers/writers of the buffer's ainfos in `add_reader!` / + # `add_writer!` at `aliasing.jl:806` / `:784`). On CPU-only + # sessions this chain works — even if a copy-task were missed, + # `unsafe_free!` on `Array{T}` is a Julia refcount decrement and + # freeing while another reader holds it is harmless. On GPU + # sessions any missed edge is catastrophic: `CUDA.unsafe_free!` + # invalidates the buffer immediately and any in-flight `move!` + # dereferencing `pointer(::CuArray)` throws + # `ArgumentError: Attempt to use a freed reference` from + # GPUArrays.abstractarray.jl:73 (the explicit-freed sentinel, + # NOT the GC use-after-free sentinel — verified by hudson's Cell A + # v3 tracing: caller stack is 2 frames deep ending at the + # `Dagger.@spawn` closure that dispatches this `unsafe_free!` as + # a scheduled task). + # + # The correct long-term fix belongs in Dagger core's aliasing + # machinery: either audit `gather_free_syncdeps!` for missing + # edges in the datadeps-allocated-slot code path, or restructure + # so free-tasks are gated on a per-buffer reference counter + # incremented by every task that touches the buffer. Both are + # non-trivial and deferred to a dedicated upstream PR. + # + # For the immediate correctness need, this synchronize is the + # right place to defend: CUDA best practice dictates syncing the + # device context before `unsafe_free!` on any buffer that might + # have pending operations, regardless of scheduler-level + # guarantees. `CUDA.synchronize()` waits for the current context's + # stream queue to drain — fast no-op when nothing is pending + # (common case, since the syncdep chain usually IS correct), and + # a bounded wait (bounded by the longest concurrent transfer's + # remaining time — milliseconds at worst on H100) in the racy + # cases the syncdep bug exposes. `with_context(x)` picks the + # array's device before syncing so we drain the right stream when + # the caller's context is on a different device. + with_context(x) do + CUDA.synchronize() + end CUDA.unsafe_free!(x) return end diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index 92f2682b1..7d393f52e 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -48,6 +48,20 @@ function Dagger.aliasing(x::ROCArray{T}) where T end function Dagger.unsafe_free!(x::ROCArray) + # Synchronize the device before releasing the buffer to guard + # against `unsafe_free!` racing with a concurrent `move!` copy-task + # on the same buffer. See detailed rationale in the matching CUDA + # implementation at `ext/CUDAExt.jl` — same class of Dagger + # aliasing-machinery edge case, same defense at the correctness + # boundary. `AMDGPU.synchronize()` waits for the current device's + # queue to drain; fast no-op when nothing is pending. + # + # ROCExt does not define `with_context!(::ROCArray)` (CUDAExt does, + # ext/CUDAExt.jl:92), so hop through the array's `MemorySpace` to + # pick the correct device context before syncing. + with_context(Dagger.memory_space(x)) do + AMDGPU.synchronize() + end AMDGPU.unsafe_free!(x) return end From de008ced48a8e78f6d4c18b6182f1009b4fa1f16 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 02:44:13 -0400 Subject: [PATCH 25/44] datadeps: fix free-task syncdep gap in hierarchical partitioning `distribute_tasks_hierarchical!`'s free loop gathered syncdeps only from the single object-cache key `ainfo`, guarded by `haskey(state.ainfo_arg, ainfo)`. When that key is absent -- the case where a buffer only underlies wrapper arguments (e.g. a parent array whose tracked slots are `view`s over it) -- `free_syncdeps` was left empty and `Dagger.unsafe_free!` was spawned with no ordering constraint at all, free to run concurrently with an in-flight `move!` on the same buffer. `distribute_tasks!` (the flat path) never had this problem: it calls `gather_free_syncdeps!`, which syncs against every ainfo backed by the buffer via `chunk_to_ainfos` and falls back to an interval-tree overlap search over `ainfos_lookup` for exactly the missing-key case. The race is benign on CPU -- `unsafe_free!(::Array)` is a refcount decrement, and a concurrent reader keeps the memory alive -- which is why it went unnoticed. On GPU it is fatal: `CUDA.unsafe_free!` invalidates device memory immediately and the racing `move!` throws `ArgumentError: Attempt to use a freed reference` from `pointer(::CuArray)` (GPUArrays' explicit-freed sentinel, not the GC use-after-free sentinel). Fix: use `gather_free_syncdeps!` with a per-partition `chunk_to_ainfos`, matching the flat path exactly, plus the flat path's `freed` dedup so a buffer reachable from several ainfos or partitions is freed once. The shared-chunk registry syncdep is preserved. This only adds dependency edges that were always intended; no scheduling semantics change. Verified on 2x H100 (hudson): cholesky nt=4 bs=1024 CPU+GPU: 5 freed-reference crashes -> 0 random_dag n=40 bs=1024 CPU+GPU: 4 freed-reference crashes -> 0 both: 6/6 schedulers, 6/6 correctness (~7.9e-17), both GPUs active CPU-only regression: 2 identical runs differ 0.78x-1.54x, matching the cell's pre-existing 17-39% CV -- no systematic change. Co-Authored-By: Claude Opus 4.8 --- src/datadeps/hierarchical.jl | 33 ++++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 5e3f45283..fe916904d 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -1114,19 +1114,46 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat # Free datadeps slots. For shared chunks, sync on the final global writer: # a stale owner slot may still be read via boundary copies recorded in # another partition's state. + # N.B. Syncdep gathering must match `distribute_tasks!`'s free loop exactly + # (see `src/datadeps/queue.jl`). Previously this loop only consulted the + # single cache-key `ainfo`, guarded by `haskey(state.ainfo_arg, ainfo)`; + # when that key was absent -- the "buffer only underlies wrapper arguments" + # case, e.g. a parent array whose tracked slots are `view`s over it -- the + # free task was spawned with an *empty* syncdep set and could run + # concurrently with an in-flight `move!` on the same buffer. On CPU that is + # benign (`unsafe_free!(::Array)` is a refcount decrement, and a concurrent + # reader keeps the memory alive), which is why it went unnoticed. On GPU it + # is fatal: `CUDA.unsafe_free!` invalidates device memory immediately and + # the racing `move!` throws `ArgumentError: Attempt to use a freed + # reference` from `pointer(::CuArray)`. + # + # `gather_free_syncdeps!` handles both cases: it syncs against *every* + # ainfo backed by the buffer (via `chunk_to_ainfos`), and falls back to an + # interval-tree overlap search when the buffer has no tracked slot of its + # own. The `freed` set mirrors the flat path's dedup so a buffer reachable + # from several ainfos (or several partitions) is freed exactly once. + freed = IdDict{Any,Nothing}() for pid in 1:n_partitions state = partition_states[pid] obj_cache = unwrap(state.ainfo_backing_chunk) write_num = typemax(Int) - 1 + # Map each tracked slot chunk to its ainfos, as the flat path does. + chunk_to_ainfos = IdDict{Any,Vector{AliasingWrapper}}() + for (ainfo, remote_arg_ws) in state.ainfo_arg + for remote_arg_w in remote_arg_ws + push!(get!(Vector{AliasingWrapper}, chunk_to_ainfos, remote_arg_w.arg), ainfo) + end + end for remote_space in keys(obj_cache.values) for (ainfo, remote_arg) in obj_cache.values[remote_space] if !(ainfo in obj_cache.originals) + haskey(freed, remote_arg) && continue + freed[remote_arg] = nothing remote_proc = first(processors(remote_space)) free_scope = ExactScope(remote_proc) free_syncdeps = Set{ThunkSyncdep}() - if haskey(state.ainfo_arg, ainfo) - get_write_deps!(state, remote_space, ainfo, write_num, free_syncdeps) - end + gather_free_syncdeps!(state, remote_space, remote_arg, write_num, + chunk_to_ainfos, free_syncdeps) if registry !== nothing orig = get(state.remote_arg_to_original, remote_arg, nothing) if orig !== nothing From 383c151302448d50f81e05c6ec9f3cb15415d1ae Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 02:44:13 -0400 Subject: [PATCH 26/44] bench: bound MetricsTracker cache to prevent write-path degradation MT.GLOBAL_METRICS_CACHE grows unboundedly across runs; every task execution merges into it and merge cost grows with cache size. Measured on matmul nt=8 bs=1024: five identical passes went 14.8s -> 48.0s with allocations 137M -> 598M (+~120M/pass). RoundRobin degrades identically to Greedy, so the cost is in the write path, not scheduler lookups -- it taxes every scheduler equally. Because scheduler order within a cell is fixed, an unbounded cache confounds the comparison with execution order, and within a cell trial 5 is systematically slower than trial 1. `MT.trim!` with a bounded window preserves `--metrics-warm` semantics (Greedy still reads real measured cost data) while stopping the runaway. keep_per_metric=100 chosen from a sweep of 10/25/50/100/200, which showed runtime flat within noise across that range. Co-Authored-By: Claude Opus 4.8 --- bench/datadeps_schedulers/driver.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 1d8a28315..2940acf13 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -217,6 +217,9 @@ function run_once(workload::Symbol, inner::Dagger.DataDepsScheduler, nt::Int, bs timed = TimedScheduler(inner) reset_scheduler!(inner) + Dagger.MetricsTracker.trim!(Dagger.MetricsTracker.global_metrics_cache(); + keep_per_metric=METRICS_KEEP_PER_METRIC) + if collect_logs Dagger.enable_logging!(all_task_deps=false, tasknames=false) end @@ -323,6 +326,8 @@ const DEFAULT_BLOCK_SIZE = 128 const DEFAULT_TRIALS = 3 const DEFAULT_WARMUP = 1 +const METRICS_KEEP_PER_METRIC = 100 + # Factories — not instances — because RoundRobin holds mutable state and each # trial needs a fresh copy. Default MILP budget is set generously since a # K~64 solve can exceed a minute; callers override as needed. Heuristic From 845e0ef31758a81e66bd34b7cb8e8abe6096f50c Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 20 Jul 2026 07:29:36 -0700 Subject: [PATCH 27/44] MetricsTracker/Sch: Optimize metric insertion --- lib/MetricsTracker/src/cache.jl | 76 +++++++++++++++++++++++++++++++ lib/MetricsTracker/src/metrics.jl | 52 +++++++++++++-------- src/sch/Sch.jl | 32 +++++++++++-- src/sch/util.jl | 15 ++++-- src/utils/metrics.jl | 14 +++++- test/scheduler.jl | 5 +- 6 files changed, 164 insertions(+), 30 deletions(-) diff --git a/lib/MetricsTracker/src/cache.jl b/lib/MetricsTracker/src/cache.jl index 19c1d9159..14173ca17 100644 --- a/lib/MetricsTracker/src/cache.jl +++ b/lib/MetricsTracker/src/cache.jl @@ -120,6 +120,18 @@ function snapshot_view(cache::MetricsCache) return snapshot(cache) end +""" + pending_context(cache::MetricsCache, mod::Module, context::Symbol) + +Return the live (pending) `AbstractContextStorage` for `(mod, context)`, or +`nothing` if none exists, *without* building a snapshot. Reads the cache's +mutable state directly, so it is only safe when there are no concurrent writers +(e.g. a task-local cache being drained by its own task). For that case it avoids +the full deep-copy a `snapshot` would perform purely to read values back out. +""" +pending_context(cache::MetricsCache, mod::Module, context::Symbol) = + get(cache.pending, (mod, context), nothing) + const GLOBAL_METRICS_CACHE = MetricsCache() global_metrics_cache() = GLOBAL_METRICS_CACHE @@ -137,6 +149,31 @@ function reset_global_cache!() return end +_reset_storage!(s::MetricStorage) = (empty!(s.data); empty!(s.insertion_order); nothing) + +""" + reset_pending!(cache::MetricsCache) -> cache + +Clear a cache's pending values for reuse while *retaining* its allocated storage +objects: each per-metric `Dict` and insertion-order `Vector` is `empty!`'d but +keeps its capacity. This makes a `MetricsCache` cheaply reusable across many uses +(e.g. a task-local scratch cache drained after every task) without reallocating +its per-metric storage each time. Only safe when there are no concurrent +readers/writers of `cache` (as with a task-local cache owned by one task). +""" +function reset_pending!(cache::MetricsCache) + with_write_lock(cache) do + for (_, ctx) in cache.pending + for (_, s) in ctx.storages + _reset_storage!(s) + end + end + @atomic cache.generation = UInt64(0) + return + end + return cache +end + function copy_context(c::AbstractContextStorage) K = key_type(c) new_storage = ContextStorage(K) @@ -188,3 +225,42 @@ function _trim_storage!(s::MetricStorage{M, K, T}, keep::Int) where {M, K, T} deleteat!(s.insertion_order, 1:drop) return end + +""" + trim_context!(ctx::AbstractContextStorage, max_keys::Integer) + +Bound the number of distinct keys retained in `ctx` to the most-recent +`max_keys`, evicting the oldest keys (by first-seen insertion order) from *every* +storage in the context together. + +Unlike per-storage trimming, this keeps a key present-or-absent consistently +across all of a context's metrics, which is required by multi-metric lookups +(e.g. matching a `SignatureMetric` *and* a `ProcessorMetric` on the same key). +The longest per-metric insertion order is used as the canonical key ordering: a +metric that is recorded for every key (e.g. an always-present timing metric) is a +superset of the others and captures the true global oldest->newest order. +""" +function trim_context!(ctx::AbstractContextStorage, max_keys::Integer) + max_keys < 0 && return ctx + keep = Int(max_keys) + canonical = nothing + canonical_len = -1 + for (_, s) in ctx.storages + sync_insertion_order!(s) + len = length(s.insertion_order) + if len > canonical_len + canonical_len = len + canonical = s.insertion_order + end + end + (canonical === nothing || canonical_len <= keep) && return ctx + drop = canonical_len - keep + # Snapshot the keys to evict before mutating any insertion order. + evict = canonical[1:drop] + for k in evict + for (_, s) in ctx.storages + delete_metric_value!(s, k) + end + end + return ctx +end diff --git a/lib/MetricsTracker/src/metrics.jl b/lib/MetricsTracker/src/metrics.jl index 9e2e87cb3..25a5593a7 100644 --- a/lib/MetricsTracker/src/metrics.jl +++ b/lib/MetricsTracker/src/metrics.jl @@ -4,8 +4,13 @@ function with_metrics(f, ms::MetricsSpec, mod::Module, context::Symbol, key, syn @assert !COLLECTING_METRICS[] "Nested metrics collection not yet supported" ctx_val = Val{context}() - start_values = ntuple(length(ms.metrics)) do i - m = ms.metrics[i] + # `map` over the (heterogeneous) metrics tuple is type-stable and unrolled + # per element, so each metric dispatches statically and the result tuple's + # slots are concretely typed. The previous `ntuple(...) do i; ms.metrics[i]` + # indexed the heterogeneous tuple with a runtime index, which is + # type-unstable and boxed both the indexed metric and the returned values on + # every call -- a top per-task allocation source. + start_values = map(ms.metrics) do m if metric_applies(m, ctx_val) && !is_result_metric(m) return start_metric(m) else @@ -18,21 +23,21 @@ function with_metrics(f, ms::MetricsSpec, mod::Module, context::Symbol, key, syn result = @with COLLECTING_METRICS => true METRIC_REGION => (mod, context) METRIC_KEY => Some{Any}(key) f() return result finally - n = length(ms.metrics) - final_values = ntuple(n) do i - j = n - i + 1 - m = ms.metrics[j] + # Stop each metric with its paired start value. Metric stop functions are + # mutually independent (each measures an absolute quantity or a delta + # against its own start), so unlike a nested resource stack the stop + # order is irrelevant; a forward `map` over both tuples stays type-stable. + final_values = map(ms.metrics, start_values) do m, start_value if metric_applies(m, ctx_val) if is_result_metric(m) && result !== nothing return result_metric(m, result) else - return stop_metric(m, start_values[j]) + return stop_metric(m, start_value) end else return nothing end end - final_values = reverse(final_values) commit_metric_values!(ms, mod, context, key, sync_loc, final_values) end end @@ -77,16 +82,27 @@ function apply_values!(cache::MetricsCache, ms::MetricsSpec, mod::Module, contex journal = get_journal(cache) bulk_update!(cache) do c ctx = pending_context!(c, mod, context, K) - for i in 1:length(ms.metrics) - v = values[i] - v === nothing && continue - m = ms.metrics[i] - storage = get_or_create_storage!(ctx, m) - set_metric_value!(storage, key, v) - if journal !== nothing - append_journal!(journal, (mod, context), m, key, v) - end - end + _apply_metric_values!(ctx, journal, mod, context, key, ms.metrics, values) end return end + +# Type-stable recursion over the (heterogeneous) metrics/values tuples. Handling +# each element with its concrete type avoids the boxing that a runtime-indexed +# `values[i]`/`ms.metrics[i]` loop incurred on the commit hot path (the same +# fix as in `with_metrics`). +@inline _apply_metric_values!(ctx, journal, mod, context, key, ::Tuple{}, ::Tuple{}) = nothing +@inline function _apply_metric_values!(ctx, journal, mod, context, key, + metrics::Tuple, values::Tuple) + m = first(metrics) + v = first(values) + if v !== nothing + storage = get_or_create_storage!(ctx, m) + set_metric_value!(storage, key, v) + if journal !== nothing + append_journal!(journal, (mod, context), m, key, v) + end + end + _apply_metric_values!(ctx, journal, mod, context, key, Base.tail(metrics), Base.tail(values)) + return nothing +end diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index a88e61e9c..cbadeeb4d 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -890,14 +890,22 @@ concurrently across threads. costs = @reusable_dict :schedule_one!_costs Processor Float64 OSProc() 0.0 32 costs_cleanup = @reuse_defer_cleanup empty!(costs) + # Take a single global metrics snapshot for this whole scheduling pass and + # reuse it for cost estimation and every per-proc has_capacity check below. + # Previously estimate_task_costs! and each has_capacity call snapshotted the + # global metrics cache independently; since other threads bump the cache's + # generation on every task completion, each of those calls rebuilt (deep- + # copied) the entire snapshot. One snapshot per pass collapses that to at + # most a single rebuild. + snap = MT.snapshot(MT.global_metrics_cache()) + # Build the per-signature index once and share it with both callers, so # each is O(1) per proc after a single O(N) build (vs two independent # end-to-end MetricsTracker scans per task on the AOT path). sig_vec_for_index = sig isa Dagger.Signature ? sig.sig : sig - runtime_index = build_signature_runtime_index( - MT.snapshot(MT.global_metrics_cache()), sig_vec_for_index) + runtime_index = build_signature_runtime_index(snap, sig_vec_for_index) estimate_task_costs!(sorted_procs, costs, state, input_procs, task; - sig, runtime_index) + sig, runtime_index, snap) input_procs_cleanup() # Select the best available processor and reserve time pressure optimistically. @@ -910,7 +918,7 @@ concurrently across threads. if can_use has_cap, est_time_util, est_alloc_util, est_occupancy = has_capacity(state, proc, gproc.pid, options.time_util, options.alloc_util, options.occupancy, sig; - runtime_index) + runtime_index, snap) if has_cap # Optimistic reservation: capture-the-ref, atomic_add! counter_ref = lock(state.worker_time_pressure) do wtp @@ -1750,6 +1758,19 @@ function do_tasks(to_proc, return_queue, tasks) @dagdebug nothing :processor "Kicked processors" end +# do_task only needs a Context to satisfy logging (`@maybelog` reads `log_sink` +# and `profile`); the proc-management machinery it carries (a ReentrantLock and a +# Condition) is never used on this path. Since `(log_sink, profile)` are constant +# for a scheduler, memoize the last-built Context instead of allocating a fresh +# one -- plus its lock and condition -- on every task. Lock-free: the hot path is +# an atomic load + key compare (a Context is only built on a config change). +# Reusable per-task scratch cache for the metrics collected during `execute!`. +# Kept task-local (each pooled scheduler worker task owns one) and reset before +# each use, so its per-metric storage Dicts/Vectors are allocated once and reused +# rather than rebuilt for every task. `do_task` runs on a pool of reusable tasks +# (see `DoTaskSpec`/`@reusable_tasks`), so the same task services many tasks and +# the cache is genuinely reused. +const DO_TASK_LOCAL_METRICS = TaskLocalValue{MT.MetricsCache}(() -> MT.MetricsCache()) """ do_task(to_proc, task::TaskSpec) -> Any @@ -1935,7 +1956,8 @@ Executes a single task specified by `task` on `to_proc`. task_sig = signature(first(data), @view data[2:end]).sig - local_metrics_cache = MT.MetricsCache() + local_metrics_cache = DO_TASK_LOCAL_METRICS[] + MT.reset_pending!(local_metrics_cache) mspec = execute_metrics_spec() @dagdebug thunk_id :execute "Executing $(typeof(f))" diff --git a/src/sch/util.jl b/src/sch/util.jl index 9986bc420..9a95dff66 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -527,10 +527,15 @@ function can_use_proc(state, task, gproc, proc, opts, scope) end function has_capacity(state, p, gp, time_util, alloc_util, occupancy, sig; - runtime_index::Union{Dagger.SignatureRuntimeIndex,Nothing}=nothing) + runtime_index::Union{Dagger.SignatureRuntimeIndex,Nothing}=nothing, + snap=nothing) T = typeof(p) sig_vec = sig isa Dagger.Signature ? sig.sig : sig - snap = MT.snapshot(MT.global_metrics_cache()) + # Reuse a caller-provided per-pass snapshot when available (see + # schedule_one!); otherwise take our own. Avoiding a fresh snapshot per + # candidate processor is what keeps the global metrics cache from being + # deep-copied once per proc on every scheduling decision. + snap = snap === nothing ? MT.snapshot(MT.global_metrics_cache()) : snap worker_id = gp isa Int ? gp : (gp isa OSProc ? gp.pid : myid()) est_time_util = if time_util !== nothing && haskey(time_util, T) round(UInt64, time_util[T] * 1000^3)::UInt64 @@ -650,7 +655,8 @@ end const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) @reuse_scope function estimate_task_costs!(sorted_procs, costs, state, procs, task; sig=nothing, - runtime_index::Union{Dagger.SignatureRuntimeIndex,Nothing}=nothing) + runtime_index::Union{Dagger.SignatureRuntimeIndex,Nothing}=nothing, + snap=nothing) # Find all Chunks chunks = @reusable_vector :estimate_task_costs_chunks Union{Chunk,Nothing} nothing 32 @@ -665,7 +671,8 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) sig = signature(task.f, task.inputs) end sig_vec = sig isa Dagger.Signature ? sig.sig : sig - snap = MT.snapshot(MT.global_metrics_cache()) + # Reuse the caller's per-pass snapshot (schedule_one!) when provided. + snap = snap === nothing ? MT.snapshot(MT.global_metrics_cache()) : snap # Build the per-signature runtime index once. Every candidate `proc` in # this call shares `sig_vec`, so `metrics_lookup_runtime_from_index` diff --git a/src/utils/metrics.jl b/src/utils/metrics.jl index 2ab5c28c9..89feeedac 100644 --- a/src/utils/metrics.jl +++ b/src/utils/metrics.jl @@ -322,8 +322,10 @@ metrics_lookup_alloc_max(snap, sig, proc) = metrics_lookup_alloc(snap, sig, proc; reducer=maximum) function extract_collected_metrics(local_cache::MT.MetricsCache, key) - snap = MT.snapshot(local_cache) - ctx = get(snap.contexts, (Dagger, :execute!), nothing) + # `local_cache` is task-local and already fully written by `with_metrics` + # (on this same task) before we drain it, so read its pending storages + # directly instead of taking a defensive deep-copy snapshot. + ctx = MT.pending_context(local_cache, Dagger, :execute!) ctx === nothing && return nothing pairs = Tuple{MT.AbstractMetric, Any}[] for (metric, storage) in ctx.storages @@ -335,6 +337,13 @@ function extract_collected_metrics(local_cache::MT.MetricsCache, key) return pairs end +# Bound the global metrics cache to the most-recent this-many tasks (distinct +# thunk_id keys) per `(mod, context)`. Without a bound the cache grows one entry +# per metric per task forever, which dominates scheduler allocations (Dict +# rehash churn) on long-running workloads. The cost model only needs recent +# samples, so we keep a rolling window. +const METRICS_CACHE_MAX_TASKS = 100 + function apply_collected_metrics!(cache::MT.MetricsCache, key::K, pairs) where K pairs === nothing && return isempty(pairs) && return @@ -345,6 +354,7 @@ function apply_collected_metrics!(cache::MT.MetricsCache, key::K, pairs) where K storage = MT.get_or_create_storage!(ctx, metric) MT.set_metric_value!(storage, key, value) end + MT.trim_context!(ctx, METRICS_CACHE_MAX_TASKS) end return end diff --git a/test/scheduler.jl b/test/scheduler.jl index ec12a64da..1240cd9a6 100644 --- a/test/scheduler.jl +++ b/test/scheduler.jl @@ -502,7 +502,10 @@ end post_snap = MetricsTracker.snapshot(cache) @test haskey(post_snap.contexts, (Dagger, :execute!)) time_values = MetricsTracker.values_for_metric(post_snap, Dagger, :execute!, MetricsTracker.TimeMetric()) - @test length(time_values) >= pre_count + 5 + # The global cache is bounded to the most-recent METRICS_CACHE_MAX_TASKS + # tasks, so the count grows by the 5 new tasks only up to that cap. + @test length(time_values) >= min(pre_count + 5, Dagger.METRICS_CACHE_MAX_TASKS) + @test length(time_values) <= Dagger.METRICS_CACHE_MAX_TASKS time_inferred = Base.return_types(MetricsTracker.lookup_value, Tuple{MetricsTracker.MetricsSnapshot, Module, Symbol, MetricsTracker.TimeMetric, Int})[1] From 9a57097c7ce11e56cdc17af209a689879d4a124a Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 20 Jul 2026 07:30:55 -0700 Subject: [PATCH 28/44] Sch: Optimize steal allocations --- src/sch/Sch.jl | 41 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index cbadeeb4d..643dc0e71 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -1352,6 +1352,25 @@ function proc_states_values(uid::UInt64=Dagger.get_tls().sch_uid) return collect(values(states.dict)) end end +""" + proc_states_values!(buf::Vector{ProcessorState}, uid) -> buf + +Refill `buf` (in place) with the current `ProcessorState`s for scheduler `uid`. +Like `proc_states_values` but reuses the caller's buffer instead of allocating a +fresh `Vector` on every call, for hot paths (e.g. the per-attempt work-stealing +scan) that would otherwise churn a short-lived list. `buf`'s concrete element +type also keeps downstream field accesses from boxing. +""" +function proc_states_values!(buf::Vector{ProcessorState}, uid::UInt64=Dagger.get_tls().sch_uid) + states = proc_states(uid) + MemPool.lock_read(states.lock) do + empty!(buf) + for (_, state) in states.dict + push!(buf, state) + end + end + return buf +end function proc_state!(f, uid::UInt64, proc::Processor) states = proc_states(uid) state = MemPool.lock_read(states.lock) do @@ -1402,6 +1421,12 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re wid = get_parent(to_proc).pid work_to_do = false + # Reused across steal attempts so the work-stealing scan below doesn't + # allocate a fresh state list + shuffle permutation on every attempt. + # Allocated once per processor runner; grows at most to the processor + # count. Concretely typed so the field accesses in the scan don't box. + steal_states = ProcessorState[] + steal_perm = Int[] while isopen(return_queue) # Wait for new tasks if !work_to_do @@ -1457,10 +1482,18 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re # Try to steal from local queues randomly # TODO: Prioritize stealing from busiest processors - states = proc_states_values(uid) - # TODO: Try to pre-allocate this - P = randperm(length(states)) - for state in getindex.(Ref(states), P) + # + # Refill the reusable state buffer from the live dict, then + # shuffle indices in place with randperm! to keep steal targets + # random. This avoids the per-attempt allocation of a fresh + # state Vector, a permutation Vector, and a broadcast Memory + # (and the boxing that the old broadcast/iterate incurred on the + # ProcessorState/ProcessorInternalState field reads). + proc_states_values!(steal_states, uid) + resize!(steal_perm, length(steal_states)) + randperm!(steal_perm) + for pidx in steal_perm + state = steal_states[pidx] other_istate = state.state if other_istate.proc === to_proc continue From 8d198c2e66f5afc690eea8baf1f35ed0fdff9310 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 20 Jul 2026 07:31:29 -0700 Subject: [PATCH 29/44] Sch: Optimize processor queue kick --- src/sch/Sch.jl | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 643dc0e71..ec29d9d8d 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -1777,16 +1777,24 @@ function do_tasks(to_proc, return_queue, tasks) end notify(istate.reschedule) - # Kick other processors to make them steal + # Kick other processors to make them steal. # TODO: Alternatively, automatically balance work instead of blindly enqueueing - states = proc_states_values(uid) - P = randperm(length(states)) - for other_state in getindex.(Ref(states), P) - other_istate = other_state.state - if other_istate.proc === to_proc - continue + # + # Iterate the processor-state dict in place under a read lock and ring each + # other processor's doorbell, instead of the previous + # collect(values)+randperm+broadcast, which allocated a fresh Vector, a + # permutation Vector, and a Memory on every task fire. Notification order is + # irrelevant for work-stealing (each idle processor independently decides + # whether to steal), so no shuffle is needed, and `notify(::Doorbell)` is a + # lock-free atomic signal that is safe to call while holding the read lock. + let states = proc_states(uid) + MemPool.lock_read(states.lock) do + for (_, other_state) in states.dict + other_istate = other_state.state + other_istate.proc === to_proc && continue + notify(other_istate.reschedule) + end end - notify(other_istate.reschedule) end @dagdebug nothing :processor "Kicked processors" end From b957755b77fda8b26f9fc136a2cb3adb36b2929b Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 20 Jul 2026 07:32:11 -0700 Subject: [PATCH 30/44] Sch: Optimize do_task move --- src/sch/Sch.jl | 156 +++++++++++++++++++++++++++++++------------------ 1 file changed, 98 insertions(+), 58 deletions(-) diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index ec29d9d8d..d3a461a25 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -1812,6 +1812,85 @@ end # (see `DoTaskSpec`/`@reusable_tasks`), so the same task services many tasks and # the cache is genuinely reused. const DO_TASK_LOCAL_METRICS = TaskLocalValue{MT.MetricsCache}(() -> MT.MetricsCache()) + +struct _DoTaskCtxMemo + key::Tuple{UInt64,Bool} + ctx::Context +end +mutable struct _DoTaskCtxMemoBox + @atomic memo::Union{Nothing,_DoTaskCtxMemo} +end +const DO_TASK_CTX = _DoTaskCtxMemoBox(nothing) +function do_task_context(log_sink, profile::Bool) + key = (objectid(log_sink), profile) + memo = @atomic DO_TASK_CTX.memo + if memo !== nothing && memo.key === key + return memo.ctx + end + ctx = Context(Processor[]; log_sink, profile) + @atomic DO_TASK_CTX.memo = _DoTaskCtxMemo(key, ctx) + return ctx +end + +# Move one argument onto `to_proc`. Factored out of `do_task` so both the inline +# (cheap local move) and spawned (potential real data movement) paths call it +# without allocating a per-argument closure. +function _move_task_arg!(ctx, thunk_id, to_proc, @nospecialize(f), arg) + value = Dagger.value(arg) + position = arg.pos + @maybelog ctx timespan_start(ctx, :move, (;thunk_id, position, processor=to_proc), (;f, data=value)) + #= FIXME: This isn't valid if x is written to + x = if x isa Chunk + value = lock(TASK_SYNC) do + if haskey(CHUNK_CACHE, x) + Some{Any}(get!(CHUNK_CACHE[x], to_proc) do + # Convert from cached value + # TODO: Choose "closest" processor of same type first + some_proc = first(keys(CHUNK_CACHE[x])) + some_x = CHUNK_CACHE[x][some_proc] + @dagdebug thunk_id :move "Cache hit for argument $id at $some_proc: $some_x" + @invokelatest move(some_proc, to_proc, some_x) + end) + else + nothing + end + end + + if value !== nothing + something(value) + else + # Fetch it + time_start = time_ns() + from_proc = processor(x) + _x = @invokelatest move(from_proc, to_proc, x) + time_finish = time_ns() + if x.handle.size !== nothing + Threads.atomic_add!(transfer_time, time_finish - time_start) + Threads.atomic_add!(transfer_size, x.handle.size) + end + + @dagdebug thunk_id :move "Cache miss for argument $id at $from_proc" + + # Update cache + lock(TASK_SYNC) do + CHUNK_CACHE[x] = Dict{Processor,Any}() + CHUNK_CACHE[x][to_proc] = _x + end + + _x + end + else + =# + new_value = @invokelatest move(to_proc, value) + #end + if new_value !== value + @dagdebug thunk_id :move "Moved argument @ $position to $to_proc: $(typeof(value)) -> $(typeof(new_value))" + end + @maybelog ctx timespan_finish(ctx, :move, (;thunk_id, position, processor=to_proc), (;f, data=new_value); tasks=[Base.current_task()]) + arg.value = new_value + return +end + """ do_task(to_proc, task::TaskSpec) -> Any @@ -1821,7 +1900,7 @@ Executes a single task specified by `task` on `to_proc`. thunk_id = task.thunk_id ctx_vars = task.ctx_vars - ctx = Context(Processor[]; log_sink=ctx_vars.log_sink, profile=ctx_vars.profile) + ctx = do_task_context(ctx_vars.log_sink, ctx_vars.profile) from_proc = OSProc() data = task.data @@ -1907,67 +1986,28 @@ Executes a single task specified by `task` on `to_proc`. else data end - fetch_tasks = map(_data) do arg - #=FIXME:REALLOC_TASKS=# - Threads.@spawn begin - value = Dagger.value(arg) - position = arg.pos - @maybelog ctx timespan_start(ctx, :move, (;thunk_id, position, processor=to_proc), (;f, data=value)) - #= FIXME: This isn't valid if x is written to - x = if x isa Chunk - value = lock(TASK_SYNC) do - if haskey(CHUNK_CACHE, x) - Some{Any}(get!(CHUNK_CACHE[x], to_proc) do - # Convert from cached value - # TODO: Choose "closest" processor of same type first - some_proc = first(keys(CHUNK_CACHE[x])) - some_x = CHUNK_CACHE[x][some_proc] - @dagdebug thunk_id :move "Cache hit for argument $id at $some_proc: $some_x" - @invokelatest move(some_proc, to_proc, some_x) - end) - else - nothing - end - end - - if value !== nothing - something(value) - else - # Fetch it - time_start = time_ns() - from_proc = processor(x) - _x = @invokelatest move(from_proc, to_proc, x) - time_finish = time_ns() - if x.handle.size !== nothing - Threads.atomic_add!(transfer_time, time_finish - time_start) - Threads.atomic_add!(transfer_size, x.handle.size) - end - - @dagdebug thunk_id :move "Cache miss for argument $id at $from_proc" - - # Update cache - lock(TASK_SYNC) do - CHUNK_CACHE[x] = Dict{Processor,Any}() - CHUNK_CACHE[x][to_proc] = _x - end - - _x - end - else - =# - new_value = @invokelatest move(to_proc, value) - #end - if new_value !== value - @dagdebug thunk_id :move "Moved argument @ $position to $to_proc: $(typeof(value)) -> $(typeof(new_value))" - end - @maybelog ctx timespan_finish(ctx, :move, (;thunk_id, position, processor=to_proc), (;f, data=new_value); tasks=[Base.current_task()]) - arg.value = new_value - return + # Move the function and arguments onto `to_proc`. Only arguments that may + # require real (possibly remote, blocking) data movement -- Chunks / + # WeakChunks -- are fetched concurrently on a spawned task; plain inline + # values (the common case for leaf tasks) just need a cheap local `move`, so + # they are moved synchronously here. This avoids spawning (and immediately + # joining) a Task per argument, which was a per-task allocation source + # (FIXME:REALLOC_TASKS). The task list itself is a reused buffer, so leaf + # tasks with no movable arguments allocate nothing here. + fetch_tasks = @reusable_vector :do_task_fetch_tasks Union{Task,Nothing} nothing 32 + fetch_tasks_cleanup = @reuse_defer_cleanup empty!(fetch_tasks) + for arg in _data + value = Dagger.value(arg) + if value isa Chunk || value isa WeakChunk + push!(fetch_tasks, Threads.@spawn _move_task_arg!(ctx, thunk_id, to_proc, f, arg)) + else + _move_task_arg!(ctx, thunk_id, to_proc, f, arg) end end for task in fetch_tasks - fetch_report(task) + task === nothing || fetch_report(task::Task) end + fetch_tasks_cleanup() f = Dagger.value(first(data)) @assert !(f isa Chunk) "Failed to unwrap thunk function" From 23fb03dae9d548c55a45b5fa5f9138bce275873f Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 11:08:51 -0400 Subject: [PATCH 31/44] datadeps: K-slot capacity-aware EFT for multi-stream processors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dagger's cost model tracked one readiness timestamp per processor, an implicit capacity of 1. That is correct for `ThreadProc` but wrong for `CuArrayDeviceProc`/`ROCArrayDeviceProc`, which multiplex N independent streams over a device and can hold N kernels in flight. EFT therefore charged K queued GPU tasks the full serial `proc_ready + task_time`, overstating completion by up to a factor of N relative to single-slot CPU processors and distorting placement. Adds `Dagger.proc_concurrency(::Processor) = 1` (src/gpu.jl) with vendor overrides returning the backend stream count, and replaces the single timestamp with a per-slot vector in `ScheduleState`. Assignment claims the earliest-free slot: start = max(data_ready, min(slots)), which is exactly the previous arithmetic when there is one slot. `proc_ready_ns` is retained as `maximum(slots)`, preserving the invariant the test suite asserts. Capacity-aware processor allocation per Sinnen & Sousa 2005 §4.3. Verified on hudson (2x EPYC 9454, 2x H100): - K=1 equivalence: 4 sequential claims identical to the old model; 8 tasks x 10ms on 4 slots -> makespan 20.0 (capacity-1 gives 80.0) - scheduler suite: 527 tests, no regressions (60/100/135+1 broken/55/70/28/11) - CPU-only regression: two identical runs span 0.78x-1.54x on this cell, K-slot lands at 0.91x-1.36x -- within the cell's own 17-39% CV - absolute gains at random_dag n=160, 98 procs: Optimizing -38%, JuMP -16%, RR/Greedy/SA ~-10% Note: this did not change relative scheduler ordering at 98 procs with 20 tasks, where nothing queues and capacity cannot bind. It is a correctness fix to the cost model, not a tuning knob. Co-Authored-By: Claude Opus 4.8 --- ext/CUDAExt.jl | 8 ++++ ext/ROCExt.jl | 6 +++ src/datadeps/scheduling.jl | 90 ++++++++++++++++++++++++++++++++------ src/gpu.jl | 24 +++++++++- 4 files changed, 113 insertions(+), 15 deletions(-) diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index b2f2289df..a6d0a8b5d 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -618,6 +618,14 @@ Dagger.scope_key_precedence(::Val{:cuda_gpus}) = 1 const DEVICES = Dict{Int, CuDevice}() const CONTEXTS = Dict{Int, CuContext}() const STREAMS = Dict{Int, Vector{CuStream}}() + +# Each `CuArrayDeviceProc` multiplexes `length(STREAMS[device])` independent +# CUDA streams, so it can have that many kernels in flight concurrently. +# Reported to the cost model so EFT does not serialise queued GPU tasks. +function Dagger.proc_concurrency(proc::CuArrayDeviceProc) + streams = get(STREAMS, proc.device, nothing) + return streams === nothing ? 1 : max(1, length(streams)) +end const SYNCDEPS = Dagger.LockedObject(Dict{Int, Tuple{Int,Int}}()) function __init__() diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index 7d393f52e..c23382b42 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -492,6 +492,12 @@ Dagger.scope_key_precedence(::Val{:rocm_gpus}) = 1 const DEVICES = Dict{Int, HIPDevice}() const CONTEXTS = Dict{Int, HIPContext}() const STREAMS = Dict{Int, Vector{HIPStream}}() + +# See `Dagger.proc_concurrency` — ROCm equivalent of the CUDA stream count. +function Dagger.proc_concurrency(proc::ROCArrayDeviceProc) + streams = get(STREAMS, proc.device_id, nothing) + return streams === nothing ? 1 : max(1, length(streams)) +end const SYNCDEPS = Dagger.LockedObject(Dict{Int, Tuple{Int,Int}}()) function __init__() diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 4e74cb7af..27709880f 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -496,22 +496,72 @@ struct GreedyScheduler <: DataDepsScheduler end mutable struct ScheduleState task_finish_ns::Dict{Int, Float64} task_proc::Dict{Int, Processor} + # Aggregate readiness: when the processor is *fully* drained, i.e. + # `maximum(proc_slots[proc])`. Retained as the schedule-level summary and + # as the makespan input; equals `proc_slots[proc][1]` when concurrency is 1. proc_ready_ns::Dict{Processor, Float64} + # Per-execution-unit readiness: one completion timestamp per concurrent + # slot (`Dagger.proc_concurrency(proc)` of them). A `ThreadProc` has one + # slot and behaves exactly as the single-timestamp model did; a multi-stream + # GPU proc has one slot per stream, so K queued tasks finish in + # ceil(K/slots) rounds rather than being charged K serial task times. + # Capacity-aware processor allocation per Sinnen & Sousa 2005 §4.3. + proc_slots::Dict{Processor, Vector{Float64}} end -ScheduleState() = ScheduleState(Dict{Int, Float64}(), Dict{Int, Processor}(), Dict{Processor, Float64}()) +ScheduleState() = ScheduleState(Dict{Int, Float64}(), Dict{Int, Processor}(), + Dict{Processor, Float64}(), Dict{Processor, Vector{Float64}}()) function Base.copy(s::ScheduleState) - return ScheduleState(copy(s.task_finish_ns), copy(s.task_proc), copy(s.proc_ready_ns)) + return ScheduleState(copy(s.task_finish_ns), copy(s.task_proc), copy(s.proc_ready_ns), + Dict{Processor, Vector{Float64}}(k => copy(v) for (k, v) in s.proc_slots)) end function Base.empty!(s::ScheduleState) empty!(s.task_finish_ns) empty!(s.task_proc) empty!(s.proc_ready_ns) + empty!(s.proc_slots) return s end +""" + _slots_for(state, proc) -> Vector{Float64} + +Per-slot completion timestamps for `proc`, allocated on first use with one +entry per concurrent execution unit. Entries start at `0.0` (idle). +""" +function _slots_for(state::ScheduleState, proc::Processor) + slots = get(state.proc_slots, proc, nothing) + if slots === nothing + slots = zeros(Float64, max(1, Dagger.proc_concurrency(proc))) + state.proc_slots[proc] = slots + end + return slots +end + +# Claim the earliest-free slot on `proc` for a task that cannot start before +# `data_ready_ns`, returning the finish time. Mutates both the slot vector and +# the aggregate `proc_ready_ns`. With one slot this is exactly +# `max(data_ready, proc_ready) + runtime`, i.e. the pre-existing behaviour. +function _claim_slot!(state::ScheduleState, proc::Processor, data_ready_ns::Float64, + runtime_ns::Float64) + slots = _slots_for(state, proc) + idx = argmin(slots) + finish = max(data_ready_ns, slots[idx]) + runtime_ns + slots[idx] = finish + state.proc_ready_ns[proc] = maximum(slots) + return finish +end + +# Finish time a task would have on `proc` without committing to it. +function _peek_slot(state::ScheduleState, proc::Processor, data_ready_ns::Float64, + runtime_ns::Float64) + slots = get(state.proc_slots, proc, nothing) + earliest = slots === nothing ? 0.0 : slots[argmin(slots)] + return max(data_ready_ns, earliest) + runtime_ns +end + # In-place copy: reuse `dst`'s dict buffers instead of allocating fresh ones. # Used by SA/IG hot loops to avoid per-iteration Dict allocations for the # candidate/best/current buffers. Equivalent to `dst = copy(src)` in state @@ -525,6 +575,20 @@ function _copy_state!(dst::ScheduleState, src::ScheduleState) for (k, v) in src.task_proc dst.task_proc[k] = v end + # Reuse slot vectors in place so the SA/IG hot loop stays allocation-free. + for (k, v) in src.proc_slots + dv = get(dst.proc_slots, k, nothing) + if dv === nothing || length(dv) != length(v) + dst.proc_slots[k] = copy(v) + else + copyto!(dv, v) + end + end + if length(dst.proc_slots) != length(src.proc_slots) + for k in collect(keys(dst.proc_slots)) + haskey(src.proc_slots, k) || delete!(dst.proc_slots, k) + end + end empty!(dst.proc_ready_ns) for (k, v) in src.proc_ready_ns dst.proc_ready_ns[k] = v @@ -553,19 +617,20 @@ function greedy_assign_task!(state::ScheduleState, snap::MT.MetricsSnapshot, best_proc = nothing best_finish = Inf + best_data_ready = 0.0 + best_runtime = 0.0 @inbounds for w in 1:n_procs cache.proc_compatible[idx, w] || continue proc = all_procs[w] target_space = cache.proc_spaces[w] data_ready_ns = _greedy_earliest_data_ready_ns_cached(snap, dag_spec, spec, target_space, state, cache, w) - proc_avail = get(state.proc_ready_ns, proc, 0.0) - start_ns = max(data_ready_ns, proc_avail) - runtime_ns = cache.task_times[idx, w] - finish = start_ns + runtime_ns + finish = _peek_slot(state, proc, data_ready_ns, runtime_ns) if finish < best_finish best_finish = finish best_proc = proc + best_data_ready = data_ready_ns + best_runtime = runtime_ns end end @@ -575,8 +640,9 @@ function greedy_assign_task!(state::ScheduleState, snap::MT.MetricsSnapshot, end state.task_proc[idx] = best_proc - state.task_finish_ns[idx] = best_finish - state.proc_ready_ns[best_proc] = best_finish + # Commit to the slot that produced `best_finish`; `_claim_slot!` recomputes + # the same value from the same earliest-free slot. + state.task_finish_ns[idx] = _claim_slot!(state, best_proc, best_data_ready, best_runtime) return best_proc end @@ -823,6 +889,7 @@ function _replay_schedule!(state::ScheduleState, snap::MT.MetricsSnapshot, # hazard from skipping the copy. empty!(state.task_finish_ns) empty!(state.proc_ready_ns) + empty!(state.proc_slots) @inbounds for idx in 1:nv(dag_spec.g) if idx in destroyed @@ -840,12 +907,7 @@ function _replay_schedule!(state::ScheduleState, snap::MT.MetricsSnapshot, data_ready_ns = _greedy_earliest_data_ready_ns_cached(snap, dag_spec, spec, target_space, state, cache, w) runtime_ns = cache.task_times[idx, w] end - proc_avail = get(state.proc_ready_ns, proc, 0.0) - start_ns = max(data_ready_ns, proc_avail) - finish = start_ns + runtime_ns - - state.task_finish_ns[idx] = finish - state.proc_ready_ns[proc] = max(get(state.proc_ready_ns, proc, 0.0), finish) + state.task_finish_ns[idx] = _claim_slot!(state, proc, data_ready_ns, runtime_ns) end end return state diff --git a/src/gpu.jl b/src/gpu.jl index fac44893d..4aeadc0f5 100644 --- a/src/gpu.jl +++ b/src/gpu.jl @@ -133,4 +133,26 @@ function write_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, to::GPUR else write_remainder!(copies, copies_offset, to[], to_ptr, n) end -end \ No newline at end of file +end +""" + proc_concurrency(proc::Processor) -> Int + +Number of tasks a processor can execute concurrently. + +Defaults to `1`: one Dagger processor maps to one execution unit (e.g. a +`ThreadProc` is a single Julia thread running one task at a time). GPU +processor types whose backends multiplex several independent streams over one +device override this with their stream count, so cost-model schedulers account +for K-way in-processor parallelism when computing earliest-finish-time. + +This matters because EFT with an implicit capacity of 1 charges each queued +task the full serial `proc_ready + task_time`, which on a multi-stream GPU +overstates completion time by up to a factor of K and distorts placement +decisions relative to single-slot CPU processors. Capacity-aware processor +allocation is the model assumed by the task-scheduling literature these +schedulers implement (Sinnen & Sousa 2005, §4.3). + +See `_slots_for` in `src/datadeps/scheduling.jl` for the K-slot readiness +tracking that consumes this. +""" +proc_concurrency(::Processor) = 1 From 060164983322d08f889810c0249ab3c8eb249408 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 12:47:10 -0400 Subject: [PATCH 32/44] bench: scope datadeps regions to all procs so GPUs reach the scheduler `spawn_datadeps` builds its candidate processor set from `get_compute_scope()`, whose default carries `DefaultEnabledTaint` and admits only processors with `default_enabled(proc) == true`. `CuArrayDeviceProc` returns `false` there -- deliberate Dagger design so accelerators are opt-in rather than used implicitly. The consequence for the benchmark is that GPUs are dropped from `all_procs` before any scheduler sees them, at the `filter!` in `distribute_tasks_hierarchical!` (src/datadeps/hierarchical.jl:969) and its twin in `distribute_tasks!` (src/datadeps/queue.jl:207). Every CPU+GPU measurement taken before this commit was therefore CPU-only in effect: the GPU was enumerated, registered and metrics-warmed, then discarded prior to scheduling. `warm_gpu_metrics_for!` was unaffected because it already passes an explicit `UnionScope` of `ExactScope`s, which bypasses the taint -- this commit applies the same mechanism to the measured region. `run_workload!` is split into a scoped outer form and `_run_workload_inner!`, so `warm_gpu_metrics_for!` can supply its own GPU-only scope without nesting two scopes. `_bench_scope_all_procs()` returns `nothing` on CPU-only sessions, leaving that path untouched. Verified on hudson (8 CPU ThreadProc + 1x H100): - AOT hook candidate set: procs=8 gpu=0 -> procs=9 gpu=1 - partition size: 2 tasks -> 20 tasks (affinity partitioning had been splitting the DAG into single-worker fragments) - CPU-only session: scope is `nothing`, matmul nt=2 bs=64 correct at 3.17e-16, default path unchanged Known remaining blocker: `_eft_runtime_ns` returns the `GREEDY_DEFAULT_RUNTIME_NS` 1000 ms fallback for every processor, so the cost model cannot yet distinguish GPU from CPU and EFT still declines the GPU (which additionally carries transfer cost). Addressed separately. Co-Authored-By: Claude Opus 4.8 --- bench/datadeps_schedulers/driver.jl | 45 ++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 2940acf13..72c893ca0 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -108,7 +108,50 @@ function build_inputs(workload::Symbol, nt::Int, bs::Int) end end +""" + _bench_scope_all_procs() -> Union{Nothing, Dagger.UnionScope} + +Scope admitting every enumerated processor, or `nothing` on a CPU-only +session. + +`spawn_datadeps` derives its candidate set from `get_compute_scope()`, whose +default carries `DefaultEnabledTaint` and therefore admits only processors +with `default_enabled(proc) == true`. `CuArrayDeviceProc` (and the ROCm +equivalent) return `false` there by design, so GPUs are dropped from +`all_procs` before any scheduler sees them -- see the `filter!` in +`distribute_tasks_hierarchical!` (src/datadeps/hierarchical.jl) and its twin +in `distribute_tasks!` (src/datadeps/queue.jl). The effect is that a datadeps +region under the default scope is CPU-only regardless of scheduler, cost +model, or available hardware. + +Passing an explicit `UnionScope` of `ExactScope`s bypasses the taint, which +is the same mechanism `warm_gpu_metrics_for!` already relies on to force its +warmup onto GPUs. Returns `nothing` when no accelerator is present so +CPU-only runs take the untouched default path. +""" +function _bench_scope_all_procs() + procs = collect(Dagger.all_processors()) + isempty(procs) && return nothing + any(!_is_cpu_proc, procs) || return nothing # CPU-only: no wrap needed + return Dagger.UnionScope([Dagger.ExactScope(p) for p in procs]) +end + +# Outer entry point: wraps the region in an all-processor scope when +# accelerators are present. `warm_gpu_metrics_for!` calls the inner form +# directly under its own GPU-only scope to avoid nesting two scopes. function run_workload!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) + scope = _bench_scope_all_procs() + if scope === nothing + _run_workload_inner!(workload, inputs, sched) + else + Dagger.with_options(; scope=scope) do + _run_workload_inner!(workload, inputs, sched) + end + end + return +end + +function _run_workload_inner!(workload::Symbol, inputs, sched::Dagger.DataDepsScheduler) if workload === :cholesky Dagger.spawn_datadeps(; scheduler=sched) do tiled_cholesky!(inputs.M) @@ -312,7 +355,7 @@ function warm_gpu_metrics_for!(workload::Symbol, nt::Int, bs::Int) warm_sched = Dagger.RoundRobinScheduler() try Dagger.with_options(; scope=gpu_scope) do - run_workload!(workload, inputs, warm_sched) + _run_workload_inner!(workload, inputs, warm_sched) end catch e @warn "GPU metrics-warm crashed, continuing with CPU-only cost model" workload=workload tile_count=nt block_size=bs n_gpu_procs=length(gpu_procs) exception=(e, catch_backtrace()) From 755688590dc91ff0c7df51d807a4d99c68a5bd82 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 13:28:44 -0400 Subject: [PATCH 33/44] bench: add tiled LU workload Right-looking tiled LU without pivoting, mirroring the tile-BLAS structure Dagger uses internally in `LinearAlgebra.lu!(::DMatrix, ::NoPivot)` (src/array/lu.jl): GETRF on the diagonal tile, TRSM down the column panel and across the row panel, GEMM on the trailing submatrix. The two TRSM calls are asymmetric, following Dagger's version -- column panel solves against U ('R','U','N','N'), row panel against L with unit diagonal ('L','L','N','U'). Pivoting is omitted so the DAG shape is static and directly comparable to `tiled_cholesky!`; partial pivoting would add a reduction over the column panel per step plus a data-dependent row swap, changing the task graph between runs. `make_lu_tiles` builds a diagonally dominant matrix (same conditioning trick as `make_spd_tiles`) so the unpivoted factorization is numerically sound. Wired into `build_inputs`, `_run_workload_inner!` and `verify_workload`. Verification reconstructs L*U from the in-place factors (unit-lower from the strict lower triangle, U from the upper triangle) and compares against the original matrix. Smoke-tested CPU-only, correctness only, no measurements: lu nt=2 bs=256 rel_err 7.54e-16 lu nt=4 bs=256 rel_err 7.26e-16 cholesky nt=4 bs=256 (regression check) rel_err 3.37e-16 Co-Authored-By: Claude Opus 4.8 --- bench/datadeps_schedulers/driver.jl | 18 ++++++++ bench/datadeps_schedulers/workloads.jl | 63 ++++++++++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 72c893ca0..a64dedefa 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -85,6 +85,9 @@ function build_inputs(workload::Symbol, nt::Int, bs::Int) return (; A, B, C, dense_A = _assemble_dense(A), dense_B = _assemble_dense(B)) + elseif workload === :lu + A = make_lu_tiles(sz, bs) + return (; A, dense_reference = _assemble_dense(A)) elseif workload === :random_dag # For random_dag, `nt` is reinterpreted as `n_tasks` (total DAG # nodes), not a tile-grid dimension — structured BLAS workloads @@ -160,6 +163,10 @@ function _run_workload_inner!(workload::Symbol, inputs, sched::Dagger.DataDepsSc Dagger.spawn_datadeps(; scheduler=sched) do tiled_matmul!(inputs.C, inputs.A, inputs.B) end + elseif workload === :lu + Dagger.spawn_datadeps(; scheduler=sched) do + tiled_lu!(inputs.A) + end elseif workload === :random_dag Dagger.spawn_datadeps(; scheduler=sched) do tiled_random_dag!(inputs.tiles, inputs.parents, @@ -204,6 +211,17 @@ function verify_workload(workload::Symbol, inputs) denom = max(norm(ref), eps(Float64)) rel = norm(L * L' - ref) / denom return (rel < 1e-8, rel) + elseif workload === :lu + # Unpivoted LU overwrites A in place: unit-lower L in the strict + # lower triangle, U in the upper triangle including the diagonal. + # Reconstruct and compare L*U against the original matrix. + F = _assemble_dense(inputs.A) + L = tril(F, -1) + LinearAlgebra.I + U = triu(F) + ref = inputs.dense_reference + denom = max(norm(ref), eps(Float64)) + rel = norm(L * U - ref) / denom + return (rel < 1e-8, rel) elseif workload === :matmul C = _assemble_dense(inputs.C) ref = inputs.dense_A * inputs.dense_B diff --git a/bench/datadeps_schedulers/workloads.jl b/bench/datadeps_schedulers/workloads.jl index adf627cc1..7929f943c 100644 --- a/bench/datadeps_schedulers/workloads.jl +++ b/bench/datadeps_schedulers/workloads.jl @@ -214,6 +214,69 @@ end return procs[mod1((i - 1) * nt + j, length(procs))] end +""" + tiled_lu!(A::AbstractMatrix) + +Right-looking tiled LU without pivoting, mirroring the tile-BLAS structure +Dagger uses internally in `LinearAlgebra.lu!(::DMatrix, ::NoPivot)` +(src/array/lu.jl): GETRF on the diagonal tile, TRSM down the column panel and +across the row panel, GEMM on the trailing submatrix. + +Pivoting is omitted so the DAG shape is static and comparable to +`tiled_cholesky!` -- partial pivoting would add a reduction over the column +panel per step and a data-dependent row-swap, changing the task graph between +runs. `make_lu_tiles` builds a diagonally dominant matrix so the unpivoted +factorization is numerically sound. + +The two TRSM calls are asymmetric and follow Dagger's internal version: + - column panel `A[m,k]` solves against U: side='R', uplo='U', non-unit diag + - row panel `A[k,n]` solves against L: side='L', uplo='L', UNIT diag +""" +function tiled_lu!(A::AbstractMatrix) + mt, nt = size(A) + for k in 1:min(mt, nt) + Dagger.@spawn LinearAlgebra.generic_lufact!(InOut(A[k, k]), + LinearAlgebra.NoPivot(); + check=false) + for m in (k+1):mt + Dagger.@spawn LinearAlgebra.BLAS.trsm!('R', 'U', 'N', 'N', + 1.0, In(A[k, k]), + InOut(A[m, k])) + end + for n in (k+1):nt + Dagger.@spawn LinearAlgebra.BLAS.trsm!('L', 'L', 'N', 'U', + 1.0, In(A[k, k]), + InOut(A[k, n])) + for m in (k+1):mt + Dagger.@spawn LinearAlgebra.BLAS.gemm!('N', 'N', -1.0, + In(A[m, k]), + In(A[k, n]), 1.0, + InOut(A[m, n])) + end + end + end + return A +end + +""" + make_lu_tiles(sz, nb) -> Matrix{Chunk} + +Diagonally dominant tiles for unpivoted LU. Adding `sz` to the diagonal makes +every pivot dominant over its row, so `generic_lufact!(.., NoPivot())` is +stable without row interchanges -- the same conditioning trick +`make_spd_tiles` uses. +""" +function make_lu_tiles(sz::Int, nb::Int) + @assert sz % nb == 0 + nt = sz ÷ nb + A = rand(sz, sz) + A[diagind(A)] .+= sz + procs = _placement_procs() + return [_distribute_tile(copy(A[(i-1)*nb+1:i*nb, (j-1)*nb+1:j*nb]), + _tile_proc(procs, nt, i, j)) + for i in 1:nt, j in 1:nt] +end + function make_spd_tiles(sz::Int, nb::Int) @assert sz % nb == 0 nt = sz ÷ nb From 4289c3b0be3b5503ecd923848775b1f7cca8d34a Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 14:08:57 -0400 Subject: [PATCH 34/44] datadeps: fix EFT signature mismatch that flattened the cost model `_eft_runtime_ns` built its lookup signature from `spec.fargs` while the annotations (`In`/`Out`/`InOut`/`Deps`) were still attached. `Sch.signature` resolves an annotated arg via `chunktype(::InOut{T}) where T = T`, a type-level unwrap that yields the `Chunk{...}` type parameter and cannot recurse. The recorded signature is built at execution time from bare `Chunk` instances, resolving through `chunktype(c::Chunk) = c.chunktype` to the materialized type. So the read side constructed Any[typeof(BLAS.trsm!), ..., Chunk{Matrix{Float64},...}, ...] against a recorded Any[typeof(BLAS.trsm!), ..., Matrix{Float64}, ...] All four tiers of `_runtime_lookup_chain` pin `LookupExact(SignatureMetric())` -- only the processor constraint relaxes -- so every lookup missed and every task fell back to `GREEDY_DEFAULT_RUNTIME_NS` (1000ms) on every processor. With a constant runtime for all (task, proc) pairs the EFT objective is flat, and every scheduler was optimizing a degenerate cost model. Unwrap annotations with `unwrap_inout` before constructing the signature, so both sides agree. Measured on cholesky nt=4 bs=1024 CPU+GPU, Greedy: 0/180 hits before, 180/180 after. Adds `DAGGER_TRACE_EFT_LOOKUP=1` to count hits/misses and dump missed signatures; gated behind a `Ref{Bool}` since this runs O(tasks x procs). Also restores the benchmark driver's `MT.trim!` call (disabled during the eviction-hypothesis test) and raises `keep_per_metric` 100 -> 10000. The cache still needs bounding -- unbounded growth took a run 14.8s -> 48.0s -- but at 100 the `unsafe_free!`/`move!` entries, ~60% of the cache by volume, evicted the compute-task samples the cost model depends on. Co-Authored-By: Claude Opus 4.8 --- bench/datadeps_schedulers/driver.jl | 6 ++++- src/datadeps/scheduling.jl | 40 ++++++++++++++++++++++++++++- 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index a64dedefa..d7c178f22 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -387,7 +387,11 @@ const DEFAULT_BLOCK_SIZE = 128 const DEFAULT_TRIALS = 3 const DEFAULT_WARMUP = 1 -const METRICS_KEEP_PER_METRIC = 100 +# Bounded to keep the metrics cache from degrading the write path (unbounded +# growth took a run from 14.8s to 48.0s), but large enough that compute-task +# samples are not evicted by the `unsafe_free!`/`move!` entries that dominate +# by volume -- at 100 they crowded out ~60% of the cache. +const METRICS_KEEP_PER_METRIC = 10000 # Factories — not instances — because RoundRobin holds mutable state and each # trial needs a fresh copy. Default MILP budget is set generously since a diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 27709880f..1b4c04bba 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -867,10 +867,48 @@ rejection. _eft_tail_args(fargs::Tuple) = Base.tail(fargs) _eft_tail_args(fargs::AbstractVector) = @view fargs[2:end] +# Diagnostics for the signature lookup: set `DAGGER_TRACE_EFT_LOOKUP=1` to +# count hits/misses and dump the first few missed signatures. Gated behind a +# `Ref{Bool}` load because `_eft_runtime_ns` runs O(tasks x procs) times. +const TRACE_EFT_LOOKUP = Ref{Union{Nothing,Bool}}(nothing) +const EFT_LOOKUP_HITS = Threads.Atomic{Int}(0) +const EFT_LOOKUP_MISSES = Threads.Atomic{Int}(0) +function trace_eft_lookup() + tr = TRACE_EFT_LOOKUP[] + tr === nothing || return tr + tr = get(ENV, "DAGGER_TRACE_EFT_LOOKUP", "0") != "0" + TRACE_EFT_LOOKUP[] = tr + return tr +end + +# `spec.fargs` carry datadeps annotations (`In`/`Out`/`InOut`/`Deps`), which are +# stripped before the task executes. `Sch.signature` resolves an annotated arg +# through `chunktype(::InOut{T}) where T = T` -- a *type-level* unwrap that +# yields the `Chunk{...}` type parameter and cannot recurse further. The +# recorded signature, built at execution time, sees a bare `Chunk` instance and +# so resolves through `chunktype(c::Chunk) = c.chunktype` to the materialized +# type (e.g. `Matrix{Float64}`). Unwrapping the annotation here hands +# `Sch.signature` the same bare `Chunk` the write side saw, so both sides +# construct an identical signature and `LookupExact(SignatureMetric(), sig)` can +# match. Without this every lookup misses all four tiers of +# `_runtime_lookup_chain` and every task falls back to +# `GREEDY_DEFAULT_RUNTIME_NS`, flattening the cost model. +_eft_unwrap_arg(arg) = with_value(arg, first(unwrap_inout(value(arg)))) + function _eft_runtime_ns(snap::MT.MetricsSnapshot, spec, proc::Processor) - sig = Sch.signature(spec.fargs[1], _eft_tail_args(spec.fargs)).sig + f = _eft_unwrap_arg(spec.fargs[1]) + tail = map(_eft_unwrap_arg, _eft_tail_args(spec.fargs)) + sig = Sch.signature(f, tail).sig worker_id = root_worker_id(proc) runtime_lookup = metrics_lookup_runtime_median(snap, sig, proc, worker_id) + if trace_eft_lookup() + if runtime_lookup === nothing + n = Threads.atomic_add!(EFT_LOOKUP_MISSES, 1) + n < 5 && @warn "EFT lookup MISS" sig proc + else + Threads.atomic_add!(EFT_LOOKUP_HITS, 1) + end + end return runtime_lookup === nothing ? Float64(GREEDY_DEFAULT_RUNTIME_NS) : Float64(runtime_lookup) end From e453fdda17334a22e35110b77bfaa0a3bd30c8a7 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 15:00:47 -0400 Subject: [PATCH 35/44] datadeps: translate sig for GPU cost lookup so post-dispatch metrics reach the scheduler MetricsTracker keys runtime samples on the post-dispatch signature. A task submitted as `BLAS.trsm!(::Matrix{Float64}, ...)` is recorded on a GPU as `CUBLAS.trsm!(::CuArray{Float64,2,DeviceMemory}, ...)`, because the `Dagger.move(::CPUProc, ::GPUProc, ::typeof(fn))` overloads substitute the vendor function at dispatch. The scheduler only ever sees the pre-dispatch form, so no GPU sample was reachable from `_eft_runtime_ns`. That did not surface as a miss: the last tier of `_runtime_lookup_chain` drops the processor constraint entirely, so a GPU lookup matched the CPU-recorded entries and returned a CPU runtime. The cost model reported the GPU as exactly as fast as an average CPU thread -- invisible in hit-rate terms, fatal for heterogeneous scheduling. Add `_translate_sig_for(sig, proc)`, generic returning `nothing` in `src/gpu.jl`, overloaded per backend. The function half is generated by piggybacking the existing reflective BLAS/LAPACK -> CUBLAS/CUSOLVER (and rocBLAS/rocSOLVER) loop that already emits the paired `move` overloads, so it stays in lockstep with dispatch and adds no maintenance surface. The array-type half is hardcoded per backend (`Matrix{T}` -> `CuArray{T,2,DeviceMemory}` / `ROCArray{T,2}`); `DeviceMemory` has no CPU-side counterpart to derive. `_eft_runtime_ns` tries the translated signature *first* where a translation exists, then falls back to the untranslated lookup. Translating only after a miss, as originally sketched, would be dead code for exactly the reason above. The fallback keeps a CPU-derived estimate when no mapping exists, still better than the 1s default. Measured, cholesky nt=4 bs=1024, GPU vs CPU for the same task: BLAS.trsm! -> CUBLAS.trsm! HIT 0.194ms GPU vs 48.008ms CPU (247x) Stopgap. The general fix is a processor-independent task identity recorded alongside the physical signature, which removes the need for translation on every backend at once. Co-Authored-By: Claude Opus 4.8 --- ext/CUDAExt.jl | 31 +++++++++++++++++++++++++++++++ ext/ROCExt.jl | 28 ++++++++++++++++++++++++++++ src/datadeps/scheduling.jl | 20 +++++++++++++++++++- src/gpu.jl | 31 +++++++++++++++++++++++++++++++ 4 files changed, 109 insertions(+), 1 deletion(-) diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index a6d0a8b5d..3a3655060 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -507,11 +507,42 @@ for lib in [BLAS, LAPACK] fn = getproperty(lib, name) cufn = getproperty(culib, name) @eval Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, ::$(typeof(fn))) = $cufn + # Companion to the `move` above: the scheduler's cost lookup + # needs the same CPU->GPU function mapping, but at the type + # level and without a value to dispatch on. + @eval Dagger._translate_fn_for(::Type{$(typeof(fn))}, ::CuArrayDeviceProc) = $(typeof(cufn)) end end end end +# Array-type half of the signature translation. Hardcoded rather than +# reflective: there are only a handful of array types to map, and the +# `DeviceMemory` parameter has no CPU-side counterpart to derive it from. +_translate_type_for(::Type{Matrix{T}}) where T = CuArray{T,2,CUDA.DeviceMemory} +_translate_type_for(::Type{Vector{T}}) where T = CuArray{T,1,CUDA.DeviceMemory} +_translate_type_for(::Type{Array{T,N}}) where {T,N} = CuArray{T,N,CUDA.DeviceMemory} +# Scalars and other non-array arguments cross unchanged. +_translate_type_for(::Type{T}) where {T<:Union{Number,Char,Symbol,Function}} = T +_translate_type_for(::Type) = nothing + +function Dagger._translate_sig_for(sig::Vector, proc::CuArrayDeviceProc) + isempty(sig) && return nothing + out = Vector{Any}(undef, length(sig)) + # Element 1 is the function type; the rest are argument types. + fn = Dagger._translate_fn_for(sig[1], proc) + fn === nothing && return nothing + out[1] = fn + for i in 2:length(sig) + t = sig[i] + t isa Type || return nothing + mapped = _translate_type_for(t) + mapped === nothing && return nothing + out[i] = mapped + end + return out +end + # Adapt RefValue Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Base.RefValue) = Dagger.GPURef(Dagger.move(from_proc, to_proc, x[]), only(Dagger.memory_spaces(to_proc))) diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index c23382b42..121cbe5d1 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -385,11 +385,39 @@ for lib in [BLAS, LAPACK] fn = getproperty(lib, name) rocfn = getproperty(roclib, name) @eval Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, ::$(typeof(fn))) = $rocfn + # Companion to the `move` above: the scheduler's cost lookup + # needs the same CPU->GPU function mapping, but at the type + # level and without a value to dispatch on. + @eval Dagger._translate_fn_for(::Type{$(typeof(fn))}, ::ROCArrayDeviceProc) = $(typeof(rocfn)) end end end end +# Array-type half of the signature translation; mirrors CUDAExt. +_translate_type_for(::Type{Matrix{T}}) where T = ROCArray{T,2} +_translate_type_for(::Type{Vector{T}}) where T = ROCArray{T,1} +_translate_type_for(::Type{Array{T,N}}) where {T,N} = ROCArray{T,N} +# Scalars and other non-array arguments cross unchanged. +_translate_type_for(::Type{T}) where {T<:Union{Number,Char,Symbol,Function}} = T +_translate_type_for(::Type) = nothing + +function Dagger._translate_sig_for(sig::Vector, proc::ROCArrayDeviceProc) + isempty(sig) && return nothing + out = Vector{Any}(undef, length(sig)) + fn = Dagger._translate_fn_for(sig[1], proc) + fn === nothing && return nothing + out[1] = fn + for i in 2:length(sig) + t = sig[i] + t isa Type || return nothing + mapped = _translate_type_for(t) + mapped === nothing && return nothing + out[i] = mapped + end + return out +end + # Adapt RefValue Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Base.RefValue) = Dagger.GPURef(Dagger.move(from_proc, to_proc, x[]), only(Dagger.memory_spaces(to_proc))) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 1b4c04bba..2de0a9993 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -873,6 +873,7 @@ _eft_tail_args(fargs::AbstractVector) = @view fargs[2:end] const TRACE_EFT_LOOKUP = Ref{Union{Nothing,Bool}}(nothing) const EFT_LOOKUP_HITS = Threads.Atomic{Int}(0) const EFT_LOOKUP_MISSES = Threads.Atomic{Int}(0) +const EFT_LOOKUP_TRANSLATED = Threads.Atomic{Int}(0) function trace_eft_lookup() tr = TRACE_EFT_LOOKUP[] tr === nothing || return tr @@ -900,13 +901,30 @@ function _eft_runtime_ns(snap::MT.MetricsSnapshot, spec, proc::Processor) tail = map(_eft_unwrap_arg, _eft_tail_args(spec.fargs)) sig = Sch.signature(f, tail).sig worker_id = root_worker_id(proc) - runtime_lookup = metrics_lookup_runtime_median(snap, sig, proc, worker_id) + # Try the accelerator-side signature *first* on procs that have a + # translation. The untranslated lookup cannot be used as the gate: the last + # tier of `_runtime_lookup_chain` drops the processor constraint entirely, + # so on a GPU it matches the CPU-recorded entries and returns a CPU runtime + # instead of missing. Translating only after a miss would therefore be dead + # code. Falling back to the untranslated lookup keeps a CPU-derived estimate + # (still better than the 1s default) when a signature has no mapping. + translated_sig = _translate_sig_for(sig, proc) + runtime_lookup = nothing + via_translation = false + if translated_sig !== nothing + runtime_lookup = metrics_lookup_runtime_median(snap, translated_sig, proc, worker_id) + via_translation = runtime_lookup !== nothing + end + if runtime_lookup === nothing + runtime_lookup = metrics_lookup_runtime_median(snap, sig, proc, worker_id) + end if trace_eft_lookup() if runtime_lookup === nothing n = Threads.atomic_add!(EFT_LOOKUP_MISSES, 1) n < 5 && @warn "EFT lookup MISS" sig proc else Threads.atomic_add!(EFT_LOOKUP_HITS, 1) + via_translation && Threads.atomic_add!(EFT_LOOKUP_TRANSLATED, 1) end end return runtime_lookup === nothing ? Float64(GREEDY_DEFAULT_RUNTIME_NS) : Float64(runtime_lookup) diff --git a/src/gpu.jl b/src/gpu.jl index 4aeadc0f5..14c082b32 100644 --- a/src/gpu.jl +++ b/src/gpu.jl @@ -156,3 +156,34 @@ See `_slots_for` in `src/datadeps/scheduling.jl` for the K-slot readiness tracking that consumes this. """ proc_concurrency(::Processor) = 1 + +""" + _translate_fn_for(::Type{F}, proc::Processor) -> Type or nothing + +Map a CPU-side function type to the accelerator-side function type that will +actually execute on `proc`. Returns `nothing` when no mapping exists. + +Generated in lockstep with the `Dagger.move(::CPUProc, ::GPUProc, ::typeof(fn))` +overloads in the vendor extensions -- the two are inherently paired, since that +`move` is precisely what performs the substitution at dispatch time. +""" +_translate_fn_for(::Type, ::Processor) = nothing + +""" + _translate_sig_for(sig::Vector, proc::Processor) -> Vector or nothing + +Rewrite a CPU-side metrics signature into the form that `proc` would record +after dispatch, or `nothing` if any element has no mapping. + +Needed because `MetricsTracker` keys runtime samples on the *post-dispatch* +signature: a task submitted as `BLAS.gemm!(::Matrix{Float64}, ...)` is recorded +on a GPU as `CUBLAS.gemm!(::CuArray{Float64,2,DeviceMemory}, ...)`. The +scheduler only ever sees the pre-dispatch form, so without translation no GPU +sample is reachable and `_runtime_lookup_chain` falls through to its +processor-agnostic tier, handing back a CPU runtime for a GPU. + +This is a stopgap. The general fix is a processor-independent task identity +recorded alongside the physical signature, which would make translation +unnecessary for every backend at once. +""" +_translate_sig_for(::Vector, ::Processor) = nothing From e490572f5d0df24445f44272fc6cbadfe82b19c3 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 16:09:28 -0400 Subject: [PATCH 36/44] metrics: make cache bound runtime-settable via metrics_cache_max_tasks! `METRICS_CACHE_MAX_TASKS` was a `const` 100, applied through `MT.trim_context!` on every metrics write. That rolling window is tuned for steady-state scheduling, where only recent samples matter, and it is far too small for benchmark scenarios that deliberately warm the cost model with high signature diversity. Concretely: a GPU warmup writes ~60 entries, then CPU warmup plus measured trials push past 100 and evict them. Heterogeneous cost lookups then silently fall back to CPU-derived estimates -- and *which* GPU signatures survive is nondeterministic, depending on task completion order, so scheduling decisions become irreproducible across reruns of an identical benchmark. Measured per-cycle demand is ~235 entries at cholesky nt=4 and ~835 at nt=8, against a bound of 100. Convert to a `Ref` with a `metrics_cache_max_tasks!` setter rather than raising the constant: Dagger's shipped default stays 100, and harnesses opt in. The benchmark driver sets 5000 -- ~6x the largest measured cycle, and ~200x below the unbounded regime that previously degraded a run from 14.8s to 48.0s. The bound is process-local; multi-worker runs must set it on each worker. Co-Authored-By: Claude Opus 4.8 --- bench/datadeps_schedulers/driver.jl | 12 ++++++++++++ src/utils/metrics.jl | 25 +++++++++++++++++++++++-- 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index d7c178f22..38272e6f6 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -393,6 +393,18 @@ const DEFAULT_WARMUP = 1 # by volume -- at 100 they crowded out ~60% of the cache. const METRICS_KEEP_PER_METRIC = 10000 +# Dagger's default rolling window (100 tasks) is tuned for steady-state +# scheduling and is far too small for a harness that warms the cost model on +# purpose: a GPU warmup writes ~60 entries, then CPU warmup plus measured +# trials push past 100 and evict them, so GPU cost lookups silently fall back +# to CPU-derived estimates -- and *which* GPU signatures survive is +# nondeterministic, making scheduling decisions irreproducible across reruns. +# Measured demand for a full warm+trials cycle is ~235 entries at cholesky +# nt=4 and ~835 at nt=8; 5000 clears the largest config by ~6x while staying +# ~200x below the unbounded regime that degraded a run from 14.8s to 48.0s. +const METRICS_CACHE_MAX_TASKS = 5000 +Dagger.metrics_cache_max_tasks!(METRICS_CACHE_MAX_TASKS) + # Factories — not instances — because RoundRobin holds mutable state and each # trial needs a fresh copy. Default MILP budget is set generously since a # K~64 solve can exceed a minute; callers override as needed. Heuristic diff --git a/src/utils/metrics.jl b/src/utils/metrics.jl index 89feeedac..779130d66 100644 --- a/src/utils/metrics.jl +++ b/src/utils/metrics.jl @@ -342,7 +342,28 @@ end # per metric per task forever, which dominates scheduler allocations (Dict # rehash churn) on long-running workloads. The cost model only needs recent # samples, so we keep a rolling window. -const METRICS_CACHE_MAX_TASKS = 100 +const METRICS_CACHE_MAX_TASKS = Ref(100) + +""" + metrics_cache_max_tasks!(n::Integer) + +Set the rolling-window bound described above, returning the previous value. + +The default of 100 is tuned for steady-state scheduling, where only recent +samples matter. It is too small for benchmark harnesses that warm the cost +model deliberately: a GPU warmup writes ~60 entries, after which CPU warmup and +measured trials push past 100 and evict the GPU samples, so heterogeneous cost +lookups silently fall back to CPU-derived estimates. Measured demand for a full +warm+trials cycle is ~235 entries at cholesky nt=4 and ~835 at nt=8. + +Note this is process-local; multi-worker runs must set it on each worker. +""" +function metrics_cache_max_tasks!(n::Integer) + n > 0 || throw(ArgumentError("metrics cache bound must be positive, got $n")) + old = METRICS_CACHE_MAX_TASKS[] + METRICS_CACHE_MAX_TASKS[] = Int(n) + return old +end function apply_collected_metrics!(cache::MT.MetricsCache, key::K, pairs) where K pairs === nothing && return @@ -354,7 +375,7 @@ function apply_collected_metrics!(cache::MT.MetricsCache, key::K, pairs) where K storage = MT.get_or_create_storage!(ctx, metric) MT.set_metric_value!(storage, key, value) end - MT.trim_context!(ctx, METRICS_CACHE_MAX_TASKS) + MT.trim_context!(ctx, METRICS_CACHE_MAX_TASKS[]) end return end From ed807239b31a7c2e7866af3746fc874816152766 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 16:24:19 -0400 Subject: [PATCH 37/44] datadeps: record FromSpace/ToSpace/MoveSize metrics at move! execution site `metrics_lookup_move_rate` queries FromSpaceMetric, ToSpaceMetric and MoveSizeMetric to estimate transfer cost for EFT. None of the three was ever written by any code path, so the lookup returned `nothing` on every call and the scheduler fell back to GREEDY_DEFAULT_TRANSFER_RATE = 1_000_000 B/s (scheduling.jl:406). At 1 MB/s an 8 MiB tile is charged 8388.6ms of transfer, a ~10,000x overestimate of PCIe, which effectively bans GPU placement: no compute advantage can offset it. The write path already existed and was simply never wired in. `instrumented_move!` wraps `move!`, derives the payload size from `source.handle.size`, and calls `_record_move_metrics!` -- but had zero callers. The four datadeps copy tasks spawned plain `Dagger.move!`. Point those four spawn sites at `instrumented_move!`. Argument order and arity are identical, and all four already pass `meta=true`, so the arguments arrive as `Chunk`s with `handle.size` populated. Measured, cholesky nt=4 bs=1024 CPU+GPU: before: FromSpace/ToSpace/MoveSize have no storage; move_rate -> nothing after: 36 entries each; MoveSize = 8388608 (one tile); move_rate -> real GPU data_ready: 8388.6ms -> 93.4ms Identified during cost-model diagnostics for CPU+GPU benchmarking. Co-Authored-By: Claude Opus 4.8 --- src/datadeps/remainders.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index cf0e562b4..e5804bfb1 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -286,7 +286,7 @@ function enqueue_remainder_copy_to!(state::DataDepsState, dest_space::MemorySpac ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true Dagger.move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true Dagger.instrumented_move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the sources and writes to the target @@ -340,7 +340,7 @@ function enqueue_remainder_copy_from!(state::DataDepsState, dest_space::MemorySp ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true Dagger.move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=remainder_syncdeps meta=true Dagger.instrumented_move!(remainder_aliasing, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the sources and writes to the target @@ -376,7 +376,7 @@ function enqueue_copy_to!(state::DataDepsState, dest_space::MemorySpace, arg_w:: ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true Dagger.move!(dep_mod, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true Dagger.instrumented_move!(dep_mod, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the source and writes to the target @@ -408,7 +408,7 @@ function enqueue_copy_from!(state::DataDepsState, dest_space::MemorySpace, arg_w ctx = Sch.eager_context() id = rand(UInt) @maybelog ctx timespan_start(ctx, :datadeps_copy, (;id), (;)) - copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true Dagger.move!(dep_mod, dest_space, source_space, arg_dest, arg_source) + copy_task = Dagger.@spawn scope=dest_scope exec_scope=dest_scope syncdeps=copy_syncdeps meta=true Dagger.instrumented_move!(dep_mod, dest_space, source_space, arg_dest, arg_source) @maybelog ctx timespan_finish(ctx, :datadeps_copy, (;id), (;thunk_id=copy_task.uid, from_space=source_space, to_space=dest_space, arg_w, from_arg=arg_source, to_arg=arg_dest)) # This copy task reads the source and writes to the target From 2f3461f587c892dddadc34ca93a7ae701f6b6293 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Mon, 20 Jul 2026 16:32:57 -0400 Subject: [PATCH 38/44] datadeps: median-of-per-move-rates aggregation in metrics_lookup_move_rate The estimator reduced `sum(size) / sum(time)` across all matching moves. That is not robust: a `move!` sample records end-to-end task time, so the first few moves carry compilation and device-context setup cost. Measured on cholesky nt=4 bs=1024 CPU+GPU, the 36 move samples ran 0.57ms to 1072ms. Five cold-start outliers (405-1072ms) contributed ~96% of total elapsed time, dragging the aggregate to ~90MB/s -- roughly 100x below the 14.8GB/s seen at steady state. Combined with the transfer charge being applied per input tile, that understated PCIe badly enough that no compute advantage could offset it, and the GPU was never selected despite measuring 247x faster than CPU on trsm!. Reduce over per-move rates with `Statistics.median` instead, matching how the compute side already reduces its samples (`metrics_lookup_runtime_median`), so both halves of the cost model use the same estimator. Exposed as a `reducer` kwarg for symmetry with `metrics_lookup_runtime`. Samples below MOVE_RATE_MIN_SIZE_BYTES (4KiB) are dropped: at that size the measurement is dominated by fixed per-task overhead rather than bandwidth, so the implied rate is noise. Measured effect, cholesky nt=4 bs=1024 CPU+GPU, GPU data_ready per tile: 8388.6ms (1MB/s fallback, pre-ed80723) -> 93.4ms (real samples, sum aggregate) -> 0.7ms (real samples, median) ~= 11.3GB/s, consistent with PCIe Greedy now places 15 of 20 cholesky tasks on the GPU; previously zero. Co-Authored-By: Claude Opus 4.8 --- src/utils/metrics.jl | 43 ++++++++++++++++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/src/utils/metrics.jl b/src/utils/metrics.jl index 779130d66..c12206728 100644 --- a/src/utils/metrics.jl +++ b/src/utils/metrics.jl @@ -416,23 +416,48 @@ metrics_lookup_move_time_min(snap, from_space, to_space) = metrics_lookup_move_time_max(snap, from_space, to_space) = metrics_lookup_move_time(snap, from_space, to_space; reducer=maximum) +# Transfers below this are dominated by fixed per-task overhead rather than +# bandwidth, so their implied rate is numeric noise. Excluded from the sample. +const MOVE_RATE_MIN_SIZE_BYTES = UInt64(4096) + +""" + metrics_lookup_move_rate(snap, from_space, to_space; reducer=Statistics.median) + +Estimated bytes/second between two memory spaces, or `nothing` if no usable +samples exist. + +Reduces over *per-move* rates rather than `sum(size) / sum(time)`. The sum form +is not robust: a `move!` sample records end-to-end task time, so the first few +moves carry compilation and device-context setup. Measured on cholesky nt=4 +bs=1024 CPU+GPU, 5 cold-start samples out of 36 (405-1072ms, against a 0.57ms +steady-state minimum) contributed ~96% of total elapsed time and dragged the +aggregate to ~90MB/s, roughly 100x below the 14.8GB/s observed at steady state. +That understates PCIe so severely that no compute advantage can offset it and +GPU placement is effectively banned. + +Taking a median over per-move rates matches how the compute side already +reduces its samples (`metrics_lookup_runtime_median`), so both halves of the +cost model use the same estimator. +""" function metrics_lookup_move_rate(snap::MT.MetricsSnapshot, - from_space::MemorySpace, to_space::MemorySpace) + from_space::MemorySpace, to_space::MemorySpace; + reducer::Function=Statistics.median) matched = _move_matching_keys(snap, from_space, to_space) isempty(matched) && return nothing - total_time = UInt64(0) - total_size = UInt64(0) + rates = Float64[] + sizehint!(rates, length(matched)) for k in matched t = MT.lookup_value(snap, Dagger, :execute!, MT.TimeMetric(), k) s = MT.lookup_value(snap, Dagger, :execute!, MoveSizeMetric(), k) - if t !== nothing && s !== nothing && t > 0 && s > 0 - total_time += t - total_size += s - end + (t === nothing || s === nothing) && continue + (t == 0 || s < MOVE_RATE_MIN_SIZE_BYTES) && continue + push!(rates, Float64(s) / (Float64(t) / 1e9)) end - (total_time == 0 || total_size == 0) && return nothing - return round(UInt64, Float64(total_size) / (Float64(total_time) / 1e9)) + isempty(rates) && return nothing + rate = reducer(rates) + (isfinite(rate) && rate > 0) || return nothing + return round(UInt64, rate) end function metrics_lookup_transfer_rate(snap::MT.MetricsSnapshot, proc::Processor, worker_id::Int) From 6382a524241953e608e84233c2f5cba6eb3467c8 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Tue, 21 Jul 2026 00:43:04 -0400 Subject: [PATCH 39/44] datadeps: stochastic reconstruction for IteratedGreedy (was deterministic == Greedy) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit IG's destroy/reconstruct loop used `greedy_assign_task!` (strict `<` tie-breaking, fully deterministic) for destroyed tasks. Since `_replay_schedule!` pins preserved tasks at their original procs and walks task ids in fixed order, the deterministic reconstruction provably reproduces the greedy seed *exactly* on every iteration, regardless of the random destroy set: by induction, if tasks 1..i-1 are all at their original procs, task i sees identical context and greedy picks the same proc. IG's neighborhood was a single point, so IG could never improve on Greedy -- contradicting the stochastic-construction step of Ruiz & Stützle (2007). Verified empirically pre-fix: across 5 RNG seeds x 200 iters x 30% destroy on cholesky nt=4 CPU+GPU, IG cost == Greedy cost exactly and 0/20 tasks differed. Fix: add `greedy_assign_task_randomized!`, an alpha-greedy (GRASP-style) variant that builds a restricted candidate list of every compatible proc whose K-slot-peeked finish is within `(1+alpha)` of the best and picks uniformly among them (reservoir sampling, no extra alloc). It uses the same `_peek_slot` / `_claim_slot!` helpers as the deterministic variant, so K-slot capacity accounting is identical. Thread an optional `rng` through `_replay_schedule!` and `iterated_greedy_step!`; only IG passes it. `rng === nothing` (every other caller, including SA's empty-destroyed replay) keeps the deterministic path byte-for-byte. `greedy_schedule!` never routes through `_replay_schedule!`, so Greedy is untouched. Strict acceptance (`new_cost < best_cost`) means IG's result is monotone from the greedy seed and can never be worse than Greedy. HONEST OUTCOME: with the default alpha=0.10 this does NOT change the benchmark numbers, because the restricted candidate list is a singleton in both regimes: in CPU+GPU the GPU is ~250x faster so its EFT is uniquely best (confirmed: no divergence even at alpha=50); in CPU-only the measured per-thread runtime spread is 35-110% so the second-best proc is already outside a 10% window. Divergence does appear at alpha>=0.5 on CPU-only, confirming the search is wired, but no strictly-better schedule was found -- consistent with the greedy EFT solution already being at the cost model's optimum for these regular tiled-BLAS DAGs. This is the "genuine flat landscape" result, not a wiring failure. `alpha` is left at 0.10 (near-greedy, conservative) and is tunable; raising it as a default is deferred pending evidence it helps on some workload. Co-Authored-By: Claude Opus 4.8 --- src/datadeps/scheduling.jl | 101 +++++++++++++++++++++++++++++++++++-- 1 file changed, 96 insertions(+), 5 deletions(-) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 2de0a9993..132b58c47 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -646,6 +646,77 @@ function greedy_assign_task!(state::ScheduleState, snap::MT.MetricsSnapshot, return best_proc end +""" + greedy_assign_task_randomized!(state, snap, dag_spec, all_procs, idx, cache, rng; alpha) + +Stochastic variant of [`greedy_assign_task!`](@ref) used only by IG's +reconstruction path. Instead of always taking the single lowest-EFT processor +(strict `<`, fully deterministic), it builds a restricted candidate list of +every compatible processor whose peeked finish time is within `(1+alpha)` of +the best, then picks one uniformly at random. + +This is what gives Iterated Greedy a real search neighborhood. With the +deterministic `greedy_assign_task!`, `_replay_schedule!` provably reproduces +the greedy seed exactly on every iteration (preserved tasks stay at their +original procs, destroyed tasks are re-derived identically), so IG could never +improve on Greedy — its neighborhood was a single point. Randomized +reconstruction is the stochastic-construction step of Ruiz & Stützle (2007); +`alpha=0` recovers the deterministic behavior. + +Uses the same K-slot peek/claim helpers as `greedy_assign_task!`, so capacity +accounting is identical; only the choice among near-best procs differs. +""" +function greedy_assign_task_randomized!(state::ScheduleState, snap::MT.MetricsSnapshot, + dag_spec::DAGSpec, all_procs::Vector{Processor}, + idx::Int, cache::EFTCostCache, + rng::Random.AbstractRNG; + alpha::Float64=IG_DEFAULT_ALPHA) + spec = dag_spec.id_to_spec[idx] + n_procs = length(all_procs) + + # Single pass: peek every compatible proc, remember what `_claim_slot!` + # will need (data_ready, runtime), and track the best finish. Peeking does + # not mutate slot state, so all candidates are evaluated against the same + # partial schedule — exactly as the deterministic variant does. + cand = Tuple{Processor,Float64,Float64,Float64}[] + best_finish = Inf + @inbounds for w in 1:n_procs + cache.proc_compatible[idx, w] || continue + proc = all_procs[w] + target_space = cache.proc_spaces[w] + data_ready_ns = _greedy_earliest_data_ready_ns_cached(snap, dag_spec, spec, target_space, state, cache, w) + runtime_ns = cache.task_times[idx, w] + finish = _peek_slot(state, proc, data_ready_ns, runtime_ns) + push!(cand, (proc, data_ready_ns, runtime_ns, finish)) + finish < best_finish && (best_finish = finish) + end + + if isempty(cand) + task_scope = @something(spec.options.compute_scope, spec.options.scope, DefaultScope()) + throw(Sch.SchedulingException("IteratedGreedyScheduler: no compatible processor for task $idx (scope: $task_scope)")) + end + + # Uniform choice among candidates within the RCL threshold, via reservoir + # sampling so no second allocation is needed. `best_finish` is always + # in-threshold, so at least one candidate is always selectable. + threshold = best_finish * (1.0 + alpha) + chosen = cand[1] + n_in = 0 + @inbounds for c in cand + if c[4] <= threshold + n_in += 1 + if rand(rng) * n_in < 1.0 # replace with probability 1/n_in + chosen = c + end + end + end + + proc, dr, rt, _ = chosen + state.task_proc[idx] = proc + state.task_finish_ns[idx] = _claim_slot!(state, proc, dr, rt) + return proc +end + function greedy_assign_task!(state::ScheduleState, snap::MT.MetricsSnapshot, dag_spec::DAGSpec, all_procs::Vector{Processor}, idx::Int) cache = _build_eft_cost_cache(snap, dag_spec, all_procs) @@ -776,6 +847,11 @@ _greedy_arg_ready_time_ns_cached(::Any, ::MT.MetricsSnapshot, ::DAGSpec, const IG_DEFAULT_N_ITERS = 32 const IG_DEFAULT_DESTROY_FRAC = 0.30 +# Restricted-candidate-list width for IG's stochastic reconstruction: a +# destroyed task is placed uniformly at random among processors whose EFT is +# within `(1+IG_DEFAULT_ALPHA)` of the best. `alpha=0` collapses to +# deterministic greedy. See `greedy_assign_task_randomized!`. +const IG_DEFAULT_ALPHA = 0.10 # `Inf` disables the wall-clock stop. const IG_DEFAULT_TIME_LIMIT_SEC = Inf @@ -936,20 +1012,33 @@ end function _replay_schedule!(state::ScheduleState, snap::MT.MetricsSnapshot, dag_spec::DAGSpec, all_procs::Vector{Processor}, destroyed::AbstractSet{Int}, - cache::EFTCostCache) + cache::EFTCostCache; + rng::Union{Nothing,Random.AbstractRNG}=nothing, + alpha::Float64=IG_DEFAULT_ALPHA) # Only the finish-time and proc-ready-time dicts need to be cleared for a # fresh replay; `task_proc` is either overwritten by `greedy_assign_task!` # (destroyed tasks) or read as the previous assignment (preserved tasks). # Since the loop walks task ids in increasing order and each iteration # only reads/writes `task_proc[idx]` for its own idx, there is no aliasing # hazard from skipping the copy. + # + # `rng === nothing` (the default) reconstructs destroyed tasks with the + # deterministic `greedy_assign_task!` — used by every non-IG caller, + # including SA (which always passes an empty `destroyed` set). When an rng + # is supplied (IG only), destroyed tasks use the randomized variant so IG + # actually searches. Greedy is untouched: `greedy_schedule!` never routes + # through here. empty!(state.task_finish_ns) empty!(state.proc_ready_ns) empty!(state.proc_slots) @inbounds for idx in 1:nv(dag_spec.g) if idx in destroyed - greedy_assign_task!(state, snap, dag_spec, all_procs, idx, cache) + if rng === nothing + greedy_assign_task!(state, snap, dag_spec, all_procs, idx, cache) + else + greedy_assign_task_randomized!(state, snap, dag_spec, all_procs, idx, cache, rng; alpha=alpha) + end else proc = state.task_proc[idx] w = get(cache.proc_to_idx, proc, 0) @@ -981,8 +1070,10 @@ end function iterated_greedy_step!(state::ScheduleState, snap::MT.MetricsSnapshot, dag_spec::DAGSpec, all_procs::Vector{Processor}, destroyed::AbstractSet{Int}, - cache::EFTCostCache) - return _replay_schedule!(state, snap, dag_spec, all_procs, destroyed, cache) + cache::EFTCostCache; + rng::Union{Nothing,Random.AbstractRNG}=nothing, + alpha::Float64=IG_DEFAULT_ALPHA) + return _replay_schedule!(state, snap, dag_spec, all_procs, destroyed, cache; rng=rng, alpha=alpha) end function iterated_greedy_step!(state::ScheduleState, snap::MT.MetricsSnapshot, @@ -1052,7 +1143,7 @@ function iterated_greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapsho _copy_state!(candidate, best_state) - iterated_greedy_step!(candidate, snap, dag_spec, all_procs, destroyed, cache) + iterated_greedy_step!(candidate, snap, dag_spec, all_procs, destroyed, cache; rng=rng) new_cost = cost_of_schedule(candidate) if new_cost < best_cost # Swap; ex-best buffers get overwritten next iter. From 4ab6d4eed9a24569160a8415a5fa915c66366673 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Tue, 21 Jul 2026 01:44:58 -0400 Subject: [PATCH 40/44] =?UTF-8?q?datadeps:=20faithful=20Ruiz-St=C3=BCtzle?= =?UTF-8?q?=20acceptance=20+=20destruction=20for=20IteratedGreedy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follows Ruiz & Stützle, "A Simple and Effective Iterated Greedy Algorithm for the Permutation Flowshop Scheduling Problem" (2005/2007). Three corrections, each verified against the paper rather than taken from the request sketch, which deviated from it in several places: 1. SA-like acceptance with a SEPARATE best-so-far (§3.2). The loop now tracks an `incumbent` that the acceptance criterion may move uphill, and a monotone `best_state` that is the returned value. The sketch overwrote `best_state` on acceptance, which would forfeit the "never worse than the greedy seed" guarantee it simultaneously claimed. Verified: never_worse=true across 6 seeds on cholesky/matmul nt=4. 2. Constant acceptance temperature per eq. (1): Temperature = T * (sum_ij p_ij) / (n * m * 10) = T * mean(p_ij) / 10 averaged over compatible (task, proc) runtimes (`_ig_acceptance_temperature`). The sketch used `T * best_cost`; best_cost is a makespan (~100x a single task time), which inflates the temperature so far that a 250x-slower GPU->CPU move is accepted ~72% of the time -- a random walk that destroys the schedule. The paper's per-task-scale temperature (0.03-0.07ms here) correctly rejects such moves. 3. T = 0.5, not 0.4. §3.4 DOE finds T=0.5 best (T=0.0 = strict acceptance stagnates). Destruction count d = 4 (§3.4 tuned optimum), applied as a cap on the frac-derived count so our K=10..64 DAGs get d=3..4, inside the paper's optimal band; the sketch's frac=0.04 would give d=1..3, too low. The randomized reconstruction (`greedy_assign_task_randomized!`, GRASP-style alpha-RCL, from 6382a52) is retained and is a deliberate ADAPTATION: the paper's construction is deterministic NEH best-insertion, whose neighborhood comes from flowshop sequence positions. Task->processor assignment has no such neighborhood (a re-derived task returns to the same proc under pinned predecessors), so construction must be randomized to give IG a search space. This is documented as an adaptation, not attributed to the paper. OUTCOME on the paper's workloads: IG still ties Greedy (improve=0.0, never_worse=true). This is now a correct, fundamental result rather than a bug: the CPU+GPU workloads are GPU-dominated (one proc ~250x faster), so the RCL collapses to a singleton and any off-GPU move is correctly rejected by the (properly scaled) acceptance criterion. No correctly-implemented metaheuristic can improve on greedy EFT when a single processor dominates -- greedy is already optimal in the cost model. The implementation is faithful; the landscape is flat. Co-Authored-By: Claude Opus 4.8 --- src/datadeps/scheduling.jl | 80 +++++++++++++++++++++++++++++++++----- 1 file changed, 71 insertions(+), 9 deletions(-) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 132b58c47..b12fec76b 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -847,14 +847,50 @@ _greedy_arg_ready_time_ns_cached(::Any, ::MT.MetricsSnapshot, ::DAGSpec, const IG_DEFAULT_N_ITERS = 32 const IG_DEFAULT_DESTROY_FRAC = 0.30 +# Ruiz & Stützle (2005/2007) §3.4 tune the destruction size to a *fixed count* +# d, with d=4 the best-performing level (not statistically different from 3 or +# 5; very low and very high d are significantly worse). Our DAGs are small +# (K=10..64), so we keep `destroy_frac` as an escape hatch but cap the actual +# count at `IG_DEFAULT_DESTROY_COUNT`, which yields d=4 for K>=14 and d=3 at +# K=10 -- both within the paper's tuned-optimal band. +const IG_DEFAULT_DESTROY_COUNT = 4 +# Ruiz & Stützle §3.2 use an SA-like acceptance criterion with a *constant* +# temperature (eq. 1, after Osman & Potts 1989): +# Temperature = T * (sum_ij p_ij) / (n * m * 10) = T * mean(p_ij) / 10 +# with T the only acceptance parameter. Their §3.4 DOE finds T=0.5 best (0.0 = +# strict/greedy acceptance stagnates). We compute the mean over compatible +# (task, proc) runtimes; see `_ig_acceptance_temperature`. +const IG_DEFAULT_TEMPERATURE = 0.5 # Restricted-candidate-list width for IG's stochastic reconstruction: a # destroyed task is placed uniformly at random among processors whose EFT is # within `(1+IG_DEFAULT_ALPHA)` of the best. `alpha=0` collapses to -# deterministic greedy. See `greedy_assign_task_randomized!`. +# deterministic greedy. NOTE: Ruiz & Stützle's construction is deterministic +# NEH best-insertion; the sequence-position neighborhood of the flowshop makes +# that construction non-trivial. Task->processor assignment has no such +# neighborhood (a re-derived task returns to the same proc under pinned +# predecessors), so we randomize the construction GRASP-style to give IG a real +# search space. This is a deliberate adaptation, not the paper's mechanism. const IG_DEFAULT_ALPHA = 0.10 # `Inf` disables the wall-clock stop. const IG_DEFAULT_TIME_LIMIT_SEC = Inf +# Ruiz & Stützle §3.2 eq. (1): constant acceptance temperature proportional to +# the mean per-task work. Averaged over compatible (task, proc) pairs only, so +# the `GREEDY_DEFAULT_RUNTIME_NS` placeholders written for incompatible pairs +# in the cost cache do not inflate the scale. +function _ig_acceptance_temperature(cache::EFTCostCache, temperature::Float64) + total = 0.0 + count = 0 + n_tasks, n_procs = size(cache.task_times) + @inbounds for k in 1:n_tasks, w in 1:n_procs + cache.proc_compatible[k, w] || continue + total += cache.task_times[k, w] + count += 1 + end + count == 0 && return 0.0 + return temperature * (total / count) / 10.0 +end + """ IteratedGreedyScheduler{S<:DataDepsScheduler} <: DataDepsScheduler @@ -1111,14 +1147,31 @@ function iterated_greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapsho cache = _build_eft_cost_cache(snap, dag_spec, all_procs) end + # Two tracked solutions, per Ruiz & Stützle §3.2 (and standard SA/ILS): + # `best_*` — best found so far; only ever improves; returned. + # `incumbent_*` — the current solution the search walks from; the SA-like + # acceptance criterion may move it *uphill* to escape + # local optima. Destruction always operates on the + # incumbent, not on best. Conflating the two (accepting a + # worse solution *into* best) would forfeit the guarantee + # that IG never returns worse than its greedy seed. best_state = copy(state) best_cost = cost_of_schedule(best_state) + incumbent = copy(state) + incumbent_cost = best_cost - # Pre-allocate candidate once; reset in place via `_copy_state!` and swap - # references on accept to avoid Dict allocations in the hot loop. - candidate = copy(best_state) + # Reset in place via `_copy_state!` to avoid Dict allocations in the loop. + candidate = copy(state) + + # Ruiz & Stützle §3.4: fixed destruction count d=4, capped at K. See + # `IG_DEFAULT_DESTROY_COUNT`. `destroy_frac` is retained as an escape hatch. n_destroy = max(1, ceil(Int, destroy_frac * n_tasks)) - n_destroy = min(n_destroy, n_tasks) + n_destroy = min(n_destroy, IG_DEFAULT_DESTROY_COUNT, n_tasks) + + # Constant acceptance temperature, eq. (1). Zero temperature recovers strict + # (greedy) acceptance. + temperature = _ig_acceptance_temperature(cache, IG_DEFAULT_TEMPERATURE) + # Shared buffers reused across iters to avoid allocations. perm_buf = collect(1:n_tasks) destroyed = Set{Int}() @@ -1141,15 +1194,24 @@ function iterated_greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapsho push!(destroyed, perm_buf[i]) end - _copy_state!(candidate, best_state) - + # Destruct + (randomized) reconstruct from the *incumbent*. + _copy_state!(candidate, incumbent) iterated_greedy_step!(candidate, snap, dag_spec, all_procs, destroyed, cache; rng=rng) new_cost = cost_of_schedule(candidate) + + # Track best-so-far (monotone). if new_cost < best_cost - # Swap; ex-best buffers get overwritten next iter. - best_state, candidate = candidate, best_state + _copy_state!(best_state, candidate) best_cost = new_cost end + + # SA-like acceptance into the incumbent, eq. (1). Better always accepted; + # worse accepted with probability exp(-ΔC / temperature). + ΔC = new_cost - incumbent_cost + if ΔC < 0 || (temperature > 0 && rand(rng) < exp(-ΔC / temperature)) + _copy_state!(incumbent, candidate) + incumbent_cost = new_cost + end end return best_state end From 6940851c331d5022242b2273eedab6c2afd381ab Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Tue, 21 Jul 2026 03:09:44 -0400 Subject: [PATCH 41/44] =?UTF-8?q?datadeps:=20Orsila=20SA=20tuning=20?= =?UTF-8?q?=E2=80=94=20k=3D2=20temperature=20range=20+=20ECP=20critical-pa?= =?UTF-8?q?th=20move?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes from Orsila, Salminen & Hämäläinen (2008), verified against the paper rather than the request sketch: 1. SA_DEFAULT_K 1.0 -> 2.0. §4.3 (Eq. 18/19) states k widens the closed-form temperature range [Tf, T0], with suitable values in [1, 9]; k=1 gives ~27% final-level worsening acceptance, k=2 ~12%. Orsila et al. (2007) use k=2 as their thoroughness/cost tradeoff, so k=2 is the paper's own cited value and is adopted here. (The sketch proposed k=3 to "fix a too-small T0 under the cost cliff"; that premise is backwards -- by Eq. 18 the GPU cliff makes t_min_sum tiny so T0 is already large -- and no in-range k lets SA accept a 250x off-GPU move anyway, so the exact value matters little; k=2 is the defensible cited choice.) 2. ECP (Enhanced Critical Path) move, Orsila §3.7.1 (Wild et al. 2003, reported "significantly better than single move with directed acyclic graphs"). Bias the SA move's task selection toward the critical path -- proxied by task_finish_ns, since the sink has the maximum finish and near-critical tasks finish late. With probability SA_ECP_BIAS_PROB=0.5 select by finish-time roulette (so every on/near-critical-path task is a likely target, the sink most likely), else uniform to keep the neighborhood ergodic. The sketch took the single max-finish task deterministically, which retries the same sink on every biased move and wastes diversity; roulette over finish time favors the whole path, not just its endpoint. The existing move already excludes the current proc (Orsila §3.6), so that correctness detail needed no change. Expected impact is small on these workloads (GPU-dominated placement and hierarchically-partitioned ~17-task sub-DAGs leave little critical-path room), but both changes make SA faithful to the paper. Regression-safe: SA still tracks a separate monotone best and returns it. Existing SA tests pass. Co-Authored-By: Claude Opus 4.8 --- src/datadeps/scheduling.jl | 45 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index b12fec76b..d2e2fafd6 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -1248,9 +1248,19 @@ end # Defaults follow Orsila, Salminen, Hämäläinen 2008 §4.3 and §5.3. const SA_DEFAULT_Q = 0.95 -const SA_DEFAULT_K = 1.0 +# Orsila §4.3: k widens the closed-form temperature range [Tf, T0] (Eq. 18/19). +# Suitable values are [1, 9]; k=1 gives ~27% final-level worsening acceptance, +# k=2 ~12%. Orsila et al. (2007) use k=2 -- "reaches a local minimum more +# likely in the end, but is more expensive than k=1" -- so k=2 is the paper's +# own thoroughness/cost tradeoff and is adopted here. (Note: k does not help +# cross the GPU cost cliff; a 250x off-GPU move is rejected at any in-range T.) +const SA_DEFAULT_K = 2.0 const SA_DEFAULT_N_RESTARTS = 1 const SA_TF_FLOOR = 1e-12 +# Probability that a proposed SA move targets a critical-path task rather than +# a uniformly-random one (Orsila §3.7.1 ECP move; Wild et al. 2003 report it +# "significantly better than single move with directed acyclic graphs"). +const SA_ECP_BIAS_PROB = 0.5 const SA_DEFAULT_TIME_LIMIT_SEC = Inf """ @@ -1398,6 +1408,37 @@ end # Single-task move per Orsila §3.6.1. The current proc is excluded from # the alternative set (Orsila §3.6) so the move is never a no-op when the # task has at least one other compatible proc. +# Orsila §3.7.1 ECP (Enhanced Critical Path) move: bias task selection toward +# the critical path -- the chain of largest cumulative compute+communication, +# which determines the makespan. `task_finish_ns` is the exact per-task finish +# under the current schedule, so a task's finish time is a direct proxy for how +# close it is to the critical-path tail (the sink has the maximum finish). +# +# Rather than always taking the single max-finish task (which would retry the +# same sink on every biased move, wasting diversity), we roulette-select with +# probability proportional to finish time: every task on or near the critical +# path is a likely target, the sink most likely, off-path tasks rarely. With +# probability `1 - SA_ECP_BIAS_PROB` we fall back to a uniform pick to keep the +# neighborhood ergodic. +function _sa_select_task_ecp(candidate::ScheduleState, n_tasks::Int, + rng::Random.AbstractRNG) + if rand(rng) < SA_ECP_BIAS_PROB + total = 0.0 + @inbounds for idx in 1:n_tasks + total += get(candidate.task_finish_ns, idx, 0.0) + end + if total > 0.0 + threshold = rand(rng) * total + acc = 0.0 + @inbounds for idx in 1:n_tasks + acc += get(candidate.task_finish_ns, idx, 0.0) + acc >= threshold && return idx + end + end + end + return rand(rng, 1:n_tasks) +end + function _sa_propose_neighbor!(candidate::ScheduleState, snap::MT.MetricsSnapshot, dag_spec::DAGSpec, all_procs::Vector{Processor}, rng::Random.AbstractRNG, @@ -1405,7 +1446,7 @@ function _sa_propose_neighbor!(candidate::ScheduleState, snap::MT.MetricsSnapsho n_tasks = nv(dag_spec.g) n_tasks == 0 && return candidate - task_idx = rand(rng, 1:n_tasks) + task_idx = _sa_select_task_ecp(candidate, n_tasks, rng) current_proc = candidate.task_proc[task_idx] n_procs = length(all_procs) From 389f222fa3a8d5fcad57c1524156933a679ff7b2 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Tue, 21 Jul 2026 05:12:43 -0400 Subject: [PATCH 42/44] datadeps: producer-finish term in greedy data-ready (critical-path awareness) Greedy's data-ready for a `Chunk` argument previously accounted only for moving the tile from its *initial* materialized location (`_greedy_arg_ready_time_ns` on `::Chunk`) -- it had no notion of a producer task, so a consumer looked ready at t=0 even while its input was still being computed. That made greedy (and its `cost_of_schedule` makespan) blind to the DAG's critical path: it optimized a precedence-free bin-packing, not a real schedule. This is the actual flattening factor for the heuristics (the missing `dag_spec.g` edges were not -- greedy never read them; controlled A/B confirmed edges vs no-edges gave identical greedy output). Add a producer-finish term in `_greedy_earliest_data_ready_ns_cached`: when a `last_writer` map is supplied, a `Chunk` consumer waits for the most recent task that wrote that exact Chunk to finish, plus the transfer of its output to the target proc. `greedy_schedule!` builds the map by object identity as it walks tasks in index order (topological), so every producer is registered before its consumers. Chunk size and the cache's per-(proc,proc) move rate give the transfer, consistent with the existing DTask-arg path. Scoped to the standalone GreedyScheduler via `producer_finish=true`; off by default. IG/SA build their seed with the flag off so it stays consistent with their `_replay_schedule!` search (untouched, per plan) -- extending producer-finish to replay + the SA proposer is the follow-up if this moves the Greedy numbers. Measured, cholesky nt=4 CPU+GPU (same DAG + metrics): greedy makespan estimate 18.164ms (off) -> 68.233ms (on); the on value reflects the true critical path. Co-Authored-By: Claude Opus 4.8 --- src/datadeps/scheduling.jl | 65 ++++++++++++++++++++++++++++++++++---- 1 file changed, 59 insertions(+), 6 deletions(-) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index d2e2fafd6..baf0b92ba 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -611,7 +611,8 @@ end function greedy_assign_task!(state::ScheduleState, snap::MT.MetricsSnapshot, dag_spec::DAGSpec, all_procs::Vector{Processor}, idx::Int, - cache::EFTCostCache) + cache::EFTCostCache; + last_writer::Union{Nothing,IdDict{Any,Int}}=nothing) spec = dag_spec.id_to_spec[idx] n_procs = length(all_procs) @@ -623,7 +624,7 @@ function greedy_assign_task!(state::ScheduleState, snap::MT.MetricsSnapshot, cache.proc_compatible[idx, w] || continue proc = all_procs[w] target_space = cache.proc_spaces[w] - data_ready_ns = _greedy_earliest_data_ready_ns_cached(snap, dag_spec, spec, target_space, state, cache, w) + data_ready_ns = _greedy_earliest_data_ready_ns_cached(snap, dag_spec, spec, target_space, state, cache, w; last_writer=last_writer) runtime_ns = cache.task_times[idx, w] finish = _peek_slot(state, proc, data_ready_ns, runtime_ns) if finish < best_finish @@ -726,13 +727,32 @@ end function greedy_schedule!(state::ScheduleState, snap::MT.MetricsSnapshot, dag_spec::DAGSpec, all_procs::Vector{Processor}; task_order::Union{Nothing, AbstractVector{Int}}=nothing, - cache::Union{EFTCostCache, Nothing}=nothing) + cache::Union{EFTCostCache, Nothing}=nothing, + producer_finish::Bool=false) if cache === nothing cache = _build_eft_cost_cache(snap, dag_spec, all_procs) end order = task_order === nothing ? (1:nv(dag_spec.g)) : task_order + # When `producer_finish` is on, track the most recent task that wrote each + # Chunk (by object identity) so a consumer's data-ready time includes the + # producer's finish (see `_greedy_earliest_data_ready_ns_cached`). Populated + # as we walk tasks in index order (a valid topological order), so every + # producer is registered before its consumers are scheduled. Off by default, + # so IG/SA seeds (which drive an untouched, producer-finish-free replay) are + # unchanged. + last_writer = producer_finish ? IdDict{Any,Int}() : nothing for idx in order - greedy_assign_task!(state, snap, dag_spec, all_procs, idx, cache) + greedy_assign_task!(state, snap, dag_spec, all_procs, idx, cache; last_writer=last_writer) + if last_writer !== nothing + spec = dag_spec.id_to_spec[idx] + for arg in spec.fargs + raw, deps = unwrap_inout(value(arg)) + raw isa Chunk || continue + for (_dep_mod, _readdep, writedep) in deps + writedep && (last_writer[raw] = idx) + end + end + end end return state end @@ -741,7 +761,7 @@ function datadeps_schedule_dag_aot!(scheduler::GreedyScheduler, schedule, dag_sp snap = MT.snapshot(MT.global_metrics_cache()) cache = _build_eft_cost_cache(snap, dag_spec, all_procs) state = ScheduleState() - greedy_schedule!(state, snap, dag_spec, all_procs; cache=cache) + greedy_schedule!(state, snap, dag_spec, all_procs; cache=cache, producer_finish=true) for idx in 1:nv(dag_spec.g) task = dag_spec.id_to_task[idx] proc = state.task_proc[idx] @@ -801,11 +821,44 @@ end function _greedy_earliest_data_ready_ns_cached(snap, dag_spec::DAGSpec, spec, target_space::MemorySpace, state::ScheduleState, - cache::EFTCostCache, target_w::Int) + cache::EFTCostCache, target_w::Int; + last_writer::Union{Nothing,IdDict{Any,Int}}=nothing) earliest_ns = 0.0 for arg in spec.fargs raw_val, _ = unwrap_inout(value(arg)) ready_ns = _greedy_arg_ready_time_ns_cached(raw_val, snap, dag_spec, target_space, state, cache, target_w) + # Producer-finish term. For a `Chunk` argument, the base helper above + # accounts only for moving the tile from its *initial* materialized + # location -- it has no notion of a producer task. So greedy is blind to + # the DAG's critical path: a consumer looks ready at t=0 even though its + # input is still being computed. When `last_writer` is supplied (built + # by `greedy_schedule!` as it walks tasks in index order), look up the + # most recent task that wrote this exact Chunk and require the consumer + # to wait for that producer to finish AND for its output to reach the + # target proc. Additive: if there is no prior writer (root/materialized + # data) the term is skipped and behavior is unchanged. `last_writer` is + # only threaded from the standalone GreedyScheduler; IG/SA seeds pass + # `nothing`, keeping them consistent with their (untouched) replay. + if last_writer !== nothing && raw_val isa Chunk + wid = get(last_writer, raw_val, nothing) + if wid !== nothing + wfinish = get(state.task_finish_ns, wid, 0.0) + wproc = get(state.task_proc, wid, nothing) + if wproc !== nothing + prod_ready = wfinish + dep_w = get(cache.proc_to_idx, wproc, 0) + if dep_w != 0 && cache.proc_spaces[dep_w] != target_space + rate = cache.move_rates[dep_w, target_w] + if rate > 0.0 + sz = raw_val.handle.size === nothing ? + GREEDY_DEFAULT_OUTPUT_SIZE : UInt64(raw_val.handle.size) + prod_ready += Float64(sz) / rate * 1e9 + end + end + ready_ns = max(ready_ns, prod_ready) + end + end + end if ready_ns > earliest_ns earliest_ns = ready_ns end From a39797bb2e53cdaef53839c420f1ebe9dff93e5a Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Tue, 21 Jul 2026 15:26:14 -0400 Subject: [PATCH 43/44] bench: paper Table-1 grid harness + MILP-objective instrumentation Adds the ROSS 2026 Table-1 benchmark and the two supporting changes it needs; required to reproduce the grid on other machines (scheduler logic is already on the branch, but the harness was not). - bench/datadeps_schedulers/paper_grid.jl: the grid (2 regimes x 2 workloads x {2,4,8} tiles x {RR,Greedy,IG,SA(IG),JuMP} x 7 seeds). Per-config metrics warmup, seed-outer scheduler-inner order, per-cell schedule-cache reset (fresh schedule per seed), records wall/aot/residual/n_tasks/milp_status/ milp_obj. IG/SA/JuMP capped at 60s (PSZ fair-comparison). SA inner is IG. - ext/JuMPExt.jl: LAST_MILP_SOLVE Ref records the last solve's termination status and proven-optimal makespan (value(t_last_end), ns) for the paper's Greedy-vs-MILP optimality-gap number. Behavior-preserving (written after the existing solve, read by the harness). - bench/datadeps_schedulers/driver.jl: type datadeps_dag_equivalent's DAGSpec args on the TimedScheduler override so it is strictly more specific than Dagger's method -- untyped args are ambiguous and MethodError once a non-empty schedule cache triggers an equivalence lookup on the hierarchical path. - bench/datadeps_schedulers/Project.toml: bench env (CUDA/Dagger/HiGHS/JuMP) so runners can use --project=bench/datadeps_schedulers + Pkg.develop the local Dagger, instead of the root Project.toml deps hack (which stays uncommitted). Co-Authored-By: Claude Opus 4.8 --- bench/datadeps_schedulers/Project.toml | 5 + bench/datadeps_schedulers/driver.jl | 7 +- bench/datadeps_schedulers/paper_grid.jl | 170 ++++++++++++++++++++++++ ext/JuMPExt.jl | 12 ++ 4 files changed, 193 insertions(+), 1 deletion(-) create mode 100644 bench/datadeps_schedulers/Project.toml create mode 100644 bench/datadeps_schedulers/paper_grid.jl diff --git a/bench/datadeps_schedulers/Project.toml b/bench/datadeps_schedulers/Project.toml new file mode 100644 index 000000000..7d7a4b6df --- /dev/null +++ b/bench/datadeps_schedulers/Project.toml @@ -0,0 +1,5 @@ +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b" +JuMP = "4076af6c-e467-56ae-b986-b466b2749572" diff --git a/bench/datadeps_schedulers/driver.jl b/bench/datadeps_schedulers/driver.jl index 38272e6f6..4bf0a65fe 100644 --- a/bench/datadeps_schedulers/driver.jl +++ b/bench/datadeps_schedulers/driver.jl @@ -58,7 +58,12 @@ Base.similar(s::TimedScheduler) = TimedScheduler(similar(s.inner), s.last_aot_ns # inner scheduler's type (TimedScheduler{Greedy} reuses Greedy entries). Dagger.datadeps_schedule_cache(sched::TimedScheduler) = Dagger.datadeps_schedule_cache(sched.inner) -Dagger.datadeps_dag_equivalent(sched::TimedScheduler, dspec1, dspec2) = +# Args typed as ::DAGSpec so this override is strictly more specific than +# Dagger's `datadeps_dag_equivalent(::DataDepsScheduler, ::DAGSpec, ::DAGSpec)`; +# with untyped args the two are ambiguous (only the scheduler arg is more +# specific), which surfaces as a MethodError the moment a non-empty cache +# triggers an equivalence lookup on the hierarchical path. +Dagger.datadeps_dag_equivalent(sched::TimedScheduler, dspec1::Dagger.DAGSpec, dspec2::Dagger.DAGSpec) = Dagger.datadeps_dag_equivalent(sched.inner, dspec1, dspec2) Dagger.datadeps_argspec_equivalent(sched::TimedScheduler, a1, a2) = Dagger.datadeps_argspec_equivalent(sched.inner, a1, a2) diff --git a/bench/datadeps_schedulers/paper_grid.jl b/bench/datadeps_schedulers/paper_grid.jl new file mode 100644 index 000000000..74c5833d1 --- /dev/null +++ b/bench/datadeps_schedulers/paper_grid.jl @@ -0,0 +1,170 @@ +# Paper Table-1 benchmark grid (ROSS 2026). Two regimes (single-node CPU+1×H100, +# single-node CPU+2×H100), 2 workloads × 3 tile counts × 5 schedulers × 7 seeds. +# Reuses driver.jl building blocks (build_inputs, run_workload!, TimedScheduler, +# verify_workload, warm_*). Env: REGIME (a/b), OUTCSV, optional SMOKE=1. +using CUDA, Dagger, HiGHS, JuMP, LinearAlgebra, Statistics, Random, Printf +include(joinpath(@__DIR__, "driver.jl")) + +# --- config --------------------------------------------------------------- +const WORKLOADS = [:cholesky, :matmul] +const NTS = [2, 4, 8] +const BS_FOR_NT = Dict(2 => 4096, 4 => 4096, 8 => 2048) # nt=8 smaller to fit VRAM +const SEEDS = 1:7 +const HEUR_TL = 60.0 # PSZ fair-comparison cap for IG/SA (matches MILP) +const MILP_TL = 60.0 +const SCHED_NAMES = ["RoundRobinScheduler", "GreedyScheduler", + "IteratedGreedyScheduler", "SimulatedAnnealingScheduler", + "JuMPScheduler"] +const RESID_GATE = Dict(:cholesky => 1e-8, :matmul => 1e-10) + +# SA's inner MUST be IG (paper §sa Algorithm 3); constructed explicitly. +function make_scheduler(name::String, seed::Int) + if name == "RoundRobinScheduler" + return Dagger.RoundRobinScheduler() + elseif name == "GreedyScheduler" + return Dagger.GreedyScheduler() + elseif name == "IteratedGreedyScheduler" + return Dagger.IteratedGreedyScheduler(; rng=MersenneTwister(seed), + time_limit_sec=HEUR_TL) + elseif name == "SimulatedAnnealingScheduler" + return Dagger.SimulatedAnnealingScheduler( + Dagger.IteratedGreedyScheduler(; rng=MersenneTwister(seed), + time_limit_sec=HEUR_TL); + rng=MersenneTwister(seed), time_limit_sec=HEUR_TL) + elseif name == "JuMPScheduler" + return Dagger.JuMPScheduler(HiGHS.Optimizer; time_limit_sec=MILP_TL) + end + error("unknown scheduler $name") +end + +milp_ext() = Base.get_extension(Dagger, :JuMPExt) +function reset_milp_info!() + ext = milp_ext(); ext === nothing || (ext.LAST_MILP_SOLVE[] = ("NONE", NaN)) +end +function read_milp_info() + ext = milp_ext(); ext === nothing && return ("NA", NaN) + return ext.LAST_MILP_SOLVE[] +end + +# n_tasks (full-DAG K) captured once per (workload, nt, bs): deterministic in +# the workload topology. Uses a capturing scheduler that records the largest +# AOT DAG (the flat CPU+GPU DAG; hierarchical partitions would be smaller). +mutable struct NTCap <: Dagger.DataDepsScheduler + inner::Dagger.GreedyScheduler + maxn::Base.RefValue{Int} +end +function Dagger.datadeps_schedule_dag_aot!(s::NTCap, sc, dag, pr, sco) + s.maxn[] = max(s.maxn[], Dagger.nv(dag.g)) + Dagger.datadeps_schedule_dag_aot!(s.inner, sc, dag, pr, sco) +end +Dagger.datadeps_schedule_task_jit!(s::NTCap, a...) = Dagger.datadeps_schedule_task_jit!(s.inner, a...) +Dagger.datadeps_schedule_cache(s::NTCap) = Dagger.datadeps_schedule_cache(s.inner) +Base.similar(s::NTCap) = NTCap(similar(s.inner), s.maxn) +function capture_ntasks(wl, nt, bs) + cap = NTCap(Dagger.GreedyScheduler(), Ref(0)) + run_workload!(wl, build_inputs(wl, nt, bs), cap) + empty!(Dagger.datadeps_schedule_cache(cap)) # don't leak into measured cells + return cap.maxn[] +end + +# One measured cell. Returns a NamedTuple row (or an :EXCEPTION row). +function run_cell(regime, wl, nt, bs, name, seed, ntk) + # Metrics cache is warmed once per config (see `main`); measured cells then + # keep enriching it. The global `metrics_cache_max_tasks!(5000)` bound set + # by driver.jl trims on every write, so no per-cell trim is needed. + inputs = build_inputs(wl, nt, bs) + sched = make_scheduler(name, seed) + # Clear the per-type schedule cache so this (scheduler, seed) computes a + # FRESH schedule -- otherwise a structurally-equivalent prior cell's cached + # schedule would be reused, collapsing the 7-seed variance for IG/SA. + empty!(Dagger.datadeps_schedule_cache(sched)) + timed = TimedScheduler(sched) + reset_milp_info!() + + t0 = time_ns() + try + run_workload!(wl, inputs, timed) + catch e + @warn "CELL EXCEPTION" regime wl nt bs name seed exception=(e, catch_backtrace()) + mstat, mobj = name == "JuMPScheduler" ? read_milp_info() : ("", NaN) + return (; regime, workload=wl, nt, bs, scheduler=name, seed, + wall_ms=NaN, aot_ms=NaN, residual=NaN, n_tasks=ntk, + milp_status="EXCEPTION:$(typeof(e))", milp_obj=mobj, failed=true) + end + wall_ms = (time_ns() - t0) / 1e6 + aot_ms = timed.last_aot_ns[] / 1e6 + _, resid = verify_workload(wl, inputs) + mstat, mobj = name == "JuMPScheduler" ? read_milp_info() : ("", NaN) + return (; regime, workload=wl, nt, bs, scheduler=name, seed, + wall_ms, aot_ms, residual=resid, n_tasks=ntk, + milp_status=mstat, milp_obj=mobj, failed=false) +end + +csv_row(r) = join((r.regime, r.workload, r.nt, r.bs, r.scheduler, r.seed, + round(r.wall_ms, digits=3), round(r.aot_ms, digits=3), + r.residual, r.n_tasks, r.milp_status, r.milp_obj), ",") +const CSV_HDR = "regime,workload,nt,bs,scheduler,seed,wall_ms,aot_ms,residual,n_tasks,milp_status,milp_obj" + +# --- time-cap verification (PSZ): IG/SA must honor time_limit_sec ---------- +function verify_time_cap() + println("SMOKE time-cap check: IG/SA at nt=8 with a 3s cap should show aot≈3s") + wl, nt, bs = :matmul, 8, BS_FOR_NT[8] + warm_gpu_metrics_for!(wl, nt, bs); warm_metrics_for!(wl, nt, bs) + for (label, s) in ( + ("IG", Dagger.IteratedGreedyScheduler(; rng=MersenneTwister(1), time_limit_sec=3.0)), + ("SA", Dagger.SimulatedAnnealingScheduler( + Dagger.IteratedGreedyScheduler(; rng=MersenneTwister(1), time_limit_sec=3.0); + rng=MersenneTwister(1), time_limit_sec=3.0))) + timed = TimedScheduler(s) + run_workload!(wl, build_inputs(wl, nt, bs), timed) + aot = timed.last_aot_ns[] / 1e9 + @printf("SMOKE %s aot=%.2fs (cap=3s; respected if <~4s and not runaway)\n", label, aot) + end +end + +# --- main ----------------------------------------------------------------- +function main() + regime = get(ENV, "REGIME", "a") + outcsv = get(ENV, "OUTCSV", "paper_grid_$(regime).csv") + smoke = get(ENV, "SMOKE", "0") != "0" + ndev = length(collect(CUDA.devices())) + @printf("paper_grid regime=%s CUDA_devices=%d smoke=%s -> %s\n", regime, ndev, smoke, outcsv) + + workloads = smoke ? [:cholesky] : WORKLOADS + nts = smoke ? [parse(Int, get(ENV,"SMOKE_NT","4"))] : NTS + seeds = smoke ? (1:1) : SEEDS + + open(outcsv, "w") do io + println(io, CSV_HDR); flush(io) + for wl in workloads, nt in nts + bs = BS_FOR_NT[nt] + ntk = capture_ntasks(wl, nt, bs) + # Per-config warmup (once): populate the metrics cache with + # ground-truth samples. Scheduler/seed-independent, so one pass + # serves all cells of this config; measured cells enrich it further. + try + warm_gpu_metrics_for!(wl, nt, bs) + catch e + @warn "GPU warmup failed (continuing on CPU cost data)" wl nt exception=e + end + warm_metrics_for!(wl, nt, bs) + @printf("[config] %s nt=%d bs=%d n_tasks=%d\n", wl, nt, bs, ntk); flush(stdout) + # Seed-outer, scheduler-inner: every scheduler sees an equally-warm + # cache at each seed (removes the "later scheduler sees richer cache" + # bias). Order: seed=1[RR,Greedy,IG,SA,MILP], seed=2[...], ... + for seed in seeds, name in SCHED_NAMES + r = run_cell(regime, wl, nt, bs, name, seed, ntk) + println(io, csv_row(r)); flush(io) + gate = get(RESID_GATE, wl, 1e-8) + flag = r.failed ? " *** FAILED ***" : + (isfinite(r.residual) && r.residual > gate ? " *** RESIDUAL GATE ***" : "") + @printf(" %-28s seed=%d wall=%.1fms aot=%.1fms resid=%.1e%s\n", + name, seed, r.wall_ms, r.aot_ms, r.residual, flag); flush(stdout) + end + end + end + (smoke && get(ENV,"TIMECAP","1")!="0") && verify_time_cap() + println("paper_grid DONE -> $outcsv") +end + +main() diff --git a/ext/JuMPExt.jl b/ext/JuMPExt.jl index d089a954e..99d388c88 100644 --- a/ext/JuMPExt.jl +++ b/ext/JuMPExt.jl @@ -50,6 +50,16 @@ end # internal solver state, which would additionally kill the whole run. const JUMP_MAX_ASSIGNMENT_VARS = 10_000 +# Benchmark instrumentation (paper optimality-gap number): records the most +# recent MILP solve's termination status and proven-optimal makespan +# (`value(t_last_end)`, in ns, comparable to Greedy's `cost_of_schedule`). +# Only OPTIMAL cells are a valid lower bound. Behavior-preserving -- written +# after the solve, read by the benchmark harness. Under hierarchical +# partitioning this holds the *last* partition's solve; the paper's +# optimality-gap cells (K ≤ 12, flat CPU+GPU) are single-solve, so it is exact +# there. +const LAST_MILP_SOLVE = Ref{Tuple{String,Float64}}(("NONE", NaN)) + function Dagger.datadeps_schedule_dag_aot!(sched::JuMPScheduler, schedule, dag_spec, all_procs, all_scope) n_tasks = nv(dag_spec.g) n_tasks == 0 && return @@ -184,6 +194,8 @@ function Dagger.datadeps_schedule_dag_aot!(sched::JuMPScheduler, schedule, dag_s optimize!(model) status = termination_status(model) + LAST_MILP_SOLVE[] = (string(status), + primal_status(model) == MOI.FEASIBLE_POINT ? Float64(value(t_last_end)) : NaN) if !(status == MOI.OPTIMAL || status == MOI.TIME_LIMIT) throw(Sch.SchedulingException("JuMPScheduler: solver returned $status; no feasible schedule found")) end From 673bff9139f3364b5f97c4115612a8dab5aa1412 Mon Sep 17 00:00:00 2001 From: Bhirud R Date: Tue, 21 Jul 2026 18:25:08 -0400 Subject: [PATCH 44/44] bench: Tuolumne MI300a one-shot benchmark (M1/M2/M4 + optgap + cache demo) Self-contained single-session harness the HPC teammate runs once on 4 MI300a APU nodes. Three node regimes (M1/M2/M4) scoped from Dagger.all_processors() grouped by hostname; 2 workloads x 3 tile counts x 5 schedulers x 7 seeds, per-config warmup, seed-major ordering. After the main grid: AOT-only optimality-gap pass on every OPTIMAL MILP cell (cost_of_schedule vs proven makespan), plus a schedule-cache demonstration (two back-to-back AOT calls per regime). Graceful per-cell failure handling, 5-min disk flush, --smoke validation, and a 4-commit tree-state check at startup. Env activates in-script (AMDGPU backend, develop-mode Dagger); launch mechanism left to the teammate. Co-Authored-By: Claude Opus 4.8 --- bench/datadeps_schedulers/TUOLUMNE_README.md | 102 ++++ .../tuolumne_env/Project.toml | 17 + bench/datadeps_schedulers/tuolumne_oneshot.jl | 498 ++++++++++++++++++ 3 files changed, 617 insertions(+) create mode 100644 bench/datadeps_schedulers/TUOLUMNE_README.md create mode 100644 bench/datadeps_schedulers/tuolumne_env/Project.toml create mode 100644 bench/datadeps_schedulers/tuolumne_oneshot.jl diff --git a/bench/datadeps_schedulers/TUOLUMNE_README.md b/bench/datadeps_schedulers/TUOLUMNE_README.md new file mode 100644 index 000000000..f0f809c6b --- /dev/null +++ b/bench/datadeps_schedulers/TUOLUMNE_README.md @@ -0,0 +1,102 @@ +# Tuolumne MI300a one-shot benchmark + +One Julia session produces every number the paper needs for Tuolumne coverage: +single-node (M1), 2-node (M2), and 4-node (M4) MI300a APU regimes, for two +workloads × three tile counts × five schedulers × seven seeds, plus a post-hoc +optimality-gap pass and a schedule-cache demonstration. + +You run it **once**. There are no re-runs — validate with `--smoke` first. + +--- + +## 1. Prerequisites (do these on a login node, with network) + +1. **Check out the paper branch** at the pinned commit. The script verifies at + startup that these four scheduler commits are in `HEAD` (matched by subject, + so a rehash is fine) and aborts with a clear message if any is missing: + - "stochastic reconstruction for IteratedGreedy" + - "faithful Ruiz-Stützle acceptance …" + - "Orsila SA tuning …" + - "producer-finish term in greedy …" + plus the MILP-objective instrumentation (`JuMPExt.LAST_MILP_SOLVE`). + +2. **ROCm + AMDGPU.jl.** ROCm ≥ 6.0 (MI300a needs a recent ROCm; match what your + AMDGPU.jl build expects — `AMDGPU.versioninfo()` should show your APUs). The + script uses `AMDGPU` as the GPU backend and Dagger's ROCExt. + +3. **Instantiate the self-contained env** (developer-mode Dagger from this repo): + ```bash + julia --project=bench/datadeps_schedulers/tuolumne_env \ + -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()' + ``` + After this, launch with **no** `--project` — the script activates the env + itself. + +--- + +## 2. Recommended allocation + +- **4 MI300a nodes**, **8 hours** wall-time (≈6 h expected; the margin covers + cold JIT and the 60 s IG/SA/MILP caps that dominate the large-K cells). +- Launch mechanism is **your call** — Flux Framework, Slurm, whatever Tuolumne + uses. The script does **not** spawn workers; it consumes whatever + `Dagger.all_processors()` reports and groups them into nodes by hostname, then + scopes M1/M2/M4 to the first 1/2/4 nodes. Start Julia with one Dagger worker + per node (each with all local CPU threads + the node's APU GPU visible), added + through Dagger's DistributedNext so Dagger can see them. +- If fewer than 4 nodes are available, the regimes that don't fit are skipped and + logged — M1 still runs on a single node. + +--- + +## 3. Run — two commands, in order + +**a) Smoke test (~30–60 s) — validate the stack before committing the full run:** +```bash +julia bench/datadeps_schedulers/tuolumne_oneshot.jl --smoke +``` +It prints the detected topology (nodes / procs / APU GPUs), runs one cholesky +nt=2 cell with RR and Greedy end-to-end, and ends with a single line: +`SMOKE: PASS` or `SMOKE: FAIL`. **Do not proceed unless it PASSes and reports at +least one APU GPU.** + +**b) Full grid (~6–8 h):** +```bash +julia bench/datadeps_schedulers/tuolumne_oneshot.jl +``` +Progress prints per config/cell so you can see forward motion. Partial CSVs are +flushed every 5 minutes, so a wall-time expiry won't lose completed work. Any +single cell that fails (correctness residual over threshold, MILP exception, +worker crash) is caught, logged, marked in the CSV, and the run continues. + +--- + +## 4. Output (lands in your invocation directory) + +| File | Contents | +|---|---| +| `tuolumne_regime_m1_1node.csv` | M1 per-cell: wall_ms, aot_ms, residual, n_tasks, milp_status, milp_obj, bs, n_nodes | +| `tuolumne_regime_m2_2node.csv` | M2 per-cell (same columns) | +| `tuolumne_regime_m4_4node.csv` | M4 per-cell | +| `tuolumne_medians.csv` | per-cell medians + stddev over 7 seeds (feeds Table 1) | +| `tuolumne_optgap.csv` | AOT-only Greedy/IG/SA `cost_of_schedule` vs MILP-optimal, with ratio (headline optimality-gap) | +| `tuolumne_cache.csv` | schedule-cache demo: `aot_ms_first_call` vs `aot_ms_cache_hit` per regime | +| `tuolumne_run_summary.txt` | git SHA, launch/elapsed times, failure count, flagged cells, full log | + +Notes: +- **bs**: 4096 for nt ∈ {2,4}; nt=8 tries 4096 first (MI300a's 128 GB unified + HBM3 may fit it) and falls back to 2048 on OOM — the actual bs is recorded per + cell and logged. +- **optgap**: only cells where MILP proved `OPTIMAL` (small K) contribute; that's + by design — a time-limited MILP is not a valid lower bound. +- **cache**: `aot_ms_cache_hit` should be ≈0 (structural-equivalence lookup). If + the two columns are similar, the cache isn't hitting on your Julia version — + flag it to us. + +--- + +## 5. What to send back + +All seven files above. If the smoke test failed, send `tuolumne_run_summary.txt` +(or the console output) and your `AMDGPU.versioninfo()` / ROCm version so we can +diagnose the stack before you spend the allocation. diff --git a/bench/datadeps_schedulers/tuolumne_env/Project.toml b/bench/datadeps_schedulers/tuolumne_env/Project.toml new file mode 100644 index 000000000..9ee2d6a6d --- /dev/null +++ b/bench/datadeps_schedulers/tuolumne_env/Project.toml @@ -0,0 +1,17 @@ +# Self-contained env for the Tuolumne MI300a one-shot benchmark. +# Dagger is added by the script via Pkg.develop(path=) so it runs the +# paper's in-tree branch, not a registry release. AMDGPU replaces CUDA (the +# hudson bench env) as the GPU backend for AMD APUs. +# +# First-time setup, on a LOGIN node with network access (see TUOLUMNE_README.md): +# julia --project=bench/datadeps_schedulers/tuolumne_env \ +# -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()' +[deps] +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b" +JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" diff --git a/bench/datadeps_schedulers/tuolumne_oneshot.jl b/bench/datadeps_schedulers/tuolumne_oneshot.jl new file mode 100644 index 000000000..ee3ac4463 --- /dev/null +++ b/bench/datadeps_schedulers/tuolumne_oneshot.jl @@ -0,0 +1,498 @@ +#!/usr/bin/env julia +# ============================================================================= +# tuolumne_oneshot.jl -- one-shot MI300a APU benchmark for the ROSS 2026 paper +# ============================================================================= +# ONE Julia session on ONE allocation covers three regimes: +# M1 single node (1 node's CPU cores + APU GPU(s)) +# M2 two nodes +# M4 four nodes +# by scoping each regime to a subset of the workers Dagger can see. It does NOT +# spawn workers itself: launch Julia with however many workers your site's +# workload manager provides (Flux Framework, Slurm, ...); this script consumes +# whatever `Dagger.all_processors()` reports and groups them into nodes. See +# TUOLUMNE_README.md. +# +# Validate first: julia bench/datadeps_schedulers/tuolumne_oneshot.jl --smoke +# Full run: julia bench/datadeps_schedulers/tuolumne_oneshot.jl +# +# Outputs land in the invocation directory: +# tuolumne_regime_m1_1node.csv / _m2_2node.csv / _m4_4node.csv +# tuolumne_medians.csv tuolumne_optgap.csv tuolumne_run_summary.txt +# ============================================================================= + +# ---- 0. Bootstrap: activate a self-contained env, no --project needed ------- +import Pkg +const SCRIPT_DIR = @__DIR__ +const REPO_ROOT = normpath(joinpath(SCRIPT_DIR, "..", "..")) +const ENV_DIR = joinpath(SCRIPT_DIR, "tuolumne_env") +try + Pkg.activate(ENV_DIR) + # `develop` the in-tree Dagger so we run the paper's branch, not a registry + # release. Idempotent; the README asks you to `instantiate` on a login node + # (with network) beforehand, so this is a no-op re-check on the compute node. + if Base.find_package("Dagger") === nothing || !haskey(Pkg.project().dependencies, "Dagger") + Pkg.develop(path=REPO_ROOT) + end + Pkg.instantiate() +catch e + println(stderr, "FATAL: environment setup failed. Run the prereq from ", + "TUOLUMNE_README.md on a login node first:") + println(stderr, " julia --project=$(ENV_DIR) -e 'using Pkg; Pkg.develop(path=\"$(REPO_ROOT)\"); Pkg.instantiate()'") + showerror(stderr, e); println(stderr) + exit(2) +end + +using Dagger, AMDGPU, HiGHS, JuMP +using LinearAlgebra, Statistics, Random, Printf, Dates +import Dagger: greedy_schedule!, iterated_greedy_schedule!, simulated_annealing_schedule!, + cost_of_schedule, ScheduleState, _build_eft_cost_cache +import Dagger.MetricsTracker as MT +# Dagger schedules over DistributedNext and re-imports its worker API into its +# own namespace (see src/Dagger.jl), so `Dagger.myid` / `Dagger.remotecall_fetch` +# are the correct handles regardless of how the site's manager spawned workers. + +# Building blocks (build_inputs, run_workload!/_run_workload_inner!, verify_workload, +# TimedScheduler, ...). driver.jl is backend-agnostic (no CUDA/AMDGPU import) and +# sets metrics_cache_max_tasks!(5000). +include(joinpath(SCRIPT_DIR, "driver.jl")) + +const SMOKE = ("--smoke" in ARGS) + +# ---- 1. Configuration ------------------------------------------------------- +const WORKLOADS = [:cholesky, :matmul] +const NTS = [2, 4, 8] +const SCHED_NAMES = ["RoundRobinScheduler", "GreedyScheduler", + "IteratedGreedyScheduler", "SimulatedAnnealingScheduler", + "JuMPScheduler"] +const SEEDS = 1:7 +const HEUR_TL = 60.0 # PSZ fair-comparison cap for IG/SA (matches MILP) +const MILP_TL = 60.0 +const RESID_GATE = Dict(:cholesky => 1e-8, :matmul => 1e-10) +# bs: 4096 for nt in {2,4}; nt=8 tries 4096 first (MI300a's 128 GB unified HBM3 +# may fit it) and falls back to 2048 if the warmup OOMs. Chosen bs is recorded +# per cell and logged. +bs_candidates(nt) = nt == 8 ? [4096, 2048] : [4096] +const REGIMES = [("m1", 1, "tuolumne_regime_m1_1node.csv"), + ("m2", 2, "tuolumne_regime_m2_2node.csv"), + ("m4", 4, "tuolumne_regime_m4_4node.csv")] +const FLUSH_INTERVAL_S = 300.0 # periodic partial-CSV flush (feature 3) + +# The four scheduler-science commits that MUST be in HEAD (matched by subject so +# a rebase/re-hash doesn't break the check). Plus the MILP-objective +# instrumentation the optimality-gap number depends on. +const REQUIRED_SUBJECTS = [ + "stochastic reconstruction for IteratedGreedy", + "Ruiz-Stützle acceptance", + "Orsila SA tuning", + "producer-finish term in greedy", +] + +# ---- 2. Scheduler construction (exactly the five, SA-over-IG) --------------- +function make_scheduler(name::String, seed::Int) + if name == "RoundRobinScheduler" + return Dagger.RoundRobinScheduler() + elseif name == "GreedyScheduler" + return Dagger.GreedyScheduler() + elseif name == "IteratedGreedyScheduler" + return Dagger.IteratedGreedyScheduler(; rng=MersenneTwister(seed), time_limit_sec=HEUR_TL) + elseif name == "SimulatedAnnealingScheduler" + return Dagger.SimulatedAnnealingScheduler( + Dagger.IteratedGreedyScheduler(; rng=MersenneTwister(seed), time_limit_sec=HEUR_TL); + rng=MersenneTwister(seed), time_limit_sec=HEUR_TL) + elseif name == "JuMPScheduler" + return Dagger.JuMPScheduler(HiGHS.Optimizer; time_limit_sec=MILP_TL) + end + error("unknown scheduler $name") +end + +# ---- 3. MILP status/objective (from the committed JuMPExt instrumentation) --- +_milp_ext() = Base.get_extension(Dagger, :JuMPExt) +function milp_instrumentation_ok() + ext = _milp_ext() + return ext !== nothing && isdefined(ext, :LAST_MILP_SOLVE) +end +reset_milp_info!() = (ext = _milp_ext(); ext === nothing || (ext.LAST_MILP_SOLVE[] = ("NONE", NaN)); nothing) +read_milp_info() = (ext = _milp_ext(); ext === nothing ? ("NA", NaN) : ext.LAST_MILP_SOLVE[]) + +# ---- 4. Node topology from the workers Dagger sees -------------------------- +struct Topology + hostnames::Vector{String} # sorted, one per node + node_procs::Dict{String,Vector{Any}} # hostname -> its Dagger procs + n_nodes::Int + total_procs::Int + total_gpus::Int +end +function detect_topology() + procs = collect(Dagger.all_processors()) + byworker = Dict{Int,Vector{Any}}() + for p in procs + push!(get!(byworker, Dagger.root_worker_id(p), Any[]), p) + end + wid_host = Dict{Int,String}() + for wid in keys(byworker) + wid_host[wid] = wid == Dagger.myid() ? gethostname() : + try Dagger.remotecall_fetch(gethostname, wid) catch; "worker$wid"; end + end + node_procs = Dict{String,Vector{Any}}() + for (wid, host) in wid_host + append!(get!(node_procs, host, Any[]), byworker[wid]) + end + hosts = sort(collect(keys(node_procs))) + ngpu = count(!_is_cpu_proc, procs) + return Topology(hosts, node_procs, length(hosts), length(procs), ngpu) +end +# Procs / scope for the first `n` nodes of a topology (M1/M2/M4). +function regime_procs(topo::Topology, n::Int) + sel = topo.hostnames[1:min(n, topo.n_nodes)] + return reduce(vcat, (topo.node_procs[h] for h in sel); init=Any[]) +end +regime_scope(procs) = Dagger.UnionScope([Dagger.ExactScope(p) for p in procs]) + +# ---- 5. Scoped run + warmup + DAG capture ---------------------------------- +# Run a workload restricted to `scope` (the regime's procs). +scoped_run!(wl, inputs, sched, scope) = + Dagger.with_options(; scope=scope) do + _run_workload_inner!(wl, inputs, sched) + end + +# Per-config warmup (once): a GPU-scoped RR pass to force real APU samples, then +# an all-regime RR pass -- both restricted to this regime's procs. Metrics are +# ground-truth and accumulate across measured cells thereafter. +function warm_config!(wl, nt, bs, rprocs) + gpu = filter(!_is_cpu_proc, rprocs) + if !isempty(gpu) + try + scoped_run!(wl, build_inputs(wl, nt, bs), Dagger.RoundRobinScheduler(), regime_scope(gpu)) + catch e + @warn "GPU warmup failed (continuing on CPU cost data)" wl nt exception=e + end + end + scoped_run!(wl, build_inputs(wl, nt, bs), Dagger.RoundRobinScheduler(), regime_scope(rprocs)) + return +end + +# Capture the largest AOT DAGSpec (+ procs + snapshot) the scheduler sees for a +# config, for n_tasks and the AOT-only optimality-gap pass. +mutable struct DagCap <: Dagger.DataDepsScheduler + inner::Dagger.GreedyScheduler + grab::Base.RefValue{Any} +end +function Dagger.datadeps_schedule_dag_aot!(s::DagCap, sc, dag, pr, sco) + cur = s.grab[] + if cur === nothing || Dagger.nv(dag.g) > Dagger.nv(cur[1].g) + s.grab[] = (dag, pr) + end + Dagger.datadeps_schedule_dag_aot!(s.inner, sc, dag, pr, sco) +end +Dagger.datadeps_schedule_task_jit!(s::DagCap, a...) = Dagger.datadeps_schedule_task_jit!(s.inner, a...) +Dagger.datadeps_schedule_cache(s::DagCap) = Dagger.datadeps_schedule_cache(s.inner) +Base.similar(s::DagCap) = DagCap(similar(s.inner), s.grab) +function capture_dag(wl, nt, bs, scope) + cap = DagCap(Dagger.GreedyScheduler(), Ref{Any}(nothing)) + scoped_run!(wl, build_inputs(wl, nt, bs), cap, scope) + empty!(Dagger.datadeps_schedule_cache(cap)) + grabbed = cap.grab[] + grabbed === nothing && return (nothing, nothing, 0) + dag, procs = grabbed + return (dag, procs, Dagger.nv(dag.g)) +end + +# ---- 6. One measured cell --------------------------------------------------- +function run_cell(regime, n_nodes, wl, nt, bs, name, seed, ntk, scope) + inputs = build_inputs(wl, nt, bs) + sched = make_scheduler(name, seed) + # Fresh schedule per (scheduler, seed): otherwise a structurally-equivalent + # prior cell's cached schedule would be reused, collapsing the 7-seed spread. + empty!(Dagger.datadeps_schedule_cache(sched)) + timed = TimedScheduler(sched) + reset_milp_info!() + t0 = time_ns() + try + scoped_run!(wl, inputs, timed, scope) + catch e + @warn "CELL EXCEPTION" regime wl nt bs name seed exception=(e, catch_backtrace()) + mstat, mobj = name == "JuMPScheduler" ? read_milp_info() : ("", NaN) + return (; regime, n_nodes, workload=wl, nt, bs, scheduler=name, seed, + wall_ms=NaN, aot_ms=NaN, residual=NaN, n_tasks=ntk, + milp_status="EXCEPTION:$(nameof(typeof(e)))", milp_obj=mobj, failed=true) + end + wall_ms = (time_ns() - t0) / 1e6 + aot_ms = timed.last_aot_ns[] / 1e6 + resid = try last(verify_workload(wl, inputs)) catch; NaN end + mstat, mobj = name == "JuMPScheduler" ? read_milp_info() : ("", NaN) + return (; regime, n_nodes, workload=wl, nt, bs, scheduler=name, seed, + wall_ms, aot_ms, residual=resid, n_tasks=ntk, + milp_status=mstat, milp_obj=mobj, failed=false) +end + +# ---- 7. CSV plumbing -------------------------------------------------------- +const MAIN_HDR = "regime,n_nodes,workload,nt,bs,scheduler,seed,wall_ms,aot_ms,residual,n_tasks,milp_status,milp_obj" +main_row(r) = join((r.regime, r.n_nodes, r.workload, r.nt, r.bs, r.scheduler, r.seed, + round(r.wall_ms, digits=3), round(r.aot_ms, digits=3), + r.residual, r.n_tasks, r.milp_status, r.milp_obj), ",") +const OPTGAP_HDR = "regime,n_nodes,workload,nt,bs,scheduler,seed,cost_of_schedule,milp_obj,ratio_to_milp" + +# ---- 8. AOT-only optimality-gap (headline Greedy/MILP, IG/MILP, SA/MILP) ----- +# For each config with at least one MILP=OPTIMAL cell, capture the DAG once and +# invoke Greedy/IG/SA in AOT-only mode (no execution) to get cost_of_schedule of +# the produced ScheduleState. Chains match the real schedulers: standalone +# Greedy uses producer_finish=true; IG/SA seed with plain greedy (as their +# actual construction does) then refine. +function aot_greedy(snap, dag, procs, cache) + st = ScheduleState() + greedy_schedule!(st, snap, dag, procs; cache=cache, producer_finish=true) + return cost_of_schedule(st) +end +function aot_ig(snap, dag, procs, cache, seed) + st = ScheduleState() + greedy_schedule!(st, snap, dag, procs; cache=cache) # seed + st = iterated_greedy_schedule!(st, snap, dag, procs; rng=MersenneTwister(seed), + cache=cache, time_limit_sec=HEUR_TL) + return cost_of_schedule(st) +end +function aot_sa(snap, dag, procs, cache, seed) + st = ScheduleState() + greedy_schedule!(st, snap, dag, procs; cache=cache) # greedy seed + st = iterated_greedy_schedule!(st, snap, dag, procs; rng=MersenneTwister(seed), + cache=cache, time_limit_sec=HEUR_TL) # IG seed for SA + simulated_annealing_schedule!(st, snap, dag, procs; rng=MersenneTwister(seed), + cache=cache, time_limit_sec=HEUR_TL) + return cost_of_schedule(st) +end +# `optimal_configs`: (regime,n_nodes,wl,nt,bs) => median milp_obj over OPTIMAL cells. +function optgap_pass!(io, optimal_configs, topo, log) + for ((regime, n_nodes, wl, nt, bs), milp_obj) in optimal_configs + rprocs = regime_procs(topo, n_nodes) + scope = regime_scope(rprocs) + dag, procs, k = capture_dag(wl, nt, bs, scope) + (dag === nothing || k == 0) && continue + snap = MT.snapshot(MT.global_metrics_cache()) + cache = _build_eft_cost_cache(snap, dag, procs) + emit(sched, seed, cost) = begin + ratio = (isfinite(milp_obj) && milp_obj > 0) ? cost / milp_obj : NaN + println(io, join((regime, n_nodes, wl, nt, bs, sched, seed, + round(cost, digits=1), round(milp_obj, digits=1), + round(ratio, digits=4)), ",")) + end + try; emit("GreedyScheduler", 0, aot_greedy(snap, dag, procs, cache)); catch e; log("optgap Greedy $wl nt=$nt failed: $e"); end + for s in SEEDS + try; emit("IteratedGreedyScheduler", s, aot_ig(snap, dag, procs, cache, s)); catch e; log("optgap IG $wl nt=$nt s=$s failed: $e"); end + try; emit("SimulatedAnnealingScheduler", s, aot_sa(snap, dag, procs, cache, s)); catch e; log("optgap SA $wl nt=$nt s=$s failed: $e"); end + end + flush(io) + end +end + +# ---- 8b. Schedule-cache demonstration (Contribution 4) ---------------------- +# For one config per regime (cholesky nt=4, seed=1, Greedy), invoke the SAME +# spawn_datadeps twice back-to-back. Call 1 populates the structural-equivalence +# cache (full AOT); call 2 hits it (equivalence lookup + return), so its aot +# should be ~0. If the two are similar, the cache isn't being hit on this Julia +# version. Fresh inputs each call -- identity differs, structure matches, which +# is exactly what the cache keys on. +function cache_demo!(io, topo, log) + wl, nt, bs, seed = :cholesky, 4, 4096, 1 + for (rtag, nnodes, _) in REGIMES + topo.n_nodes < nnodes && continue + rprocs = regime_procs(topo, nnodes); scope = regime_scope(rprocs) + try + warm_config!(wl, nt, bs, rprocs) + # Clean cache, then two back-to-back runs. GreedyScheduler shares one + # per-type cache, so call 2 sees call 1's entry. + empty!(Dagger.datadeps_schedule_cache(Dagger.GreedyScheduler())) + t1 = TimedScheduler(Dagger.GreedyScheduler()) + scoped_run!(wl, build_inputs(wl, nt, bs), t1, scope) + aot1 = t1.last_aot_ns[] / 1e6 + t2 = TimedScheduler(Dagger.GreedyScheduler()) + scoped_run!(wl, build_inputs(wl, nt, bs), t2, scope) + aot2 = t2.last_aot_ns[] / 1e6 + println(io, join((rtag, wl, nt, bs, "GreedyScheduler", seed, + round(aot1, digits=3), round(aot2, digits=3)), ",")) + flush(io) + log(@sprintf(" cache-demo %s: first=%.3fms hit=%.3fms (%.1fx)", + rtag, aot1, aot2, aot1 / max(aot2, 1e-6))) + catch e + @warn "cache-demo failed" rtag exception=e + println(io, join((rtag, wl, nt, bs, "GreedyScheduler", seed, "NaN", "NaN"), ",")) + flush(io) + end + end +end + +# ---- 9. Tree verification (feature 5) --------------------------------------- +function verify_tree(log) + gitlog = try readchomp(`git -C $REPO_ROOT log --oneline -60`) catch; "" end + missing = filter(s -> !occursin(s, gitlog), REQUIRED_SUBJECTS) + sha = try readchomp(`git -C $REPO_ROOT rev-parse --short HEAD`) catch; "unknown" end + if !isempty(missing) + log("FATAL: HEAD ($sha) is missing required commits:") + foreach(s -> log(" - \"$s\""), missing) + log("Check out the paper branch at the pinned SHA (see TUOLUMNE_README.md).") + return (false, sha) + end + if !milp_instrumentation_ok() + log("FATAL: JuMPExt.LAST_MILP_SOLVE not found -- the MILP-objective") + log("instrumentation commit is missing; optimality-gap numbers unavailable.") + return (false, sha) + end + log("tree OK -- HEAD=$sha, all required commits + MILP instrumentation present") + return (true, sha) +end + +# ---- 10. Smoke: one tiny cell end-to-end, PASS/FAIL (feature 1) ------------- +function run_smoke() + println("=== TUOLUMNE SMOKE (validates env + branch + ROCM/AMDGPU stack) ===") + ok, sha = verify_tree(println) + ok || (println("SMOKE: FAIL (tree state)"); return false) + topo = detect_topology() + @printf("detected: %d node(s), %d procs, %d APU GPU(s) [%s]\n", + topo.n_nodes, topo.total_procs, topo.total_gpus, join(topo.hostnames, ", ")) + if topo.total_gpus == 0 + println("SMOKE: WARNING -- no APU GPUs visible to Dagger. AMDGPU.jl/ROCm may not") + println(" be initialised, or workers lack GPU access. Fix before the full run.") + end + rprocs = regime_procs(topo, 1); scope = regime_scope(rprocs) + wl, nt, bs = :cholesky, 2, 4096 + try + warm_config!(wl, nt, bs, rprocs) + _, _, k = capture_dag(wl, nt, bs, scope) + for name in ("RoundRobinScheduler", "GreedyScheduler") + r = run_cell("smoke", 1, wl, nt, bs, name, 1, k, scope) + gate = RESID_GATE[wl] + status = r.failed ? "FAILED" : (r.residual > gate ? "RESID>gate" : "ok") + @printf(" %-22s wall=%.1fms aot=%.1fms resid=%.2e -> %s\n", + name, r.wall_ms, r.aot_ms, r.residual, status) + (r.failed || r.residual > gate) && (println("SMOKE: FAIL"); return false) + end + catch e + println("SMOKE: FAIL (exception)"); showerror(stdout, e, catch_backtrace()); println() + return false + end + println("SMOKE: PASS -- environment is ready for the full run.") + return true +end + +# ---- 11. Full grid ---------------------------------------------------------- +function main() + started = now() + summary = String[] + log(msg) = (s = "[$(Dates.format(now(), "HH:MM:SS"))] $msg"; println(s); push!(summary, s); flush(stdout)) + + ok, sha = verify_tree(log) + ok || (write_summary(summary, sha, started, 0); exit(1)) + topo = detect_topology() + log(@sprintf("topology: %d node(s), %d procs, %d APU GPU(s): %s", + topo.n_nodes, topo.total_procs, topo.total_gpus, join(topo.hostnames, ", "))) + + optgap_io = open("tuolumne_optgap.csv", "w"); println(optgap_io, OPTGAP_HDR); flush(optgap_io) + optimal_configs = Dict{Any,Float64}() + all_rows = Vector{Any}() + fail_count = 0 + last_flush = time() + + for (rtag, nnodes, outcsv) in REGIMES + if topo.n_nodes < nnodes + log("SKIP regime $rtag ($nnodes nodes) -- only $(topo.n_nodes) node(s) allocated") + continue + end + rprocs = regime_procs(topo, nnodes); scope = regime_scope(rprocs) + rgpu = count(!_is_cpu_proc, rprocs) + log(@sprintf("REGIME %s: %d node(s), %d procs (%d APU GPU) -> %s", + uppercase(rtag), nnodes, length(rprocs), rgpu, outcsv)) + rstart = time() + io = open(outcsv, "w"); println(io, MAIN_HDR); flush(io) + for wl in WORKLOADS, nt in NTS + # pick bs (nt=8 tries 4096 then 2048) via the warmup + bs = 0 + for cand in bs_candidates(nt) + try + warm_config!(wl, nt, cand, rprocs); bs = cand; break + catch e + log(" $wl nt=$nt bs=$cand warmup failed ($(nameof(typeof(e)))); trying smaller") + end + end + if bs == 0 + log(" $wl nt=$nt: no bs fit -- config skipped"); continue + end + _, _, ntk = capture_dag(wl, nt, bs, scope) + log(@sprintf(" [config] %s nt=%d bs=%d n_tasks=%d", wl, nt, bs, ntk)) + optimal_objs = Float64[] + for seed in SEEDS, name in SCHED_NAMES # seed-major, scheduler-inner + r = run_cell(rtag, nnodes, wl, nt, bs, name, seed, ntk, scope) + println(io, main_row(r)); push!(all_rows, r) + r.failed && (fail_count += 1) + gate = RESID_GATE[wl] + if !r.failed && isfinite(r.residual) && r.residual > gate + log(@sprintf(" *** RESIDUAL GATE %s %s nt=%d seed=%d resid=%.2e", rtag, name, nt, seed, r.residual)) + fail_count += 1 + end + if name == "JuMPScheduler" && r.milp_status == "OPTIMAL" && isfinite(r.milp_obj) + push!(optimal_objs, r.milp_obj) + end + if time() - last_flush > FLUSH_INTERVAL_S + flush(io); flush(optgap_io); last_flush = time() + end + end + isempty(optimal_objs) || (optimal_configs[(rtag, nnodes, wl, nt, bs)] = median(optimal_objs)) + end + close(io) + log(@sprintf("REGIME %s done in %.1f min", uppercase(rtag), (time()-rstart)/60)) + end + + log("schedule-cache demonstration (Contribution 4)") + open("tuolumne_cache.csv", "w") do cio + println(cio, "regime,workload,nt,bs,scheduler,seed,aot_ms_first_call,aot_ms_cache_hit") + cache_demo!(cio, topo, log) + end + + log("optimality-gap AOT pass over $(length(optimal_configs)) OPTIMAL config(s)") + optgap_pass!(optgap_io, optimal_configs, topo, log); close(optgap_io) + + write_medians(all_rows) + write_summary(summary, sha, started, fail_count) + log(@sprintf("ALL DONE. %d cells, %d flagged/failed. Elapsed %.1f min.", + length(all_rows), fail_count, (now()-started).value/60000)) +end + +# ---- 12. Aggregates --------------------------------------------------------- +function write_medians(rows) + open("tuolumne_medians.csv", "w") do io + println(io, "regime,n_nodes,workload,nt,bs,scheduler,wall_ms_median,wall_ms_stddev,", + "aot_ms_median,aot_ms_stddev,residual_max,n_tasks,milp_status,milp_obj") + groups = Dict{Any,Vector{Any}}() + for r in rows + push!(get!(groups, (r.regime, r.n_nodes, r.workload, r.nt, r.bs, r.scheduler), Any[]), r) + end + med(xs) = (ys = filter(isfinite, xs); isempty(ys) ? NaN : median(ys)) + sdv(xs) = (ys = filter(isfinite, xs); length(ys) < 2 ? 0.0 : std(ys)) + for (key, rs) in sort(collect(groups), by = k -> string(k[1])) + wl = [r.wall_ms for r in rs]; ao = [r.aot_ms for r in rs] + resid = maximum(filter(isfinite, [r.residual for r in rs]); init=NaN) + mstat = any(r -> r.milp_status == "OPTIMAL", rs) ? "OPTIMAL" : + (isempty(rs) ? "" : rs[1].milp_status) + mobj = med([r.milp_obj for r in rs]) + println(io, join((key..., round(med(wl), digits=3), round(sdv(wl), digits=3), + round(med(ao), digits=3), round(sdv(ao), digits=3), + resid, rs[1].n_tasks, mstat, mobj), ",")) + end + end +end + +function write_summary(summary, sha, started, fail_count) + open("tuolumne_run_summary.txt", "w") do io + println(io, "Tuolumne one-shot benchmark -- run summary") + println(io, "git SHA (HEAD): ", sha) + println(io, "launch time: ", started) + println(io, "end time: ", now()) + println(io, "flagged/failed cells: ", fail_count) + println(io, "\n--- log ---") + foreach(l -> println(io, l), summary) + end +end + +# ---- entry ------------------------------------------------------------------ +if SMOKE + exit(run_smoke() ? 0 : 1) +else + main() +end