diff --git a/.buildkite/pipeline-julia.yml b/.buildkite/pipeline-julia.yml index ea0dbaad7..284569e90 100644 --- a/.buildkite/pipeline-julia.yml +++ b/.buildkite/pipeline-julia.yml @@ -128,6 +128,23 @@ steps: env: CI_USE_OPENCL: "1" + # MPI × OpenCL SPMD datadeps suite. Runs on the same macOS/aarch64 agent as the + # non-MPI OpenCL job, but drives test/mpi_opencl.jl under MPI.jl's bundled + # `mpiexec` via test/run_mpi.jl (no `julia-test` plugin, no system MPI). + - label: Julia 1.11 (OpenCL, MPI) + <<: *gputest + timeout_in_minutes: 20 + plugins: + - JuliaCI/julia#v1: + version: "1.11" + agents: + queue: "juliaecosystem" + os: macos + arch: aarch64 + command: | + julia --project=test/openclenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + julia --project=test/openclenv test/run_mpi.jl 2 2 test/mpi_opencl.jl + - label: Julia 1 - TimespanLogging <<: *test timeout_in_minutes: 20 diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 86e678cbd..b3f72a3c9 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -2,6 +2,16 @@ timeout_in_minutes: 60 if: build.message !~ /\[skip tests\]/ +# MPI × GPU jobs launch the SPMD datadeps suite (test/mpi_.jl) under +# `mpiexec`. They use only the `julia` plugin (not `julia-test`, which would run +# the non-MPI `test/runtests.jl`), then instantiate the backend's MPI test env +# and drive it with `test/run_mpi.jl`, which uses MPI.jl's bundled `mpiexec` +# (no system MPI required). 2 ranks share whatever GPU(s) the agent exposes; +# the data plane is host-staged P2P, so a single device is sufficient. +.mpigputest: &mpigputest + timeout_in_minutes: 60 + if: build.message !~ /\[skip tests\]/ + steps: - label: Julia 1.11 (CUDA) <<: *gputest @@ -56,5 +66,50 @@ steps: env: CI_USE_METAL: "1" + - label: Julia 1.11 (CUDA, MPI) + <<: *mpigputest + plugins: + - JuliaCI/julia#v1: + version: "1.11" + agents: + queue: "cuda" + command: | + julia --project=test/cudaenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + julia --project=test/cudaenv test/run_mpi.jl 2 2 test/mpi_cuda.jl + + - label: Julia 1.11 (ROCm, MPI) + <<: *mpigputest + plugins: + - JuliaCI/julia#v1: + version: "1.11" + agents: + queue: "rocm" + command: | + julia --project=test/rocmenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + julia --project=test/rocmenv test/run_mpi.jl 2 2 test/mpi_rocm.jl + + - label: Julia 1.11 (oneAPI, MPI) + <<: *mpigputest + plugins: + - JuliaCI/julia#v1: + version: "1.11" + agents: + queue: "oneapi" + command: | + julia --project=test/oneapienv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + julia --project=test/oneapienv test/run_mpi.jl 2 2 test/mpi_oneapi.jl + + - label: Julia 1.11 (Metal, MPI) + <<: *mpigputest + timeout_in_minutes: 20 + plugins: + - JuliaCI/julia#v1: + version: "1.11" + agents: + queue: "metal" + command: | + julia --project=test/metalenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + julia --project=test/metalenv test/run_mpi.jl 2 2 test/mpi_metal.jl + # env: # SECRET_CODECOV_TOKEN: "" diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index eed72ba58..0b25e5576 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -231,3 +231,35 @@ jobs: } else { await github.rest.issues.createComment({ owner, repo, issue_number, body }); } + + mpi-cpu: + name: MPI CPU - Julia ${{ matrix.version }} / ${{ matrix.ranks }} ranks + runs-on: ubuntu-latest + timeout-minutes: 60 + if: ${{ !contains(github.event.head_commit.message, '[skip tests]') }} + strategy: + fail-fast: false + matrix: + version: + - '1.10' # LTS + - '1.11' + - '1' # latest stable + ranks: + - 2 + - 4 + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + arch: x64 + - uses: julia-actions/cache@v3 + # test/mpienv pulls MPI (MPICH_jll by default) plus the checked-out Dagger; + # no system MPI installation is required. We `develop` the checkout rather + # than relying on the `[sources]` entry because `[sources]` is a Julia 1.11+ + # feature and is silently ignored on 1.10 (LTS), which would otherwise + # resolve a stale registered Dagger lacking the MPI API (`accelerate!`). + - name: Instantiate CPU MPI environment + run: julia --project=test/mpienv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + - name: Run CPU MPI tests + run: julia --project=test/mpienv test/run_mpi.jl ${{ matrix.ranks }} 2 test/mpi.jl diff --git a/Project.toml b/Project.toml index 68b410991..6b4a2e9f2 100644 --- a/Project.toml +++ b/Project.toml @@ -40,6 +40,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" +MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" @@ -55,6 +56,7 @@ GraphVizSimpleExt = "Colors" IntelExt = "oneAPI" JSON3Ext = "JSON3" LinuxPerfExt = "LinuxPerf" +MPIExt = "MPI" MetalExt = "Metal" OpenCLExt = "OpenCL" PlotsExt = ["DataFrames", "Plots"] @@ -78,6 +80,7 @@ Graphs = "1" JSON3 = "1" KernelAbstractions = "0.9" LinuxPerf = "0.4.2" +MPI = "0.20.22" MacroTools = "0.5" MemPool = "0.4.17" Metal = "1.1" diff --git a/contrib/mpi/check_comm_asymmetry.jl b/contrib/mpi/check_comm_asymmetry.jl new file mode 100644 index 000000000..684240ec5 --- /dev/null +++ b/contrib/mpi/check_comm_asymmetry.jl @@ -0,0 +1,111 @@ +#!/usr/bin/env julia +# Parse MPI+Dagger logs and report communication decision asymmetry per tag. +# Asymmetry: for the same tag, one rank decides to send (local+bcast, sender+communicated, etc.) +# and another rank decides to infer (inferred, uninvolved) and never recv → deadlock. +# +# Usage: julia check_comm_asymmetry.jl < logfile +# Or: mpiexec -n 10 julia ... run_matmul.jl 2>&1 | tee matmul.log; julia check_comm_asymmetry.jl < matmul.log + +const SEND_DECISIONS = Set([ + "local+bcast", "sender+communicated", "sender+inferred", "receiver+bcast", + "aliasing", # when followed by local+bcast we already capture local+bcast +]) +const RECV_DECISIONS = Set([ + "communicated", "receiver", "sender+communicated", # received data +]) +const INFER_DECISIONS = Set([ + "inferred", "uninvolved", # did not recv (uses inferred type) +]) + +function parse_line(line) + # Match [rank X][tag Y] then any [...] and capture the last bracket pair before space or end + rank = nothing + tag = nothing + decision = nothing + category = nothing # aliasing, execute!, remotecall_endpoint + for m in eachmatch(r"\[rank\s+(\d+)\]", line) + rank = parse(Int, m.captures[1]) + end + for m in eachmatch(r"\[tag\s+(\d+)\]", line) + tag = parse(Int, m.captures[1]) + end + for m in eachmatch(r"\[(execute!|aliasing|remotecall_endpoint)\]", line) + category = m.captures[1] + end + # Decision is usually in last [...] that looks like [word] or [word+word] + for m in eachmatch(r"\]\[([^\]]+)\]", line) + candidate = m.captures[1] + # Normalize: "communicated" "inferred" "local+bcast" "sender+inferred" "receiver" etc. + if occursin("inferred", candidate) && !occursin("communicated", candidate) + decision = "inferred" + break + elseif occursin("communicated", candidate) + decision = "communicated" + break + elseif occursin("local+bcast", candidate) + decision = "local+bcast" + break + elseif occursin("sender+", candidate) + decision = startswith(candidate, "sender+inferred") ? "sender+inferred" : "sender+communicated" + break + elseif candidate == "receiver" + decision = "receiver" + break + elseif candidate == "receiver+bcast" + decision = "receiver+bcast" + break + elseif candidate == "inplace_move" + decision = "inplace_move" + break + end + end + return rank, tag, category, decision +end + +function main() + # tag => Dict(rank => decision) + by_tag = Dict{Int, Dict{Int, String}}() + for line in eachline(stdin) + rank, tag, category, decision = parse_line(line) + isnothing(rank) && continue + isnothing(tag) && continue + isnothing(decision) && continue + if !haskey(by_tag, tag) + by_tag[tag] = Dict{Int, String}() + end + by_tag[tag][rank] = decision + end + + # For each tag, check: is there at least one sender and one inferrer (non-receiver)? + send_keys = Set(["local+bcast", "sender+communicated", "sender+inferred", "receiver+bcast"]) + infer_keys = Set(["inferred", "sender+inferred"]) # sender+inferred means sender didn't need to recv + recv_keys = Set(["communicated", "receiver", "sender+communicated"]) + + asymmetries = [] + for (tag, ranks) in sort(collect(by_tag), by = first) + senders = [r for (r, d) in ranks if d in send_keys] + inferrers = [r for (r, d) in ranks if d in infer_keys || d == "uninvolved"] + receivers = [r for (r, d) in ranks if d in recv_keys] + # Asymmetry: someone sends (bcast) so will send to ALL other ranks; someone chose infer and won't recv. + if !isempty(senders) && !isempty(inferrers) + push!(asymmetries, (tag, senders, inferrers, receivers, ranks)) + end + end + + if isempty(asymmetries) + println("No communication decision asymmetry found (no tag has both sender and inferrer).") + return + end + + println("=== Communication decision asymmetry (can cause deadlock) ===\n") + for (tag, senders, inferrers, receivers, ranks) in asymmetries + println("Tag $tag:") + println(" Senders (will bcast to all others): $senders") + println(" Inferrers (did not recv): $inferrers") + println(" Receivers: $receivers") + println(" All decisions: $ranks") + println() + end +end + +main() diff --git a/contrib/mpi/check_comm_asymmetry.py b/contrib/mpi/check_comm_asymmetry.py new file mode 100644 index 000000000..31a117442 --- /dev/null +++ b/contrib/mpi/check_comm_asymmetry.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +""" +Parse MPI+Dagger logs and report communication decision asymmetry per tag. +Asymmetry: for the same tag, one rank decides to send (local+bcast, etc.) +and another decides to infer (inferred) and never recv → deadlock. + +Usage: + # Capture full log (all ranks' Core.println from mpi.jl go to stdout): + mpiexec -n 10 julia --project=/path/to/Dagger.jl benchmarks/run_matmul.jl 2>&1 | tee matmul.log + # Then look for asymmetry (same tag: one rank sends, another infers → deadlock): + python3 check_comm_asymmetry.py < matmul.log +""" + +import re +import sys +from collections import defaultdict + +SEND_DECISIONS = {"local+bcast", "sender+communicated", "sender+inferred", "receiver+bcast"} +RECV_DECISIONS = {"communicated", "receiver", "sender+communicated"} +INFER_DECISIONS = {"inferred", "uninvolved", "sender+inferred"} + + +def parse_line(line: str): + rank = tag = category = decision = None + m = re.search(r"\[rank\s+(\d+)\]", line) + if m: + rank = int(m.group(1)) + m = re.search(r"\[tag\s+(\d+)\]", line) + if m: + tag = int(m.group(1)) + m = re.search(r"\[(execute!|aliasing|remotecall_endpoint)\]", line) + if m: + category = m.group(1) + # Capture decision from [...] blocks + for m in re.finditer(r"\]\[([^\]]+)\]", line): + candidate = m.group(1) + if "inferred" in candidate and "communicated" not in candidate: + decision = "inferred" + break + if "communicated" in candidate: + decision = "communicated" + break + if "local+bcast" in candidate: + decision = "local+bcast" + break + if candidate.startswith("sender+"): + decision = "sender+inferred" if "inferred" in candidate else "sender+communicated" + break + if candidate == "receiver": + decision = "receiver" + break + if candidate == "receiver+bcast": + decision = "receiver+bcast" + break + if candidate == "inplace_move": + decision = "inplace_move" + break + return rank, tag, category, decision + + +def main(): + by_tag = defaultdict(dict) # tag -> {rank: decision} + for line in sys.stdin: + rank, tag, category, decision = parse_line(line) + if rank is None or tag is None or decision is None: + continue + by_tag[tag][rank] = decision + + send_keys = {"local+bcast", "sender+communicated", "sender+inferred", "receiver+bcast"} + infer_keys = {"inferred", "sender+inferred", "uninvolved"} + recv_keys = {"communicated", "receiver", "sender+communicated"} + + asymmetries = [] + for tag in sorted(by_tag.keys()): + ranks = by_tag[tag] + senders = [r for r, d in ranks.items() if d in send_keys] + inferrers = [r for r, d in ranks.items() if d in infer_keys] + receivers = [r for r, d in ranks.items() if d in recv_keys] + if senders and inferrers: + asymmetries.append((tag, senders, inferrers, receivers, ranks)) + + if not asymmetries: + print("No communication decision asymmetry found (no tag has both sender and inferrer).") + return + + print("=== Communication decision asymmetry (can cause deadlock) ===\n") + for tag, senders, inferrers, receivers, ranks in asymmetries: + print(f"Tag {tag}:") + print(f" Senders (will bcast to all others): {senders}") + print(f" Inferrers (did not recv): {inferrers}") + print(f" Receivers: {receivers}") + print(f" All decisions: {dict(ranks)}") + print() + + +if __name__ == "__main__": + main() diff --git a/contrib/mpi/mpi_transfer_bench.jl b/contrib/mpi/mpi_transfer_bench.jl new file mode 100644 index 000000000..f475086c0 --- /dev/null +++ b/contrib/mpi/mpi_transfer_bench.jl @@ -0,0 +1,183 @@ +# MPI data-plane benchmark: transfer bandwidth and datadeps macro timings. +# +# Run (CPU): mpiexec -n 2 julia --project --threads=2 benchmarks/mpi_transfer_bench.jl +# Run (GPU): mpiexec -n 2 julia --project=../gpuenv --threads=2 benchmarks/mpi_transfer_bench.jl +# +# Note on GPU-direct validation: the default MPICH_jll is not CUDA-aware +# (MPI.has_cuda() == false), so device buffers host-stage. To validate the +# device-direct path for real, switch to a CUDA-aware library, e.g.: +# using MPIPreferences; MPIPreferences.use_system_binary() # UCX/CUDA OpenMPI +# and force the path on with DAGGER_MPI_GPU_DIRECT=1. + +using Dagger, MPI, LinearAlgebra, Random, Printf + +Dagger.accelerate!(:mpi) +const comm = MPI.COMM_WORLD +const rank = MPI.Comm_rank(comm) +const nranks = MPI.Comm_size(comm) +@assert nranks == 2 "transfer bench expects exactly 2 ranks" + +const HAVE_CUDA = try + @eval using CUDA + CUDA.functional() +catch + false +end + +function bench_transfer(label, make; iters=10) + if rank == 0 + @printf("%-14s %10s %12s\n", label, "size", "bandwidth") + end + for p in 16:2:26 + nbytes = 2^p + n = nbytes ÷ sizeof(Float64) + src = make(n) + dst = make(n) + tag = 20_000 + p + # Warmup (includes compile and any pool/pinning setup) + for _ in 1:3 + if rank == 0 + Dagger.send_yield!(src, comm, 1, tag) + Dagger.recv_yield!(dst, comm, 1, tag) + else + Dagger.recv_yield!(dst, comm, 0, tag) + Dagger.send_yield!(src, comm, 0, tag) + end + end + MPI.Barrier(comm) + t0 = time_ns() + for _ in 1:iters + if rank == 0 + Dagger.send_yield!(src, comm, 1, tag) + Dagger.recv_yield!(dst, comm, 1, tag) + else + Dagger.recv_yield!(dst, comm, 0, tag) + Dagger.send_yield!(src, comm, 0, tag) + end + end + dt = (time_ns() - t0) / 1e9 + # Each iteration moves nbytes in each direction (round trip) + bw = 2 * iters * nbytes / dt / 1e9 + if rank == 0 + @printf("%-14s %7.2f MiB %9.3f GB/s\n", "", nbytes / 2^20, bw) + end + MPI.Barrier(comm) + end +end + +function time_region(f; reps=2) + f() # warmup (compile) + best = Inf + for _ in 1:reps + MPI.Barrier(comm) + t0 = time_ns() + f() + dt = (time_ns() - t0) / 1e9 + best = min(best, dt) + end + return best +end + +function bench_macro_cpu() + Random.seed!(1) + N, B = 512, 128 + A = rand(N, N); A = A * A'; A[diagind(A)] .+= N + t_chol = time_region() do + DA = DArray(A, Blocks(B, B)) + cholesky(DA) + nothing + end + M1 = rand(N, N); M2 = rand(N, N) + t_mul = time_region() do + DA = DArray(M1, Blocks(B, B)) + DB = DArray(M2, Blocks(B, B)) + DC = zeros(Blocks(B, B), N, N) + mul!(DC, DA, DB) + nothing + end + rank == 0 && @printf("macro CPU: cholesky %7.3f s mul! %7.3f s (N=%d, B=%d)\n", t_chol, t_mul, N, B) +end + +function bench_macro_gpu() + CUDAExt = Base.get_extension(Dagger, :CUDAExt) + CuProc = CUDAExt.CuArrayDeviceProc + procs = sort(collect(Dagger.get_processors(Dagger.MPIClusterProc(comm))); + by=p->(p.rank, Dagger.short_name(p))) + gpu_scope = Dagger.UnionScope([Dagger.ExactScope(p) for p in procs if p.innerProc isa CuProc]...) + + Random.seed!(2) + N, B = 512, 128 + A = rand(N, N); A = A * A'; A[diagind(A)] .+= N + t_chol = time_region() do + DA = DArray(A, Blocks(B, B)) + Dagger.with_options(;scope=gpu_scope) do + cholesky(DA) + end + nothing + end + M1 = rand(N, N); M2 = rand(N, N) + t_mul = time_region() do + DA = DArray(M1, Blocks(B, B)) + DB = DArray(M2, Blocks(B, B)) + DC = zeros(Blocks(B, B), N, N) + Dagger.with_options(;scope=gpu_scope) do + mul!(DC, DA, DB) + end + nothing + end + rank == 0 && @printf("macro GPU: cholesky %7.3f s mul! %7.3f s (N=%d, B=%d)\n", t_chol, t_mul, N, B) +end + +rank == 0 && println("== MPI transfer bench (nranks=$nranks, has_cuda=$(MPI.has_cuda()), CUDA=$(HAVE_CUDA)) ==") + +bench_transfer("Array", n->rand(n)) +if HAVE_CUDA + bench_transfer("CuArray", n->CUDA.rand(Float64, n)) +end + +# Same-node device IPC protocol round trip (handle + DtoD, no host hop) +function bench_ipc(; iters=10) + rank == 0 && @printf("%-14s %10s %12s\n", "CuArray-IPC", "size", "bandwidth") + for p in 16:2:26 + nbytes = 2^p + n = nbytes ÷ sizeof(Float64) + src = CUDA.rand(Float64, n) + dst = CUDA.rand(Float64, n) + tag = 21_000 + p + for _ in 1:3 + if rank == 0 + Dagger.mpi_ipc_send(src, comm, 1, tag) + Dagger.mpi_ipc_recv!(dst, comm, 1, tag) + else + Dagger.mpi_ipc_recv!(dst, comm, 0, tag) + Dagger.mpi_ipc_send(src, comm, 0, tag) + end + end + MPI.Barrier(comm) + t0 = time_ns() + for _ in 1:iters + if rank == 0 + Dagger.mpi_ipc_send(src, comm, 1, tag) + Dagger.mpi_ipc_recv!(dst, comm, 1, tag) + else + Dagger.mpi_ipc_recv!(dst, comm, 0, tag) + Dagger.mpi_ipc_send(src, comm, 0, tag) + end + end + dt = (time_ns() - t0) / 1e9 + bw = 2 * iters * nbytes / dt / 1e9 + rank == 0 && @printf("%-14s %7.2f MiB %9.3f GB/s\n", "", nbytes / 2^20, bw) + MPI.Barrier(comm) + end +end +if HAVE_CUDA && Dagger.same_node(Dagger.current_acceleration(), 0, 1) + bench_ipc() +end + +bench_macro_cpu() +if HAVE_CUDA + bench_macro_gpu() +end + +MPI.Barrier(comm) +rank == 0 && println("== done ==") diff --git a/contrib/mpi/run_distribute_fetch.jl b/contrib/mpi/run_distribute_fetch.jl new file mode 100644 index 000000000..822e1ad2c --- /dev/null +++ b/contrib/mpi/run_distribute_fetch.jl @@ -0,0 +1,42 @@ +#!/usr/bin/env julia +# Create a matrix with a fixed reproducible pattern, distribute it with an +# MPI procgrid, then on each rank fetch and println the chunk(s) it owns. +# Usage (from repo root, use full path to Dagger.jl): +# mpiexec -n 4 julia --project=/path/to/Dagger.jl benchmarks/run_distribute_fetch.jl + +using MPI +using Dagger + +if !isdefined(Dagger, :accelerate!) + error("Dagger.accelerate! not found. Run with the local Dagger project: julia --project=/path/to/Dagger.jl ...") +end +Dagger.accelerate!(:mpi) + +const comm = MPI.COMM_WORLD +const rank = MPI.Comm_rank(comm) +const nranks = MPI.Comm_size(comm) + +# Fixed reproducible pattern: 6×6 matrix, M[i,j] = 10*i + j (same on all ranks) +const N = 6 +const BLOCK = 2 +A = [10 * i + j for i in 1:N, j in 1:N] + +# Procgrid: use Dagger's compatible processors so the procgrid passes validation +availprocs = collect(Dagger.compatible_processors()) +nblocks = (cld(N, BLOCK), cld(N, BLOCK)) +procgrid = reshape( + [availprocs[mod(i - 1, length(availprocs)) + 1] for i in 1:prod(nblocks)], + nblocks, +) + +# Distribute so chunk (i,j) is computed on procgrid[i,j] +D = distribute(A, Blocks(BLOCK, BLOCK), procgrid) +D_fetched = fetch(D) + +# On each rank: fetch and print only the chunk(s) this rank owns +for (idx, ch) in enumerate(D_fetched.chunks) + if ch isa Dagger.Chunk && ch.handle isa Dagger.MPIRef && ch.handle.rank == rank + data = fetch(ch) + println("rank $rank chunk $idx: ", data) + end +end diff --git a/contrib/mpi/run_matmul.jl b/contrib/mpi/run_matmul.jl new file mode 100644 index 000000000..b0a870168 --- /dev/null +++ b/contrib/mpi/run_matmul.jl @@ -0,0 +1,105 @@ +#!/usr/bin/env julia +# N×N matmul benchmark (Float32); block size scales with number of ranks. +# Usage (use the full path to Dagger.jl, not "..."): +# mpiexec -n 10 julia --project=/home/felipetome/dagger-dev/mpi/Dagger.jl benchmarks/run_matmul.jl +# Set CHECK_CORRECTNESS=true to collect and compare against GPU baseline: +# CHECK_CORRECTNESS=true mpiexec -n 10 julia --project=/home/felipetome/dagger-dev/mpi/Dagger.jl benchmarks/run_matmul.jl + +using MPI +using Dagger +using LinearAlgebra + +if !isdefined(Dagger, :accelerate!) + error("Dagger.accelerate! not found. Run with the local Dagger project: julia --project=/path/to/Dagger.jl ...") +end +Dagger.accelerate!(:mpi) + +const N = 2_000 +const comm = MPI.COMM_WORLD +const rank = MPI.Comm_rank(comm) +const nranks = MPI.Comm_size(comm) +# Block size proportional to ranks: ~nranks blocks in 2D => side blocks ≈ √nranks +const BLOCK = max(1, ceil(Int, N / ceil(Int, sqrt(nranks)))) + +const CHECK_CORRECTNESS = parse(Bool, get(ENV, "CHECK_CORRECTNESS", "false")) + +if rank == 0 + println("Benchmark: ", nranks, " ranks, N=", N, ", block size ", BLOCK, "×", BLOCK, " (matmul)") +end + +# Allocate and fill matrices in blocks (Float32) +A = rand(Blocks(BLOCK, BLOCK), Float32, N, N) +B = rand(Blocks(BLOCK, BLOCK), Float32, N, N) + +# Matrix multiply C = A * B +t_matmul = @elapsed begin + C = A * B +end + +if rank == 0 + println("Matmul time: ", round(t_matmul; digits=4), " s") +end + +# Optional: tree collect via datadeps. All ranks participate in the region. +if CHECK_CORRECTNESS + t_collect = @elapsed begin + A_full = collect(A) + B_full = collect(B) + C_dagger = collect(C) + end + if rank == 0 + println("Collecting result and computing baseline for correctness check (GPU)...") + using CUDA + CUDA.functional() || error("CUDA not functional; cannot compute GPU baseline. Check CUDA driver and device.") + t_upload = @elapsed begin + A_g = CUDA.cu(A_full) + B_g = CUDA.cu(B_full) + end + println("Collect + upload time: ", round(t_collect + t_upload; digits=4), " s") + + t_baseline = @elapsed begin + C_ref_g = A_g * B_g + end + println("Baseline (GPU/CUDA) time: ", round(t_baseline; digits=4), " s") + + # Require all elements within 100× machine epsilon relative error (componentwise) + C_dagger_cpu = C_dagger + C_ref_cpu = Array(C_ref_g) + eps_f = eps(Float32) + rtol = 50.0f0 * eps_f + diff = C_dagger_cpu .- C_ref_cpu + # rel_ij = |diff|/|C_ref|, denominator at least eps to avoid div by zero + denom = max.(abs.(C_ref_cpu), eps_f) + rel_err = abs.(diff) ./ denom + max_rel_err = Float32(maximum(rel_err)) + ok = max_rel_err <= rtol + if ok + println("Correctness: OK (max rel_err = ", max_rel_err, " <= 100×eps = ", rtol, ")") + else + println("Correctness: FAIL (max rel_err = ", max_rel_err, " > 100×eps = ", rtol, ")") + end + + # Per-block: which blocks have any element with rel_err > 100×eps + n_bi = ceil(Int, N / BLOCK) + n_bj = ceil(Int, N / BLOCK) + bad_blocks = Tuple{Int,Int,Float32}[] + for bi in 1:n_bi, bj in 1:n_bj + ri = (bi - 1) * BLOCK + 1 : min(bi * BLOCK, N) + rj = (bj - 1) * BLOCK + 1 : min(bj * BLOCK, N) + block_rel = Float32(maximum(@view(rel_err[ri, rj]))) + if block_rel > rtol + push!(bad_blocks, (bi, bj, block_rel)) + end + end + if isempty(bad_blocks) + println("Per-block: all ", n_bi * n_bj, " blocks within 100×eps rel_err.") + else + println("Per-block: ", length(bad_blocks), " block(s) exceed 100×eps rel_err (block size ", BLOCK, "×", BLOCK, "):") + sort!(bad_blocks; by = x -> -x[3]) + for (bi, bj, block_rel) in bad_blocks + println(" block [", bi, ",", bj, "] rows ", (bi - 1) * BLOCK + 1, ":", min(bi * BLOCK, N), + ", cols ", (bj - 1) * BLOCK + 1, ":", min(bj * BLOCK, N), " max rel_err = ", block_rel) + end + end + end +end diff --git a/contrib/mpi/run_mpi_tests.sh b/contrib/mpi/run_mpi_tests.sh new file mode 100755 index 000000000..61fef80dd --- /dev/null +++ b/contrib/mpi/run_mpi_tests.sh @@ -0,0 +1,111 @@ +#!/usr/bin/env bash +# Run the full Dagger MPI test suite (CPU; GPU optional). +# +# This script lives at contrib/mpi/run_mpi_tests.sh and locates the repo root +# relative to its own path, so it can be run from anywhere: +# contrib/mpi/run_mpi_tests.sh +# N=4 THREADS=2 contrib/mpi/run_mpi_tests.sh +# SKIP_GPU=1 contrib/mpi/run_mpi_tests.sh # skip mpi_cuda.jl and mpi_rocm.jl +# SKIP_CUDA=1 contrib/mpi/run_mpi_tests.sh # skip mpi_cuda.jl only (e.g. AMD/ROCm systems) +# SKIP_ROCM=1 contrib/mpi/run_mpi_tests.sh # skip mpi_rocm.jl only (e.g. NVIDIA/CUDA systems) +# CUDA_GPU_PROJECT=test/cudaenv contrib/mpi/run_mpi_tests.sh +# ROCM_GPU_PROJECT=test/rocmenv contrib/mpi/run_mpi_tests.sh +# +# Override the Julia binary with JULIA_BIN if needed. + +set -euo pipefail + +# Repo root is two levels up from this script (contrib/mpi/). +DAGGER_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +N="${N:-4}" +N_GPU="${N_GPU:-2}" +THREADS="${THREADS:-2}" +CUDA_GPU_PROJECT="${CUDA_GPU_PROJECT:-${DAGGER_DIR}/test/cudaenv}" +ROCM_GPU_PROJECT="${ROCM_GPU_PROJECT:-${DAGGER_DIR}/test/rocmenv}" + +if [[ -z "${JULIA_BIN:-}" ]]; then + JULIA_BIN="$(ls -d "${HOME}"/.julia/juliaup/julia-*+0.x64.linux.gnu/bin/julia 2>/dev/null | sort -V | tail -1)" +fi + +if [[ ! -x "${JULIA_BIN}" ]]; then + echo "ERROR: Julia binary not found. Set JULIA_BIN to your julia executable." >&2 + exit 1 +fi + +if ! command -v mpiexec >/dev/null 2>&1; then + echo "ERROR: mpiexec not found on PATH." >&2 + exit 1 +fi + +run_mpi() { + local label="$1" + local nproc="$2" + local project="$3" + local script="$4" + shift 4 + + echo "" + echo "======================================================================" + echo "MPI suite: ${label}" + echo " ranks: ${nproc}" + echo " project: ${project}" + echo " script: ${script}" + echo "======================================================================" + mpiexec -n "${nproc}" "${JULIA_BIN}" --project="${project}" --threads="${THREADS}" "${script}" "$@" +} + +failed=0 +pass() { echo "PASS: $1"; } +fail() { echo "FAIL: $1" >&2; failed=1; } + +echo "Dagger MPI test runner" +echo " julia: ${JULIA_BIN}" +echo " mpiexec: $(command -v mpiexec)" +echo " N: ${N}" +echo " threads: ${THREADS}" + +# --- Dagger CPU MPI suite (mpiexec) --- +if run_mpi "mpi" "${N}" "${DAGGER_DIR}" "${DAGGER_DIR}/test/mpi.jl"; then + pass "mpi" +else + fail "mpi" +fi + +# --- MPI GPU (CUDA, optional) --- +if [[ "${SKIP_GPU:-0}" == "1" || "${SKIP_CUDA:-0}" == "1" ]]; then + echo "" + echo "SKIP: mpi_cuda (SKIP_GPU=${SKIP_GPU:-0}, SKIP_CUDA=${SKIP_CUDA:-0})" +elif [[ ! -f "${CUDA_GPU_PROJECT}/Project.toml" ]]; then + echo "" + echo "SKIP: mpi_cuda (no Project.toml at CUDA_GPU_PROJECT=${CUDA_GPU_PROJECT})" +else + if run_mpi "mpi_cuda" "${N_GPU}" "${CUDA_GPU_PROJECT}" "${DAGGER_DIR}/test/mpi_cuda.jl"; then + pass "mpi_cuda" + else + fail "mpi_cuda" + fi +fi + +# --- MPI GPU (ROCm, optional) --- +if [[ "${SKIP_GPU:-0}" == "1" || "${SKIP_ROCM:-0}" == "1" ]]; then + echo "" + echo "SKIP: mpi_rocm (SKIP_GPU=${SKIP_GPU:-0}, SKIP_ROCM=${SKIP_ROCM:-0})" +elif [[ ! -f "${ROCM_GPU_PROJECT}/Project.toml" ]]; then + echo "" + echo "SKIP: mpi_rocm (no Project.toml at ROCM_GPU_PROJECT=${ROCM_GPU_PROJECT})" +else + if run_mpi "mpi_rocm" "${N_GPU}" "${ROCM_GPU_PROJECT}" "${DAGGER_DIR}/test/mpi_rocm.jl"; then + pass "mpi_rocm" + else + fail "mpi_rocm" + fi +fi + +echo "" +if [[ "${failed}" -eq 0 ]]; then + echo "All MPI suites passed." + exit 0 +else + echo "One or more MPI suites failed." >&2 + exit 1 +fi diff --git a/contrib/mpi/run_qr.jl b/contrib/mpi/run_qr.jl new file mode 100644 index 000000000..c5915db2a --- /dev/null +++ b/contrib/mpi/run_qr.jl @@ -0,0 +1,46 @@ +#!/usr/bin/env julia +# 10k×10k QR + matmul benchmark; block size scales with number of ranks. +# Usage: mpiexec -n 100 julia --project=/path/to/Dagger.jl benchmarks/bench_100rank_qr_matmul.jl +# Or: bash benchmarks/run_100rank_qr_matmul.sh . + +using MPI +using Dagger +using LinearAlgebra + +Dagger.accelerate!(:mpi) + +const N = 10_000 +const comm = MPI.COMM_WORLD +const rank = MPI.Comm_rank(comm) +const nranks = MPI.Comm_size(comm) +# Block size proportional to ranks: ~nranks blocks in 2D => side blocks ≈ √nranks +const BLOCK = max(1, ceil(Int, N / ceil(Int, sqrt(nranks)))) + +if rank == 0 + println("Benchmark: ", nranks, " ranks, N=", N, ", block size ", BLOCK, "×", BLOCK, " (QR + matmul)") +end + +# Allocate and fill 10k×10k matrix in 1k×1k blocks +A = rand(Blocks(BLOCK, BLOCK), Float64, N, N) +MPI.Barrier(comm) + +# QR factorization (computing Q runs the full factorization) +t_qr = @elapsed begin + qr!(A) +end +MPI.Barrier(comm) + +if rank == 0 + println("QR time: ", round(t_qr; digits=4), " s") +end + +# Matrix multiply A * A +t_matmul = @elapsed begin + C = A * A +end +MPI.Barrier(comm) + +if rank == 0 + println("Matmul time: ", round(t_matmul; digits=4), " s") +end + diff --git a/contrib/tuolumne/dagger_mpi_1gpu.flux b/contrib/tuolumne/dagger_mpi_1gpu.flux new file mode 100644 index 000000000..7d6c61a29 --- /dev/null +++ b/contrib/tuolumne/dagger_mpi_1gpu.flux @@ -0,0 +1,25 @@ +#!/bin/bash +# Full Dagger MPI test suite on Tuolumne — single GPU (1 MPI rank for ROCm tests). +# +# Before submit: +# 1. source ~/tuolumne_dagger.profile +# 2. Replace YOUR_BANK below (or set FLUX_BANK in profile and use it here) +# 3. flux batch scripts/tuolumne/dagger_mpi_1gpu.flux +# +#flux: -N 1 +#flux: -q pdebug +#flux: -t 2h +#flux: --setattr=bank=YOUR_BANK +#flux: --job-name=dagger-mpi-1gpu +#flux: --output=dagger-mpi-1gpu.o{{id}} +#flux: --error=dagger-mpi-1gpu.e{{id}} + +set -euo pipefail +source "${HOME}/tuolumne_dagger.profile" +cd "${DAGGER_DIR}" + +# One Flux task on an exclusive node; mpiexec inside run_mpi_tests.sh. +# N=4 -> CPU MPI suites (mpi_datadeps, cholesky, lu_smoke) +# N_GPU=1 -> mpi_rocm.jl with 1 rank (single MI300A APU) +flux run -N1 -x -n 1 -- \ + env SKIP_CUDA=1 N=4 N_GPU=1 THREADS=4 contrib/mpi/run_mpi_tests.sh diff --git a/contrib/tuolumne/dagger_mpi_4gpu.flux b/contrib/tuolumne/dagger_mpi_4gpu.flux new file mode 100644 index 000000000..2ea2b2d0d --- /dev/null +++ b/contrib/tuolumne/dagger_mpi_4gpu.flux @@ -0,0 +1,23 @@ +#!/bin/bash +# Full Dagger MPI test suite on Tuolumne — single node, all 4 MI300A APUs. +# +# Before submit: +# 1. source ~/tuolumne_dagger.profile +# 2. Replace YOUR_BANK below +# 3. flux batch scripts/tuolumne/dagger_mpi_4gpu.flux +# +#flux: -N 1 +#flux: -q pdebug +#flux: -t 2h +#flux: --setattr=bank=YOUR_BANK +#flux: --job-name=dagger-mpi-4gpu +#flux: --output=dagger-mpi-4gpu.o{{id}} +#flux: --error=dagger-mpi-4gpu.e{{id}} + +set -euo pipefail +source "${HOME}/tuolumne_dagger.profile" +cd "${DAGGER_DIR}" + +# N_GPU=4 -> one MPI rank per APU (Tuolumne has 4 MI300A per node) +flux run -N1 -x -n 1 -- \ + env SKIP_CUDA=1 N=4 N_GPU=4 THREADS=4 contrib/mpi/run_mpi_tests.sh diff --git a/contrib/tuolumne/tuolumne_dagger.profile.example b/contrib/tuolumne/tuolumne_dagger.profile.example new file mode 100644 index 000000000..093db76b4 --- /dev/null +++ b/contrib/tuolumne/tuolumne_dagger.profile.example @@ -0,0 +1,33 @@ +# Tuolumne (LLNL) environment for Dagger MPI testing. +# Copy to $HOME/tuolumne_dagger.profile and edit FLUX_BANK / module versions. +# +# cp scripts/tuolumne/tuolumne_dagger.profile.example ~/tuolumne_dagger.profile +# vi ~/tuolumne_dagger.profile # set FLUX_BANK, JULIA_BIN, DAGGER_DIR +# source ~/tuolumne_dagger.profile + +# --- Required: set your Flux bank (flux banks) --- +export FLUX_BANK=YOUR_BANK + +# --- ROCm (check: module avail rocm) --- +module load rocm/6.4.1 + +# MI300A compilation / runtime arch hint +export AMREX_AMD_ARCH=gfx942 + +# --- Julia (install to Lustre; do not use NFS $HOME for depots) --- +export JULIA_BIN=/p/lustre5/${USER}/julia/bin/julia +export PATH="$(dirname "${JULIA_BIN}"):${PATH}" + +# Keep Julia depot on Lustre for compute jobs +export JULIA_DEPOT_PATH=/p/lustre5/${USER}/.julia_depot:${JULIA_DEPOT_PATH:-} + +# --- Dagger checkout (Lustre) --- +export DAGGER_DIR=/p/lustre5/${USER}/dagger-mpi/Dagger.jl + +# --- Test runner defaults (Tuolumne is ROCm-only) --- +export SKIP_CUDA=1 +export ROCM_GPU_PROJECT=${DAGGER_DIR}/test/rocmenv + +# Optional: GPU-aware MPI (not required for current host-staged MPI GPU tests) +# export MPICH_GPU_SUPPORT_ENABLED=1 +# export MPICH_OFI_NIC_POLICY=GPU diff --git a/docs/make.jl b/docs/make.jl index 1e05647d8..611f7301d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -32,6 +32,7 @@ makedocs(; "Stencils" => "stencils.md", ], "GPUs" => "gpu.md", + "MPI" => "mpi.md", "Option Propagation" => "propagation.md", "Logging and Visualization" => [ "Logging: Basics" => "logging.md", diff --git a/docs/src/datadeps.md b/docs/src/datadeps.md index a4b6fe6bd..0c303c0a7 100644 --- a/docs/src/datadeps.md +++ b/docs/src/datadeps.md @@ -22,6 +22,10 @@ Dagger.spawn_datadeps() do end ``` +Datadeps also runs across distributed memory on top of MPI, using the same API +shown here; see the [MPI](mpi.md) page for how to launch and scope datadeps +regions across cluster ranks. + In this example, we have two Dagger tasks being launched, one adding `A` into `B`, and the other copying `B` into `C`. The `add!` task is specifying that `A` is being only read from (`In` for "input"), and that `B` is being read diff --git a/docs/src/gpu.md b/docs/src/gpu.md index 626de7f56..e13d5ec61 100644 --- a/docs/src/gpu.md +++ b/docs/src/gpu.md @@ -4,6 +4,10 @@ Dagger supports GPU acceleration for CUDA, ROCm (AMD), Intel oneAPI, Metal (Appl Dagger's GPU support is built on top of the [KernelAbstractions.jl](https://github.com/JuliaGPU/KernelAbstractions.jl) package, as well as the specific GPU-specific packages for each backend (e.g. [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl), [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl), [oneAPI.jl](https://github.com/JuliaGPU/oneAPI.jl), [Metal.jl](https://github.com/JuliaGPU/Metal.jl), and [OpenCL.jl](https://github.com/JuliaGPU/OpenCL.jl)). Dagger's GPU support is designed to be fully interoperable with the Julia GPU ecosystem, allowing you to use Dagger to distribute your GPU computations across multiple devices. +GPU computing also composes with Dagger's MPI backend for multi-node, +distributed-memory GPU execution; see the [MPI](mpi.md) page (in particular the +GPUs-with-MPI section) for details. + There are a few ways to use Dagger's GPU support: 1. **KernelAbstractions**: Use the `KernelAbstractions.jl` interface to write GPU kernels, and then use `Dagger.Kernel` and `Dagger.@spawn` to execute them. @@ -359,4 +363,25 @@ results = map(workers_with_gpus) do scope fetch(Dagger.@spawn sum(DA)) end end -``` \ No newline at end of file +``` + +## MPI, staging, and same-node IPC + +All GPU backends expose the same Dagger hooks used by the MPI acceleration and +host staging pools: + +| Hook | Purpose | +|------|---------| +| `mpi_remap_space` / `value_memory_space` | Stamp owning MPI rank on VRAM spaces | +| `gpu_memory_kind` / `pin_buffer!` | Key and page-lock host staging buffers | +| `mpi_device_sync` / `mpi_device_direct` | Sync and optional GPU-aware MPI sends | +| `ipc_eligible` / `ipc_export` / `ipc_copyto!` | Same-node on-device IPC (when available) | + +### Per-backend capabilities + +- **CUDA**: Full stack — pinned staging, `cuIpc*` same-node IPC, GPU-direct when `MPI.has_cuda()` (override with `DAGGER_MPI_GPU_DIRECT=0/1`). +- **ROCm (AMD)**: Pinned staging, `hipIpc*` same-node IPC, GPU-direct when `MPI.has_rocm()` (same env override). +- **oneAPI (Intel)**: Level Zero `zeMem*IpcHandle` same-node IPC. GPU-direct is **opt-in** because MPI.jl has no `has_ze`: set `DAGGER_MPI_GPU_DIRECT=1`, or use Intel MPI with `I_MPI_OFFLOAD≥1` (optional `JULIA_MPI_HAS_ONEAPI=true`). +- **Metal / OpenCL**: Same remap / sync / staging-pool hooks; **no** CUDA-style process IPC and **no** GPU-aware MPI path in Julia — transfers host-stage. `DAGGER_IPC=0` disables IPC on backends that support it. + +Dedicated MPI smoke suites live under `test/mpi_{cuda,rocm,oneapi,opencl,metal}.jl` with matching `test/*env` projects. diff --git a/docs/src/mpi.md b/docs/src/mpi.md new file mode 100644 index 000000000..431e24239 --- /dev/null +++ b/docs/src/mpi.md @@ -0,0 +1,377 @@ +# MPI (Distributed-Memory Execution) + +Dagger can run on top of [MPI](https://juliaparallel.org/MPI.jl/stable/) to +scale across the nodes of a cluster or supercomputer, while keeping the same +programming model you already use with the Distributed backend: ordinary +tasks, [`DArray`s](darray.md), and [Datadeps](datadeps.md) regions. This is implemented as a +package extension (`MPIExt`) that loads automatically when `MPI` is imported +alongside `Dagger`, and is activated by switching Dagger's *acceleration* to MPI: + +```julia +using Dagger, MPI +Dagger.accelerate!(:mpi) +``` + +The remainder of this page explains the execution model, how to launch and +write MPI programs with Dagger, how data ownership and scheduling work across +ranks, how to target specific ranks (and GPUs) with scopes, and how to test +MPI code. + +## Distributed vs. MPI: which backend? + +Dagger has two execution backends ("accelerations"): + +- **Distributed** (the default): a single "head" process drives scheduling and + farms tasks out to worker processes added with `addprocs`. Workers can be + added and removed dynamically, and the head process holds the global view of + the task graph. This is ideal for interactive use, heterogeneous clusters, + and elastic workloads. +- **MPI**: every rank runs the *same* program (SPMD — Single Program, Multiple + Data). There is no central head process; instead, all ranks cooperatively and + deterministically make the same scheduling decisions, and exchange data + point-to-point over MPI. This is the model expected by most HPC schedulers + (Slurm, Flux, PBS) and is the right choice when you need to run across many + nodes launched by `mpiexec`/`srun`/`flux run`. + +The key practical difference is *how you launch* and *how data is owned* — the +task and array APIs are shared. Code written against `spawn_datadeps` and +`DArray` will generally run under either backend. + +## Requirements and setup + +1. Add `MPI` to your project and configure it as usual for MPI.jl (the default + MPICH JLL works out of the box; use + [`MPIPreferences`](https://juliaparallel.org/MPI.jl/stable/configuration/) + to bind a system MPI on a cluster). +2. Import both packages and switch acceleration to MPI: + +```julia +using Dagger, MPI +Dagger.accelerate!(:mpi) +``` + +`accelerate!(:mpi)` initializes MPI if it has not been initialized yet +(`MPI.Init` with thread level `:multiple`), registers one Dagger processor set +per rank, and builds the per-node topology used for local/IPC fast paths. It +uses `MPI.COMM_WORLD` by default. + +!!! note "Thread support" + Dagger uses multiple threads within each rank, so MPI must be initialized + with `MPI.THREAD_MULTIPLE`. `accelerate!(:mpi)` requests this for you when it + initializes MPI; if you call `MPI.Init` yourself beforehand, request + `threadlevel=:multiple`. + +## Launching an MPI program + +An MPI Dagger program is a normal Julia script launched under `mpiexec` (or your +site's launcher), with one Julia process per rank. Give each rank multiple +threads so Dagger can exploit intra-rank parallelism: + +```sh +mpiexec -n 4 julia --project --threads=4 my_program.jl +``` + +If you do not have a system `mpiexec` handy, Dagger's test helper +`test/run_mpi.jl` launches a script through the `mpiexec` that MPI.jl itself +provides: + +```sh +# julia --project= test/run_mpi.jl