examples: Reframe the MPI data-indexing benchmark

The first version pitted distributed devito against single-process NumPy
as if competing, which is the wrong comparison. Reframe into two studies,
auto-selected by the launch size, each making one point:

* serial (no mpirun): the data[idx] wrapper adds only a small overhead
  over NumPy, with no per-element Python loop;
* parallel (mpirun -n N): a routed assignment's cost scales with the
  volume communicated, while the NumPy reference (pure compute) is
  unaffected by the rank count.

Sweeps the number of indexed/communicated points and prints a table per
study.

Author: mloubout

examples/mpi/benchmark_data_indexing.py

@@ -1,137 +1,168 @@
 """
-Benchmark Devito's MPI-distributed data indexing against single-process NumPy.
+Benchmark the overhead of Devito's ``data[idx]`` get/set over plain NumPy.
 
-Devito turns ``data[idx]`` get/set on an MPI-distributed array into the minimal
-point-to-point exchange (see ``devito/data/distributed``). This script measures
-that machinery against the equivalent operation on a single-process NumPy array
-holding the full domain, for three representative cases:
+The script runs one of two studies, picked automatically from the MPI launch
+size, each making a single point:
 
-* a routed scatter assignment ``data[idx] = values`` (global-index labelled,
-  the path that moves data across ranks);
-* the matching routed read ``data[idx]``;
-* a basic strided slice ``data[::2]`` (local/induced, no communication).
+* **Serial** (run without ``mpirun``) -- with MPI off, ``data[idx]`` is a thin
+  wrapper over NumPy indexing. The benchmark shows the overhead is small: a
+  fixed cost plus a low per-point cost (vectorised, with no per-element Python
+  loop), so for realistic operation sizes it stays negligible. ::
 
-For the routed cases the distributed timing is reported both *cold* (first call,
-including plan construction) and *warm* (cached plan -- the steady state of a
-time loop, e.g. sparse injection every step). The NumPy column is the
-single-process reference for the same logical operation on the full domain.
+      python examples/mpi/benchmark_data_indexing.py
 
-Run, for example::
+* **Parallel** (run under ``mpirun -n N`` with ``N > 1``) -- a routed
+  assignment ``data[idx] = values`` moves data point-to-point. The benchmark
+  sweeps the number of communicated points and shows the added cost scales with
+  that volume (as expected for communication), while the per-rank NumPy work is
+  not slowed down -- the NumPy reference time is independent of the rank count. ::
 
-    DEVITO_MPI=1 mpirun -n 4 python examples/mpi/benchmark_data_indexing.py
-    DEVITO_MPI=1 mpirun -n 4 python examples/mpi/benchmark_data_indexing.py \\
-        --shape 1024 1024 --reps 50
+      DEVITO_MPI=1 mpirun -n 4 python examples/mpi/benchmark_data_indexing.py
+
+Pass ``--reps`` to change the number of timed repetitions.
 """
 
 import argparse
+import os
 from time import perf_counter
 
 import numpy as np
 
 from devito import Function, Grid, configuration
-from devito.data.distributed import exchange as _exchange
-from devito.mpi import MPI
+
+# Global axis length used by both studies; large enough to distribute, small
+# enough to stay in memory on one process.
+AXIS = 1 << 20
+
+# Numbers of indexed/communicated points to sweep.
+POINTS = [16, 128, 1024, 8192, 65536]
 
 
-def _measure(call, reps, comm, clear=None):
+def _set(target, idx, vals):
+    """Assign ``target[idx] = vals`` (``target`` a NumPy array or ``f.data``)."""
+    target[idx] = vals
+
+
+def _median_time(func, args, reps, comm=None):
     """
-    Median wall time of ``call`` over ``reps`` repetitions.
+    Median wall time of ``func(*args)`` over ``reps`` repetitions, after a
+    warm-up call.
 
-    The wall time of one repetition is the slowest rank (an ``MPI.MAX`` reduce
-    of the per-rank time). With ``clear`` given it is invoked before every timed
-    call to drop the plan cache, measuring the cold path; otherwise one warm-up
-    call is made first and the cached path is measured.
+    Under MPI the per-repetition time is the slowest rank (an ``MPI.MAX``
+    reduction), since a collective operation finishes only when every rank does.
     """
