Distributed arrays and halo cells

Download notebookView on GitHubRun when the site was built.

A tour of mpiarray: how an array is split over MPI ranks, what the halo cells hold, and how reductions and operators work. This page runs on a single rank; the same code runs unchanged under mpiexec -n N.

import numpy as np

import mpiarray as mpa

a = mpa.array([1, 2, 3, 4])
a  # printing shows this rank's block and never communicates

DistributedArray(shape=(4,), dtype=int64, split=(0,), rank 0 of 1 holds [0:4]: [1, 2, 3, 4])

Splitting an array

Arrays are split along their first axis by default; split chooses other axes, and split=None gives every rank the whole array. mpa.process_grid shows how the ranks are shared for any rank count, without MPI:

for size in (4, 6, 12, 7):
    print(size, "ranks ->", mpa.process_grid(size, 2, (0, 1)))

4 ranks -> (2, 2)
6 ranks -> (2, 3)
12 ranks -> (3, 4)
7 ranks -> (1, 7)

A prime number of ranks cannot be split over two axes, so it all goes to the last one.

Every array has a layout that tells each rank its place: the process grid, its neighbours and the index range it owns.

b = mpa.zeros((64, 48), split=(0, 1), periodic=(True, False))
layout = b.layout
print("process grid:", layout.process_grid)
print("neighbours (left, right) per axis:", layout.neighbours)
print("owned index ranges:", layout.index_bounds)

process grid: (1, 1)
neighbours (left, right) per axis: ((0, 0), (-2, -2))
owned index ranges: ((0, 64), (0, 48))

On one rank, the periodic axis has this rank as its own neighbour, and the wall axis has MPI.PROC_NULL on both sides.

Halo cells

With halo=1, each rank stores its block with one extra cell on each side. local is the block, local_with_halos the whole storage:

u = mpa.arange(8.0, halo=1, periodic=True)
print("block:      ", u.local)
print("with halos: ", u.local_with_halos)

block:       [0. 1. 2. 3. 4. 5. 6. 7.]
with halos:  [0. 0. 1. 2. 3. 4. 5. 6. 7. 0.]

update_halos() copies the neighbours’ boundary values into the halo cells, as a finite-difference stencil needs. On this periodic axis they wrap around:

u.update_halos()
print(u.local_with_halos)

# a centred difference only needs the local block with its halos
v = u.local_with_halos
print("second difference:", v[:-2] - 2 * v[1:-1] + v[2:])

[7. 0. 1. 2. 3. 4. 5. 6. 7. 0.]
second difference: [ 8.  0.  0.  0.  0.  0.  0. -8.]

accumulate_halos() goes the other way: values written into halo cells, for example charge deposited by particles near the edge of a block, are added into the neighbour’s block, and the halo cells are zeroed:

rho = mpa.zeros(8, halo=1, periodic=True)
rho.local_with_halos[0] = 1.0  # deposited left of the first cell
rho.local_with_halos[-1] = 2.0  # deposited right of the last one
rho.accumulate_halos()
print(rho.local_with_halos)

[0. 2. 0. 0. 0. 0. 0. 0. 1. 0.]

Reductions and operators

Reductions are collective over the ranks, skip the halo cells and return the same host scalar on every rank. Arithmetic and NumPy ufuncs return new distributed arrays with the same layout:

x = mpa.linspace(0.0, 1.0, 8)
y = 2 * x + np.sin(x)
print("sum:", y.sum(), " max:", y.max(), " norm:", y.norm())
print("element 3:", y[3])
print("gathered on every rank:", y.gather())

sum: 11.633144897399504  max: 2.8414709848078967  norm: 4.887094305532312
element 3: 1.272714712135909
gathered on every rank: [0.         0.42808602 0.85327142 1.27271471 1.68369136 2.08364933
 2.47026108 2.84147098]

On several ranks

Save the cells above as a script and run it with mpiexec -n 4 python tutorial.py: each rank then holds its own block, update_halos and accumulate_halos exchange messages with the neighbour ranks, and sum, norm, y[3] and gather communicate, so every rank must call them. On the CuPy backend (CUNUMPY_BACKEND=cupy), call xp.mpi.set_mpi_cuda_aware(...) or xp.mpi.mpi_is_cuda_aware(comm) once first.