diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 490fd9965..5f9b2a108 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, { @@ -237,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 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 diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 9b28fdb1d..c2b359678 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,79 @@ 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(state.cond) do + 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) +# 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) do + notify(state.cond) 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) + 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) + return lock(state.cond) do + 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 +1300,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 +1346,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 +1359,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 +1394,15 @@ 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) do + notify(state.cond, + ConcurrencyViolationError("MPI broadcast relay stopped while a consumer was waiting"); + error=true) + end return nothing end @@ -1301,7 +1424,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 @@ -1714,11 +1837,70 @@ 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}, dims::Dims) where {T} = + Dagger.alloc_uninit(space.innerSpace, T, dims) + +# Allocate the slot on the destination rank without shipping the payload. Only +# the dimensions travel, which is the entire saving: the buffer's contents are +# established later by Datadeps' copy-to phase. +# +# 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, size(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(size(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 + dims = recv_yield(accel.comm, from_rank, tag)::Dims + return tochunk(Dagger.alloc_uninit(to_space, T_dest, dims), 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 @@ -1746,11 +1928,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 diff --git a/src/array/stencil.jl b/src/array/stencil.jl index a02bcac8f..6bbc86dca 100644 --- a/src/array/stencil.jl +++ b/src/array/stencil.jl @@ -798,7 +798,9 @@ 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) @@ -827,9 +829,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 +880,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 +914,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 +947,50 @@ 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. +""" +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}[] + for idx in CartesianIndices(src_chunks) + region_metadata, neighbor_chunks = select_neighborhood_info(src_chunks, idx, neigh_dist, boundary) + target_space = memory_space(write_chunks[idx]) + for i in eachindex(region_metadata) + region_code, is_boundary, boundary_dims, _ = region_metadata[i] + neighbor = neighbor_chunks[i] + memory_space(neighbor) == target_space && continue + region_metadata[i] = (region_code, is_boundary, boundary_dims, true) + task = Dagger.@spawn name="stencil_halo" scope=memory_space_scope(memory_space(neighbor)) 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 +1322,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 +1332,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/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index c60f4af4c..a80059189 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) @@ -1182,6 +1275,19 @@ 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 dimensions 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}, dims::Dims) where T + wid = root_worker_id(to_proc) + if wid == myid() + return tochunk(alloc_uninit(to_space, T, dims), to_proc, to_space) + end + return remotecall_fetch(wid, to_proc, to_space, T, dims) do to_proc, to_space, T, dims + return tochunk(alloc_uninit(to_space, T, dims), 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 +1344,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, size(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/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/src/memory-spaces.jl b/src/memory-spaces.jl index cfd5a8df6..f1c4af59f 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -53,6 +53,70 @@ 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. +""" +memory_space_scope(space::MemorySpace) = ExactScope(first(processors(space))) + +### 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_uninit(space::MemorySpace, ::Type{T}, dims::Dims) -> T + +Allocate an uninitialized `T` of shape `dims` in `space`. 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. +""" +slot_may_be_uninit(from_space::MemorySpace, to_space::MemorySpace, ::Type{T}) where {T} = + from_space != to_space && T <: DenseArray && can_alloc_uninit(to_space, T) + ### In-place Data Movement unwrap(x::Chunk) = unwrap(x.handle) 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 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