MPI: Final merge sprint of the whole MPI infrastructure - #720
Merged
Conversation
jpsamaroo
force-pushed
the
fda/mpi-spmd-datadeps
branch
2 times, most recently
from
July 19, 2026 20:29
9c17cdb to
a5416c1
Compare
Contributor
Dagger benchmarks:
|
| master | dirty | master / dirty | |
|---|---|---|---|
| array/dagger/N=1024 (block 512)/add (X + X) | 4.55 ± 0.18 ms | 4.54 ± 0.11 ms | 1 ± 0.047 |
| array/dagger/N=1024 (block 512)/alloc (rand) | 4.3 ± 0.69 ms | 3.56 ± 0.92 ms | 1.21 ± 0.37 |
| array/dagger/N=1024 (block 512)/broadcast (X .+ 1) | 2.93 ± 0.64 ms | 3.01 ± 0.29 ms | 0.973 ± 0.23 |
| array/dagger/N=1024 (block 512)/map (sin.(X)) | 7.19 ± 0.49 ms | 7.31 ± 1.5 ms | 0.983 ± 0.21 |
| array/dagger/N=1024 (block 512)/norm | 1.16 ± 0.042 ms | 1.31 ± 0.019 ms | 0.888 ± 0.034 |
| array/dagger/N=1024 (block 512)/reduce (sum) | 2.61 ± 0.32 ms | 2.56 ± 0.83 ms | 1.02 ± 0.35 |
| array/dagger/N=1024 (block 512)/transpose (permutedims) | 7.79 ± 0.43 ms | 7.18 ± 0.12 ms | 1.08 ± 0.063 |
| array/dagger/N=256 (block 256)/add (X + X) | 1.36 ± 0.15 ms | 1.38 ± 0.19 ms | 0.986 ± 0.17 |
| array/dagger/N=256 (block 256)/alloc (rand) | 1.25 ± 0.049 ms | 1.09 ± 0.11 ms | 1.15 ± 0.12 |
| array/dagger/N=256 (block 256)/broadcast (X .+ 1) | 0.591 ± 3.2 ms | 0.62 ± 0.27 ms | 0.953 ± 5.1 |
| array/dagger/N=256 (block 256)/map (sin.(X)) | 1.21 ± 0.18 ms | 1.3 ± 0.036 ms | 0.935 ± 0.14 |
| array/dagger/N=256 (block 256)/norm | 0.471 ± 0.08 ms | 0.604 ± 0.076 ms | 0.78 ± 0.16 |
| array/dagger/N=256 (block 256)/reduce (sum) | 0.997 ± 0.64 ms | 1.14 ± 0.076 ms | 0.875 ± 0.57 |
| array/dagger/N=256 (block 256)/transpose (permutedims) | 0.831 ± 0.13 ms | 1.04 ± 0.18 ms | 0.799 ± 0.19 |
| linalg/dagger/N=1024 (block 512)/cholesky | 19.5 ± 3.4 ms | 25.1 ± 1.3 ms | 0.778 ± 0.14 |
| linalg/dagger/N=1024 (block 512)/lu | 0.0478 ± 0.0033 s | 0.0453 ± 0.0029 s | 1.05 ± 0.1 |
| linalg/dagger/N=1024 (block 512)/matmul (A*A) | 0.0452 ± 0.0019 s | 0.0718 ± 0.027 s | 0.629 ± 0.24 |
| linalg/dagger/N=1024 (block 512)/matvec (A*x) | 3.26 ± 0.045 ms | 3.19 ± 0.18 ms | 1.02 ± 0.059 |
| linalg/dagger/N=1024 (block 512)/qr | 0.367 ± 0.0065 s | 0.382 ± 0.02 s | 0.963 ± 0.053 |
| linalg/dagger/N=1024 (block 512)/solve (A\b via lu) | 0.0532 ± 0.0012 s | 0.0547 ± 0.0029 s | 0.973 ± 0.055 |
| linalg/dagger/N=1024 (block 512)/syrk (A'*A) | 0.037 ± 0.002 s | 0.0364 ± 0.00071 s | 1.02 ± 0.057 |
| linalg/dagger/N=256 (block 256)/cholesky | 3.08 ± 0.42 ms | 3.36 ± 0.78 ms | 0.917 ± 0.25 |
| linalg/dagger/N=256 (block 256)/lu | 4.7 ± 0.82 ms | 4.85 ± 0.51 ms | 0.969 ± 0.2 |
| linalg/dagger/N=256 (block 256)/matmul (A*A) | 2.26 ± 0.14 ms | 2.83 ± 0.5 ms | 0.796 ± 0.15 |
| linalg/dagger/N=256 (block 256)/matvec (A*x) | 1.43 ± 0.21 ms | 1.55 ± 0.051 ms | 0.924 ± 0.14 |
| linalg/dagger/N=256 (block 256)/qr | 7.39 ± 0.24 ms | 7.66 ± 0.38 ms | 0.965 ± 0.057 |
| linalg/dagger/N=256 (block 256)/solve (A\b via lu) | 8.94 ± 1.3 ms | 9.83 ± 1.3 ms | 0.909 ± 0.18 |
| linalg/dagger/N=256 (block 256)/syrk (A'*A) | 4.08 ± 0.31 ms | 4.21 ± 1.3 ms | 0.968 ± 0.31 |
| stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) | 0.0621 ± 0.0013 s | 27 ± 2.8 ms | 2.3 ± 0.24 |
| stencil/dagger/N=1024 (block 512)/assign (const) | 1.34 ± 0.11 ms | 1.54 ± 0.042 ms | 0.867 ± 0.074 |
| stencil/dagger/N=1024 (block 512)/multi-expr | 3.55 ± 0.26 ms | 4.06 ± 0.44 ms | 0.873 ± 0.11 |
| stencil/dagger/N=1024 (block 512)/neighbors (Clamp) | 30.8 ± 7.5 ms | 25.1 ± 1 ms | 1.23 ± 0.3 |
| stencil/dagger/N=1024 (block 512)/neighbors (Pad) | 24.2 ± 7.5 ms | 25.1 ± 0.62 ms | 0.968 ± 0.3 |
| stencil/dagger/N=1024 (block 512)/neighbors (Reflect) | 30.2 ± 4.8 ms | 24.9 ± 1.3 ms | 1.21 ± 0.2 |
| stencil/dagger/N=1024 (block 512)/neighbors (Wrap) | 0.0318 ± 0.0022 s | 24.4 ± 4.2 ms | 1.3 ± 0.24 |
| stencil/dagger/N=1024 (block 512)/update (+) | 2.11 ± 0.45 ms | 2.99 ± 0.37 ms | 0.704 ± 0.17 |
| stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) | 6.79 ± 0.28 ms | 5.62 ± 0.31 ms | 1.21 ± 0.083 |
| stencil/dagger/N=256 (block 256)/assign (const) | 0.677 ± 0.11 ms | 0.758 ± 0.054 ms | 0.893 ± 0.16 |
| stencil/dagger/N=256 (block 256)/multi-expr | 1.44 ± 0.21 ms | 1.54 ± 0.029 ms | 0.937 ± 0.14 |
| stencil/dagger/N=256 (block 256)/neighbors (Clamp) | 4.54 ± 0.64 ms | 5.16 ± 0.63 ms | 0.881 ± 0.16 |
| stencil/dagger/N=256 (block 256)/neighbors (Pad) | 4.65 ± 0.63 ms | 4.87 ± 0.27 ms | 0.956 ± 0.14 |
| stencil/dagger/N=256 (block 256)/neighbors (Reflect) | 4.58 ± 0.35 ms | 4.89 ± 0.18 ms | 0.937 ± 0.079 |
| stencil/dagger/N=256 (block 256)/neighbors (Wrap) | 5 ± 0.8 ms | 5.03 ± 0.32 ms | 0.994 ± 0.17 |
| stencil/dagger/N=256 (block 256)/update (+) | 0.745 ± 0.062 ms | 0.918 ± 0.11 ms | 0.812 ± 0.12 |
| time_to_load | 1.03 ± 0.0034 s | 1.09 ± 0.004 s | 0.949 ± 0.0047 |
Plots
⚠️ Regressions (> 25.0%)
linalg/dagger/N=1024 (block 512)/matmul (A*A): +59.0%stencil/dagger/N=1024 (block 512)/update (+): +42.0%linalg/dagger/N=1024 (block 512)/cholesky: +28.5%array/dagger/N=256 (block 256)/norm: +28.2%linalg/dagger/N=256 (block 256)/matmul (A*A): +25.6%array/dagger/N=256 (block 256)/transpose (permutedims): +25.1%
Improvements (> 25.0% faster)
stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap): -56.4%
Full results and plots (download the benchmark-results artifact).
jpsamaroo
force-pushed
the
fda/mpi-spmd-datadeps
branch
from
July 19, 2026 21:05
a5416c1 to
0dab20f
Compare
jpsamaroo
marked this pull request as ready for review
July 20, 2026 16:37
Adds MPI acceleration for Dagger with a focus on SPMD Datadeps: MPIAcceleration/MPIRef, deterministic rank-uniform task placement, processor grids, multi-span (KA) copies, move_rewrap header+children protocol, and per-backend GPU MPI/IPC hooks. All MPI.jl-dependent code lives in `ext/MPIExt.jl` (MPI is a weakdep). Scheduler integration is ported onto the new NG scheduler: `schedule_one!` filters/orders processors via the acceleration (`compatible_processors(accel, ...)`, `select_processors_uniform!`) and ignores rank-local capacity/occupancy under uniform execution; `do_task` sets the task acceleration, sequences argument moves (`schedule_argument_move`), and binds moved values via `bind_moved_argument`. This collapses the earlier MPI branch (which evolved over many commits: MPIRPC added then dropped, `src/mpi.jl` moved to the extension, etc.) into its end state, rebased onto current master. The datadeps buffer-free path keeps the upstream `gather_free_syncdeps!` while adding the MPI-uniform task `tag`; `system_uuid` and other changes that landed upstream separately are taken from master. Co-authored-by: Felipe Tome <tomefelipe0@gmail.com> Co-authored-by: yanzin00 <ylsg0303@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Adds the CPU MPI test suite (`test/mpi.jl`) and a shared parametrized GPU MPI suite (`test/mpi_gpu_suite.jl`) with thin per-backend wrappers (CUDA/ROCm/OpenCL/Metal/oneAPI). `test/mpienv` + `test/run_mpi.jl` launch MPI tests via MPI.jl's `mpiexec` (no system MPI needed), wired into a new `mpi-cpu` GHA job matrix (Julia 1.10/1, 2/4 ranks) alongside the existing CI jobs. Adds per-backend GPU test envs and processor-grid tests. Co-authored-by: Felipe Tome <tomefelipe0@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Moves the MPI benchmark and helper scripts out of the (now AirspeedVelocity-based) `benchmark/` tree into `contrib/mpi/`, since they don't fit the new benchmarking system. Adds the Tuolumne agent testing guide and Flux job scripts. Co-authored-by: Felipe Tome <tomefelipe0@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Add SPMD MPI x GPU datadeps jobs to the Buildkite GPU pipelines, covering CUDA, ROCm, oneAPI, and Metal (in .buildkite/pipeline.yml) plus OpenCL (in .buildkite/pipeline-julia.yml, on the existing macOS/aarch64 agent). Each job uses only the `julia` plugin (not `julia-test`, which runs the non-MPI test/runtests.jl), instantiates the backend's MPI test env (test/<backend>env), and launches the corresponding test/mpi_<backend>.jl suite with 2 ranks / 2 threads via test/run_mpi.jl. That helper uses MPI.jl's bundled `mpiexec`, so no system MPI is required, and the GPU data plane is host-staged P2P so a single device per agent suffices. Also fix contrib/mpi/run_mpi_tests.sh, which computed DAGGER_DIR from its own directory and so pointed at contrib/mpi/ after the script was relocated there; resolve the repo root two levels up instead. Co-authored-by: Cursor <cursoragent@cursor.com>
Add a comprehensive "MPI" manual page (docs/src/mpi.md) and wire it into the documentation navigation. The page covers: - When to choose the MPI backend vs. Distributed - Requirements, setup, and `Dagger.accelerate!(:mpi)` - Launching under mpiexec (and the test/run_mpi.jl helper) - The SPMD execution model and `Dagger.check_uniformity!` - A first DArray + datadeps example - Data ownership (MPIRef), processors, and memory spaces - Targeting ranks with `Dagger.scope(mpi_rank=...)` / `mpi_ranks` - Deterministic processor grids for array block placement - Datadeps under MPI - GPUs with MPI (host-staged P2P, GPU scopes, GPU-aware MPI) - collect/fetch semantics, testing, CI, and limitations/tips Also add brief cross-links to the new page from the GPUs and Datadeps pages. Co-authored-by: Cursor <cursoragent@cursor.com>
jpsamaroo
force-pushed
the
fda/mpi-spmd-datadeps
branch
from
July 23, 2026 21:20
402d0dc to
d38cb7e
Compare
This was referenced Jul 25, 2026
Closed
Closed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added the whole SPMD-uniform structure and it's EXT companions for accelerators, as well as the MPI RPC infrastructure that was needed.