numba-mpi

Overview

numba-mpi provides Python wrappers to the C MPI API callable from within Numba JIT-compiled code (@jit mode). For an outline of the project, rationale, architecture, and features, refer to: numba-mpi paper in SoftwareX (open access) (please cite if numba-mpi is used in your research).

Support is provided for a subset of MPI routines covering: size/rank, send/recv, allreduce, reduce, bcast, scatter/gather & allgather, barrier, wtime and basic asynchronous communication with isend/irecv (only for contiguous arrays); for request handling including wait/waitall/waitany and test/testall/testany.

The API uses NumPy and supports both numeric and character datatypes (e.g., broadcast). Auto-generated docstring-based API docs are published on the web: https://numba-mpi.github.io/numba-mpi

Packages can be obtained from PyPI, Conda Forge, Arch Linux or by invoking pip install git+https://github.com/numba-mpi/numba-mpi.git.

numba-mpi is a pure-Python package. The codebase includes a test suite used through the GitHub Actions workflows (thanks to mpi4py's setup-mpi!) for automated testing on: Linux (MPICH, OpenMPI & Intel MPI), macOS (MPICH & OpenMPI) and Windows (MS MPI). Note, that some of those combinations may not be fully supported yet - see Known Issues for more information.

Features that are not implemented yet include (help welcome!): - support for non-default communicators - support for MPI_IN_PLACE in [all]gather/scatter and allreduce - support for MPI_Type_create_struct (Numpy structured arrays) - ...

Hello world send/recv example:

```python import numba, numba_mpi, numpy

@numba.jit() def hello(): src = numpy.array([1., 2., 3., 4., 5.]) dsttst = numpy.emptylike(src)

if numba_mpi.rank() == 0:
    numba_mpi.send(src, dest=1, tag=11)
elif numba_mpi.rank() == 1:
    numba_mpi.recv(dst_tst, source=0, tag=11)

hello() ```

Example comparing numba-mpi vs. mpi4py performance:

The example below compares Numba+mpi4py vs. Numba+numba-mpi performance. The sample code estimates $\pi$ by numerical integration of $\int0^1 (4/(1+x^2))dx=\pi$ dividing the workload into `nintervalshandled by separate MPI processes and then obtaining a sum usingallreduce(see, e.g., analogous [Matlab docs example](https://www.mathworks.com/help/parallel-computing/numerical-estimation-of-pi-using-message-passing.html)). The computation is carried out in a JIT-compiled functiongetpipart()and is repeated NTIMES. The repetitions and the MPI-handled reduction are done outside or inside of the JIT-compiled block formpi4pyandnumba-mpi, respectively. Timing is repeatedNREPEATtimes and the minimum time is reported. The generated plot shown below depicts the speedup obtained by replacingmpi4py withnumbampi, plotted as a function ofNTIMES / nintervals- the number of MPI calls per interval. The speedup, which stems from avoiding roundtrips between JIT-compiled and Python code is significant (150%-300%) in all cases. The more often communication is needed (smallernintervals), the larger the measured speedup. Note that nothing in the actual number crunching (within thegetpipart()function) or in the employed communication logic (handled by the same MPI library) differs between thempi4pyornumba-mpisolutions. These are the overhead ofmpi4pyhigher-level abstractions and the overhead of repeatedly entering and leaving the JIT-compiled block if usingmpi4py, which can be eliminated by usingnumba-mpi, and which the measured differences in execution time stem from. ``python import timeit, mpi4py, numba, numpy as np, numba_mpi

N_TIMES = 10000 RTOL = 1e-3

@numba.jit def getpipart(nintervals=1000000, rank=0, size=1): h = 1 / nintervals partialsum = 0.0 for i in range(rank + 1, nintervals, size): x = h * (i - 0.5) partialsum += 4 / (1 + x**2) return h * partialsum

@numba.jit def pinumbampi(nintervals): pi = np.array([0.]) part = np.emptylike(pi) for _ in range(NTIMES): part[0] = getpipart(nintervals, numbampi.rank(), numbampi.size()) numbampi.allreduce(part, pi, numbampi.Operator.SUM) assert abs(pi[0] - np.pi) / np.pi < RTOL

def pimpi4py(nintervals): pi = np.array([0.]) part = np.emptylike(pi) for _ in range(NTIMES): part[0] = getpipart(nintervals, mpi4py.MPI.COMMWORLD.rank, mpi4py.MPI.COMMWORLD.size) mpi4py.MPI.COMMWORLD.Allreduce(part, (pi, mpi4py.MPI.DOUBLE), op=mpi4py.MPI.SUM) assert abs(pi[0] - np.pi) / np.pi < RTOL

plotx = [x for x in range(1, 11)] ploty = {'numbampi': [], 'mpi4py': []} for x in plotx: for impl in ploty: ploty[impl].append(min(timeit.repeat( f"pi{impl}(nintervals={N_TIMES // x})", globals=locals(), number=1, repeat=10 )))

if numbampi.rank() == 0: from matplotlib import pyplot pyplot.figure(figsize=(8.3, 3.5), tightlayout=True) pyplot.plot(plotx, np.array(ploty['mpi4py'])/np.array(ploty['numbampi']), marker='o') pyplot.xlabel('number of MPI calls per interval') pyplot.ylabel('mpi4py/numba-mpi wall-time ratio') pyplot.title(f'mpiexec -np {numbampi.size()}') pyplot.grid() pyplot.savefig('readmeplot.svg') ```

Known Issues

NOTE: Issues listed below only relate to combinations of platforms and MPI distributions that we target to support, but due to various reason are currently not working and are temporarily excluded from automated testing:

• tests on Ubuntu 2024.4 that use MPICH are not run due to failures caused by newer version of MPICH (4.2.0); note, that previous tests ran using version 4.0.2 of MPICH (that is installed by default on Ubuntu 2022.4 using apt) were passing (see related issue - TODO #162),
• tests on MacOS (both Intel- and ARM-based) that use OpenMPI are currently not run due to failures being under investigation (see related issue - TODO #163).

MPI resources on the web:

• MPI standard and general information:
  • https://www.mpi-forum.org/docs
  • https://en.wikipedia.org/wiki/MessagePassingInterface
• MPI implementations:
  • OpenMPI: https://www.open-mpi.org
  • MPICH: https://www.mpich.org
  • MS MPI: https://learn.microsoft.com/en-us/message-passing-interface
  • Intel MPI: https://intel.com/content/www/us/en/developer/tools/oneapi/mpi-library-documentation.html
• MPI bindings:
  • Python: https://mpi4py.readthedocs.io
  • Python/JAX: https://mpi4jax.readthedocs.io
  • Julia: https://juliaparallel.org/MPI.jl
  • Rust: https://docs.rs/mpi
  • C++: https://boost.org/doc/html/mpi.html
  • R: https://cran.r-project.org/web/packages/Rmpi

Acknowledgements:

We thank all contributors and users who reported feedback to the project through GitHub issues.

Development of numba-mpi has been supported by the Polish National Science Centre (grant no. 2020/39/D/ST10/01220), the Max Planck Society and the European Union (ERC, EmulSim, 101044662). We further acknowledge Poland’s high-performance computing infrastructure PLGrid (HPC Centers: ACK Cyfronet AGH) for providing computer facilities and support within computational grant no. PLG/2023/016369.