Halo exchange

Deep dive

How it works

Parts:

• Cartesian rank grid — px × py ranks, periodic or fixed edges (MPI_Cart_create, MPI_Cart_shift).
• Padded local array — (ny+2h) × (nx+2h), interior at [h, h+n).
• Pack/unpack buffers — contiguous copies of strided faces (or an MPI subarray datatype).

Per time step:

1. Post MPI_Irecvs. Phase 1: pack the left/right interior columns, send, unpack into the halo columns.
2. Phase 2: top/bottom rows, full width including the halo columns just received: corners ride along, no diagonal messages (MPI_Isend/MPI_Waitall per phase).
3. Update every interior cell, swap buffers.
4. Overlap: after posting sends, update cells farther than the stencil radius from the edge, wait, then update the outer ring.
5. Deep halo (width k): exchange once, run k steps; step s updates the array shrunk by s+1 cells per side, eating the halo layer by layer.

 phase 1 (▒ columns)   phase 2 (▓ rows, full width)
 ┌─┬──────┬─┐         ┌────────────┐
 │▒│ own  │▒│         │▓▓▓▓▓▓▓▓▓▓▓▓│ corners included
 └─┴──────┴─┘         └────────────┘

Code

Ranks simulated in one process, messages are Vec<f64>. For h dividing T the result equals the single-domain run bit for bit.

fn jacobi(c: f64, n: f64, s: f64, w: f64, e: f64) -> f64 { c + 0.2 * (n + s + w + e - 4.0 * c) } // one expression => bit-identical

type R = (usize, usize, usize, usize); // rows r0..r1, cols c0..c1

/// One rank: (ny+2h) x (nx+2h) cells, interior at [h, h+n).
pub struct Rank { nx: usize, ny: usize, h: usize, u: Vec<f64>, v: Vec<f64> }

impl Rank {
    fn w(&self) -> usize { self.nx + 2 * self.h }
    fn pack(&self, (r0, r1, c0, c1): R) -> Vec<f64> { // strided rectangle -> contiguous message
        let (w, mut b) = (self.w(), Vec::with_capacity((r1 - r0) * (c1 - c0)));
        for i in r0..r1 { b.extend_from_slice(&self.u[i * w + c0..i * w + c1]); }
        b
    }
    fn unpack(&mut self, (r0, r1, c0, c1): R, b: &[f64]) {
        let w = self.w();
        for (k, i) in (r0..r1).enumerate() { self.u[i * w + c0..i * w + c1].copy_from_slice(&b[k * (c1 - c0)..][..c1 - c0]); }
    }
    /// k <= h steps; step s only trusts cells more than s from the edge.
    fn steps(&mut self, k: usize) {
        let (w, rows) = (self.w(), self.ny + 2 * self.h);
        for s in 0..k {
            for i in 1 + s..rows - 1 - s {
                for p in i * w + 1 + s..i * w + w - 1 - s {
                    self.v[p] = jacobi(self.u[p], self.u[p - w], self.u[p + w], self.u[p - 1], self.u[p + 1]);
                }
            }
            std::mem::swap(&mut self.u, &mut self.v);
        }
    }
}

pub struct Decomp { px: usize, py: usize, ranks: Vec<Rank>, pub bytes: usize }

impl Decomp {
    pub fn new(gx: usize, gy: usize, px: usize, py: usize, h: usize, init: impl Fn(usize, usize) -> f64) -> Self {
        let (nx, ny) = (gx / px, gy / py);
        let ranks = (0..px * py).map(|id| {
            let n = (nx + 2 * h) * (ny + 2 * h);
            let mut r = Rank { nx, ny, h, u: vec![0.0; n], v: vec![0.0; n] };
            for i in 0..ny { for j in 0..nx { let w = r.w(); r.u[(i + h) * w + j + h] = init(id / px * ny + i, id % px * nx + j); } }
            r
        }).collect();
        Decomp { px, py, ranks, bytes: 0 }
    }
    /// Columns of interior rows first, then FULL-width rows: corners arrive in phase 2.
    pub fn exchange(&mut self) {
        let (nx, ny, h) = (self.ranks[0].nx, self.ranks[0].ny, self.ranks[0].h);
        let (w, rows) = (nx + 2 * h, ny + 2 * h);
        let plans: [[(R, R, isize, isize); 2]; 2] = [
            [((h, h + ny, h, 2 * h), (h, h + ny, h + nx, w), -1, 0), ((h, h + ny, nx, nx + h), (h, h + ny, 0, h), 1, 0)],
            [((h, 2 * h, 0, w), (h + ny, rows, 0, w), 0, -1), ((ny, ny + h, 0, w), (0, h, 0, w), 0, 1)],
        ];
        for plan in plans {
            let mut inflight = Vec::new(); // like Isend: pack all, then unpack
            for id in 0..self.ranks.len() {
                for (src, dst, dx, dy) in plan {
                    let to = (((id / self.px) as isize + dy).rem_euclid(self.py as isize) as usize) * self.px
                        + ((id % self.px) as isize + dx).rem_euclid(self.px as isize) as usize;
                    let buf = self.ranks[id].pack(src);
                    self.bytes += buf.len() * 8;
                    inflight.push((to, dst, buf));
                }
            }
            for (to, dst, buf) in inflight { self.ranks[to].unpack(dst, &buf); } // like Waitall
        }
    }
    /// h steps per message round.
    pub fn compute(&mut self) { for r in &mut self.ranks { r.steps(r.h); } }
    pub fn gather(&self, gx: usize) -> Vec<f64> {
        let mut g = vec![0.0; gx * self.py * self.ranks[0].ny];
        for (id, r) in self.ranks.iter().enumerate() {
            for i in 0..r.ny { for j in 0..r.nx { g[(id / self.px * r.ny + i) * gx + id % self.px * r.nx + j] = r.u[(i + r.h) * r.w() + j + r.h]; } }
        }
        g
    }
}

