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