From 207aed95eeb7d7dbddfb5ed5737057c367d021fb Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 13 Aug 2026 13:59:53 -0700 Subject: [PATCH 01/10] DArray/stencil: Avoid xfer of full chunks for boundary across different spaces --- src/array/stencil.jl | 126 +++++++++++++++++++++++++++++++++---------- src/memory-spaces.jl | 29 ++++++++++ 2 files changed, 127 insertions(+), 28 deletions(-) diff --git a/src/array/stencil.jl b/src/array/stencil.jl index a02bcac8f..6c4accbdb 100644 --- a/src/array/stencil.jl +++ b/src/array/stencil.jl @@ -798,13 +798,19 @@ function select_neighborhood_chunks(chunks, idx, neigh_dist, boundary) end # Returns (region_metadata, neighbor_chunk_dtasks) without spawning intermediate load tasks. -# region_metadata: Vector of (region_code, is_boundary, boundary_dims). +# region_metadata: Vector of (region_code, is_boundary, boundary_dims, is_prematerialized). +# `is_prematerialized` is set by `stencil_region_info`; it is always `false` here, +# meaning the region is still to be taken from the neighboring chunk itself. # neighbor_chunk_dtasks: Vector of raw chunk DTasks (resolved to arrays when build_halo_new runs). function select_neighborhood_info(chunks, idx, neigh_dist, boundary) validate_neigh_dist(neigh_dist) N = ndims(chunks) chunk_dist = 1 - region_metadata = Tuple[] + # Concretely typed: every element has this exact shape, and an abstract + # `Tuple[]` would box each one on `push!` and box each field again on the + # destructure in `stencil_region_info`, which reads all 3^N-1 of them per + # chunk. + region_metadata = Tuple{NTuple{N,Int},Bool,NTuple{N,Bool},Bool}[] neighbor_chunks = Any[] for i in 0:(3^N - 1) @@ -827,9 +833,9 @@ function select_neighborhood_info(chunks, idx, neigh_dist, boundary) else new_idx = idx end - push!(region_metadata, (region_code, true, boundary_dims)) + push!(region_metadata, (region_code, true, boundary_dims, false)) else - push!(region_metadata, (region_code, false, ntuple(_ -> false, N))) + push!(region_metadata, (region_code, false, ntuple(_ -> false, N), false)) end push!(neighbor_chunks, chunks[new_idx]) end @@ -878,22 +884,31 @@ struct CopyHalos end return view(arr, ranges...) end -@inline function _build_fused_halos(::ViewHalos, neigh_dist, boundary, region_metadata, +# One halo region, taken from the neighboring chunk `chunk` that supplies it. +@inline halo_region(::ViewHalos, neigh_dist, boundary, region_code, is_boundary, boundary_dims, chunk) = + neighbor_region_view(chunk, region_code, neigh_dist) +@inline halo_region(::CopyHalos, neigh_dist, boundary, region_code, is_boundary, boundary_dims, chunk) = + is_boundary ? load_boundary_region(boundary, chunk, region_code, neigh_dist, boundary_dims) : + load_neighbor_region(chunk, region_code, neigh_dist) + +# Same region, but always materialized: used when the region is computed in the +# neighbor's memory space and then shipped, where a view would drag its parent along. +@inline halo_region_materialized(::ViewHalos, neigh_dist, boundary, region_code, is_boundary, boundary_dims, chunk) = + load_neighbor_region(chunk, region_code, neigh_dist) +@inline halo_region_materialized(style::CopyHalos, neigh_dist, boundary, region_code, is_boundary, boundary_dims, chunk) = + halo_region(style, neigh_dist, boundary, region_code, is_boundary, boundary_dims, chunk) + +@inline function _build_fused_halos(style, neigh_dist, boundary, region_metadata, neighbor_chunks::NTuple{NH,Any}) where NH return ntuple(Val(NH)) do i - region_code, _, _ = region_metadata[i] - neighbor_region_view(neighbor_chunks[i], region_code, neigh_dist) - end -end -@inline function _build_fused_halos(::CopyHalos, neigh_dist, boundary, region_metadata, - neighbor_chunks::NTuple{NH,Any}) where NH - return ntuple(Val(NH)) do i - region_code, is_boundary, boundary_dims = region_metadata[i] - chunk = neighbor_chunks[i] - if is_boundary - load_boundary_region(boundary, chunk, region_code, neigh_dist, boundary_dims) + region_code, is_boundary, boundary_dims, is_prematerialized = region_metadata[i] + if is_prematerialized + # `stencil_region_info` already built this region in its owner's memory + # space; `neighbor_chunks[i]` is the region itself, not a chunk to slice. + neighbor_chunks[i] else - load_neighbor_region(chunk, region_code, neigh_dist) + halo_region(style, neigh_dist, boundary, region_code, is_boundary, + boundary_dims, neighbor_chunks[i]) end end end @@ -903,8 +918,10 @@ end Wraps `center` (used in place, not copied) plus halo regions taken from `neighbor_chunks` into a `HaloArray`. `region_metadata` is the per-region -`(region_code, is_boundary, boundary_dims)` triple precomputed on the submitting -task by `select_neighborhood_info`. +`(region_code, is_boundary, boundary_dims, is_prematerialized)` tuple precomputed +on the submitting task by `stencil_region_info`. Where `is_prematerialized` is +set, the corresponding entry of `neighbor_chunks` is the halo region itself +rather than the chunk to slice it out of. """ function build_fused_halo(neigh_dist, boundary, region_metadata, center::AbstractArray{T,N}, neighbor_chunks::Vararg{Any,NH}) where {T,N,NH} @@ -934,6 +951,64 @@ function stencil_source_chunks(read_chunks, write_chunks) return map(task -> fetch(task; raw=true), snapshot_tasks) end +""" + stencil_region_info(src_chunks, write_chunks, neigh_dist, boundary) -> table + +Per-chunk `(region_metadata, neighbor_values)` for a neighborhood read, laid out +like `src_chunks`. `neighbor_values[i]` is what the sweeping task needs in order +to obtain halo region `i`, and the region's `is_prematerialized` flag says which +of the two forms it takes: + +- a neighboring chunk, sliced by the sweeping task itself (`false`), or +- the halo region itself, already extracted (`true`). + +The first form is what we want when the neighbor is in the same memory space as +the chunk being written, because slicing it there is free. When the neighbor lives +elsewhere it is emphatically not: passing the neighbor as a task argument makes +Datadeps copy the *entire* chunk across the space boundary to serve a halo of a +few elements, which for one MPI rank per chunk means shipping the whole +neighborhood every sweep. So in that case the region is extracted by a task +scoped to the neighbor's own space and only the region crosses the boundary. +""" +# `memory_space` of a chunk `DTask` is not free -- it `fetch`es the task's result +# reference, boxing it -- and this loop asks the same question about the same +# chunks repeatedly: every chunk is some neighbor's neighbor 3^N-1 times over. +# Memoize by identity so each distinct chunk is resolved once per sweep. +function _memoized_memory_space!(memo::IdDict{Any,Any}, @nospecialize(chunk)) + space = get(memo, chunk, nothing) + space === nothing || return space + space = memory_space(chunk) + memo[chunk] = space + return space +end + +function stencil_region_info(src_chunks, write_chunks, neigh_dist, boundary) + style = halo_build_style(boundary) + table = Array{Any}(undef, size(src_chunks)) + # (index, region) => extraction task, resolved after all of them are spawned so + # that the extractions run concurrently rather than one chunk at a time. + extractions = Pair{Tuple{CartesianIndex{ndims(src_chunks)},Int},Any}[] + space_memo = IdDict{Any,Any}() + for idx in CartesianIndices(src_chunks) + region_metadata, neighbor_chunks = select_neighborhood_info(src_chunks, idx, neigh_dist, boundary) + target_space = _memoized_memory_space!(space_memo, write_chunks[idx]) + for i in eachindex(region_metadata) + region_code, is_boundary, boundary_dims, _ = region_metadata[i] + neighbor = neighbor_chunks[i] + neighbor_space = _memoized_memory_space!(space_memo, neighbor) + neighbor_space == target_space && continue + region_metadata[i] = (region_code, is_boundary, boundary_dims, true) + task = Dagger.@spawn name="stencil_halo" scope=memory_space_scope(neighbor_space) halo_region_materialized(style, neigh_dist, boundary, region_code, is_boundary, boundary_dims, neighbor) + push!(extractions, (idx, i) => task) + end + table[idx] = (Tuple(region_metadata), neighbor_chunks) + end + for ((idx, i), task) in extractions + table[idx][2][i] = fetch(task; raw=true) + end + return table +end + @inline load_neighborhood(arr::HaloArray{T,N}, idx) where {T,N} = StencilNeighborhood(arr, idx, arr.halo_width) @inline load_neighborhood(arr::HaloInterior{T,N}, idx) where {T,N} = @@ -1265,7 +1340,8 @@ macro stencil(orig_ex) # 2a. For each neighborhood read_var, pre-compute on the main task (so no DArray # is ever passed into @spawn) the chunk array to read neighbors from and, per - # chunk, the region metadata plus the neighboring chunks themselves. + # chunk, the region metadata plus what each halo region is to be built from + # (see `stencil_region_info`). # # `stencil_source_chunks` substitutes a snapshot when the expression writes back # into the chunks it reads (`A[idx] = f(@neighbors(A[idx]))`); everywhere else it @@ -1274,17 +1350,11 @@ macro stencil(orig_ex) for read_var in read_vars if read_var in keys(neighborhoods) neigh_dist, boundary = neighborhoods[read_var] - @gensym region_info_table src_chunks region_meta neighbor_cks + @gensym region_info_table src_chunks neigh_sym_map[read_var] = (; region_info_table, src_chunks) push!(final_ex.args, :($validate_neigh_dist($neigh_dist, ndims($read_var)))) push!(final_ex.args, :($src_chunks = $stencil_source_chunks($chunks($read_var), $chunks($write_var)))) - push!(final_ex.args, :($region_info_table = Array{Any}(undef, size($src_chunks)))) - push!(final_ex.args, quote - for $chunk_idx in $CartesianIndices($src_chunks) - ($region_meta, $neighbor_cks) = $select_neighborhood_info($src_chunks, $chunk_idx, $neigh_dist, $boundary) - $region_info_table[$chunk_idx] = (tuple($region_meta...), $neighbor_cks) - end - end) + push!(final_ex.args, :($region_info_table = $stencil_region_info($src_chunks, $chunks($write_var), $neigh_dist, $boundary))) end end diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index cfd5a8df6..65f176133 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -53,6 +53,35 @@ function processors(space::CPURAMMemorySpace) end end +""" + memory_space_scope(space::MemorySpace) -> AbstractScope + +A scope restricting execution to `space`, for tasks that must run where their data +already lives. Picks a single processor, like Datadeps does for its copy and free +tasks: under uniform (SPMD) execution every rank must pick the same one, and +`processors` is ordered deterministically while a `UnionScope` of all of them +would not be. + +Memoized per space so the returned scope is *identity*-stable, not merely equal. +`Sch.compatible_processors_cached` keys on `objectid(scope)` and only takes a hit +when the stored scope is `===` the query, so handing out a fresh `ExactScope` per +call would miss that cache on every spawn, pay a full `compatible_processors` +scan, and evict live entries from its fixed-size LFU. Datadeps' exec scopes are +identity-stable for exactly this reason; callers here (e.g. `stencil_halo` +spawns) need the same property. +""" +const MEMORY_SPACE_SCOPE_CACHE = LockedObject(Dict{MemorySpace,ExactScope}()) +function memory_space_scope(space::MemorySpace) + @safe_lock1 MEMORY_SPACE_SCOPE_CACHE cache begin + value = get(cache, space, nothing) + if value === nothing + value = ExactScope(first(processors(space))) + cache[space] = value + end + return value + end +end + ### In-place Data Movement unwrap(x::Chunk) = unwrap(x.handle) From 5464f5328501aefdf5626b83fb80fa271472f735 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 17 Aug 2026 15:34:30 -0700 Subject: [PATCH 02/10] CI/benchmarks: Add Distributed and MPI benchmarks --- .github/workflows/CI.yml | 197 +++++++++++++++++++++++++++++++++--- benchmark/benchmarks.jl | 92 ++++++++++++++++- benchmark/ci.jl | 6 ++ benchmark/common.jl | 33 +++++- benchmark/suites/array.jl | 47 ++++----- benchmark/suites/linalg.jl | 87 ++++++++-------- benchmark/suites/stencil.jl | 105 +++++++++---------- benchmark/worker_mpi.jl | 154 ++++++++++++++++++++++++++++ 8 files changed, 582 insertions(+), 139 deletions(-) create mode 100644 benchmark/worker_mpi.jl diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 490fd9965..f81b1098e 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -153,9 +153,6 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 120 if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && !contains(github.event.head_commit.message, '[skip benchmarks]') }} - permissions: - contents: read - pull-requests: write env: # `benchmark/ci.jl` compares the working tree ("dirty") against this base # revision and fails the job if any benchmark regresses past the threshold. @@ -167,6 +164,14 @@ jobs: # 4 vCPU / ~16 GB, and noisier than dedicated benchmarking hardware. The # orchestrator is OOM-robust, but smaller scales keep wall-clock bounded. BENCHMARK_SCALE: '[256, 1024]' + # Tile size matters as much as N here, and one tile is not enough. At 512 + # alone, N=256 is a *single* chunk (`square_block` returns N when N <= tile) + # and N=1024 is only 2x2, so anything touching data distribution or halo + # exchange is barely exercised. But small tiles are not simply "better": + # tile size sets the ratio of per-tile task overhead to bytes moved per + # transfer, and a data-movement change can be a large win at one tile size + # and a regression at the other. Measure both regimes. + BENCHMARK_BLOCKSIZE: '[128, 512]' BENCHMARK_SECONDS: '3' BENCHMARK_SAMPLES: '5' steps: @@ -198,30 +203,192 @@ jobs: if: always() uses: actions/upload-artifact@v7 with: - name: benchmark-results + name: benchmark-results-default + path: ${{ env.BENCHMARK_OUTPUT_DIR }}/ + if-no-files-found: ignore + # PR comment posting is handled by the `benchmarks-report` job below, + # once this job and its Distributed/MPI siblings have all finished -- + # a single writer avoids three jobs racing to read-modify-write the + # same PR comment. + + benchmarks-distributed: + name: Benchmarks (Distributed, vs master) + runs-on: ubuntu-latest + timeout-minutes: 120 + if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && !contains(github.event.head_commit.message, '[skip benchmarks]') }} + env: + BENCHMARK_BASE_REV: master + # Multiple OS processes on a shared runner are noisier than the + # single-process default job, so allow more slack before flagging a + # regression. + BENCHMARK_REGRESSION_THRESHOLD: '0.35' + # Parallelism comes from BENCHMARK_PROCS (processes), not threads: 3 + # extra Distributed workers + the driver = 4 processes x 1 thread, + # matching the 4 vCPUs on ubuntu-latest. + BENCHMARK_CI_THREADS: '1' + BENCHMARK_PROCS: '3:1' + BENCHMARK_OUTPUT_DIR: benchmark_results_distributed + BENCHMARK_SCALE: '[256, 1024]' + # See the `benchmarks` job: both tiles, so neither the many-small-chunks + # nor the few-large-chunks regime goes unmeasured. + BENCHMARK_BLOCKSIZE: '[128, 512]' + BENCHMARK_SECONDS: '3' + BENCHMARK_SAMPLES: '5' + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + - name: Fetch base revision + run: git fetch -f origin "$BENCHMARK_BASE_REV:$BENCHMARK_BASE_REV" || git fetch origin "$BENCHMARK_BASE_REV" || true + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + arch: x64 + - uses: julia-actions/cache@v3 + - name: Run benchmarks (current vs master, Distributed) + run: julia --color=yes benchmark/ci.jl + - name: Write job summary + if: always() + run: | + if [ -f "${BENCHMARK_OUTPUT_DIR}/report.md" ]; then + grep -v 'artifact://' "${BENCHMARK_OUTPUT_DIR}/report.md" >> "$GITHUB_STEP_SUMMARY" + else + echo "No benchmark report was produced." >> "$GITHUB_STEP_SUMMARY" + fi + - name: Upload benchmark artifacts + if: always() + uses: actions/upload-artifact@v7 + with: + name: benchmark-results-distributed path: ${{ env.BENCHMARK_OUTPUT_DIR }}/ if-no-files-found: ignore - - name: Comment results on PR + + benchmarks-mpi: + name: Benchmarks (MPI, vs master) + runs-on: ubuntu-latest + timeout-minutes: 120 + if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && !contains(github.event.head_commit.message, '[skip benchmarks]') }} + env: + BENCHMARK_BASE_REV: master + # See benchmarks-distributed: shared-runner multi-process noise. + BENCHMARK_REGRESSION_THRESHOLD: '0.35' + # 4 MPI ranks x 1 thread, matching the 4 vCPUs on ubuntu-latest. + # BENCHMARK_MPI_RANKS switches benchmark/benchmarks.jl to spawn the + # worker under mpiexec and use worker_mpi.jl (SPMD) automatically. + BENCHMARK_CI_THREADS: '1' + BENCHMARK_MPI_RANKS: '4' + BENCHMARK_OUTPUT_DIR: benchmark_results_mpi + BENCHMARK_SCALE: '[256, 1024]' + # See the `benchmarks` job: both tiles, so neither the many-small-chunks + # nor the few-large-chunks regime goes unmeasured. + BENCHMARK_BLOCKSIZE: '[128, 512]' + BENCHMARK_SECONDS: '3' + BENCHMARK_SAMPLES: '5' + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + - name: Fetch base revision + run: git fetch -f origin "$BENCHMARK_BASE_REV:$BENCHMARK_BASE_REV" || git fetch origin "$BENCHMARK_BASE_REV" || true + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + arch: x64 + - uses: julia-actions/cache@v3 + - name: Run benchmarks (current vs master, MPI) + run: julia --color=yes benchmark/ci.jl + - name: Write job summary + if: always() + run: | + if [ -f "${BENCHMARK_OUTPUT_DIR}/report.md" ]; then + grep -v 'artifact://' "${BENCHMARK_OUTPUT_DIR}/report.md" >> "$GITHUB_STEP_SUMMARY" + else + echo "No benchmark report was produced." >> "$GITHUB_STEP_SUMMARY" + fi + - name: Upload benchmark artifacts + if: always() + uses: actions/upload-artifact@v7 + with: + name: benchmark-results-mpi + path: ${{ env.BENCHMARK_OUTPUT_DIR }}/ + if-no-files-found: ignore + + benchmarks-report: + name: Benchmarks (PR comment) + needs: [benchmarks, benchmarks-distributed, benchmarks-mpi] + runs-on: ubuntu-latest + timeout-minutes: 10 + if: ${{ always() && github.event_name == 'pull_request' }} + permissions: + contents: read + pull-requests: write + steps: + - name: Download default benchmark results + if: always() + continue-on-error: true + uses: actions/download-artifact@v7 + with: + name: benchmark-results-default + path: results-default + - name: Download Distributed benchmark results + if: always() + continue-on-error: true + uses: actions/download-artifact@v7 + with: + name: benchmark-results-distributed + path: results-distributed + - name: Download MPI benchmark results + if: always() + continue-on-error: true + uses: actions/download-artifact@v7 + with: + name: benchmark-results-mpi + path: results-mpi + - name: Comment combined results on PR # Best-effort: fork PRs get a read-only token, so the comment will 403; # don't fail the job over it. - if: ${{ always() && github.event_name == 'pull_request' }} + if: always() continue-on-error: true uses: actions/github-script@v9 with: script: | const fs = require('fs'); - const reportPath = `${process.env.BENCHMARK_OUTPUT_DIR}/report.md`; - if (!fs.existsSync(reportPath)) { - core.info('No report.md found; skipping PR comment.'); - return; + + function section(dir) { + const reportPath = `${dir}/report.md`; + if (!fs.existsSync(reportPath)) { + return null; + } + return fs.readFileSync(reportPath, 'utf8') + .split('\n') + .filter(line => !line.includes('artifact://')) + .join('\n'); } + + function collapsible(summary, content) { + const body = content === null + ? '_Results unavailable (job did not produce a report)._' + : content; + return `
\n${summary}\n\n${body}\n\n
`; + } + const marker = ''; - let body = fs.readFileSync(reportPath, 'utf8') - .split('\n') - .filter(line => !line.includes('artifact://')) - .join('\n'); const runUrl = `${process.env.GITHUB_SERVER_URL}/${process.env.GITHUB_REPOSITORY}/actions/runs/${process.env.GITHUB_RUN_ID}`; - body = `${marker}\n${body}\n\n[Full results and plots](${runUrl}) (download the \`benchmark-results\` artifact).`; + + const parts = [ + marker, + '## Dagger benchmarks: `dirty` vs `master`', + '', + collapsible('Multi-threaded benchmarks (4 threads)', section('results-default')), + '', + collapsible('Distributed benchmarks (4 processes)', section('results-distributed')), + '', + collapsible('MPI benchmarks (4 ranks)', section('results-mpi')), + '', + `[Full results and plots](${runUrl}) (download the \`benchmark-results-*\` artifacts).`, + ]; + const body = parts.join('\n'); + const { owner, repo } = context.repo; const issue_number = context.issue.number; const comments = await github.paginate(github.rest.issues.listComments, { diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 8ba874af5..1f5f45273 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -70,8 +70,17 @@ # estimated peak allocation may use before a size is skipped. Defaults to # "0.2" (conservative). Raise it to allow larger sizes, lower it for more # headroom. -# - BENCHMARK_BLOCKSIZE: Target square tile size (elements per side) for dense -# suites. Defaults to "512". +# - BENCHMARK_BLOCKSIZE: Target square tile size(s) (elements per side) for dense +# suites, given like BENCHMARK_SCALE: an integer or an iterable of integers +# (e.g. "512", "[128, 512]"). Each requested tile produces its own benchmark +# group per N, keyed by the resulting block; tiles that collapse to the same +# block at a given N (both 512 and 256 give 256 at N=256) are emitted once. +# Defaults to "512". +# +# Sweeping this matters more than it looks: tile size sets the ratio between +# per-tile task overhead and bytes moved per transfer, so a change to data +# movement can be a large win at one tile size and a regression at another. +# A single tile can only ever see one of those regimes. # - BENCHMARK_SPARSE_BLOCKS: Number of blocks per dimension for sparse/banded # operators (keeps tile counts bounded at large N). Defaults to "16". # - BENCHMARK_PROCS: Worker/thread topology, as "numprocs:numthreads". This @@ -80,6 +89,15 @@ # runs on the worker process, which uses the orchestrator's thread count). # - BENCHMARK_REMOTES: Remote hosts on which to start workers, in the format # accepted by `Distributed.addprocs` (colon-separated). Optional. +# - BENCHMARK_MPI_RANKS: When set to a positive integer, runs the benchmark +# worker under `mpiexec -n ` (via MPI.jl's bundled mpiexec, so no +# system MPI install is required) instead of as a plain subprocess, and the +# worker calls `Dagger.accelerate!(:mpi)`. Because MPI ranks are SPMD, this +# also switches the default worker script from `worker.jl` to +# `worker_mpi.jl` (a single flat pass over every benchmark, run identically +# on every rank, instead of `worker.jl`'s incremental per-benchmark +# 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_PROC_TIMEOUT: Wall-clock seconds to wait for a single benchmark @@ -121,7 +139,15 @@ BenchmarkTools.tune!(b::PrecomputedTrial, args...; kwargs...) = b const WORKDIR = let d = get(ENV, "BENCHMARK_WORKDIR", "") isempty(d) ? mktempdir() : (mkpath(d); abspath(d)) end -const WORKER_SCRIPT = get(ENV, "BENCHMARK_WORKER_SCRIPT", joinpath(@__DIR__, "worker.jl")) +const MPI_RANKS = parse(Int, get(ENV, "BENCHMARK_MPI_RANKS", "0")) +if MPI_RANKS > 0 + using MPI +end +# Defaults to worker_mpi.jl (the SPMD worker) when BENCHMARK_MPI_RANKS is set, +# otherwise worker.jl (the incremental request/response worker); overridable +# either way via BENCHMARK_WORKER_SCRIPT. +const WORKER_SCRIPT = get(ENV, "BENCHMARK_WORKER_SCRIPT", + joinpath(@__DIR__, MPI_RANKS > 0 ? "worker_mpi.jl" : "worker.jl")) const JULIA_BIN = first(Base.julia_cmd().exec) const NTHREADS = Threads.nthreads() const PROC_TIMEOUT = parse(Float64, get(ENV, "BENCHMARK_PROC_TIMEOUT", "3600")) @@ -143,7 +169,23 @@ function spawn_worker() rm(joinpath(WORKDIR, "ready"); force=true) rm(joinpath(WORKDIR, "request.json"); force=true) rm(joinpath(WORKDIR, "request.json.tmp"); force=true) - return run(worker_cmd(); wait=false) + rm(joinpath(WORKDIR, "done"); force=true) + rm(joinpath(WORKDIR, "results_mpi_manifest.json"); force=true) + 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 + # callback (mirrors test/run_mpi.jl) -- so the actual `run` call must + # happen *inside* the `do` block. Extracting the bare executable via + # e.g. `MPI.mpiexec(identity)` and running it afterwards drops that + # environment and the spawned process dies immediately. + local proc + MPI.mpiexec() do mpiexec + proc = run(`$mpiexec -n $MPI_RANKS $(worker_cmd())`; wait=false) + end + return proc + else + return run(worker_cmd(); wait=false) + end end function wait_ready(proc; timeout=600) @@ -246,10 +288,50 @@ function run_all_external() return results end +# MPI counterpart to `run_all_external`: worker_mpi.jl is SPMD (every rank +# must run the identical benchmark at the same time), so there is no +# per-benchmark request/response round trip here -- the worker runs its +# entire flat pass unattended and signals completion via a `done` sentinel, +# which this just waits for (bounded by `PROC_TIMEOUT`, same as a single +# `await_response` would be) before loading whatever it produced. +function run_all_mpi() + results = Dict{Vector{String},BenchmarkTools.Trial}() + + proc = spawn_worker() + donepath = joinpath(WORKDIR, "done") + 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 + @warn "MPI benchmark run exceeded $(PROC_TIMEOUT)s; killing and returning partial results." + kill(proc) + break + end + sleep(POLL) + end + if process_running(proc) + try; wait(proc); catch; end + end + + manifestpath = joinpath(WORKDIR, "results_mpi_manifest.json") + isfile(manifestpath) || return results + manifest = JSON3.read(read(manifestpath, String)) + for entry in manifest + kp = String[string(k) for k in entry.keypath] + resultpath = joinpath(WORKDIR, String(entry.file)) + isfile(resultpath) || continue + results[kp] = BenchmarkTools.load(resultpath)[1] + end + return results +end + # --- Assemble SUITE from the externally-measured trials --------------------- const SUITE = BenchmarkGroup() -for (keypath, trial) in run_all_external() +for (keypath, trial) in (MPI_RANKS > 0 ? run_all_mpi() : run_all_external()) SUITE[keypath] = PrecomputedTrial(trial) end diff --git a/benchmark/ci.jl b/benchmark/ci.jl index 8dea14f4c..6af535478 100644 --- a/benchmark/ci.jl +++ b/benchmark/ci.jl @@ -55,6 +55,12 @@ const EXTRA_PKGS = String[ # with Dagger; the orchestrator/worker need it for their file-based IPC. #"JSON3", ] +# MPI is only needed (and only installed) when the caller requests an MPI +# benchmark run (see benchmark/benchmarks.jl's BENCHMARK_MPI_RANKS), so the +# plain/Distributed CI runs don't pay for pulling in MPICH_jll. +if get(ENV, "BENCHMARK_MPI_RANKS", "0") != "0" + push!(EXTRA_PKGS, "MPI") +end mkpath(OUTPUT_DIR) diff --git a/benchmark/common.jl b/benchmark/common.jl index f5f65e152..e988be090 100644 --- a/benchmark/common.jl +++ b/benchmark/common.jl @@ -64,13 +64,44 @@ sparse_bytes(N; nmats=1, density=0.0, T=Float64, nvecs=2) = """Whether an estimated allocation of `bytes` fits the conservative budget.""" fits_budget(bytes) = bytes <= MEM_BUDGET +"""Target square tile sizes (elements per side) for the dense suites. Given as a +Julia expression evaluating to an integer or an iterable of integers, exactly +like `BENCHMARK_SCALE`. + +Tile size is not a detail that can be pinned to one value and forgotten: it sets +the ratio between per-tile task overhead and bytes moved per transfer, and a +change to data movement can be a large win at one tile size and a regression at +another. Sweeping it is what lets a comparison tell those two regimes apart.""" +const blocksizes = let s = get(ENV, "BENCHMARK_BLOCKSIZE", "") + if isempty(s) + [512] + else + parsed = eval(Meta.parse(s)) + parsed isa Integer ? [parsed] : collect(parsed) + end +end + """Square tile size (elements per side) for a dense N×N matrix, targeting `tile` elements per side but never exceeding N.""" -function square_block(N; tile=parse(Int, get(ENV, "BENCHMARK_BLOCKSIZE", "512"))) +function square_block(N; tile=first(blocksizes)) N <= tile && return N return cld(N, cld(N, tile)) end +"""The distinct tile sizes `square_block` actually yields at dimension `N`. + +Different requested tiles collapse to the same block once clamped to `N` (at +N=256 both 512 and 256 give 256), and the suites key their benchmark groups by +the resulting block, so emitting a group per *requested* tile would collide.""" +function blocks_for(N) + seen = Int[] + for tile in blocksizes + b = square_block(N; tile) + b in seen || push!(seen, b) + end + return seen +end + """Block size for sparse/banded N×N operators, keeping the per-dimension block count bounded (so we don't create an enormous number of mostly-empty tiles).""" function banded_block(N; maxblocks=parse(Int, get(ENV, "BENCHMARK_SPARSE_BLOCKS", "16"))) diff --git a/benchmark/suites/array.jl b/benchmark/suites/array.jl index a881fbdbc..7d93c4889 100644 --- a/benchmark/suites/array.jl +++ b/benchmark/suites/array.jl @@ -17,37 +17,38 @@ function array_suite(ctx; method, accels) for N in scales # Elementwise ops hold at most the input plus a same-size result. fits_budget(dense_bytes(N; nmats=2, T=T)) || continue - b = square_block(N) - sub = BenchmarkGroup() + for b in blocks_for(N) + sub = BenchmarkGroup() - sub["alloc (rand)"] = @benchmarkable(wait(rand(Blocks($b, $b), $T, $N, $N)), - teardown = (@everywhere GC.gc())) + sub["alloc (rand)"] = @benchmarkable(wait(rand(Blocks($b, $b), $T, $N, $N)), + teardown = (@everywhere GC.gc())) - sub["broadcast (X .+ 1)"] = @benchmarkable(wait(X .+ 1), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["broadcast (X .+ 1)"] = @benchmarkable(wait(X .+ 1), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); 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)), - 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)), + 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)), - 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)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["transpose (permutedims)"] = @benchmarkable(wait(permutedims(X)), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); 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)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["reduce (sum)"] = @benchmarkable(sum(X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); 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)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["norm"] = @benchmarkable(norm(X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["norm"] = @benchmarkable(norm(X), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + teardown = (X = nothing; @everywhere GC.gc())) - suite["N=$N (block $b)"] = sub + suite["N=$N (block $b)"] = sub + end end suite diff --git a/benchmark/suites/linalg.jl b/benchmark/suites/linalg.jl index 3ddd000db..fd42a530d 100644 --- a/benchmark/suites/linalg.jl +++ b/benchmark/suites/linalg.jl @@ -23,58 +23,59 @@ function linalg_suite(ctx; method, accels) suite = BenchmarkGroup() for N in scales - b = square_block(N) - sub = BenchmarkGroup() + for b in blocks_for(N) + sub = BenchmarkGroup() - # 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)), - teardown = (A = nothing; @everywhere GC.gc())) + # 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)), + 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)), - 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)), + 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)), - 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)), + teardown = (A = nothing; @everywhere GC.gc())) - sub["qr"] = @benchmarkable(wait(qr(A).factors), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); 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)), + 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)), - teardown = (A = nothing; b = nothing; @everywhere GC.gc())) - end + 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)), + 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)), - teardown = (A = 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)), + 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)) - sub["svd"] = @benchmarkable(wait(svd(A).U), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), - 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)) + sub["svd"] = @benchmarkable(wait(svd(A).U), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); 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)), - teardown = (A = nothing; x = 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)), + teardown = (A = nothing; x = nothing; @everywhere GC.gc())) + end - isempty(sub) || (suite["N=$N (block $b)"] = sub) + isempty(sub) || (suite["N=$N (block $b)"] = sub) + end end suite diff --git a/benchmark/suites/stencil.jl b/benchmark/suites/stencil.jl index 021430863..764c58b88 100644 --- a/benchmark/suites/stencil.jl +++ b/benchmark/suites/stencil.jl @@ -105,63 +105,64 @@ function stencil_suite(ctx; method, accels) end for N in scales - b = square_block(N) - sub = BenchmarkGroup() - - # In-place stencils hold at most the input plus a same-size result. - 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)), - teardown = (B = nothing; @everywhere GC.gc())) + for b in blocks_for(N) + sub = BenchmarkGroup() + + # In-place stencils hold at most the input plus a same-size result. + 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)), + 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)), + 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)), + 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)), + 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)), + 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)), + 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)), + teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + end 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)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + # 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)), + teardown = (A = 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)), - 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)), - 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)), - 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)), - 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)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) - end + isempty(sub) || (suite["N=$N (block $b)"] = sub) end - - # 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)), - teardown = (A = nothing; @everywhere GC.gc())) - end - - isempty(sub) || (suite["N=$N (block $b)"] = sub) end suite diff --git a/benchmark/worker_mpi.jl b/benchmark/worker_mpi.jl new file mode 100644 index 000000000..de0cbaead --- /dev/null +++ b/benchmark/worker_mpi.jl @@ -0,0 +1,154 @@ +# Full-SPMD, single-shot MPI benchmark worker. +# +# Spawned by the orchestrator (benchmarks.jl) via `mpiexec -n N julia +# worker_mpi.jl workdir` when BENCHMARK_MPI_RANKS is set. Unlike worker.jl (a +# plain OS subprocess driven one benchmark at a time over a file-based +# request/response protocol), MPI ranks are SPMD: every rank must call the +# same collective Dagger operations at the same time, so there is no +# orchestrator-in-the-loop per benchmark here. Instead every rank builds the +# identical benchmark suite (mirroring test/mpi.jl's bootstrap) and runs a +# single flat pass over every benchmark in lockstep; rank 0 alone talks to +# the filesystem, recording each Trial and finishing with a manifest plus a +# `done` sentinel that the orchestrator polls for. +# +# Protocol (rank 0 only, all writes atomic via tmp-then-rename): +# - `result_mpi_.json`: a BenchmarkTools.save'd Trial for the i'th leaf +# that completed (in sorted-keypath order). +# - `results_mpi_manifest.json`: [{keypath, file}, ...] pairing each +# completed leaf's key path to its result file. +# - `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. + +using BenchmarkTools +using Distributed +using Dates, Random, Statistics, LinearAlgebra, InteractiveUtils +import JSON3 +using MPI + +include(joinpath(@__DIR__, "common.jl")) + +const WORKDIR = abspath(ARGS[1]) + +function atomic_write(path, data) + tmp = path * ".tmp" + open(io -> write(io, data), tmp, "w") + mv(tmp, path; force=true) + return nothing +end + +using Dagger +Dagger.accelerate!(:mpi) +Dagger.check_uniformity!(true) +const comm = MPI.COMM_WORLD +const rank = MPI.Comm_rank(comm) + +# --- Load acceleration backends (only if requested) ------------------------- +# No Distributed workers exist under MPI (each rank is its own OS process), +# so a plain `using` suffices here (worker.jl uses `@everywhere using` because +# it may have addprocs'd extra Distributed workers). + +for accel in accelerations + if accel == "cuda" + try + using DaggerGPU, CUDA + catch err + error("Failed to load CUDA acceleration; ensure DaggerGPU and CUDA " * + "are available (e.g. `benchpkg ... -a DaggerGPU,CUDA`)\n$err") + end + elseif accel == "amdgpu" + try + using DaggerGPU, AMDGPU + catch err + error("Failed to load AMDGPU acceleration; ensure DaggerGPU and " * + "AMDGPU are available (e.g. `benchpkg ... -a DaggerGPU,AMDGPU`)\n$err") + end + else + error("Unknown acceleration: $accel") + end +end + +# --- Build the benchmark suites --------------------------------------------- +# Every rank builds the identical suite (deterministic; no per-rank +# branching), so BenchmarkTools.leaves(SUITE) enumerates the same benchmarks +# everywhere. `ctx` is accepted-but-unused by every suite (verified: array, +# linalg, sparse, stencil), so it's passed as `nothing` here rather than a +# Distributed-flavored `Context()` -- test/mpi.jl doesn't set +# `Dagger.Sch.EAGER_CONTEXT[]` either, and doing so here would risk pinning +# the scheduler to a single-process view instead of the MPI-aware processor +# set that `Dagger.accelerate!(:mpi)` establishes. + +const suite_setup = Dict{String,Function}() +for suite in suites + suite_setup[suite] = include(joinpath(@__DIR__, "suites", suite * ".jl")) +end + +const SUITE = BenchmarkGroup() +for (suite_name, bench_list) in benches + suite_group = BenchmarkGroup() + for bench in bench_list + method_key = isempty(bench.accels) ? bench.method : + "$(bench.method)+$(join(bench.accels, "+"))" + rank == 0 && @info "[worker_mpi] Creating benchmarks for suite=$suite_name method=$method_key" + suite_group[method_key] = + suite_setup[suite_name](nothing; method=bench.method, accels=bench.accels) + end + SUITE[suite_name] = suite_group +end + +# Apply consistent run parameters to every benchmark in the tree. `evals=1` is +# required because the suites use `setup`/`teardown`. +for (_, b) in BenchmarkTools.leaves(SUITE) + b.params.seconds = bench_seconds + b.params.samples = bench_samples + b.params.evals = 1 + b.params.gcsample = true +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 +# iteration order. + +leaves = sort(collect(BenchmarkTools.leaves(SUITE)); by = kv -> String[string(k) for k in kv[1]]) + +const results = Vector{Tuple{Vector{String},BenchmarkTools.Trial}}() # rank 0 only + +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) + rank == 0 && push!(results, (kp, trial)) + catch err + 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()) + end + end +end + +# --- Rank 0: persist results and advertise completion ----------------------- + +if rank == 0 + manifest = Vector{Any}() + for (i, (kp, trial)) in enumerate(results) + fname = "result_mpi_$(i).json" + BenchmarkTools.save(joinpath(WORKDIR, fname), trial) + push!(manifest, (; keypath=kp, file=fname)) + end + atomic_write(joinpath(WORKDIR, "results_mpi_manifest.json"), JSON3.write(manifest)) + atomic_write(joinpath(WORKDIR, "done"), "1") + @info "[worker_mpi] Done; $(length(results))/$(length(leaves)) benchmark(s) succeeded." +end From bbdf2ed36e7f879aed1306c0ab4a6886c36e9a52 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 19 Aug 2026 02:40:41 -0700 Subject: [PATCH 03/10] datadeps: Allocate cross-space slots instead of transferring them A Datadeps slot's contents are never assumed to be current. `generate_slot!` says so outright -- it deliberately does not sync with the owner -- and `compute_remainder_for_arg!` decides what to copy purely from `arg_history` / `arg_owner`, never from the buffer. A slot in a space that does not yet appear in the argument's history is therefore always filled by a copy-to (`FullCopy`, or a span-exact `MultiRemainderAliasing` once part of it is current) before any task can read it. So for a dense, isbits payload crossing a space boundary, `move_rewrap` was sending bytes that were guaranteed to be overwritten. It now allocates the destination buffer and sends nothing: only the element type and dimensions travel (broadcast like a header under MPI, closure arguments under Distributed). The allocation is Libc-backed straight from `alloc_libc_array` rather than allocated and then copied into Libc memory by `libc_backed`, since there is nothing to copy. Under MPI that also fixes a latent no-op -- `libc_backed` was being applied to a `Chunk`, which hits the identity method. `from_space != to_space` is the load-bearing guard. When the two coincide the leaf transfer is the identity, so the "slot" *is* the original data and `compute_remainder_for_arg!` returns `NoAliasing()`; handing back a fresh buffer there would silently discard the argument. Allocation also stays routed through `aliased_object!`, so an already-present parent is still reused rather than shadowed by a detached buffer -- which is what keeps an array and its view sharing one destination allocation. N.B. This depends on the preceding free-syncdeps fix. Transferring the slot was masking a use-after-free: a recycled Libc block was immediately overwritten with real data, so a racing read returned plausible values. Without the transfer the same read hits uninitialized memory, which is how that bug first became visible (NaN out of a distributed Cholesky). Measured on 4 MPI ranks, a 1024x1024 `@stencil` sweep with `@neighbors`: 689ms/sweep at 128x128 tiles (872ms before) and 316ms at 512x512 tiles (509ms before). `test/mpi.jl` 398/403/388/376 passing; Distributed datadeps 1272 passing. Test contract updated accordingly: cross-space slots are no longer asserted to hold a copy of the source, only to have the right shape, type, and parent sharing; same-space slots must still be the real data. Two new testsets cover the allocation itself and the partial-currency invariant it rests on -- a task reading spans never written in its own space. --- ext/MPIExt.jl | 65 ++++++++++++++++++++++++++++++- src/datadeps/aliasing.jl | 27 ++++++++++++- src/memory-spaces.jl | 83 ++++++++++++++++++++++++++++++++++++++++ test/mpi.jl | 81 +++++++++++++++++++++++++++++++++------ 4 files changed, 243 insertions(+), 13 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 9b28fdb1d..3945268e9 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -1714,11 +1714,74 @@ function mpi_endpoint_transfer(accel::MPIAcceleration, from_proc, to_proc, from_ end end +# Uninitialized slot allocation is a property of the rank-local storage, so +# defer to the inner space. Rank identity is already handled by `!=` on the +# spaces themselves (`MPIMemorySpace` compares innerSpace/comm/rank). +Dagger.can_alloc_uninit(space::MPIMemorySpace, ::Type{T}) where {T} = + Dagger.can_alloc_uninit(space.innerSpace, T) +Dagger.alloc_uninit(space::MPIMemorySpace, ::Type{T}, header) where {T} = + Dagger.alloc_uninit(space.innerSpace, T, header) + +# Allocate the slot on the destination rank without shipping the payload. Only +# the `alloc_header` travels -- dimensions, for an array -- which is the entire +# saving: the buffer's contents are established later by Datadeps' copy-to +# phase. The header is whatever the value type declares, so a type whose shape +# is not its `size` participates without special-casing here. +# +# N.B. Point-to-point, deliberately *not* `bcast_yield`, even though the payload +# is header-sized. Only the destination rank needs the dimensions, so this keeps +# exactly the communication pattern of the `mpi_endpoint_transfer` it replaces: +# one message between the same two ranks. Broadcasting instead makes every rank +# a participant, and an intermediate rank cannot forward while it is blocked +# inside a task awaiting dispatch -- the cross-rank wait cycle described above +# `bcast_tree_children`, seen as a hang on Julia 1.10/1.11. +function mpi_endpoint_alloc(accel::MPIAcceleration, to_proc, to_space, ::Type{T_dest}, w::MPIWireValue) where T_dest + local_rank = MPI.Comm_rank(accel.comm) + from_rank = w.space.rank + if from_rank == to_space.rank + # Same rank (e.g. CPU -> GPU on one rank): nothing to communicate + if local_rank == to_space.rank + value = Dagger.alloc_uninit(to_space, T_dest, Dagger.alloc_header(wire_value(w))) + return tochunk(value, to_proc, to_space; type=T_dest) + end + return tochunk(nothing, to_proc, to_space; type=T_dest) + end + # Every rank takes a tag so the sequence stays rank-uniform, exactly as + # `mpi_endpoint_transfer` does before its own send/recv. + tag = to_tag() + if local_rank == from_rank + send_yield(Dagger.alloc_header(wire_value(w)), accel.comm, to_space.rank, tag) + return tochunk(nothing, to_proc, to_space; type=T_dest) + elseif local_rank == to_space.rank + # Typed from `T_dest` alone: the receiving rank holds no value to derive + # it from, and a type-level answer is rank-uniform by construction. + header = recv_yield(accel.comm, from_rank, tag)::Dagger.alloc_header_type(T_dest) + return tochunk(Dagger.alloc_uninit(to_space, T_dest, header), to_proc, to_space; type=T_dest) + else + return tochunk(nothing, to_proc, to_space; type=T_dest) + end +end + # Generic / wrapper MPIWireValue: leaf transfer, or header+children rebuild function move_rewrap(accel::MPIAcceleration, cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, w::MPIWireValue{T}) where T child_types = move_rewrap_child_types(T) if child_types === nothing - # Leaf: transfer to the destination, sharing via the aliased-object cache + # Leaf: materialize the slot on the destination rank, sharing via the + # aliased-object cache. A dense isbits payload crossing a rank boundary + # need not be sent at all -- Datadeps' copy-to phase fills the slot, and + # for a wrapper (e.g. a ChunkView) it sends only the wrapper's own spans + # instead of the entire backing array. See `slot_may_be_uninit`. + # + # N.B. The predicate depends only on the two spaces and the destination + # type, all rank-uniform, so every rank takes the same branch and the + # collectives below stay matched. + T_dest = move_type(mpi_inner_proc(as_mpi_proc(from_proc, from_space)), + mpi_inner_proc(as_mpi_proc(to_proc, to_space)), T) + if Dagger.slot_may_be_uninit(from_space, to_space, T_dest) + return aliased_object!(cache, w) do w + return mpi_endpoint_alloc(accel, to_proc, to_space, T_dest, w) + end + end return aliased_object!(cache, w) do w return Dagger.libc_backed(mpi_endpoint_transfer(accel, from_proc, to_proc, from_space, to_space, w)) end diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index c60f4af4c..9754e4360 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1182,6 +1182,20 @@ function remotecall_endpoint_toplevel(f, accel::DistributedAcceleration, cache:: return f(accel, cache, from_proc, to_proc, from_space, to_space, unwrap(data))::Chunk end end +# Allocate the slot directly on the destination worker. Only the element type +# and the `alloc_header` (dimensions, for an array) cross the wire; `data` +# deliberately does not, which is the whole point (a `remotecall_fetch` closing +# over it would serialize the payload). +function remotecall_endpoint_alloc(accel::DistributedAcceleration, to_proc, to_space, ::Type{T}, header) where T + wid = root_worker_id(to_proc) + if wid == myid() + return tochunk(alloc_uninit(to_space, T, header), to_proc, to_space) + end + return remotecall_fetch(wid, to_proc, to_space, T, header) do to_proc, to_space, T, header + return tochunk(alloc_uninit(to_space, T, header), to_proc, to_space) + end +end + function remotecall_endpoint_transfer(f, accel::DistributedAcceleration, from_proc, to_proc, from_space, to_space, data) wid = root_worker_id(to_proc) if wid == myid() @@ -1238,7 +1252,18 @@ move_rewrap(accel, cache::AliasedObjectCache, from_proc::Processor, to_proc::Pro function move_rewrap(accel, cache::AliasedObjectCache, from_proc::Processor, to_proc::Processor, from_space::MemorySpace, to_space::MemorySpace, data) parts = move_rewrap_parts(data) if parts === nothing - # Leaf: transfer the value, sharing via the aliased-object cache + # Leaf: materialize the slot in the destination space, sharing via the + # aliased-object cache. For a dense isbits payload crossing a space + # boundary the bytes need not travel at all -- Datadeps' copy-to phase + # establishes the contents, and for a wrapper (a view, say) it copies + # only the wrapper's own spans rather than the whole backing array. + # See `slot_may_be_uninit` for why this is safe. + T_dest = move_type(from_proc, to_proc, typeof(data)) + if slot_may_be_uninit(from_space, to_space, T_dest) + return aliased_object!(cache, data) do data + return remotecall_endpoint_alloc(accel, to_proc, to_space, T_dest, alloc_header(data)) + end + end return aliased_object!(cache, data) do data return remotecall_endpoint_transfer(accel, from_proc, to_proc, from_space, to_space, data) do accel, from_proc, to_proc, from_space, to_space, data return tochunk(libc_backed(move(from_proc, to_proc, data)), to_proc, to_space) diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 65f176133..2796b6636 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -82,6 +82,89 @@ function memory_space_scope(space::MemorySpace) end end +### Uninitialized slot allocation + +""" + can_alloc_uninit(space::MemorySpace, ::Type{T}) -> Bool + +Whether a value of type `T` can be conjured in `space` without moving any bytes +into it. Only dense, `isbits`-element arrays qualify: a pointer-ful element type +would leave the buffer holding garbage references rather than merely garbage +numbers, and a non-dense type is not described by its dimensions alone. + +Must depend only on `space` and `T`, both of which are rank-uniform under SPMD +execution, because the decision is taken before any collective (see +`slot_may_be_uninit`). +""" +can_alloc_uninit(::MemorySpace, ::Type) = false +can_alloc_uninit(::CPURAMMemorySpace, ::Type{<:Array{T}}) where {T} = isbitstype(T) + +""" + alloc_header(x) -> header + +The value-derived metadata [`alloc_uninit`](@ref) needs in order to conjure an +uninitialized stand-in for `x`. For a dense array that is its `size`; a type +whose shape is not described by dimensions alone supplies whatever its own +`alloc_uninit` method consumes. + +Kept distinct from `alloc_uninit` because the two run in different places: the +header is derived where the *value* lives, while the allocation happens in the +destination space. Under SPMD execution that separation is the point -- the +header is exactly what crosses the wire, so it should be as small as the type +allows (see `mpi_endpoint_alloc` in `ext/MPIExt.jl`). +""" +alloc_header(x::DenseArray) = size(x) + +""" + alloc_header_type(::Type{T}) -> Type + +The type of `alloc_header(::T)`. Needed wherever the header is obtained before +any value of `T` exists locally -- the MPI slot path types its `recv` with it -- +so it must follow from `T` alone, which also makes it rank-uniform. +""" +alloc_header_type(::Type{<:DenseArray{T,N}}) where {T,N} = Dims{N} + +""" + alloc_uninit(space::MemorySpace, ::Type{T}, header) -> T + +Allocate an uninitialized `T` in `space`, shaped by `header` (see +[`alloc_header`](@ref)). Only called when [`can_alloc_uninit`](@ref) returned +`true` for the same `space`/`T`. + +The CPU allocation is Libc-backed from the start rather than allocated and then +copied into Libc memory (what `libc_backed` would do), since there is nothing to +copy: the buffer's contents are established later by Datadeps' copy-to phase. +""" +alloc_uninit(::CPURAMMemorySpace, ::Type{<:Array{T,N}}, dims::Dims{N}) where {T,N} = + alloc_libc_array(T, dims) + +""" + slot_may_be_uninit(from_space, to_space, ::Type{T}) -> Bool + +Whether a Datadeps slot of type `T` being created in `to_space` for data living +in `from_space` may be left uninitialized, skipping the transfer entirely. + +This is sound because a slot's contents are never assumed to be current. +`generate_slot!` explicitly does not sync with the owner, and +`compute_remainder_for_arg!` decides what to copy purely from `arg_history` / +`arg_owner` — never from the buffer. A slot in a space that does not yet appear +in the argument's history is therefore always filled by a copy-to (`FullCopy`, +or a span-exact `MultiRemainderAliasing` once part of it is current) before any +task can read it. + +`from_space != to_space` is the load-bearing guard. When the two coincide, +`move_rewrap`'s leaf transfer is the identity, so the "slot" *is* the original +data, and `compute_remainder_for_arg!` returns `NoAliasing()` (owner space == +target space, empty history) — no copy is scheduled, and handing back a fresh +buffer there would silently discard the argument's contents. + +What may be conjured is decided entirely by `can_alloc_uninit`, so a type opts +in by defining that plus [`alloc_header`](@ref) / [`alloc_uninit`](@ref) — there +is deliberately no `DenseArray` bound here to be extended around. +""" +slot_may_be_uninit(from_space::MemorySpace, to_space::MemorySpace, ::Type{T}) where {T} = + from_space != to_space && can_alloc_uninit(to_space, T) + ### In-place Data Movement unwrap(x::Chunk) = unwrap(x.handle) diff --git a/test/mpi.jl b/test/mpi.jl index 78cff3fb5..cda70f5f4 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -231,9 +231,16 @@ end # must never alias, even with coinciding SPMD addresses @test !Dagger.will_alias(src_ainfo, dst_ainfo) end - # The slot holds a copy of the source data on the destination + # A same-space slot wraps (aliases) the original data. A + # cross-rank slot is deliberately left *uninitialized*: nothing + # is sent, and Datadeps' copy-to phase establishes the contents + # later (see `slot_may_be_uninit`). Only shape/type is promised. if rank == dst - @test Dagger.unwrap(slot) == A + if src == dst + @test Dagger.unwrap(slot) == A + else + @test size(Dagger.unwrap(slot)) == size(A) + end end end @@ -261,9 +268,16 @@ end # ... while the views themselves remain disjoint @test !Dagger.will_alias(aA, aB) if rank == dst - @test Dagger.unwrap(slotA) == vA - @test Dagger.unwrap(slotB) == vB + # Parent sharing holds regardless: both views must resolve + # to one destination allocation, initialized or not. @test parent(Dagger.unwrap(slotA)) === parent(Dagger.unwrap(slotB)) + if src == dst + @test Dagger.unwrap(slotA) == vA + @test Dagger.unwrap(slotB) == vB + else + @test size(Dagger.unwrap(slotA)) == size(vA) + @test size(Dagger.unwrap(slotB)) == size(vB) + end end end @@ -277,7 +291,11 @@ end @test Dagger.check_uniform(slot.handle) @test Dagger.chunktype(slot) <: UpperTriangular if rank == dst - @test Dagger.unwrap(slot) == U + if src == dst + @test Dagger.unwrap(slot) == U + else + @test size(Dagger.unwrap(slot)) == size(U) + end end end @@ -315,9 +333,14 @@ end @test aA.base_ptr == aB.base_ptr @test !Dagger.will_alias(aA, aB) if rank == dst - @test Dagger.unwrap(slotA) == view(A, 1:4, 1:8) - @test Dagger.unwrap(slotB) == view(A, 5:8, 1:8) @test parent(Dagger.unwrap(slotA)) === parent(Dagger.unwrap(slotB)) + if src == dst + @test Dagger.unwrap(slotA) == view(A, 1:4, 1:8) + @test Dagger.unwrap(slotB) == view(A, 5:8, 1:8) + else + @test size(Dagger.unwrap(slotA)) == (4, 8) + @test size(Dagger.unwrap(slotB)) == (4, 8) + end end # Nested ChunkView flattens to the same slices as a direct view @@ -330,8 +353,10 @@ end slot_direct = make_slot(fresh_cache(space_for_rank(dst)), src, dst, cv_direct) @test Dagger.chunktype(slot_nested) <: SubArray if rank == dst - @test Dagger.unwrap(slot_nested) == Dagger.unwrap(slot_direct) - @test Dagger.unwrap(slot_nested) == view(A, 2:3, 1:4) + @test size(Dagger.unwrap(slot_nested)) == size(Dagger.unwrap(slot_direct)) + if src == dst + @test Dagger.unwrap(slot_nested) == view(A, 2:3, 1:4) + end end end @@ -355,8 +380,12 @@ end end if rank == dst H2 = Dagger.unwrap(slot) - @test H2.center == H.center - @test all(h2 == h for (h2, h) in zip(H2.halos, H.halos)) + @test size(H2.center) == size(H.center) + @test all(size(h2) == size(h) for (h2, h) in zip(H2.halos, H.halos)) + if src == dst + @test H2.center == H.center + @test all(h2 == h for (h2, h) in zip(H2.halos, H.halos)) + end end end end @@ -436,6 +465,36 @@ end @test fetch(cv_top) ≈ ref_blk[1:2, :] end +@testset "Uninitialized cross-rank slots" begin + # Cross-space slots for dense isbits payloads are allocated on the + # destination rather than transferred (`slot_may_be_uninit`), and are + # Libc-backed straight from the allocator so Datadeps can free them eagerly + # without the allocate-then-copy that `libc_backed` would otherwise do. + A = rand(4, 4) + for (src, dst) in rank_pairs + obj = Dagger.tochunk(A, proc_for_rank(src), space_for_rank(src)) + slot = make_slot(fresh_cache(space_for_rank(dst)), src, dst, obj) + if rank == dst + val = Dagger.unwrap(slot) + if src == dst + # No copy-to is scheduled when owner space == target space, so + # the slot must still be the real data + @test val == A + else + @test Dagger.is_libc_allocated(val) + end + end + end + + # The predicate itself, including the guard that makes it safe + s0, s1 = space_for_rank(0), space_for_rank(min(1, nranks-1)) + @test !Dagger.slot_may_be_uninit(s0, s0, Matrix{Float64}) + @test !Dagger.slot_may_be_uninit(s0, s1, Matrix{String}) + if nranks > 1 + @test Dagger.slot_may_be_uninit(s0, s1, Matrix{Float64}) + end +end + @testset "Partial slot currency" begin # The load-bearing invariant behind uninitialized slots: a task may read # spans that were never written in its own space, and those must be copied From 9832121d274fcc33b11d22f4a43c89f676f1b60c Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 20 Aug 2026 10:41:05 -0700 Subject: [PATCH 04/10] CI: Give the 4-rank MPI CPU job more timeout headroom test/mpi.jl at 4 ranks passes cleanly (398/398) in ~12 minutes on an idle workstation, matching the ~12m benchmark in c13465a8's commit message -- no reproducible deadlock. CI reportedly runs ~40 minutes and occasionally exceeds the 60-minute timeout, consistent with slower/shared runners eating the margin rather than a hang. Bump 4-rank's budget to 90 minutes; 2-rank stays at 60 since it isn't reported as tight. --- .github/workflows/CI.yml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index f81b1098e..5f9b2a108 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -404,7 +404,14 @@ jobs: mpi-cpu: name: MPI CPU - Julia ${{ matrix.version }} / ${{ matrix.ranks }} ranks runs-on: ubuntu-latest - timeout-minutes: 60 + # 4 ranks has less headroom than 2: a clean local run of test/mpi.jl at 4 + # ranks takes ~12-13 minutes on an idle 12-core workstation (matching the + # ~12m historical benchmark in c13465a8's commit message, with all 398 + # tests passing and no hang), but shared/oversubscribed `ubuntu-latest` + # runners reportedly take ~40 minutes and occasionally exceed 60 -- a + # timeout-margin problem on slower CI hardware, not a reproducible + # deadlock. Give the 4-rank job more slack than 2-rank needs. + timeout-minutes: ${{ matrix.ranks == 4 && 90 || 60 }} if: ${{ !contains(github.event.head_commit.message, '[skip tests]') }} strategy: fail-fast: false From a351639e976a729db35c5555a0412f9f8b6106c1 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 20 Aug 2026 14:45:45 -0700 Subject: [PATCH 05/10] MPIExt: Broadcast move_rewrap's header via the dispatch-decoupled relay move_rewrap runs from generate_slot! inside the per-task planning loop, and each task is handed to eager_launch! (and so starts executing, concurrently) before the loop moves on to plan the next task's header broadcast. That means this broadcast is not at the "sequential, non-overlapping point" the raw bcast_yield tree assumes -- it can race an already-dispatched task's own execute!/poolget activity, closing the same forwarder wait-cycle documented above bcast_meta_yield. Confirmed: a 4-rank mpiexec run of test/mpi.jl hung almost immediately with the old bcast_yield header broadcast under load from a concurrent MPI job on the same machine, and passed cleanly (16m44s, 398/398, no warnings) after switching to bcast_meta_yield. Co-Authored-By: Claude Sonnet 5 --- ext/MPIExt.jl | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 3945268e9..3869790c2 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -1809,11 +1809,20 @@ function move_rewrap(accel::MPIAcceleration, cache::AliasedObjectCache, from_pro mode = move_rewrap_header_mode(T) header = if mode === :broadcast tag = to_tag() + # `move_rewrap` is reached from `generate_slot!` inside the per-task + # planning loop, and each task is handed to `eager_launch!` (and so + # begins executing, concurrently, on this same rank) before the loop + # moves on to plan the next task. That means this broadcast is NOT at + # a "sequential, non-overlapping point" the way the doc comment above + # `bcast_meta_yield` assumes -- it can race an already-dispatched + # task's own `execute!`/`poolget` activity on this or another rank. + # Use the dispatch-decoupled relay so a rank busy running a task's + # `execute!` still forwards this header on message arrival. if local_rank == w.space.rank _, hdr = move_rewrap_parts(wire_value(w)) - bcast_yield(accel.comm, w.space.rank, tag, hdr) + bcast_meta_yield(accel.comm, w.space.rank, tag, hdr) else - bcast_yield(accel.comm, w.space.rank, tag) + bcast_meta_yield(accel.comm, w.space.rank, tag) end elseif mode === :none nothing From 2466a46253a1527d4681ea724a1a495fbecae4da Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 20 Aug 2026 20:10:08 -0700 Subject: [PATCH 06/10] MPIExt: Drop the blocking barrier that deadlocks the 4-rank MPI job The 4-rank CPU MPI job has been running out its CI timeout on Julia 1.10 and 1.11 with nothing in the log to point at. It reproduces once the run is given CI's CPU budget rather than a workstation's: pinned to 4 CPUs (what ubuntu-latest provides), Julia 1.10 at HEAD deadlocks outright, and SIGUSR1 backtrace dumps catch all four ranks inside `MPI_Barrier`, called from `check_uniform`. Nothing here is version-specific -- it is a race whose odds turn on GC and scheduling timing, which is why 1.12 slips through. Three defects, in decreasing order of severity: * `check_uniform` ended with `MPI.Barrier`. That is a blocking `ccall`, so it parks the OS thread inside MPI: the thread runs no other Julia task -- notably not `bcast_relay_loop`, which other ranks depend on to have a broadcast forwarded -- and it never reaches a GC safepoint, so any other thread requesting a collection stalls the process until the barrier returns. Rank r parks here waiting for rank s while s waits on a broadcast only r's relay can forward, and neither can move. The barrier bought nothing either: `compare_all` directly above it is already an arrival barrier, since it cannot return until every rank has entered it and sent. * `bcast_slot_wait` kept a one-shot slot keyed by tag alone. A tag does not identify a single broadcast -- `to_tag()` returns the planning task's thunk id, so every `move_rewrap` header broadcast issued while generating that task's slots shares it, one per wrapper level and one per argument. The second delivery overwrote the first before its consumer read it; that consumer then took the wrong value and the next one waited forever on a payload that had already arrived. Now a FIFO keyed by (root, tag), drained in the rank-uniform order MPI's non-overtaking guarantee establishes. * The relay busy-spun on `MPI_Improbe`, taking 107 of one thread's 127 profile samples -- a core per rank, which a 4-vCPU runner hosting 4 ranks cannot spare. It now spins hot through a burst and backs off after it, keeping the added latency off all but the first message of the next burst. Both waits that had no deadlock detection now have it, so a future regression here reports itself instead of silently spending the CI budget. Verified on Julia 1.10, 4 ranks x 2 threads pinned to 4 CPUs. HEAD deadlocks in the Stencils testset after ~35 minutes: ranks 1-3 time out on a `compare_all` recv from rank 0 on tag 1073741823 (`MPI.tag_ub()`) while rank 0, sitting in the undetected barrier, reports nothing at all, and `mpiexec` never exits. With this change the suite passes on all four ranks twice over, in 18m46s and 18m19s, with the expected 398/403/376/388 test counts. Co-Authored-By: Claude Opus 5 --- ext/MPIExt.jl | 187 ++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 157 insertions(+), 30 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 3869790c2..7ec2ea697 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -65,7 +65,26 @@ function check_uniform(value::Integer, original=value) Core.print("[$rank] value=$value, original=$original\n") throw(ArgumentError("Non-uniform value")) end - MPI.Barrier(comm) + # N.B. No `MPI.Barrier` here. `compare_all` above already *is* an arrival + # barrier -- it cannot return until every other rank has entered it and + # sent -- so a barrier afterwards synchronizes nothing extra, while a + # blocking one is actively harmful under Dagger's multithreaded runtime: + # + # * `MPI.Barrier` is a blocking `ccall`, so it parks the calling OS thread + # inside MPI until every rank arrives. That thread runs no other Julia + # task meanwhile -- in particular not `bcast_relay_loop`, the progress + # engine other ranks depend on to have a broadcast forwarded. + # * Julia cannot see into a plain `ccall`, so the parked thread never + # reaches a GC safepoint: any other thread that requests a collection + # stalls there until the barrier returns, which takes the relay down with + # it even when a second thread was free to run it. + # + # Those two close a real cross-rank cycle. Rank r parks here waiting for + # rank s; rank s is waiting on a broadcast whose next hop is r's relay, + # which cannot run; so s never reaches this barrier and r never leaves + # it. `MPI.Barrier` has no deadlock detection either, so it surfaces as a + # CI job silently exhausting its timeout. The P2P inside `compare_all` is + # covered by `mpi_deadlock_detect` and yields, so it has neither problem. return matched end @@ -1152,18 +1171,37 @@ function bcast_tree_children(rank::Int, root::Int, sz::Int) return children end -# One-shot promise for a delivered broadcast payload (one per tag). +# Delivered broadcast payloads for one (root, tag) pair, in arrival order. +# +# N.B. A FIFO, not a one-shot promise, because a tag does NOT identify a single +# broadcast. `to_tag()` returns the planning task's thunk id, so every +# `move_rewrap` header broadcast issued while generating that task's slots +# shares one tag -- including one per level of a nested wrapper, and one per +# argument. A one-shot slot keeps only the last of those: the earlier payload is +# overwritten before its consumer reads it, the consumer then takes the wrong +# value and the next consumer waits forever on a delivery that already happened. +# That is a silent hang, since nothing downstream of the relay is covered by +# `mpi_deadlock_detect`. +# +# Queue order matches sender intent: MPI's non-overtaking guarantee delivers +# same-(source, tag) messages in send order, each relay forwards in that same +# order, and SPMD planning issues and consumes them in one rank-uniform order. +# Keying on the root as well as the tag keeps two roots that happen to share a +# tag from interleaving into a single queue, where no ordering rule applies. mutable struct BcastSlot - ev::Base.Event - value::Any - done::Bool - BcastSlot() = new(Base.Event(), nothing, false) + values::Vector{Any} + BcastSlot() = new(Any[]) end mutable struct BcastState bcast_comm::MPI.Comm - slots::Dict{UInt32,BcastSlot} - lock::Threads.ReentrantLock + 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. + cond::Threads.Condition running::Threads.Atomic{Bool} relay::Union{Task,Nothing} end @@ -1181,30 +1219,85 @@ bcast_state_for(comm::MPI.Comm) = bcast_serialize(x) = (io = IOBuffer(); Serialization.serialize(io, x); take!(io)) -function bcast_deliver!(state::BcastState, tag::UInt32, value) - slot = lock(state.lock) do - get!(BcastSlot, state.slots, tag) +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.value = value - slot.done = true - notify(slot.ev) return end -function bcast_slot_wait(state::BcastState, tag::UInt32) - slot = lock(state.lock) do - get!(BcastSlot, state.slots, tag) - end - # `Base.Event` latches: a `notify` that already fired makes `wait` return - # immediately, so relay-before-consumer and consumer-before-relay both work. - wait(slot.ev) - value = slot.value - lock(state.lock) do - delete!(state.slots, tag) +# Wake every consumer so each can re-check its own elapsed time. Called by the +# relay while it is idle, which makes the relay the heartbeat that drives +# 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) + return +end + +function bcast_slot_wait(state::BcastState, root::Int, tag::UInt32) + key = (root, tag) + # Enqueueing under the same lock means relay-before-consumer and + # consumer-before-relay both work: a payload that arrived first is simply + # already queued when we first look. + # + # The wait is routed through `mpi_deadlock_detect` like every other + # cross-rank wait in this extension. Before, this was a bare `wait` on a + # one-shot `Event`, so a delivery that was lost (see `BcastSlot`) stalled + # 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) + 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`. + @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 + end + warn_period = mpi_deadlock_detect(detect, time_start, warn_period, timeout_period, + rank, tag, "bcast_meta delivery", root) + wait(state.cond) + end end - return value end +# How many consecutive empty probes to spin through before the relay starts +# sleeping between them, and how long it then sleeps. +# +# A bare `Improbe`/`yield()` loop never blocks, so it keeps a whole OS thread at +# 100% forever: a profile of a 4-rank run spent 107 of one thread's 127 samples +# inside `MPI_Improbe`. With one relay per rank that burns one core per rank, +# which is survivable on a workstation but not on a 4-vCPU CI runner hosting 4 +# ranks -- there the relays alone claim the entire machine and the threads doing +# real work are left to fight over what is left. The cost scales with rank +# count, so it hits the 4-rank job far harder than the 2-rank one. +# +# Broadcast traffic is bursty -- planning issues them in clusters -- so spin hot +# first and only back off once the burst is clearly over. That keeps the added +# latency off the messages inside a burst and pays it at most once, on the first +# message of the next one. +const BCAST_RELAY_SPIN_ITERS = 4096 +const BCAST_RELAY_IDLE_SLEEP = 1e-4 +# Heartbeat cadence once idle, counted in backoff iterations: at ~150us apiece +# this wakes blocked consumers roughly every 150ms, which is fine granularity +# for a 10s warn / 120s error timer and costs nothing when nobody is waiting. +const BCAST_RELAY_HEARTBEAT_ITERS = 1000 + # Always-running per-rank progress engine: receive a broadcast on `bcast_comm`, # forward the raw bytes to this rank's tree children, then deliver the payload. # Forwarding is NOT gated by datadeps task dispatch — that is what breaks the @@ -1213,14 +1306,41 @@ function bcast_relay_loop(state::BcastState) comm = state.bcast_comm rank = MPI.Comm_rank(comm) sz = MPI.Comm_size(comm) + idle = 0 + # Backing off means blocking the OS thread, which only gives time back to + # anyone else if this rank has another thread to run them on. On a + # single-threaded rank it would instead starve the very tasks the relay + # exists to serve, so there we keep the pure `yield()` loop. + may_backoff = Threads.nthreads() > 1 while state.running[] try MPI.Finalized() && break got, msg, stat = MPI.Improbe(MPI.ANY_SOURCE, MPI.ANY_TAG, comm, MPI.Status) if !got - yield() + idle += 1 + if idle <= BCAST_RELAY_SPIN_ITERS || !may_backoff + yield() + else + # N.B. `Libc.systemsleep`, not `sleep`: `sleep` parks the + # task on a libuv timer, which only fires when some thread + # reaches the scheduler's idle loop. A rank whose other + # thread is inside a blocking `ccall` has no such thread, so + # the relay could stay asleep exactly when it is needed + # most. Sleeping the thread outright cannot be starved that + # way, and 100us is short enough that the GC-unsafe window + # it opens is irrelevant. + Libc.systemsleep(BCAST_RELAY_IDLE_SLEEP) + yield() + end + # Nothing is arriving, which is exactly when a stalled consumer + # needs its deadlock timer re-checked -- and the relay is the + # one thing guaranteed to still be running to do it. + if idle % BCAST_RELAY_HEARTBEAT_ITERS == 0 + bcast_heartbeat!(state) + end continue end + idle = 0 src = MPI.Get_source(stat) tag = UInt32(MPI.Get_tag(stat)) count = MPI.Get_count(stat, UInt8) @@ -1232,7 +1352,7 @@ function bcast_relay_loop(state::BcastState) sreq = MPI.Isend(buf, comm; dest=child, tag=tag) __wait_for_request(sreq, comm, rank, child, tag, "bcast_relay", "send") end - bcast_deliver!(state, tag, value) + bcast_deliver!(state, Int(root), tag, value) catch err # Comm torn down (disable/finalize) races the loop: stop quietly. # A genuine mid-session fault surfaces via errormonitor. @@ -1245,8 +1365,8 @@ end function start_bcast_relay!(accel_comm::MPI.Comm) bcast_comm = MPI.Comm_dup(accel_comm) - state = BcastState(bcast_comm, Dict{UInt32,BcastSlot}(), - Threads.ReentrantLock(), Threads.Atomic{Bool}(true), nothing) + state = BcastState(bcast_comm, Dict{Tuple{Int,UInt32},BcastSlot}(), + Threads.Condition(), Threads.Atomic{Bool}(true), nothing) lock(BCAST_STATES) do states states[accel_comm] = state end @@ -1280,6 +1400,13 @@ function stop_bcast_relay!(accel_comm::MPI.Comm) catch end end + # The relay is the heartbeat that lets a blocked consumer re-check its + # 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) return nothing end @@ -1301,7 +1428,7 @@ function bcast_meta_yield(comm::MPI.Comm, root::Integer, tag, value=nothing) end return value else - return bcast_slot_wait(state, utag) + return bcast_slot_wait(state, Int(root), utag) end end From 8f315fb08b94fd03b71e39f2a1b43a3722658b77 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 5 Sep 2026 12:07:06 -0700 Subject: [PATCH 07/10] datadeps: Compute free syncdeps from the buffer's own recorded aliasing `gather_free_syncdeps!`'s uniform-execution fallback synced on the object cache's `key` ainfo, guarded on that key being tracked: haskey(state.ainfos_overlaps, wrapped) && get_write_deps!(state, space, wrapped, write_num, syncdeps) That guard is close to never true. `ainfos_overlaps` is populated by `populate_ainfo!` only for ainfos tracked as a direct task dependency, and the buffers reaching this branch are exactly the ones that are not -- they merely underlie wrapper arguments (the parent array shared by several `view`s, whose tracked slots are the views rather than the buffer). Worse, the key ainfo describes the original object in its *source* space: every span carries its space, and `intersect` skips candidates sharing no space index (memory-spaces.jl:428), so a source-space ainfo cannot overlap a destination-space one even when it is tracked. The branch therefore emitted an `unsafe_free!` with an *empty* syncdep set, racing every task still reading the buffer. What the free loop actually needs to know is which tasks touch memory *inside* the buffer, and that is the tracked ainfos in the buffer's own space overlapping its extent -- which requires the buffer's *destination-space* aliasing. That cannot be recomputed here: `aliasing` is a collective under uniform (SPMD) execution, so reaching it from the free loop costs one broadcast per freed buffer at a sequential point. It need not be recomputed, though. `set_stored!` already computes exactly this ainfo, on every rank, when it allocates the buffer -- it needs it to register the buffer in `derived` -- and then discards it. Retain it in a new `value_ainfos` field and hand it to `gather_free_syncdeps!`, which now runs one exact, space-filtered overlap search for both execution modes. The uniform/non-uniform split disappears, as does the non-uniform path's own free-time `aliasing` call. The result is exact rather than conservative. Syncing against every tracked ainfo in the space instead would hold every buffer until that space's work had drained -- the opposite of what deferring frees is meant to buy. A missing record is now an error rather than a quiet `return`: degrading to an empty syncdep set is precisely the use-after-free this function exists to prevent. Also adopts, from 8469a619 on jps/datadeps-region-async (which closes the same hole by taking the collective instead), the parts independent of that choice: `gather_overlap_syncdeps!` as a shared helper, and `DATADEPS_ASSERT_FREE_SYNCDEPS`, a debug-gated invariant check that re-derives the answer by linear scan + `will_alias` rather than the interval tree, so a bug in the fast path cannot also hide from the assertion meant to catch it. Distributed datadeps suite (4 workers) 2068 passing, 2 broken, no failures -- including the new end-to-end regression test, which asserts every `unsafe_free!` a shared-parent-views region emits has non-empty syncdeps and runs with the invariant check forced on. 4-rank `test/mpi.jl` 443/470/455/466 passing, no failures or errors. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019Pp6hFZNHR44pLhDrR6aXq --- src/datadeps/aliasing.jl | 179 ++++++++++++++++++++++++++--------- src/datadeps/hierarchical.jl | 18 ++-- src/datadeps/queue.jl | 3 +- test/datadeps.jl | 131 +++++++++++++++++++++++++ 4 files changed, 279 insertions(+), 52 deletions(-) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 9754e4360..a8ba0d2ab 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -302,6 +302,14 @@ struct AliasedObjectCacheStore derived::Dict{AbstractAliasing,AbstractAliasing} stored::Dict{MemorySpace,Set{AbstractAliasing}} values::Dict{MemorySpace,Dict{AbstractAliasing,Chunk}} + # The aliasing of each buffer in `values`, *in the space that holds it*. + # `set_stored!` computes this anyway (it registers the buffer in `derived`), + # so retaining it is free -- and it is the only rank-uniform description of + # a buffer's own extent available after planning, since recomputing it means + # calling `aliasing`, a collective under SPMD. See `gather_free_syncdeps!`. + # Only Datadeps-allocated copies get an entry; the user's originals (added + # by `set_key_stored!`) do not, as they are never freed. + value_ainfos::Dict{MemorySpace,Dict{AbstractAliasing,AbstractAliasing}} # The `(space, key)` pairs identifying the user's original data, as opposed # to Datadeps-allocated copies. A `key` ainfo is recorded here at the space # where it is first registered (its source space, see `set_key_stored!`), @@ -315,6 +323,7 @@ AliasedObjectCacheStore(accel::Acceleration) = Dict{AbstractAliasing,AbstractAliasing}(), Dict{MemorySpace,Set{AbstractAliasing}}(), Dict{MemorySpace,Dict{AbstractAliasing,Chunk}}(), + Dict{MemorySpace,Dict{AbstractAliasing,AbstractAliasing}}(), Set{Tuple{MemorySpace,AbstractAliasing}}()) """ @@ -354,8 +363,29 @@ function set_stored!(cache::AliasedObjectCacheStore, dest_space::MemorySpace, va 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 + +""" + 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). + +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. +""" +function stored_value_ainfo(cache::AliasedObjectCacheStore, space::MemorySpace, key::AbstractAliasing) + ainfos = get(cache.value_ainfos, space, nothing) + ainfos === nothing && return nothing + return get(ainfos, key, nothing) +end function set_key_stored!(cache::AliasedObjectCacheStore, space::MemorySpace, ainfo::AbstractAliasing, value::Chunk) @check_uniform(value) push!(cache.keys, ainfo) @@ -935,65 +965,128 @@ function _get_read_deps!(state::DataDepsState, dest_space::MemorySpace, ainfo::A end end """ - gather_free_syncdeps!(state, space, key_ainfo, remote_arg, write_num, chunk_to_ainfos, syncdeps) + gather_overlap_syncdeps!(state, ainfo, write_num, syncdeps) + +Register `ainfo` into `state.ainfos_lookup` and add to `syncdeps` the owner and +readers of every *other* tracked ainfo overlapping it. Used by +`gather_free_syncdeps!` for buffers that are not themselves a directly-tracked +task dependency, so their overlaps are not already precomputed in +`state.ainfos_overlaps`. + +We reuse the lookup's interval-tree overlap search (which prunes most +`will_alias` comparisons via bounding spans) rather than scanning every tracked +ainfo. `intersect` requires its query ainfo to be registered, so we `push!` it +in first; this is safe because the free loop is the final step of +`distribute_tasks!`, after which the lookup is no longer consulted for +scheduling. The search is space-filtered for free: every span carries its space, +so only ainfos in `ainfo`'s own space can be returned. + +Safe to call with an `ainfo` that has a content-equal (same `.hash`) entry +already registered elsewhere in the lookup: `AliasingWrapper` equality (and +`Dict` lookups keyed by it) are by content hash, not object identity, so a +redundant push here costs only a duplicate interval-tree entry, never a missed +or double-counted syncdep -- `syncdeps` is a `Set`, and the `other_ainfo === +ainfo` self-skip below excludes only the literal instance just pushed (which, +being brand new, was never itself recorded as anyone's owner or reader). +""" +function gather_overlap_syncdeps!(state::DataDepsState, ainfo::AliasingWrapper, write_num::Int, syncdeps) + ainfo.inner isa NoAliasing && return + ainfo_idx = push!(state.ainfos_lookup, ainfo) + for other_ainfo in intersect(state.ainfos_lookup, ainfo; ainfo_idx) + other_ainfo === ainfo && continue + 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)) + end + for (reader_task, reader_write_num) in get(state.ainfos_readers, other_ainfo, ()) + reader_write_num != write_num && push!(syncdeps, ThunkSyncdep(reader_task)) + end + end +end + +# Debug-only invariant: freeing a buffer must never produce an empty syncdep +# set while the state still records a writer or readers overlapping it. Gated +# behind this `Ref` (following the same pattern as `CHECK_UNIFORMITY`, +# acceleration.jl:75) so it costs nothing -- not even a function call, since +# `assert_free_syncdeps!` itself is only ever invoked from within +# `gather_free_syncdeps!`, which is already off the hot path -- in normal +# operation. Flip it on in tests that want this checked. +# +# Deliberately re-derives the answer independently of `gather_free_syncdeps!` +# / `gather_overlap_syncdeps!` (a linear scan over every tracked ainfo plus +# `will_alias`, instead of the interval-tree search they use) so that a bug in +# the fast path can't also hide from the assertion meant to catch it. +const DATADEPS_ASSERT_FREE_SYNCDEPS = Ref(false) +function assert_free_syncdeps!(state::DataDepsState, ainfo::AliasingWrapper, write_num::Int, syncdeps) + DATADEPS_ASSERT_FREE_SYNCDEPS[] || return + ainfo.inner isa NoAliasing && return + for (other_ainfo, owner) in state.ainfos_owner + owner === nothing && continue + 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" + end + for (other_ainfo, readers) in state.ainfos_readers + for (reader_task, reader_write_num) in readers + reader_write_num == write_num && continue + will_alias(ainfo, other_ainfo) || continue + @assert ThunkSyncdep(reader_task) in syncdeps "gather_free_syncdeps! omitted a live reader $reader_task ($other_ainfo) for buffer overlapping $ainfo" + end + end +end + +""" + gather_free_syncdeps!(state, space, buf_ainfo, remote_arg, write_num, chunk_to_ainfos, syncdeps) Collect into `syncdeps` every task that must complete before the backing buffer -`remote_arg` (a Datadeps-allocated copy in `space`) can be freed. +`remote_arg` (a Datadeps-allocated copy in `space`) can be freed. `buf_ainfo` is +the buffer's own aliasing *in `space`*, from `stored_value_ainfo`. If `remote_arg` is itself a tracked slot (the common case -- whole-object arguments), its ainfos are in `chunk_to_ainfos` and we reuse their precomputed overlap sets. Otherwise the buffer only underlies wrapper arguments (e.g. it is the parent array shared by several `view`s, whose tracked slots are the views -rather than this buffer); in that case we compute the buffer's own aliasing and -sync with every tracked ainfo that overlaps its memory. Under uniform (SPMD) -execution that computation is unavailable, and we fall back to the rank-uniform -cache key ainfo `key_ainfo`. +rather than this buffer), and the tasks to wait on are exactly those touching +memory *inside* the buffer -- i.e. the tracked ainfos in `space` that overlap +`buf_ainfo`. + +N.B. The buffer's aliasing has to be supplied rather than computed here. The +object cache's `key` ainfo describes the original object in its *source* space, +and every span carries its space, so a key ainfo can never overlap a +destination-space ainfo -- keying this lookup on it yields an empty syncdep set +and an `unsafe_free!` that races the tasks still reading the buffer. Nor can we +recompute it: `aliasing` is a collective under uniform (SPMD) execution (the +owner computes and broadcasts), and only the owning rank holds the data, so +reaching it from here would broadcast into a collective no other rank enters -- +consuming a tag on one rank (desyncing every subsequent `to_tag`) and then +deadlocking. `set_stored!` already computed this ainfo, on every rank, when it +allocated the buffer; we simply use what it recorded. """ -function gather_free_syncdeps!(state::DataDepsState, space::MemorySpace, key_ainfo, remote_arg, write_num::Int, chunk_to_ainfos, syncdeps) +function gather_free_syncdeps!(state::DataDepsState, space::MemorySpace, buf_ainfo, remote_arg, write_num::Int, chunk_to_ainfos, syncdeps) ainfos = get(chunk_to_ainfos, remote_arg, nothing) if ainfos !== nothing for ainfo in ainfos get_write_deps!(state, space, ainfo, write_num, syncdeps) + assert_free_syncdeps!(state, ainfo, write_num, syncdeps) end return end - # Under uniform (SPMD) execution we cannot compute the buffer's aliasing here: - # `aliasing` is collective (the owner computes and broadcasts), and only the - # owning rank holds the data, so having it alone reach the call below would - # broadcast into a collective no other rank enters -- consuming a tag on one - # rank (desyncing every subsequent `to_tag`) and then deadlocking. Fall back - # to the rank-uniform cache key ainfo, which is metadata available identically - # on every rank. The cache stores raw ainfos, so wrap it to match the - # `AliasingWrapper` keys used by the overlap tracking. - if uniform_execution() - wrapped = key_ainfo isa AliasingWrapper ? key_ainfo : AliasingWrapper(key_ainfo) - haskey(state.ainfos_overlaps, wrapped) && - get_write_deps!(state, space, wrapped, write_num, syncdeps) - return - end - - # Buffer underlies wrapper arguments: find all tracked ainfos overlapping - # it. We reuse the lookup's interval-tree overlap search (which prunes most - # `will_alias` comparisons via bounding spans) rather than scanning every - # tracked ainfo. `intersect` requires its query ainfo to be registered, so - # we `push!` the buffer's aliasing in first; this is safe because the free - # loop is the final step of `distribute_tasks!`, after which the lookup is - # no longer consulted for scheduling. - buf_ainfo = AliasingWrapper(aliasing(remote_arg)) - buf_ainfo.inner isa NoAliasing && return - ainfo_idx = push!(state.ainfos_lookup, buf_ainfo) - for other_ainfo in intersect(state.ainfos_lookup, buf_ainfo; ainfo_idx) - other_ainfo === buf_ainfo && continue - 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)) - end - for (reader_task, reader_write_num) in get(state.ainfos_readers, other_ainfo, ()) - reader_write_num != write_num && push!(syncdeps, ThunkSyncdep(reader_task)) - end - end + # Every freeable buffer was allocated by `set_stored!`, which records its + # aliasing; only the user's originals (never freed) lack one. Returning + # quietly here would emit an `unsafe_free!` with an empty syncdep set -- the + # exact use-after-free this function exists to prevent -- so treat a missing + # entry as the bookkeeping bug it is. + buf_ainfo === nothing && + error("No recorded aliasing for Datadeps buffer in $space; cannot determine free syncdeps") + + # Buffer underlies wrapper arguments: sync with the owner and readers of + # every tracked ainfo overlapping it. + wrapped = buf_ainfo isa AliasingWrapper ? buf_ainfo : AliasingWrapper(buf_ainfo) + gather_overlap_syncdeps!(state, wrapped, write_num, syncdeps) + 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) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 2a1f5dd2a..94d7e4617 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -1488,13 +1488,14 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat write_num = typemax(Int) - 1 # Map each tracked slot chunk to its ainfos, exactly as flat - # `distribute_tasks!` does. A slot's object-cache *key* ainfo is computed - # from the source object, so it is frequently absent from `ainfo_arg` - # (which is keyed by destination-space ainfos). Keying the syncdep lookup - # on the key ainfo alone therefore yields an empty syncdep set, and the - # resulting `unsafe_free!` races the very tasks still reading that slot - # -- freeing e.g. the copy-in buffer for an `In(::DTask)` argument out - # from under its consumer. + # `distribute_tasks!` does. A buffer that is not itself a tracked slot + # is not covered by this map, and its object-cache *key* ainfo cannot + # stand in for it (that ainfo describes the source object, in the source + # space); `gather_free_syncdeps!` handles those from the buffer's own + # recorded destination-space aliasing instead. Getting this wrong yields + # an empty syncdep set, and an `unsafe_free!` that races the very tasks + # still reading that slot -- freeing e.g. the copy-in buffer for an + # `In(::DTask)` argument out from under its consumer. chunk_to_ainfos = IdDict{Any,Vector{AliasingWrapper}}() for (ainfo, remote_arg_ws) in state.ainfo_arg for remote_arg_w in remote_arg_ws @@ -1513,7 +1514,8 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat haskey(freed, remote_arg) && continue freed[remote_arg] = nothing free_syncdeps = Set{ThunkSyncdep}() - gather_free_syncdeps!(state, remote_space, ainfo, remote_arg, + buf_ainfo = stored_value_ainfo(obj_cache, remote_space, ainfo) + gather_free_syncdeps!(state, remote_space, buf_ainfo, remote_arg, write_num, chunk_to_ainfos, free_syncdeps) if registry !== nothing orig = get(state.remote_arg_to_original, remote_arg, nothing) diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 1b6bc815c..2db311336 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -231,7 +231,8 @@ function distribute_tasks!(queue::DataDepsTaskQueue) haskey(freed, remote_arg) && continue freed[remote_arg] = nothing free_syncdeps = Set{ThunkSyncdep}() - gather_free_syncdeps!(state, remote_space, ainfo, remote_arg, write_num, chunk_to_ainfos, free_syncdeps) + buf_ainfo = stored_value_ainfo(obj_cache, remote_space, ainfo) + gather_free_syncdeps!(state, remote_space, buf_ainfo, remote_arg, write_num, chunk_to_ainfos, free_syncdeps) # `tag` keeps the free task rank-uniform under MPI/uniform execution. Dagger.@spawn scope=free_scope syncdeps=free_syncdeps tag=datadeps_task_tag() Dagger.unsafe_free!(remote_arg) end diff --git a/test/datadeps.jl b/test/datadeps.jl index 1fb6562f3..edf99ad07 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -182,6 +182,33 @@ function test_move_rewrap_aliasing(obj, dest_space) Dagger.subtract_spans!(tree, [Dagger.ManyMemorySpan{N}((Dagger.LocalMemorySpan(ss), Dagger.LocalMemorySpan(ds))) for (ss, ds) in zip(src_spans, dst_spans)]) @test isempty(tree) + + # Every Datadeps-allocated buffer must carry its destination-space aliasing. + # `gather_free_syncdeps!` needs it to find the tasks still using a buffer + # that is not itself a tracked slot (the parent of a `view`, say), and it + # cannot recompute it: `aliasing` is a collective under SPMD execution. A + # buffer that lost this entry would be freed with an empty syncdep set, + # racing its readers. + store = Dagger.unwrap(dummy_backing) + dest_buffers = get(store.values, dest_space, nothing) + @test dest_buffers !== nothing + if dest_buffers !== nothing + for (key, buf) in dest_buffers + Dagger.is_original(store, dest_space, key) && continue + buf_ainfo = Dagger.stored_value_ainfo(store, dest_space, key) + @test buf_ainfo !== nothing + buf_ainfo === nothing && continue + # It must describe the buffer itself, not the source object it was + # copied from -- that is the distinction the free path turns on. + recorded = Dagger.memory_spans(buf_ainfo) + actual = Dagger.memory_spans(Dagger.aliasing(buf, identity)) + @test length(recorded) == length(actual) + for (r, a) in zip(recorded, actual) + @test Dagger.span_start(r) == Dagger.span_start(a) + @test Dagger.span_len(r) == Dagger.span_len(a) + end + end + end end @testset "Aliased Object Copying" begin nw = nprocs() @@ -1102,3 +1129,107 @@ end end end end + +# Regression tests for the free-syncdeps hole: `gather_free_syncdeps!` +# (src/datadeps/aliasing.jl) must never hand back an empty syncdep set while +# some task could still be reading or writing the buffer about to be freed. +# Previously, the region barrier at the end of `spawn_datadeps` hid this -- +# every task had already retired by the time the free loop ran -- but that +# cover disappears once frees can be deferred, so it needs to be correct on +# its own merits. +@testset "gather_free_syncdeps!" begin + @testset "overlap search finds a live writer" begin + # A buffer that is not itself a tracked slot (the shared parent behind + # a view, say) is not in `chunk_to_ainfos`, and its object-cache *key* + # ainfo describes the source object in the source space, so it can + # never overlap a destination-space ainfo. Driving the lookup from the + # buffer's own recorded aliasing must find the overlapping writer. + state = Dagger.DataDepsState() + + # A directly-tracked ainfo (as if some task took `In`/`Out` of a view + # over part of `A`) with a live writer recorded against it. + A = zeros(8) + view_ainfo = Dagger.AliasingWrapper(Dagger.aliasing(view(A, 1:4))) + push!(state.ainfos_lookup, view_ainfo) + state.ainfos_readers[view_ainfo] = Pair{Dagger.DTask,Int}[] + writer_task = Dagger.@spawn 1 + 1 + fetch(writer_task) + state.ainfos_owner[view_ainfo] = writer_task => 1 + + # The buffer being freed represents *all* of `A` (as the shared parent + # backing a view would), so it overlaps `view_ainfo` but is not + # content-identical to it -- `state.ainfos_overlaps` has no entry for + # it, matching the case that was silently mishandled. + buf_ainfo = Dagger.AliasingWrapper(Dagger.aliasing(A)) + @test !haskey(state.ainfos_overlaps, buf_ainfo) + + remote_arg = Dagger.tochunk(A) + space = Dagger.memory_space(remote_arg) + chunk_to_ainfos = IdDict{Any,Vector{Dagger.AliasingWrapper}}() + syncdeps = Set{Dagger.ThunkSyncdep}() + Dagger.gather_free_syncdeps!(state, space, buf_ainfo, remote_arg, 2, + chunk_to_ainfos, syncdeps) + + @test !isempty(syncdeps) + @test Dagger.ThunkSyncdep(writer_task) in syncdeps + end + + @testset "a buffer with no recorded aliasing is an error, not a silent free" begin + # Losing the recorded aliasing must not degrade to an empty syncdep + # set: that is exactly the use-after-free this path exists to prevent. + state = Dagger.DataDepsState() + remote_arg = Dagger.tochunk(zeros(4)) + space = Dagger.memory_space(remote_arg) + chunk_to_ainfos = IdDict{Any,Vector{Dagger.AliasingWrapper}}() + @test_throws ErrorException Dagger.gather_free_syncdeps!( + state, space, nothing, remote_arg, 2, chunk_to_ainfos, + Set{Dagger.ThunkSyncdep}()) + end + + @testset "buffer underlying shared views" begin + # End-to-end version of the same hole: two views into disjoint slices + # of one parent array, each written by a task on the same remote + # worker. The parent array is moved there exactly once (`move_rewrap` + # dedups children of both views onto the same object-cache entry), and + # that shared parent buffer is never itself a direct task argument -- + # it only "underlies" the two view arguments -- so its free syncdeps + # must be computed via the aliasing-overlap search, not a direct + # `chunk_to_ainfos` hit. Confirm every `unsafe_free!` task the region + # emits has a non-empty syncdep set. + if nprocs() >= 2 + w = workers()[1] + A = rand(4, 4) + v1 = view(A, 1:2, :) + v2 = view(A, 3:4, :) + old_assert = Dagger.DATADEPS_ASSERT_FREE_SYNCDEPS[] + Dagger.DATADEPS_ASSERT_FREE_SYNCDEPS[] = true + logs = try + with_logs() do + Dagger.spawn_datadeps() do + Dagger.@spawn scope=Dagger.scope(worker=w) mut_V!(Out(v1)) + Dagger.@spawn scope=Dagger.scope(worker=w) mut_V!(Out(v2)) + end + end + finally + Dagger.DATADEPS_ASSERT_FREE_SYNCDEPS[] = old_assert + end + @test all(==(1), A) + + free_tids = Int[] + for wl in keys(logs) + _logs = logs[wl] + for idx in 1:length(_logs[:core]) + core_log = _logs[:core][idx] + if core_log.category == :add_thunk && core_log.kind == :start && + _logs[:taskfuncnames][idx] == "unsafe_free!" + push!(free_tids, _logs[:id][idx].thunk_id::Int) + end + end + end + @test !isempty(free_tids) + for tid in free_tids + @test !isempty(taskdeps_for_task(logs, tid)) + end + end + end +end From c32ea2776c8e590b77ed819d5ec0b94cc7e3b18f Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 6 Sep 2026 07:15:54 -0700 Subject: [PATCH 08/10] datadeps/hierarchical: Record the free-syncdep constraint on cross-space planning `_hierarchical_copy_from_and_free!` derives each buffer's free syncdeps from a single partition's tracking. That is exact today, but only incidentally: a buffer is reachable solely from the partition whose cache allocated it, because every partition builds its own slots through its own `AliasedObjectCacheStore` and the one mechanism that hands a slot across a partition boundary (`_sync_incoming_ownership!`) is currently dead -- the parallel per-partition path requires a single exec memory space, which makes `build_shared_chunk_registry` return `nothing`. Re-enabling parallel planning across memory spaces breaks that premise: a consumer partition's boundary copy-to reads `owner_slot`, a buffer living in the producer's cache and absent from the consumer's, so per-partition tracking alone cannot see that reader. The free loop already anticipates this by also syncing on the chunk's final global writer, which transitively covers any reader that is an ancestor of it. Note the case that argument does not obviously reach -- a consumer partition that only reads and never commits ownership, whose boundary copy-to is not an ancestor of the final writer -- so whoever revives that path settles it rather than rediscovering it. Comments only; no behaviour change. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019Pp6hFZNHR44pLhDrR6aXq --- src/datadeps/hierarchical.jl | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 94d7e4617..bb997eb65 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -140,7 +140,33 @@ end # space-keyed currency tracking (`arg_current` / `arg_owner` / per-space slots), # which is what actually breaks. It also keys on the argument object, so a # `ChunkView` and the `Chunk` it views are tracked separately and their sharing -# is missed. Anyone reviving this must fix both before flipping the gate. +# is missed. +# +# A third area needs re-checking, in the free loop rather than in planning. +# `_hierarchical_copy_from_and_free!` walks one partition's object cache at a +# time and derives each buffer's free syncdeps from *that* partition's +# `ainfos_owner` / `ainfos_readers` / `ainfos_lookup`. That is exact today only +# because a buffer is reachable solely from the partition whose cache allocated +# it: every partition builds its own slots through its own +# `AliasedObjectCacheStore`, and the one mechanism that hands a slot across a +# partition boundary is `_sync_incoming_ownership!` below, which never runs. +# Re-enable it and a consumer partition's boundary copy-to reads `owner_slot`, a +# buffer living in the *producer's* cache and absent from the consumer's, so +# per-partition tracking alone cannot see that reader. (The `freed` dedup does +# not help: the buffer is in exactly one cache, so there is nothing to +# deduplicate.) +# +# The free loop already anticipates this by also syncing on `entry.owner_task`, +# the chunk's final global writer, which transitively covers every reader that +# is an ancestor of it -- each cross-partition hand-off makes the consumer's +# copy-to wait on the producer, and the next writer waits on that copy. What +# that argument does not obviously reach is a consumer partition that only +# *reads* the chunk and never commits ownership: its boundary copy-to is not an +# ancestor of the final writer, so nothing orders it against the free. Confirm +# that case (or extend the syncdeps to every partition the registry entry +# names) before flipping the gate. +# +# Anyone reviving this must settle all three before flipping the gate. # # Each partition schedules with its own `DataDepsState`, so `arg_owner` / # `arg_history` / physical slots are per-partition. When a backing chunk is @@ -1481,6 +1507,13 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat # several ainfos, or recorded in more than one partition's object cache, must # be freed exactly once. A double `unsafe_free!` is harmless on CPU (refcount # decrement) but releases device memory twice on GPU backends. + # + # N.B. The per-partition syncdep derivation below is exact only while a + # buffer is reachable from the single partition whose cache allocated it, + # which holds because `_sync_incoming_ownership!` (the only cross-partition + # slot hand-off) is currently dead code. See the constraint recorded with + # `SharedChunkRegistry` above before re-enabling parallel planning across + # memory spaces. freed = IdDict{Any,Nothing}() for pid in 1:n_partitions state = partition_states[pid] From f4235b830a4732cea7783543e897619d8580931e Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 6 Sep 2026 07:57:45 -0700 Subject: [PATCH 09/10] datadeps/hierarchical: Rank by written args separately during partition_dag Hoist the two affinity vectors out of the per-vertex loop. Splitting one affinity count into separate write/read counts turned one array allocation per task into two, on the planning path. Both are scratch and dead across iterations, so allocate them once per `partition_dag` and `fill!` them at the point of use. --- src/datadeps/hierarchical.jl | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index bb997eb65..43e2128f2 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -810,6 +810,11 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe if multi_owner owner_to_partition = Dict(o => i for (i, o) in enumerate(owners)) default_scope = DefaultScope() + # Hoisted out of the per-vertex loop and refilled: these are scratch, + # dead across iterations, and allocating them per task put two arrays + # per task on the planning path. + write_affinity = zeros(Int, n_owners) + read_affinity = zeros(Int, n_owners) for v in 1:n meta = task_metas[v] task_scope = @something(meta.pair.spec.options.compute_scope, meta.pair.spec.options.scope, default_scope) @@ -851,8 +856,8 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe # N.B. Argument *counts*, not byte counts: `datasize` of a chunk is # only known on its owning rank (see `datasize(::MPIRef)`), and this # decision must come out identical on every rank under SPMD. - write_affinity = zeros(Int, n_owners) - read_affinity = zeros(Int, n_owners) + fill!(write_affinity, 0) + fill!(read_affinity, 0) for dep in meta.deps arg_space = memory_space(dep.arg_w.arg) arg_oid = partition_affinity_id(arg_space) From f56d9b49818f6cc6e9171ab18fff22d68ee1e41d Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 6 Sep 2026 07:57:46 -0700 Subject: [PATCH 10/10] Sch/MPI: Cut fixed per-task overhead in placement and dispatch Stop paying twice for the inference-memo key, and stop boxing in the lookup. `mpi_execute_bcast_plan` asks two questions about one dispatch, per task, per rank. `cached_nothrow` built its own `Tuple{typeof(f), arg_types...}` key alongside the one `cached_return_type` had already built -- splatting arg types into a `Type` allocates. The key construction (and the `Union{}` check that decides whether a key exists at all) moves into `call_signature_key`, both memos take a precomputed key, and the MPI path builds it once. The memo bodies also moved from `lock(...) do` to `@safe_lock1`: the closure captured the key and reassigned its result across the get/infer/store sequence, putting it in a heap `Box` on a per-spawn path. `cached_return_type` gets the same treatment, since it shares the body and the caller. The two-argument forms are kept, so callers that ask only one question are unchanged. --- ext/MPIExt.jl | 9 +++++--- src/submission.jl | 54 +++++++++++++++++++++++++++++++++-------------- 2 files changed, 44 insertions(+), 19 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 7ec2ea697..e8831f8f3 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -2180,11 +2180,14 @@ mpi_result_space(result, proc::MPIProcessor) = # the task bodies it is deciding about. function mpi_execute_bcast_plan(f, args, proc::MPIProcessor) arg_types = map(chunktype, args) + # One key for both memo lookups: building it splats `arg_types` into a + # `Type` and allocates, and this runs per task on every rank. + key = Dagger.call_signature_key(f, arg_types) if !(proc.innerProc isa ThreadProc) - inferred = Dagger.cached_return_type(f, arg_types) + inferred = Dagger.cached_return_type(key, f, arg_types) return (; need_type_bcast=true, nothrow=false, inferred) end - inferred = Dagger.cached_return_type(f, arg_types) + inferred = Dagger.cached_return_type(key, f, arg_types) # `Nothing` is a concrete type and is deliberately NOT forced onto the # broadcast path: a `nothing` return (the common in-place / mutating task) # is fully known on every rank (all ranks stamp `Chunk{Nothing}` and @@ -2196,7 +2199,7 @@ function mpi_execute_bcast_plan(f, args, proc::MPIProcessor) if need_type_bcast return (; need_type_bcast=true, nothrow=false, inferred) end - nothrow = Dagger.cached_nothrow(f, arg_types) + nothrow = Dagger.cached_nothrow(key, f, arg_types) return (; need_type_bcast=false, nothrow, inferred) end diff --git a/src/submission.jl b/src/submission.jl index 2583decf6..aa8eb7b3f 100644 --- a/src/submission.jl +++ b/src/submission.jl @@ -421,16 +421,37 @@ end # repeated `(f, arg-types)` shape on the submission hot path. const RETURN_TYPE_CACHE = LockedObject(Dict{Type,Type}()) -function cached_return_type(@nospecialize(f), @nospecialize(arg_types::Tuple)) - # `Union{}` (the bottom type) is a legal inferred arg type — it arises when an - # upstream task is inferred never to return — but it cannot appear as a tuple - # field, so we can't form a `Tuple{...}` cache key for it. Such calls are rare, - # so infer them directly rather than caching. +""" + call_signature_key(f, arg_types) -> Type or nothing + +The cache key shared by [`cached_return_type`](@ref) and [`cached_nothrow`](@ref). +`nothing` when no key can be formed, which is the caller's cue to answer without +consulting a cache. + +`Union{}` (the bottom type) is a legal inferred arg type — it arises when an +upstream task is inferred never to return — but it cannot appear as a tuple +field, so no `Tuple{...}` key exists for it. Such calls are rare, so they are +answered directly rather than cached. + +Split out so a caller asking both questions about one dispatch builds it once: +forming it splats `arg_types` into a `Type` and allocates, and the MPI path runs +this per task, per rank (see `mpi_execute_bcast_plan`). +""" +function call_signature_key(@nospecialize(f), @nospecialize(arg_types::Tuple)) for T in arg_types - T === Union{} && return Base.promote_op(f, arg_types...) + T === Union{} && return nothing end - key = Tuple{typeof(f), arg_types...} - return lock(RETURN_TYPE_CACHE) do cache + return Tuple{typeof(f), arg_types...} +end + +cached_return_type(@nospecialize(f), @nospecialize(arg_types::Tuple)) = + cached_return_type(call_signature_key(f, arg_types), f, arg_types) +function cached_return_type(@nospecialize(key), @nospecialize(f), @nospecialize(arg_types::Tuple)) + key === nothing && return Base.promote_op(f, arg_types...) + # `@safe_lock1`, not `lock(...) do`: the closure captures `key`/`f`/`arg_types` + # and reassigns `rt` across the get/infer/store sequence, which puts it in a + # heap `Box` on a per-spawn path. + @safe_lock1 RETURN_TYPE_CACHE cache begin rt = get(cache, key, nothing) rt === nothing || return rt rt = Base.promote_op(f, arg_types...) @@ -451,14 +472,15 @@ end # grew to dominate exactly as chunks got smaller with more ranks. const NOTHROW_CACHE = LockedObject(Dict{Type,Bool}()) -function cached_nothrow(@nospecialize(f), @nospecialize(arg_types::Tuple)) - # `Union{}` cannot appear as a tuple field, so no cache key can be formed - # (see `cached_return_type`); assume the worst rather than infer. - for T in arg_types - T === Union{} && return false - end - key = Tuple{typeof(f), arg_types...} - return lock(NOTHROW_CACHE) do cache +cached_nothrow(@nospecialize(f), @nospecialize(arg_types::Tuple)) = + cached_nothrow(call_signature_key(f, arg_types), f, arg_types) +function cached_nothrow(@nospecialize(key), @nospecialize(f), @nospecialize(arg_types::Tuple)) + # No key means `Union{}` among the arg types (see `call_signature_key`); + # assume the worst rather than infer. + key === nothing && return false + # `@safe_lock1` rather than `lock(...) do`, for the reason in + # `cached_return_type`: the closure would box `nothrow`. + @safe_lock1 NOTHROW_CACHE cache begin nothrow = get(cache, key, nothing) nothrow === nothing || return nothrow nothrow = Core.Compiler.is_nothrow(Base.infer_effects(f, arg_types))