/// Single-domain periodic reference.
pub fn reference(gx: usize, gy: usize, steps: usize, init: impl Fn(usize, usize) -> f64) -> Vec<f64> {
    let mut u: Vec<f64> = (0..gx * gy).map(|p| init(p / gx, p % gx)).collect();
    let mut v = u.clone();
    for _ in 0..steps {
        let at = |i: usize, j: usize| u[i % gy * gx + j % gx];
        for i in 0..gy { for j in 0..gx {
            v[i * gx + j] = jacobi(at(i, j), at(i + gy - 1, j), at(i + 1, j), at(i, j + gx - 1), at(i, j + 1));
        } }
        std::mem::swap(&mut u, &mut v);
    }
    u
}

Test: T/h times exchange(); compute();; gather equals reference by to_bits() for h = 1, 2, 3, 6 on several rank grids (compiled, run).

Complexity and constants

• Per rank, n×n block, width h: n² updates, 4hn + 4h² doubles sent; surface/volume ∝ 4/n (3D: 6/n). 1024² grid, 16 ranks, h = 1: 4×4 blocks send 8224 B per step and rank, 16×1 strips 17440 B (counted by the code above).
• Message time α + bytes/β: small blocks are latency-bound (α order of µs on HPC fabrics); the stencil itself is memory-bound (roofline).
• Deep halo k: messages ÷ k, bytes per step nearly equal (8224 → 8448 B at k = 8, measured), extra cells ≈ 5.6 % at n = 256, k = 8 (computed).

Efficiency

• Cost: one step of a 256² f64 block on one Apple M4 Pro core: about 25 µs compute plus 2.7 µs copying four faces in-process, 8.2 KB sent (measured, best of several runs; the network adds α + 8.2 KB/β). Bound: memory bandwidth for large blocks, latency for small ones.
• Levers:
  • Squarer blocks — 4×4 vs 16×1: half the halo bytes, about 3× less pack time (measured) — when the rank grid is free.
  • Contiguous faces — a strided column face cost about 5× a row face at equal bytes (16×1 vs 1×16, measured) — keep the busiest direction contiguous or use a datatype.
  • Overlap — hides up to min(compute, comm) per step — needs asynchronous progress.
  • Deep halo k — messages ÷ k for ≈ (k−1)/n extra work per dimension (computed; ghost cell expansion, Ding & He 2001) — latency-bound, small blocks.
  • GPU-aware MPI or NVSHMEM — device buffers go into the calls, host staging can be skipped (Open MPI, NVSHMEM docs) — GPU-resident grids.
• Naive → standard → state of the art:
  • Per-cell blocking Send/Recv · one message per cell
  • Per-face nonblocking, dimension-ordered, overlapped · 2d messages per step, hidden behind compute
  • Deep halos plus GPU-initiated puts or persistent neighbourhood collectives · messages ÷ k, host out of the path

Implementation notes

• Allocate buffers once (the sketch allocates per face); post receives first (eager vs rendezvous); MPI_Send_init persistent requests cut per-step setup.
• MPI_Neighbor_alltoallw with per-face subarray datatypes packs and sends in one call; with MPI_Dist_graph_create_adjacent listing all 8 (2D) or 26 (3D) neighbours corners need no second phase, but messages are more and smaller.

Pitfalls

• Phase 2 before phase 1 is unpacked corrupts corners (a star stencil needs none).
• Overlap that does not overlap: many MPI libraries progress nonblocking transfers only inside MPI calls; call MPI_Test between tiles.
• Halo narrower than the reach (higher order, k substeps) gives silently wrong edges: assert h >= radius·k.

Variants and alternatives

In production

PETSc DMDA, AMReX FillBoundary, Cabana; most MPI stencil and climate codes.

Implementations

Sources

• MPI Forum, MPI: A Message-Passing Interface Standard, Version 4.1, 2023. standard
• C. Ding, Y. He, A ghost cell expansion method for reducing communications in solving PDE problems, SC 2001. paper
• T. Hoefler, J. L. Träff, Sparse collective operations for MPI, IPDPS 2009. paper
• Docs: PETSc DMDACreate2d, AMReX Basics, Open MPI CUDA, NVSHMEM.