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/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/driver.jl b/bench/datadeps_schedulers/driver.jl index 3d86125b0..4bf0a65fe 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,11 +45,25 @@ 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) = 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) @@ -63,12 +90,76 @@ 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 + # 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 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) @@ -77,16 +168,37 @@ 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 === :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, + inputs.initial_A, inputs.initial_B) + end end 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 @@ -104,12 +216,46 @@ 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 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 @@ -137,6 +283,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 @@ -179,20 +328,104 @@ 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()) + # 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() + try + Dagger.with_options(; scope=gpu_scope) do + _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()) + 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 const DEFAULT_WARMUP = 1 +# 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 + +# 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. -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,14 +434,33 @@ 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 +# 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, @@ -218,52 +470,130 @@ 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 + # 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 + @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 @@ -349,14 +679,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 @@ -365,15 +706,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) 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/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 diff --git a/bench/datadeps_schedulers/workloads.jl b/bench/datadeps_schedulers/workloads.jl index 9d0452e6f..7929f943c 100644 --- a/bench/datadeps_schedulers/workloads.jl +++ b/bench/datadeps_schedulers/workloads.jl @@ -1,8 +1,20 @@ using LinearAlgebra +using Random 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 +38,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 +52,405 @@ function tiled_matmul!(C::Matrix{<:Matrix}, A::Matrix{<:Matrix}, B::Matrix{<:Mat return C end +""" + _placement_procs() -> Vector{Dagger.Processor} + +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 +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. + +## 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()) + 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(cpu_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(cpu_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 + +# 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 + +Ships `tile` (currently master-resident) to `target_proc`'s worker via +`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 +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 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 (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 + +# 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 + +""" + 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 + 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 + +# ─── 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 diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index e4edce813..3a3655060 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -59,6 +59,61 @@ 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 — 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 @@ -452,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))) @@ -563,6 +649,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/JuMPExt.jl b/ext/JuMPExt.jl index a139cdfeb..99d388c88 100644 --- a/ext/JuMPExt.jl +++ b/ext/JuMPExt.jl @@ -11,9 +11,13 @@ 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 -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}) @@ -21,12 +25,58 @@ function _milp_compatible_procs(spec, all_procs::Vector{Dagger.Processor}) return filter(p -> proc_in_scope(p, task_scope), all_procs) end +# 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 + +# 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 - 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 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()) # Cost matrix uses the same EFT fallback as the heuristics so MILP # optimises against an identical cost model. @@ -144,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 @@ -158,7 +210,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/ext/ROCExt.jl b/ext/ROCExt.jl index 92f2682b1..121cbe5d1 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 @@ -371,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))) @@ -478,6 +520,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/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/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index b4e2f0734..fe916904d 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,17 @@ 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)) \ +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 @@ -396,14 +334,21 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr all_results_lock = ReentrantLock() @sync for (wid, worker_args) in by_worker Threads.@spawn begin + # `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)) \ +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 @@ -436,8 +381,8 @@ function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWr 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}) @@ -463,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}, @@ -553,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}) @@ -590,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 @@ -691,21 +631,100 @@ 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 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}, + 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)) + + 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. + 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 + + 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, 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, @@ -733,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 @@ -754,19 +767,14 @@ end schedule_partition_full!(queue, queue_lock, partition_id, partition_verts, dag, seen_tasks, local_procs, vertex_to_partition, task_submitted, - region_uids) -> DataDepsState - -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!`. + region_uids, registry; + precomputed_schedule=nothing) + -> (DataDepsState, Dict{DTask,Processor}) + +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, @@ -778,16 +786,15 @@ 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)) - # 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 @@ -802,42 +809,30 @@ 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] - _pdag, schedule = datadeps_build_schedule!(temp_queue.scheduler, partition_pairs, - local_procs, local_scope; region_uids) - - # 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. + + if precomputed_schedule === nothing + _pdag, schedule = datadeps_build_schedule!(temp_queue.scheduler, partition_pairs, + local_procs, local_scope; region_uids) + else + 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 + + # `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) @@ -846,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, @@ -866,7 +854,7 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, end end - return state + return state, schedule end struct LockedEnqueueQueue <: AbstractTaskQueue @@ -880,9 +868,71 @@ function enqueue!(leq::LockedEnqueueQueue, pairs::Vector{DTaskPair}) @lock leq.lock enqueue!(leq.inner, pairs) end -# Parallel partition scheduling wraps errors in TaskFailedException / -# CompositeException via `@sync`/`Threads.@spawn`. Unwrap so callers and tests -# see the root `SchedulingException` / scheduler error. +""" + _hierarchical_dag_spec_and_cache_lookup(scheduler, seen_tasks) + -> (dag_spec::DAGSpec, precomputed::Union{Dict{DTask,Processor},Nothing}) + +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}) + 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) + +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, + 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 + +# 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) @@ -898,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 @@ -919,7 +958,9 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) return end - # Get the set of all processors + dag_spec, precomputed_schedule = + _hierarchical_dag_spec_and_cache_lookup(queue.scheduler, seen_tasks) + all_procs = Processor[] scope = get_compute_scope() for w in procs() @@ -931,29 +972,48 @@ 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 3: Partition the DAG - vertex_to_partition, n_partitions, partition_procs, _multi_worker = - partition_dag(dag, task_metas, all_procs) + # 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 + push!(data_workers, root_worker_id(memory_space(dep.arg_w.arg))) + end + end + data_is_distributed = length(data_workers) >= 2 + + # 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}() + datadeps_schedule_dag_aot!(queue.scheduler, trial, dag_spec, all_procs, all_scope) + if !isempty(trial) + precomputed_schedule = trial + end + end + + 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 - # 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] @@ -964,23 +1024,21 @@ 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, 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 +1047,10 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) end _hierarchical_copy_from_and_free!(partition_states, n_partitions, registry) + + if !cache_hit + _hierarchical_persist_schedule!(queue.scheduler, dag_spec, partition_schedules) + end end function _hierarchical_max_write_num(state::DataDepsState, arg_w::ArgumentWrapper) @@ -1017,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) @@ -1033,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] @@ -1052,23 +1111,49 @@ 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. + # 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 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 diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index c8d200968..baf0b92ba 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -380,6 +380,28 @@ 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 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 + 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 + +# 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) const GREEDY_DEFAULT_OUTPUT_SIZE = UInt64(1_048_576) @@ -391,11 +413,26 @@ 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) + +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) + 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, @@ -459,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 @@ -488,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 @@ -510,25 +611,27 @@ 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) 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) - + 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 = 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 @@ -538,11 +641,83 @@ 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 +""" + 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) @@ -552,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 @@ -567,10 +761,13 @@ 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] - 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 @@ -624,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 @@ -670,6 +900,49 @@ _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. 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 @@ -691,6 +964,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,23 +979,32 @@ 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 IteratedGreedyScheduler(; kwargs...) = IteratedGreedyScheduler(GreedyScheduler(); kwargs...) -# 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. +# 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, + destroy_frac=s.destroy_frac, + rng=copy(s.rng), + time_limit_sec=s.time_limit_sec) + +# 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) = @@ -740,34 +1027,107 @@ Mutates and returns `state`. The caller is responsible for providing a freshly-`copy(prev_state)` if the previous state must be preserved on rejection. """ +# `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] + +# 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) +const EFT_LOOKUP_TRANSLATED = 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], @view spec.fargs[2:end]).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) + # 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) 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}, - 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) @@ -781,12 +1141,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 @@ -804,8 +1159,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, @@ -831,33 +1188,56 @@ 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 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 + + # Reset in place via `_copy_state!` to avoid Dict allocations in the loop. + candidate = copy(state) - # 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) + # 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}() sizehint!(destroyed, n_destroy) + # `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 - # 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. + (time_ns() - start_ns) >= budget_ns && break + # 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] @@ -867,18 +1247,24 @@ 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) + # 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 - # 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). - 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 @@ -896,6 +1282,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, @@ -903,7 +1290,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 @@ -912,9 +1301,20 @@ 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 """ SimulatedAnnealingScheduler{S<:DataDepsScheduler, R<:Random.AbstractRNG} <: DataDepsScheduler @@ -938,6 +1338,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, 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 @@ -953,22 +1357,31 @@ 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 SimulatedAnnealingScheduler(; kwargs...) = SimulatedAnnealingScheduler(IteratedGreedyScheduler(); kwargs...) +Base.similar(s::SimulatedAnnealingScheduler) = + SimulatedAnnealingScheduler(similar(s.inner); + q=s.q, k=s.k, n_restarts=s.n_restarts, + rng=copy(s.rng), + time_limit_sec=s.time_limit_sec) + datadeps_schedule_cache(sched::SimulatedAnnealingScheduler) = datadeps_schedule_cache(sched.inner) datadeps_dag_equivalent(sched::SimulatedAnnealingScheduler, dspec1::DAGSpec, dspec2::DAGSpec) = @@ -1048,6 +1461,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, @@ -1055,7 +1499,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) @@ -1114,10 +1558,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 @@ -1151,7 +1597,13 @@ 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 + # 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) + 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) @@ -1165,6 +1617,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) @@ -1222,18 +1675,22 @@ 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] - 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 @@ -1300,6 +1757,9 @@ struct JuMPScheduler <: DataDepsScheduler end end +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 @@ -1324,13 +1784,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`. -- `sa_q` — forwarded to `SimulatedAnnealingScheduler`. -- `sa_k` — forwarded to `SimulatedAnnealingScheduler`. -- `sa_n_restarts` — forwarded to `SimulatedAnnealingScheduler`. -- `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 @@ -1339,9 +1795,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(; @@ -1351,9 +1809,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")) @@ -1361,9 +1821,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, @@ -1372,14 +1834,30 @@ 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 end +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, + 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)) + """ opt_uses_milp(sched::OptimizingScheduler, n_tasks::Integer) -> Bool @@ -1413,12 +1891,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/src/gpu.jl b/src/gpu.jl index fac44893d..14c082b32 100644 --- a/src/gpu.jl +++ b/src/gpu.jl @@ -133,4 +133,57 @@ 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 + +""" + _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 diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 4e2c59174..d3a461a25 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 @@ -888,7 +889,23 @@ 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) + + # 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(snap, sig_vec_for_index) + estimate_task_costs!(sorted_procs, costs, state, input_procs, task; + sig, runtime_index, snap) input_procs_cleanup() # Select the best available processor and reserve time pressure optimistically. @@ -900,7 +917,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, snap) if has_cap # Optimistic reservation: capture-the-ref, atomic_add! counter_ref = lock(state.worker_time_pressure) do wtp @@ -1334,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 @@ -1384,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 @@ -1439,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 @@ -1726,20 +1777,120 @@ 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 +# 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()) + +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 @@ -1749,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 @@ -1835,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" @@ -1925,7 +2037,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 62f280edd..9a95dff66 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -526,15 +526,27 @@ 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, + 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 else - runtime = metrics_lookup_runtime(snap, sig_vec, p, worker_id) + # 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 + 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 +653,10 @@ 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, + snap=nothing) # Find all Chunks chunks = @reusable_vector :estimate_task_costs_chunks Union{Chunk,Nothing} nothing 32 @@ -656,7 +671,23 @@ 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` + # 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. + # + # 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) @@ -664,15 +695,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) + 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) 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() diff --git a/src/utils/metrics.jl b/src/utils/metrics.jl index 571491a55..c12206728 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), @@ -192,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 @@ -205,6 +337,34 @@ 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 = 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 isempty(pairs) && return @@ -215,6 +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[]) end return end @@ -255,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) diff --git a/test/datadeps/scheduling.jl b/test/datadeps/scheduling.jl index c12f2f363..668db45e1 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 && @@ -2220,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 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]