From 3a88db02bf6637d43ddaddf72c39db710315fa79 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 17 Aug 2026 15:34:30 -0700 Subject: [PATCH] CI/benchmarks: Add Distributed and MPI benchmarks --- .github/workflows/CI.yml | 197 +++++++++++++++++++++++++++++++++--- benchmark/benchmarks.jl | 92 ++++++++++++++++- benchmark/ci.jl | 6 ++ benchmark/common.jl | 33 +++++- benchmark/suites/array.jl | 47 ++++----- benchmark/suites/linalg.jl | 87 ++++++++-------- benchmark/suites/stencil.jl | 105 +++++++++---------- benchmark/worker_mpi.jl | 154 ++++++++++++++++++++++++++++ 8 files changed, 582 insertions(+), 139 deletions(-) create mode 100644 benchmark/worker_mpi.jl diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 490fd9965..f81b1098e 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -153,9 +153,6 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 120 if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && !contains(github.event.head_commit.message, '[skip benchmarks]') }} - permissions: - contents: read - pull-requests: write env: # `benchmark/ci.jl` compares the working tree ("dirty") against this base # revision and fails the job if any benchmark regresses past the threshold. @@ -167,6 +164,14 @@ jobs: # 4 vCPU / ~16 GB, and noisier than dedicated benchmarking hardware. The # orchestrator is OOM-robust, but smaller scales keep wall-clock bounded. BENCHMARK_SCALE: '[256, 1024]' + # Tile size matters as much as N here, and one tile is not enough. At 512 + # alone, N=256 is a *single* chunk (`square_block` returns N when N <= tile) + # and N=1024 is only 2x2, so anything touching data distribution or halo + # exchange is barely exercised. But small tiles are not simply "better": + # tile size sets the ratio of per-tile task overhead to bytes moved per + # transfer, and a data-movement change can be a large win at one tile size + # and a regression at the other. Measure both regimes. + BENCHMARK_BLOCKSIZE: '[128, 512]' BENCHMARK_SECONDS: '3' BENCHMARK_SAMPLES: '5' steps: @@ -198,30 +203,192 @@ jobs: if: always() uses: actions/upload-artifact@v7 with: - name: benchmark-results + name: benchmark-results-default + path: ${{ env.BENCHMARK_OUTPUT_DIR }}/ + if-no-files-found: ignore + # PR comment posting is handled by the `benchmarks-report` job below, + # once this job and its Distributed/MPI siblings have all finished -- + # a single writer avoids three jobs racing to read-modify-write the + # same PR comment. + + benchmarks-distributed: + name: Benchmarks (Distributed, vs master) + runs-on: ubuntu-latest + timeout-minutes: 120 + if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && !contains(github.event.head_commit.message, '[skip benchmarks]') }} + env: + BENCHMARK_BASE_REV: master + # Multiple OS processes on a shared runner are noisier than the + # single-process default job, so allow more slack before flagging a + # regression. + BENCHMARK_REGRESSION_THRESHOLD: '0.35' + # Parallelism comes from BENCHMARK_PROCS (processes), not threads: 3 + # extra Distributed workers + the driver = 4 processes x 1 thread, + # matching the 4 vCPUs on ubuntu-latest. + BENCHMARK_CI_THREADS: '1' + BENCHMARK_PROCS: '3:1' + BENCHMARK_OUTPUT_DIR: benchmark_results_distributed + BENCHMARK_SCALE: '[256, 1024]' + # See the `benchmarks` job: both tiles, so neither the many-small-chunks + # nor the few-large-chunks regime goes unmeasured. + BENCHMARK_BLOCKSIZE: '[128, 512]' + BENCHMARK_SECONDS: '3' + BENCHMARK_SAMPLES: '5' + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + - name: Fetch base revision + run: git fetch -f origin "$BENCHMARK_BASE_REV:$BENCHMARK_BASE_REV" || git fetch origin "$BENCHMARK_BASE_REV" || true + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + arch: x64 + - uses: julia-actions/cache@v3 + - name: Run benchmarks (current vs master, Distributed) + run: julia --color=yes benchmark/ci.jl + - name: Write job summary + if: always() + run: | + if [ -f "${BENCHMARK_OUTPUT_DIR}/report.md" ]; then + grep -v 'artifact://' "${BENCHMARK_OUTPUT_DIR}/report.md" >> "$GITHUB_STEP_SUMMARY" + else + echo "No benchmark report was produced." >> "$GITHUB_STEP_SUMMARY" + fi + - name: Upload benchmark artifacts + if: always() + uses: actions/upload-artifact@v7 + with: + name: benchmark-results-distributed path: ${{ env.BENCHMARK_OUTPUT_DIR }}/ if-no-files-found: ignore - - name: Comment results on PR + + benchmarks-mpi: + name: Benchmarks (MPI, vs master) + runs-on: ubuntu-latest + timeout-minutes: 120 + if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && !contains(github.event.head_commit.message, '[skip benchmarks]') }} + env: + BENCHMARK_BASE_REV: master + # See benchmarks-distributed: shared-runner multi-process noise. + BENCHMARK_REGRESSION_THRESHOLD: '0.35' + # 4 MPI ranks x 1 thread, matching the 4 vCPUs on ubuntu-latest. + # BENCHMARK_MPI_RANKS switches benchmark/benchmarks.jl to spawn the + # worker under mpiexec and use worker_mpi.jl (SPMD) automatically. + BENCHMARK_CI_THREADS: '1' + BENCHMARK_MPI_RANKS: '4' + BENCHMARK_OUTPUT_DIR: benchmark_results_mpi + BENCHMARK_SCALE: '[256, 1024]' + # See the `benchmarks` job: both tiles, so neither the many-small-chunks + # nor the few-large-chunks regime goes unmeasured. + BENCHMARK_BLOCKSIZE: '[128, 512]' + BENCHMARK_SECONDS: '3' + BENCHMARK_SAMPLES: '5' + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + - name: Fetch base revision + run: git fetch -f origin "$BENCHMARK_BASE_REV:$BENCHMARK_BASE_REV" || git fetch origin "$BENCHMARK_BASE_REV" || true + - uses: julia-actions/setup-julia@v3 + with: + version: '1' + arch: x64 + - uses: julia-actions/cache@v3 + - name: Run benchmarks (current vs master, MPI) + run: julia --color=yes benchmark/ci.jl + - name: Write job summary + if: always() + run: | + if [ -f "${BENCHMARK_OUTPUT_DIR}/report.md" ]; then + grep -v 'artifact://' "${BENCHMARK_OUTPUT_DIR}/report.md" >> "$GITHUB_STEP_SUMMARY" + else + echo "No benchmark report was produced." >> "$GITHUB_STEP_SUMMARY" + fi + - name: Upload benchmark artifacts + if: always() + uses: actions/upload-artifact@v7 + with: + name: benchmark-results-mpi + path: ${{ env.BENCHMARK_OUTPUT_DIR }}/ + if-no-files-found: ignore + + benchmarks-report: + name: Benchmarks (PR comment) + needs: [benchmarks, benchmarks-distributed, benchmarks-mpi] + runs-on: ubuntu-latest + timeout-minutes: 10 + if: ${{ always() && github.event_name == 'pull_request' }} + permissions: + contents: read + pull-requests: write + steps: + - name: Download default benchmark results + if: always() + continue-on-error: true + uses: actions/download-artifact@v7 + with: + name: benchmark-results-default + path: results-default + - name: Download Distributed benchmark results + if: always() + continue-on-error: true + uses: actions/download-artifact@v7 + with: + name: benchmark-results-distributed + path: results-distributed + - name: Download MPI benchmark results + if: always() + continue-on-error: true + uses: actions/download-artifact@v7 + with: + name: benchmark-results-mpi + path: results-mpi + - name: Comment combined results on PR # Best-effort: fork PRs get a read-only token, so the comment will 403; # don't fail the job over it. - if: ${{ always() && github.event_name == 'pull_request' }} + if: always() continue-on-error: true uses: actions/github-script@v9 with: script: | const fs = require('fs'); - const reportPath = `${process.env.BENCHMARK_OUTPUT_DIR}/report.md`; - if (!fs.existsSync(reportPath)) { - core.info('No report.md found; skipping PR comment.'); - return; + + function section(dir) { + const reportPath = `${dir}/report.md`; + if (!fs.existsSync(reportPath)) { + return null; + } + return fs.readFileSync(reportPath, 'utf8') + .split('\n') + .filter(line => !line.includes('artifact://')) + .join('\n'); } + + function collapsible(summary, content) { + const body = content === null + ? '_Results unavailable (job did not produce a report)._' + : content; + return `
\n${summary}\n\n${body}\n\n
`; + } + const marker = ''; - let body = fs.readFileSync(reportPath, 'utf8') - .split('\n') - .filter(line => !line.includes('artifact://')) - .join('\n'); const runUrl = `${process.env.GITHUB_SERVER_URL}/${process.env.GITHUB_REPOSITORY}/actions/runs/${process.env.GITHUB_RUN_ID}`; - body = `${marker}\n${body}\n\n[Full results and plots](${runUrl}) (download the \`benchmark-results\` artifact).`; + + const parts = [ + marker, + '## Dagger benchmarks: `dirty` vs `master`', + '', + collapsible('Multi-threaded benchmarks (4 threads)', section('results-default')), + '', + collapsible('Distributed benchmarks (4 processes)', section('results-distributed')), + '', + collapsible('MPI benchmarks (4 ranks)', section('results-mpi')), + '', + `[Full results and plots](${runUrl}) (download the \`benchmark-results-*\` artifacts).`, + ]; + const body = parts.join('\n'); + const { owner, repo } = context.repo; const issue_number = context.issue.number; const comments = await github.paginate(github.rest.issues.listComments, { diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 8ba874af5..1f5f45273 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -70,8 +70,17 @@ # estimated peak allocation may use before a size is skipped. Defaults to # "0.2" (conservative). Raise it to allow larger sizes, lower it for more # headroom. -# - BENCHMARK_BLOCKSIZE: Target square tile size (elements per side) for dense -# suites. Defaults to "512". +# - BENCHMARK_BLOCKSIZE: Target square tile size(s) (elements per side) for dense +# suites, given like BENCHMARK_SCALE: an integer or an iterable of integers +# (e.g. "512", "[128, 512]"). Each requested tile produces its own benchmark +# group per N, keyed by the resulting block; tiles that collapse to the same +# block at a given N (both 512 and 256 give 256 at N=256) are emitted once. +# Defaults to "512". +# +# Sweeping this matters more than it looks: tile size sets the ratio between +# per-tile task overhead and bytes moved per transfer, so a change to data +# movement can be a large win at one tile size and a regression at another. +# A single tile can only ever see one of those regimes. # - BENCHMARK_SPARSE_BLOCKS: Number of blocks per dimension for sparse/banded # operators (keeps tile counts bounded at large N). Defaults to "16". # - BENCHMARK_PROCS: Worker/thread topology, as "numprocs:numthreads". This @@ -80,6 +89,15 @@ # runs on the worker process, which uses the orchestrator's thread count). # - BENCHMARK_REMOTES: Remote hosts on which to start workers, in the format # accepted by `Distributed.addprocs` (colon-separated). Optional. +# - BENCHMARK_MPI_RANKS: When set to a positive integer, runs the benchmark +# worker under `mpiexec -n ` (via MPI.jl's bundled mpiexec, so no +# system MPI install is required) instead of as a plain subprocess, and the +# worker calls `Dagger.accelerate!(:mpi)`. Because MPI ranks are SPMD, this +# also switches the default worker script from `worker.jl` to +# `worker_mpi.jl` (a single flat pass over every benchmark, run identically +# on every rank, instead of `worker.jl`'s incremental per-benchmark +# request/response protocol -- see `run_all_mpi` vs `run_all_external`). +# Defaults to "0" (disabled; plain subprocess, worker.jl). # - BENCHMARK_SECONDS: Time budget (seconds) per benchmark. Defaults to "30". # - BENCHMARK_SAMPLES: Max samples per benchmark. Defaults to "5". # - BENCHMARK_PROC_TIMEOUT: Wall-clock seconds to wait for a single benchmark @@ -121,7 +139,15 @@ BenchmarkTools.tune!(b::PrecomputedTrial, args...; kwargs...) = b const WORKDIR = let d = get(ENV, "BENCHMARK_WORKDIR", "") isempty(d) ? mktempdir() : (mkpath(d); abspath(d)) end -const WORKER_SCRIPT = get(ENV, "BENCHMARK_WORKER_SCRIPT", joinpath(@__DIR__, "worker.jl")) +const MPI_RANKS = parse(Int, get(ENV, "BENCHMARK_MPI_RANKS", "0")) +if MPI_RANKS > 0 + using MPI +end +# Defaults to worker_mpi.jl (the SPMD worker) when BENCHMARK_MPI_RANKS is set, +# otherwise worker.jl (the incremental request/response worker); overridable +# either way via BENCHMARK_WORKER_SCRIPT. +const WORKER_SCRIPT = get(ENV, "BENCHMARK_WORKER_SCRIPT", + joinpath(@__DIR__, MPI_RANKS > 0 ? "worker_mpi.jl" : "worker.jl")) const JULIA_BIN = first(Base.julia_cmd().exec) const NTHREADS = Threads.nthreads() const PROC_TIMEOUT = parse(Float64, get(ENV, "BENCHMARK_PROC_TIMEOUT", "3600")) @@ -143,7 +169,23 @@ function spawn_worker() rm(joinpath(WORKDIR, "ready"); force=true) rm(joinpath(WORKDIR, "request.json"); force=true) rm(joinpath(WORKDIR, "request.json.tmp"); force=true) - return run(worker_cmd(); wait=false) + rm(joinpath(WORKDIR, "done"); force=true) + rm(joinpath(WORKDIR, "results_mpi_manifest.json"); force=true) + if MPI_RANKS > 0 + # `MPI.mpiexec` only sets up the environment mpiexec needs (library + # paths for the bundled MPICH_jll, etc.) for the dynamic extent of the + # callback (mirrors test/run_mpi.jl) -- so the actual `run` call must + # happen *inside* the `do` block. Extracting the bare executable via + # e.g. `MPI.mpiexec(identity)` and running it afterwards drops that + # environment and the spawned process dies immediately. + local proc + MPI.mpiexec() do mpiexec + proc = run(`$mpiexec -n $MPI_RANKS $(worker_cmd())`; wait=false) + end + return proc + else + return run(worker_cmd(); wait=false) + end end function wait_ready(proc; timeout=600) @@ -246,10 +288,50 @@ function run_all_external() return results end +# MPI counterpart to `run_all_external`: worker_mpi.jl is SPMD (every rank +# must run the identical benchmark at the same time), so there is no +# per-benchmark request/response round trip here -- the worker runs its +# entire flat pass unattended and signals completion via a `done` sentinel, +# which this just waits for (bounded by `PROC_TIMEOUT`, same as a single +# `await_response` would be) before loading whatever it produced. +function run_all_mpi() + results = Dict{Vector{String},BenchmarkTools.Trial}() + + proc = spawn_worker() + donepath = joinpath(WORKDIR, "done") + t0 = time() + while !isfile(donepath) + if !process_running(proc) + @error "MPI benchmark worker exited before signaling completion; producing partial results." + break + end + if time() - t0 > PROC_TIMEOUT + @warn "MPI benchmark run exceeded $(PROC_TIMEOUT)s; killing and returning partial results." + kill(proc) + break + end + sleep(POLL) + end + if process_running(proc) + try; wait(proc); catch; end + end + + manifestpath = joinpath(WORKDIR, "results_mpi_manifest.json") + isfile(manifestpath) || return results + manifest = JSON3.read(read(manifestpath, String)) + for entry in manifest + kp = String[string(k) for k in entry.keypath] + resultpath = joinpath(WORKDIR, String(entry.file)) + isfile(resultpath) || continue + results[kp] = BenchmarkTools.load(resultpath)[1] + end + return results +end + # --- Assemble SUITE from the externally-measured trials --------------------- const SUITE = BenchmarkGroup() -for (keypath, trial) in run_all_external() +for (keypath, trial) in (MPI_RANKS > 0 ? run_all_mpi() : run_all_external()) SUITE[keypath] = PrecomputedTrial(trial) end diff --git a/benchmark/ci.jl b/benchmark/ci.jl index 8dea14f4c..6af535478 100644 --- a/benchmark/ci.jl +++ b/benchmark/ci.jl @@ -55,6 +55,12 @@ const EXTRA_PKGS = String[ # with Dagger; the orchestrator/worker need it for their file-based IPC. #"JSON3", ] +# MPI is only needed (and only installed) when the caller requests an MPI +# benchmark run (see benchmark/benchmarks.jl's BENCHMARK_MPI_RANKS), so the +# plain/Distributed CI runs don't pay for pulling in MPICH_jll. +if get(ENV, "BENCHMARK_MPI_RANKS", "0") != "0" + push!(EXTRA_PKGS, "MPI") +end mkpath(OUTPUT_DIR) diff --git a/benchmark/common.jl b/benchmark/common.jl index f5f65e152..e988be090 100644 --- a/benchmark/common.jl +++ b/benchmark/common.jl @@ -64,13 +64,44 @@ sparse_bytes(N; nmats=1, density=0.0, T=Float64, nvecs=2) = """Whether an estimated allocation of `bytes` fits the conservative budget.""" fits_budget(bytes) = bytes <= MEM_BUDGET +"""Target square tile sizes (elements per side) for the dense suites. Given as a +Julia expression evaluating to an integer or an iterable of integers, exactly +like `BENCHMARK_SCALE`. + +Tile size is not a detail that can be pinned to one value and forgotten: it sets +the ratio between per-tile task overhead and bytes moved per transfer, and a +change to data movement can be a large win at one tile size and a regression at +another. Sweeping it is what lets a comparison tell those two regimes apart.""" +const blocksizes = let s = get(ENV, "BENCHMARK_BLOCKSIZE", "") + if isempty(s) + [512] + else + parsed = eval(Meta.parse(s)) + parsed isa Integer ? [parsed] : collect(parsed) + end +end + """Square tile size (elements per side) for a dense N×N matrix, targeting `tile` elements per side but never exceeding N.""" -function square_block(N; tile=parse(Int, get(ENV, "BENCHMARK_BLOCKSIZE", "512"))) +function square_block(N; tile=first(blocksizes)) N <= tile && return N return cld(N, cld(N, tile)) end +"""The distinct tile sizes `square_block` actually yields at dimension `N`. + +Different requested tiles collapse to the same block once clamped to `N` (at +N=256 both 512 and 256 give 256), and the suites key their benchmark groups by +the resulting block, so emitting a group per *requested* tile would collide.""" +function blocks_for(N) + seen = Int[] + for tile in blocksizes + b = square_block(N; tile) + b in seen || push!(seen, b) + end + return seen +end + """Block size for sparse/banded N×N operators, keeping the per-dimension block count bounded (so we don't create an enormous number of mostly-empty tiles).""" function banded_block(N; maxblocks=parse(Int, get(ENV, "BENCHMARK_SPARSE_BLOCKS", "16"))) diff --git a/benchmark/suites/array.jl b/benchmark/suites/array.jl index a881fbdbc..7d93c4889 100644 --- a/benchmark/suites/array.jl +++ b/benchmark/suites/array.jl @@ -17,37 +17,38 @@ function array_suite(ctx; method, accels) for N in scales # Elementwise ops hold at most the input plus a same-size result. fits_budget(dense_bytes(N; nmats=2, T=T)) || continue - b = square_block(N) - sub = BenchmarkGroup() + for b in blocks_for(N) + sub = BenchmarkGroup() - sub["alloc (rand)"] = @benchmarkable(wait(rand(Blocks($b, $b), $T, $N, $N)), - teardown = (@everywhere GC.gc())) + sub["alloc (rand)"] = @benchmarkable(wait(rand(Blocks($b, $b), $T, $N, $N)), + teardown = (@everywhere GC.gc())) - sub["broadcast (X .+ 1)"] = @benchmarkable(wait(X .+ 1), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["broadcast (X .+ 1)"] = @benchmarkable(wait(X .+ 1), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["add (X + X)"] = @benchmarkable(wait(X + X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["add (X + X)"] = @benchmarkable(wait(X + X), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["map (sin.(X))"] = @benchmarkable(wait(sin.(X)), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["map (sin.(X))"] = @benchmarkable(wait(sin.(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["transpose (permutedims)"] = @benchmarkable(wait(permutedims(X)), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["transpose (permutedims)"] = @benchmarkable(wait(permutedims(X)), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["reduce (sum)"] = @benchmarkable(sum(X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["reduce (sum)"] = @benchmarkable(sum(X), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + teardown = (X = nothing; @everywhere GC.gc())) - sub["norm"] = @benchmarkable(norm(X), - setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), - teardown = (X = nothing; @everywhere GC.gc())) + sub["norm"] = @benchmarkable(norm(X), + setup = (X = rand(Blocks($b, $b), $T, $N, $N); wait(X)), + teardown = (X = nothing; @everywhere GC.gc())) - suite["N=$N (block $b)"] = sub + suite["N=$N (block $b)"] = sub + end end suite diff --git a/benchmark/suites/linalg.jl b/benchmark/suites/linalg.jl index 3ddd000db..fd42a530d 100644 --- a/benchmark/suites/linalg.jl +++ b/benchmark/suites/linalg.jl @@ -23,58 +23,59 @@ function linalg_suite(ctx; method, accels) suite = BenchmarkGroup() for N in scales - b = square_block(N) - sub = BenchmarkGroup() + for b in blocks_for(N) + sub = BenchmarkGroup() - # gemm needs A and the result resident; factorizations copy internally. - if fits_budget(dense_bytes(N; nmats=3, T=T)) - sub["matmul (A*A)"] = @benchmarkable(wait(A * A), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; @everywhere GC.gc())) + # gemm needs A and the result resident; factorizations copy internally. + if fits_budget(dense_bytes(N; nmats=3, T=T)) + sub["matmul (A*A)"] = @benchmarkable(wait(A * A), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; @everywhere GC.gc())) - sub["syrk (A'*A)"] = @benchmarkable(wait(A' * A), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; @everywhere GC.gc())) + sub["syrk (A'*A)"] = @benchmarkable(wait(A' * A), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; @everywhere GC.gc())) - sub["lu"] = @benchmarkable(wait(lu(A, RowMaximum()).factors), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; @everywhere GC.gc())) + sub["lu"] = @benchmarkable(wait(lu(A, RowMaximum()).factors), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; @everywhere GC.gc())) - sub["qr"] = @benchmarkable(wait(qr(A).factors), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; @everywhere GC.gc())) + sub["qr"] = @benchmarkable(wait(qr(A).factors), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; @everywhere GC.gc())) - sub["solve (A\\b via lu)"] = @benchmarkable(wait(lu(A, RowMaximum()) \ b), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); - b = rand(Blocks($b), $T, $N); wait(A)), - teardown = (A = nothing; b = nothing; @everywhere GC.gc())) - end + sub["solve (A\\b via lu)"] = @benchmarkable(wait(lu(A, RowMaximum()) \ b), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); + b = rand(Blocks($b), $T, $N); wait(A)), + teardown = (A = nothing; b = nothing; @everywhere GC.gc())) + end - # Cholesky additionally holds the SPD-construction temporary. - if fits_budget(dense_bytes(N; nmats=4, T=T)) - sub["cholesky"] = @benchmarkable(wait(cholesky(A).factors), - setup = (A = _spd($T, $N, $b)), - teardown = (A = nothing; @everywhere GC.gc())) - end + # Cholesky additionally holds the SPD-construction temporary. + if fits_budget(dense_bytes(N; nmats=4, T=T)) + sub["cholesky"] = @benchmarkable(wait(cholesky(A).factors), + setup = (A = _spd($T, $N, $b)), + teardown = (A = nothing; @everywhere GC.gc())) + end - # SVD (tiled one-sided Jacobi) additionally holds the internally-copied - # scratch matrix, the accumulated V factor, and (across multiple - # workers) a restaged copy of A, on top of the resident input. - if fits_budget(dense_bytes(N; nmats=5, T=T)) - sub["svd"] = @benchmarkable(wait(svd(A).U), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; @everywhere GC.gc())) - end + # SVD (tiled one-sided Jacobi) additionally holds the internally-copied + # scratch matrix, the accumulated V factor, and (across multiple + # workers) a restaged copy of A, on top of the resident input. + if fits_budget(dense_bytes(N; nmats=5, T=T)) + sub["svd"] = @benchmarkable(wait(svd(A).U), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; @everywhere GC.gc())) + end - # gemv is cheap (one matrix + two vectors). - if fits_budget(dense_bytes(N; nmats=1, T=T)) - sub["matvec (A*x)"] = @benchmarkable(wait(A * x), - setup = (A = rand(Blocks($b, $b), $T, $N, $N); - x = rand(Blocks($b), $T, $N); wait(A)), - teardown = (A = nothing; x = nothing; @everywhere GC.gc())) - end + # gemv is cheap (one matrix + two vectors). + if fits_budget(dense_bytes(N; nmats=1, T=T)) + sub["matvec (A*x)"] = @benchmarkable(wait(A * x), + setup = (A = rand(Blocks($b, $b), $T, $N, $N); + x = rand(Blocks($b), $T, $N); wait(A)), + teardown = (A = nothing; x = nothing; @everywhere GC.gc())) + end - isempty(sub) || (suite["N=$N (block $b)"] = sub) + isempty(sub) || (suite["N=$N (block $b)"] = sub) + end end suite diff --git a/benchmark/suites/stencil.jl b/benchmark/suites/stencil.jl index 021430863..764c58b88 100644 --- a/benchmark/suites/stencil.jl +++ b/benchmark/suites/stencil.jl @@ -105,63 +105,64 @@ function stencil_suite(ctx; method, accels) end for N in scales - b = square_block(N) - sub = BenchmarkGroup() - - # In-place stencils hold at most the input plus a same-size result. - if fits_budget(dense_bytes(N; nmats=2, T=T)) - if assign_ok - sub["assign (const)"] = @benchmarkable(stencil_assign!(B, $T), - setup = (B = zeros(Blocks($b, $b), $T, $N, $N); wait(B)), - teardown = (B = nothing; @everywhere GC.gc())) + for b in blocks_for(N) + sub = BenchmarkGroup() + + # In-place stencils hold at most the input plus a same-size result. + if fits_budget(dense_bytes(N; nmats=2, T=T)) + if assign_ok + sub["assign (const)"] = @benchmarkable(stencil_assign!(B, $T), + setup = (B = zeros(Blocks($b, $b), $T, $N, $N); wait(B)), + teardown = (B = nothing; @everywhere GC.gc())) + end + + if wrap_ok + sub["neighbors (Wrap)"] = @benchmarkable(stencil_neighbors_wrap!(A, B), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + end + + if pad_ok + sub["neighbors (Pad)"] = @benchmarkable(stencil_neighbors_pad!(A, B), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + end + + if clamp_ok + sub["neighbors (Clamp)"] = @benchmarkable(stencil_neighbors_clamp!(A, B), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + end + + if reflect_ok + sub["neighbors (Reflect)"] = @benchmarkable(stencil_neighbors_reflect!(A, B), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + end + + if update_ok + sub["update (+)"] = @benchmarkable(stencil_update_plus!(A, B), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + end + + if multi_ok + sub["multi-expr"] = @benchmarkable(stencil_multi_expr!(A, B, $T), + setup = (A = zeros(Blocks($b, $b), $T, $N, $N); + B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + end end - if wrap_ok - sub["neighbors (Wrap)"] = @benchmarkable(stencil_neighbors_wrap!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) + # Functional allocation syntax also materializes an output DArray. + if alloc_ok && fits_budget(dense_bytes(N; nmats=3, T=T)) + sub["alloc (neighbors Wrap)"] = @benchmarkable(wait(stencil_alloc_neighbors_wrap(A)), + setup = (A = ones(Blocks($b, $b), $T, $N, $N); wait(A)), + teardown = (A = nothing; @everywhere GC.gc())) end - if pad_ok - sub["neighbors (Pad)"] = @benchmarkable(stencil_neighbors_pad!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) - end - - if clamp_ok - sub["neighbors (Clamp)"] = @benchmarkable(stencil_neighbors_clamp!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) - end - - if reflect_ok - sub["neighbors (Reflect)"] = @benchmarkable(stencil_neighbors_reflect!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) - end - - if update_ok - sub["update (+)"] = @benchmarkable(stencil_update_plus!(A, B), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) - end - - if multi_ok - sub["multi-expr"] = @benchmarkable(stencil_multi_expr!(A, B, $T), - setup = (A = zeros(Blocks($b, $b), $T, $N, $N); - B = zeros(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; B = nothing; @everywhere GC.gc())) - end + isempty(sub) || (suite["N=$N (block $b)"] = sub) end - - # Functional allocation syntax also materializes an output DArray. - if alloc_ok && fits_budget(dense_bytes(N; nmats=3, T=T)) - sub["alloc (neighbors Wrap)"] = @benchmarkable(wait(stencil_alloc_neighbors_wrap(A)), - setup = (A = ones(Blocks($b, $b), $T, $N, $N); wait(A)), - teardown = (A = nothing; @everywhere GC.gc())) - end - - isempty(sub) || (suite["N=$N (block $b)"] = sub) end suite diff --git a/benchmark/worker_mpi.jl b/benchmark/worker_mpi.jl new file mode 100644 index 000000000..de0cbaead --- /dev/null +++ b/benchmark/worker_mpi.jl @@ -0,0 +1,154 @@ +# Full-SPMD, single-shot MPI benchmark worker. +# +# Spawned by the orchestrator (benchmarks.jl) via `mpiexec -n N julia +# worker_mpi.jl workdir` when BENCHMARK_MPI_RANKS is set. Unlike worker.jl (a +# plain OS subprocess driven one benchmark at a time over a file-based +# request/response protocol), MPI ranks are SPMD: every rank must call the +# same collective Dagger operations at the same time, so there is no +# orchestrator-in-the-loop per benchmark here. Instead every rank builds the +# identical benchmark suite (mirroring test/mpi.jl's bootstrap) and runs a +# single flat pass over every benchmark in lockstep; rank 0 alone talks to +# the filesystem, recording each Trial and finishing with a manifest plus a +# `done` sentinel that the orchestrator polls for. +# +# Protocol (rank 0 only, all writes atomic via tmp-then-rename): +# - `result_mpi_.json`: a BenchmarkTools.save'd Trial for the i'th leaf +# that completed (in sorted-keypath order). +# - `results_mpi_manifest.json`: [{keypath, file}, ...] pairing each +# completed leaf's key path to its result file. +# - `done`: written last, once every benchmark has been attempted. +# +# Note this worker does not have worker.jl's per-scale OOM isolation: a +# caught `OutOfMemoryError` aborts the *whole* MPI job (`MPI.Abort`) rather +# than exiting only the local rank and letting the orchestrator retry smaller +# scales, since a partially-alive rank set can't make collective progress. +# Other exceptions are caught per-rank and just skip that one benchmark: +# since every rank runs the identical deterministic computation, ranks are +# expected to fail together at the same call, so no cross-rank coordination +# is needed for the common case. + +using BenchmarkTools +using Distributed +using Dates, Random, Statistics, LinearAlgebra, InteractiveUtils +import JSON3 +using MPI + +include(joinpath(@__DIR__, "common.jl")) + +const WORKDIR = abspath(ARGS[1]) + +function atomic_write(path, data) + tmp = path * ".tmp" + open(io -> write(io, data), tmp, "w") + mv(tmp, path; force=true) + return nothing +end + +using Dagger +Dagger.accelerate!(:mpi) +Dagger.check_uniformity!(true) +const comm = MPI.COMM_WORLD +const rank = MPI.Comm_rank(comm) + +# --- Load acceleration backends (only if requested) ------------------------- +# No Distributed workers exist under MPI (each rank is its own OS process), +# so a plain `using` suffices here (worker.jl uses `@everywhere using` because +# it may have addprocs'd extra Distributed workers). + +for accel in accelerations + if accel == "cuda" + try + using DaggerGPU, CUDA + catch err + error("Failed to load CUDA acceleration; ensure DaggerGPU and CUDA " * + "are available (e.g. `benchpkg ... -a DaggerGPU,CUDA`)\n$err") + end + elseif accel == "amdgpu" + try + using DaggerGPU, AMDGPU + catch err + error("Failed to load AMDGPU acceleration; ensure DaggerGPU and " * + "AMDGPU are available (e.g. `benchpkg ... -a DaggerGPU,AMDGPU`)\n$err") + end + else + error("Unknown acceleration: $accel") + end +end + +# --- Build the benchmark suites --------------------------------------------- +# Every rank builds the identical suite (deterministic; no per-rank +# branching), so BenchmarkTools.leaves(SUITE) enumerates the same benchmarks +# everywhere. `ctx` is accepted-but-unused by every suite (verified: array, +# linalg, sparse, stencil), so it's passed as `nothing` here rather than a +# Distributed-flavored `Context()` -- test/mpi.jl doesn't set +# `Dagger.Sch.EAGER_CONTEXT[]` either, and doing so here would risk pinning +# the scheduler to a single-process view instead of the MPI-aware processor +# set that `Dagger.accelerate!(:mpi)` establishes. + +const suite_setup = Dict{String,Function}() +for suite in suites + suite_setup[suite] = include(joinpath(@__DIR__, "suites", suite * ".jl")) +end + +const SUITE = BenchmarkGroup() +for (suite_name, bench_list) in benches + suite_group = BenchmarkGroup() + for bench in bench_list + method_key = isempty(bench.accels) ? bench.method : + "$(bench.method)+$(join(bench.accels, "+"))" + rank == 0 && @info "[worker_mpi] Creating benchmarks for suite=$suite_name method=$method_key" + suite_group[method_key] = + suite_setup[suite_name](nothing; method=bench.method, accels=bench.accels) + end + SUITE[suite_name] = suite_group +end + +# Apply consistent run parameters to every benchmark in the tree. `evals=1` is +# required because the suites use `setup`/`teardown`. +for (_, b) in BenchmarkTools.leaves(SUITE) + b.params.seconds = bench_seconds + b.params.samples = bench_samples + b.params.evals = 1 + b.params.gcsample = true +end + +# --- Run every benchmark once, identically on every rank -------------------- +# Sorted by keypath (not Dict/BenchmarkGroup insertion order) so ranks agree +# even if suite construction were ever to introduce nondeterministic +# iteration order. + +leaves = sort(collect(BenchmarkTools.leaves(SUITE)); by = kv -> String[string(k) for k in kv[1]]) + +const results = Vector{Tuple{Vector{String},BenchmarkTools.Trial}}() # rank 0 only + +for (keypath, bench) in leaves + kp = String[string(k) for k in keypath] + rank == 0 && @info "[worker_mpi] Running: $(join(kp, " / "))" + try + trial = BenchmarkTools.run(bench) + rank == 0 && push!(results, (kp, trial)) + catch err + if err isa OutOfMemoryError + rank == 0 && @warn "[worker_mpi] OutOfMemoryError; aborting MPI job" benchmark = join(kp, " / ") + flush(stdout); flush(stderr) + MPI.Abort(comm, 137) + exit(137) # unreachable unless MPI.Abort fails to terminate us + else + rank == 0 && @warn "[worker_mpi] Benchmark errored (skipped)" benchmark = join(kp, " / ") exception = (err, catch_backtrace()) + end + end +end + +# --- Rank 0: persist results and advertise completion ----------------------- + +if rank == 0 + manifest = Vector{Any}() + for (i, (kp, trial)) in enumerate(results) + fname = "result_mpi_$(i).json" + BenchmarkTools.save(joinpath(WORKDIR, fname), trial) + push!(manifest, (; keypath=kp, file=fname)) + end + atomic_write(joinpath(WORKDIR, "results_mpi_manifest.json"), JSON3.write(manifest)) + atomic_write(joinpath(WORKDIR, "done"), "1") + @info "[worker_mpi] Done; $(length(results))/$(length(leaves)) benchmark(s) succeeded." +end