MPI and halos

Download notebook View on GitHubRun when the site was built.

Split a field over MPI ranks: blocks, halo nodes, deposits that cross block edges, and gathers near them.

With comm and domain_decomposition, each rank stores its own block of the grid and num_ghost_points halo nodes on every side (the layout is mpiarray’s). A deposit into a rank’s block can reach into its halo nodes; exchange_halos adds them to the neighbours’ nodes. Before a gather, fill_halos copies the neighbours’ nodes into the halos. Everything global (sum, to_numpy, evaluate, export) is collective. The same script runs on one process, where nothing is split.

import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

A script for several ranks

Every rank keeps the markers inside its block: from its first node up to its last node plus one cell (the cell between its last node and the next rank’s first). The deposit then completes on the neighbours after exchange_halos. The script compares the result with the serial deposit of all markers.

script = r"""
import numpy as np
from maybempi import MPI

from mpifield import Field

comm = MPI.COMM_WORLD
kwargs = {"ndim": 2, "lower_bound": (0.0, 0.0), "upper_bound": (1.0, 1.0)}
positions = np.random.default_rng(0).random((2, 100_000))  # the same on every rank

density = Field(num_gridpoints=(129, 129), comm=comm,
                domain_decomposition=[True, True], num_ghost_points=1, **kwargs)
grid = density.block_grid(cpu=True)  # the stored block, halo nodes included
first = grid.bounds.lower_bound + grid.grid_spacing
end = first + np.array(density.local.shape[:2]) * grid.grid_spacing
mine = np.ones(positions.shape[1], dtype=bool)
for d in range(2):
    last = density.proc_index_bounds[d][1] == density.num_gridpoints[d]
    below = positions[d] <= end[d] if last else positions[d] < end[d]
    mine &= (positions[d] >= first[d]) & below
local = np.ascontiguousarray(positions[:, mine])

density.scatter(local, np.ones((1, local.shape[1])))
density.exchange_halos()  # the halo contributions go to the neighbours

serial = Field(num_gridpoints=(129, 129), **kwargs)
serial.scatter(positions, np.ones((1, positions.shape[1])))
# collective calls (to_numpy, sum, gather) run on every rank, outside the if
error = np.abs(density.to_numpy() - serial.to_numpy()).max()
total = density.sum()
blocks = comm.gather((comm.Get_rank(), density.proc_index_bounds, int(mine.sum())))
if comm.Get_rank() == 0:
    for rank, bounds, count in blocks:
        print(f"rank {rank}: nodes {bounds}, {count} markers")
    print(f"total {total:.0f}, largest difference to serial {error:.1e}")
"""
path = Path(tempfile.mkdtemp()) / "deposit.py"
path.write_text(script)

Running it

Serially it is the plain deposit; under mpiexec -n 4 the grid is split into a 2 x 2 process grid. The result is the same to round-off.

print(
    subprocess.run(
        [sys.executable, str(path)], capture_output=True, text=True, check=True
    ).stdout
)

rank 0: nodes [(0, 129), (0, 129)], 100000 markers
total 100000, largest difference to serial 0.0e+00

mpiexec = shutil.which("mpiexec")
if mpiexec is None:
    print("mpiexec is not installed here; run: mpiexec -n 4 python deposit.py")
else:
    run = subprocess.run(
        [mpiexec, "-n", "4", "--oversubscribe", sys.executable, str(path)],
        capture_output=True,
        text=True,
        check=True,
        timeout=600,
    )
    print(run.stdout)

rank 0: nodes [(0, 65), (0, 65)], 25998 markers
rank 1: nodes [(0, 65), (65, 129)], 24902 markers
rank 2: nodes [(65, 129), (0, 65)], 25024 markers
rank 3: nodes [(65, 129), (65, 129)], 24076 markers
total 100000, largest difference to serial 3.6e-15

What to call when

After	Call	Why
`scatter`	`exchange_halos()`	the halo nodes hold contributions for the neighbours’ nodes
changing the nodes, before `gather`	`fill_halos()`	markers in a block’s last cell read the next rank’s first node
`smooth`, `gradient`, `evaluate`, `export`	nothing	they exchange or gather themselves

Field.on_grid(grid) makes a field split like a simgeometry grid, with one halo node when the grid is decomposed: the setup a PIC code with a decomposed solver grid needs. Halo nodes on a periodic axis hold the wrapped-around nodes, so the kernels see a bounded block (block_grid().periodic is all False there).

A weighting says how many halo nodes it needs (LINEAR.halo is 1, QUADRATIC.halo 2), and scatter and gather raise on a split field with fewer. They leave out the halo nodes beyond a bounded edge of the domain, which no exchange fills: markers there deposit on the edge nodes, as on an unsplit field (block_grid(inside_domain=True) is the grid they use).