MPI

With MPI enabled, Devito decomposes the grid over the ranks and each rank owns a slab of it. A rank only ever needs the part of the model that covers its own slab, and WaveModel is built so that this is all it reads: a model can be given as a file that is read on demand rather than as an array, and no rank then reads it in full nor holds it in memory. This is what makes models that do not fit on a single node tractable.

1. Enabling MPI

MPI has to be on before the first Grid is created, either through the environment or in the script:

DEVITO_MPI=diag2 mpiexec -n 4 python my_script.py
from devito import configuration

if not configuration['mpi']:
    configuration['mpi'] = 'basic'

DEVITO_MPI selects the mode: basic, diag2 (single node) or overlap2 (multi node). Everything below runs unchanged on a single process.

2. A model read on demand

Any physical parameter of a WaveModel accepts, in place of an array, an object that exposes a shape and can be sliced:

import numpy as np
from recipes import WaveModel

vp = np.load('vp.npy', mmap_mode='r')          # or an HDF5 dataset, or a reader

model = WaveModel(origin, spacing, vp.shape, space_order, vp, nbl=40,
                  bcs="damp")

Devito asks such a value only for the window the calling rank owns, so with four ranks the file is read as four disjoint windows, one per rank, and never as a whole. Points of the absorbing layer repeat the closest model edge point, exactly as for an array.

A memory mapped array and an HDF5 dataset qualify as they are; anything else needs two members:

Member Purpose
shape Model shape, without the absorbing layer
__getitem__(window) window is a tuple of slices, one per axis; returns that block as an array

Parameters derived from the input ones are computed on the windows that are read, so they are not formed in full either. This covers the Lamé parameters, built from vp, vs and b, and delta, capped below epsilon:

model = WaveModel(origin, spacing, vp.shape, space_order, vp, nbl=40,
                  vs=np.load('vs.npy', mmap_mode='r'),
                  b=np.load('b.npy', mmap_mode='r'))

Quantization is the one thing that still needs a full pass: Pdtype=np.float16 derives the quantization range from the parameter, which is read in chunks when it is not in memory.

3. Reading a SEG-Y model

examples/mpi_modeling.py models a shot on the SEG/Chevron 2014 2D benchmark, reading the velocity straight from its SEG-Y file. Traces are the unit of I/O of a SEG-Y file, which is all the window interface needs:

from pysegy import segy_scan


class SegyVelocity:

    def __init__(self, path, scale=1e-3):
        scan = segy_scan(path, keys=['GroupX'])
        if len(scan) != 1:
            raise ValueError(f"{path} scans as {len(scan)} gathers, this "
                             "expects a model")

        self.record = scan[0]
        self.scale = scale

        header = scan.fileheader.bfh
        x0, x1 = self.record.summary['GroupX']

        self.shape = (self.record.ntraces, header.ns)
        self.spacing = ((x1 - x0) / (self.shape[0] - 1), header.dt / 1000)
        self.origin = (x0, 0.)

    def __getitem__(self, window):
        traces, samples = window
        # The record holds its traces as columns, and reads those indexed
        return self.scale * self.record.data[samples, traces].T

A model has no gathers, so it scans as a single record holding every trace, and indexing its data reads the traces of the window and nothing else. The velocities come out in the file’s m/s and are scaled to the km/s the recipes use.

segy_scan indexes the file without reading any trace data, and gives the grid geometry along the way: the number of traces, the sample count and interval, and the range of the trace coordinates. Reading a window is then one call, and the model is built straight from it:

velocity = SegyVelocity(path)
model = WaveModel(velocity.origin, velocity.spacing, velocity.shape,
                  space_order, velocity, nbl=40, bcs="damp")

Running it on four ranks shows what each of them took off the file:

Rank 0: reading traces 0:1910 and samples 0:600 of (3820, 1200)
Rank 1: reading traces 0:1910 and samples 600:1200 of (3820, 1200)
Rank 2: reading traces 1910:3820 and samples 0:600 of (3820, 1200)
Rank 3: reading traces 1910:3820 and samples 600:1200 of (3820, 1200)
(3820, 1200) model, (12.5, 5.0) m spacing, on a (3900, 1280) grid over 4 rank(s)

4. Sources, receivers and results

Sparse objects are distributed too, by the rank that owns the grid point they sit on. Coordinates are set through the global index space, so the same assignment runs on every rank and only the owner stores anything:

solver.src.coordinates.data[0, :] = np.array([source_x, depth])

A tuple is not accepted there under MPI, the right hand side has to be an array. Results are collected with data_gather, which returns the assembled array on the chosen rank and None on the others:

shot = solver.rec.data_gather(rank=0)
if shot is not None:
    ...

The same applies to a wavefield: Function.data_gather(rank=0) assembles it on one rank, and Function.data_local is the rank-local view when each rank is to write its own slab instead.

5. Reductions

Under MPI, Function.data is a rank-local view, so a numpy reduction of it only sees what the calling rank owns. A decision taken on such a reduction makes the ranks diverge, and diverging on a collective call — a halo exchange, an operator that all ranks must enter — hangs the run. recipes.utils provides the global versions, used wherever the model takes such a decision:

from recipes.utils import global_min, global_max

vmax = global_max(model.vp)     # reduced across all ranks

Devito’s mmin/mmax builtins are global too; they run an operator, which is worth it for a reduction inside a kernel but not for a one-off decision on the host.

Back to top