-    if clear is None:
-        call()
+    func(*args)
     times = []
     for _ in range(reps):
-        if clear is not None:
-            clear()
-        comm.Barrier()
+        if comm is not None:
+            comm.Barrier()
         start = perf_counter()
-        call()
+        func(*args)
         elapsed = perf_counter() - start
-        times.append(comm.allreduce(elapsed, op=MPI.MAX))
+        if comm is not None:
+            from devito.mpi import MPI
+            elapsed = comm.allreduce(elapsed, op=MPI.MAX)
+        times.append(elapsed)
     return float(np.median(times))
 
 
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument('--shape', type=int, nargs='+', default=[512, 512],
-                        help='global grid shape (default: 512 512)')
-    parser.add_argument('--reps', type=int, default=30,
-                        help='timed repetitions per operation (default: 30)')
-    parser.add_argument('--npoint', type=int, default=None,
-                        help='routed points (default: an eighth of axis 0)')
-    args = parser.parse_args()
+def _distinct_indices(npoint):
+    """``npoint`` distinct global indices, evenly spread over the axis."""
+    step = max(1, AXIS // npoint)
+    return (np.arange(npoint) * step) % AXIS
 
-    configuration['mpi'] = True
 
-    shape = tuple(args.shape)
-    # Creating the Grid initialises the MPI distributor (and MPI itself)
-    grid = Grid(shape=shape)
-    comm = grid.distributor.comm
-    rank, nprocs = grid.distributor.myrank, grid.distributor.nprocs
+def serial_study(reps):
+    """Overhead of ``data[idx]`` over NumPy with MPI off (single process)."""
+    configuration['mpi'] = False
+    grid = Grid(shape=(AXIS,))
     f = Function(name='f', grid=grid, dtype=np.float32)
-
-    # Full replicated reference array; also the NumPy single-process baseline
-    full = np.arange(int(np.prod(shape)), dtype=f.dtype).reshape(shape)
+    full = np.arange(AXIS, dtype=f.dtype)
     f.data[:] = full
 
-    # Global row indices to route (reversed, so they cross rank boundaries);
-    # each rank contributes the slice it would hold in an external layout
-    npoint = args.npoint or max(1, shape[0] // 8)
-    glb_idx = np.arange(shape[0])[::-1][:npoint]
-    loc_idx = np.array_split(glb_idx, nprocs)[rank]
-    loc_vals = full[loc_idx]
-    ref = full.copy()
-
-    clear = _exchange._build.cache_clear
-
-    rows = []
-
-    # 1. Routed scatter assignment: data[idx] = values
-    rows.append((
-        'scatter  data[idx] = v',
-        _measure(lambda: ref.__setitem__(glb_idx, full[glb_idx]), args.reps, comm),
-        _measure(lambda: f.data.__setitem__(loc_idx, loc_vals), args.reps, comm,
-                 clear=clear),
-        _measure(lambda: f.data.__setitem__(loc_idx, loc_vals), args.reps, comm),
-    ))
-
-    # 2. Routed gather read: data[idx]
-    rows.append((
-        'gather   data[idx]',
-        _measure(lambda: ref[glb_idx], args.reps, comm),
-        _measure(lambda: f.data[loc_idx], args.reps, comm, clear=clear),
-        _measure(lambda: f.data[loc_idx], args.reps, comm),
-    ))
-
-    # 3. Basic strided slice (local/induced, no communication)
-    rows.append((
-        'slice    data[::2]',
-        _measure(lambda: ref[::2, ::2], args.reps, comm),
-        None,
-        _measure(lambda: f.data[::2, ::2], args.reps, comm),
-    ))
+    print(f"\nSerial overhead of data[idx] = v over NumPy "
+          f"(axis={AXIS}, reps={reps})\n")
+    header = ('points'.rjust(10) + 'numpy [us]'.rjust(14)
+              + 'devito [us]'.rjust(14) + 'overhead [us]'.rjust(16))
+    print(header)
+    print('-' * len(header))
+    for npoint in POINTS:
+        idx = _distinct_indices(npoint)
+        vals = full[idx]
+        ref = full.copy()
+        t_np = _median_time(_set, (ref, idx, vals), reps)
+        t_dv = _median_time(_set, (f.data, idx, vals), reps)
+        print(str(npoint).rjust(10)
+              + format(t_np * 1e6, '.1f').rjust(14)
+              + format(t_dv * 1e6, '.1f').rjust(14)
+              + format((t_dv - t_np) * 1e6, '.1f').rjust(16))
+    print("\noverhead = devito - numpy; a small fixed cost plus a low per-point "
+          "cost (no\nper-element Python loop) -- tens of microseconds at "
+          "realistic sizes.")
+
+
+def mpi_study(comm, nprocs, reps):
+    """Communication cost of a routed assignment vs the volume exchanged."""
+    configuration['mpi'] = True
+    grid = Grid(shape=(AXIS,))
+    f = Function(name='f', grid=grid, dtype=np.float32)
+    full = np.arange(AXIS, dtype=f.dtype)
+    f.data[:] = full
+    rank = grid.distributor.myrank
 
     if rank == 0:
-        print(f"\nMPI distributed data indexing vs NumPy "
-              f"(shape={shape}, ranks={nprocs}, "
-              f"routed points={npoint}, reps={args.reps})\n")
-        header = ('operation'.ljust(24) + 'numpy [us]'.rjust(14)
-                  + 'devito cold [us]'.rjust(18) + 'devito warm [us]'.rjust(18))
+        print(f"\nRouted assignment data[idx] = v under MPI "
+              f"(axis={AXIS}, ranks={nprocs}, reps={reps})\n")
+        header = ('points'.rjust(10) + 'KiB sent'.rjust(12)
+                  + 'numpy [us]'.rjust(14) + 'devito [us]'.rjust(14)
+                  + 'comm [us]'.rjust(14))
         print(header)
         print('-' * len(header))
-        for name, np_t, cold, warm in rows:
-            cold_s = '-' if cold is None else format(cold * 1e6, '.1f')
-            print(name.ljust(24)
-                  + format(np_t * 1e6, '.1f').rjust(14)
-                  + cold_s.rjust(18)
-                  + format(warm * 1e6, '.1f').rjust(18))
-        print("\ncold = first call (plan construction); "
-              "warm = cached plan (steady state of a time loop).")
+
+    for npoint in POINTS:
+        global_idx = _distinct_indices(npoint)
+        # Each rank supplies the points it would hold in an external layout
+        # (reversed, so the values cross rank boundaries when routed).
+        loc_idx = np.array_split(global_idx[::-1], nprocs)[rank]
+        loc_vals = full[loc_idx]
+        ref = full.copy()
+
+        # NumPy reference: the whole operation on a single process. It is the
+        # pure-compute cost and is unaffected by the rank count.
+        t_np = _median_time(_set, (ref, global_idx, full[global_idx]), reps, comm)
+        # Devito: routed, point-to-point, with the plan cached (steady state).
+        t_dv = _median_time(_set, (f.data, loc_idx, loc_vals), reps, comm)
+
+        if rank == 0:
+            kib = npoint * full.dtype.itemsize / 1024
+            print(str(npoint).rjust(10)
+                  + format(kib, '.1f').rjust(12)
+                  + format(t_np * 1e6, '.1f').rjust(14)
+                  + format(t_dv * 1e6, '.1f').rjust(14)
+                  + format((t_dv - t_np) * 1e6, '.1f').rjust(14))
+
+    if rank == 0:
+        print("\ncomm = devito - numpy; grows with the volume communicated, as "
+              "expected,\nwhile the numpy reference (pure compute) is unchanged.")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--reps', type=int, default=20,
+                        help='timed repetitions per point count (default: 20)')
+    args = parser.parse_args()
+
+    # Detect the launch size without initialising MPI (so the serial study can
+    # keep MPI off and exercise the pure-NumPy path).
+    nprocs = int(os.environ.get('OMPI_COMM_WORLD_SIZE',
+                                os.environ.get('PMI_SIZE', '1')))
+
+    if nprocs == 1:
+        serial_study(args.reps)
+    else:
+        configuration['mpi'] = True
+        grid = Grid(shape=(8,))  # initialise the distributor (and MPI)
+        mpi_study(grid.distributor.comm, nprocs, args.reps)
 
 
 if __name__ == '__main__':