Fix stencils over-transfer of neighbor data

.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 `<details>\n<summary>${summary}</summary>\n\n${body}\n\n</details>`;
+            }
+
             const marker = '<!-- dagger-benchmarks -->';
-            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

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 <ranks>` (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
 

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)