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