serial mpi

CHANGELOG.md

@@ -41,6 +41,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `KernelCatalog.from_package(..., include_dirs=None)` is now an explicit keyword; by default the source root of the top-level package (the directory containing it) is an include directory of every CUDA kernel, in addition to the kernel's own folder, so kernels can `#include "my_pkg/common.cuh"`.
 
 ### Added
+- `xp.mpi.get_mpi(use_mpi=None)`, `launched_under_mpi()`, `SerialMPI` and `SerialComm`: MPI only when it is used. `launched_under_mpi()` tells from the launcher's environment (Open MPI, MPICH, Intel MPI, PMIx/`srun`, MVAPICH2, Hydra, Cray ALPS/PALS; `CUNUMPY_MPI=1/0` overrides), without importing mpi4py, whether the process was started by `mpirun`/`mpiexec`/`srun`. `get_mpi()` returns `mpi4py.MPI` then, and otherwise a serial stand-in whose `COMM_WORLD` is a communicator of size 1: object collectives return what one rank gets (`allreduce(x)` is `x`, `gather(x)` is `[x]`), buffer collectives copy (`Allreduce`, `Allgather(v)`, `Gather(v)`, `Scatter(v)`, `Sendrecv` to self; nothing with `IN_PLACE`), on NumPy and CuPy buffers, and unsupported methods raise `AttributeError`. The stand-in has the module attributes serial code needs (`IN_PLACE`, reduction ops, datatypes, `PROC_NULL`, `ROOT`, `COMM_NULL`, `Request.Waitall`, `Prequest.Startall`, `Wtime`, ...). It replaces feectools' `MockComm`/`MockMPI`, which returned `None` for every call (`MockComm().allreduce(5)` was `None`, `MockComm().rank` a function). Kept in `cunumpy/_mpi_serial.py` without dependencies on the rest of cunumpy, so that it could become a package of its own.
 - `cunumpy/morton.cuh` and `xp.morton_keys`, `morton_encode`, `morton_decode`, `morton_scales`, `MAX_MORTON_LEVELS`: Morton (Z-order) keys of 2D and 3D points (`uint64`, up to 32 and 21 bits per axis), the same in a kernel and on the host (bit for bit), for sorting particles along a space-filling curve and building quadtrees and octrees from sorted keys.
 - `xp.sort_by_key(keys, *arrays)`: one stable argsort of `keys` applied to every array; returns the sorted keys, the order and the sorted arrays.
 - `Kernel.from_folder(package, ...)`: the kernel of one kernel folder, with the options of `KernelCatalog.from_package` (`host_suffix`, `compile_host`, `dispatch`, `include_dirs`, CUDA options...). Every version in the folder is an implementation: `<name><host_suffix>.py` (`"pyccel"`, compiled with `compile_host`, and `"python"`, uncompiled), `<name>_numba.py`, `<name>_numpy.py` and `<name>_cuda.cu`. A folder's own `__init__.py` can declare `kernel = xp.Kernel.from_folder(__name__, ...)`, so code imports the kernel from where it is written; `from_package` now builds each kernel with it. `Kernel.implementations` lists them, `Kernel.selected(device=False)` tells which one a call runs.

docs/source/api.md

@@ -429,6 +429,48 @@ mapping differs:
 device_id = xp.cuda.set_device_for_rank(mpi_rank)
 ```
 
+### `mpi.get_mpi(use_mpi=None)`, `mpi.launched_under_mpi()`
+
+Importing `mpi4py.MPI` starts MPI (`MPI_Init`), which takes time and makes
+every collective cost something even on one process. `launched_under_mpi()`
+tells, from the environment that `mpirun`/`mpiexec`/`srun` set up and without
+importing mpi4py, whether the process belongs to an MPI job
+(`CUNUMPY_MPI=1`/`0` overrides it). `get_mpi()` returns `mpi4py.MPI` then, and
+otherwise the serial stand-in, so that the same code runs with and without
+MPI:
+
+```python
+MPI = xp.mpi.get_mpi()           # decided once per process
+comm = MPI.COMM_WORLD
+comm.Allreduce(MPI.IN_PLACE, rho, op=MPI.SUM)   # nothing to do on one process
+n_total = comm.allreduce(n_local)               # n_local itself
+if isinstance(MPI, xp.mpi.SerialMPI):
+    ...                                          # a serial run
+```
+
+`get_mpi(True)` imports mpi4py (`ImportError` if missing), `get_mpi(False)`
+returns the stand-in. Launched under MPI without mpi4py installed, `get_mpi()`
+warns (`RuntimeWarning`) and returns the stand-in.
+
+### `mpi.SerialMPI`, `mpi.SerialComm`
+
+The stand-in for `mpi4py.MPI` and its communicators of size 1.
+`SerialComm` returns (object methods) or copies (buffer methods) what one rank
+gets: `bcast`, `allreduce`, `reduce` and `scan` return their argument,
+`gather`/`allgather` return `[x]`, `scatter([x])` returns `x`;
+`Allreduce`/`Reduce`/`Allgather`/`Gather`/`Scatter`/`Alltoall`/`Scan` copy the
+send buffer into the receive buffer (nothing with `IN_PLACE`), the vector
+forms (`Allgatherv`, `Gatherv`, `Scatterv`) use the displacement of rank 0,
+`Bcast` and `Barrier` do nothing, and `sendrecv`/`Sendrecv` work to and from
+rank 0 or `PROC_NULL`. Buffers are NumPy or CuPy arrays, or mpi4py buffer
+specifications (`[array, MPI.DOUBLE]`). Non-blocking versions return completed
+requests. Any other method raises `AttributeError`, so a missing feature shows
+instead of doing nothing. `SerialMPI` has `COMM_WORLD`, `COMM_SELF`,
+`COMM_NULL` (false), `IN_PLACE`, the reduction operations, the common
+datatypes (`isinstance(MPI.DOUBLE, MPI.Datatype)` holds), `PROC_NULL`, `ROOT`,
+`ANY_SOURCE`, `UNDEFINED`, `Comm`/`Intracomm`, `Request`, `Prequest`, `Status`,
+`Wtime()` and `Is_initialized()` (False).
+
 ### `mpi.local_rank()`
 
 The rank of the process within its node, read from the environment variables

docs/source/guides/mpi.md

@@ -51,6 +51,17 @@ rank and assumes ranks are numbered contiguously per node. Prefer
 
 Importing `mpi4py.MPI` initializes MPI by default. Do it after step 1.
 
+To run the same program serially without starting MPI, get the module from
+`xp.mpi.get_mpi()` instead of importing it: it returns `mpi4py.MPI` when the
+process was started by `mpirun`/`mpiexec`/`srun`, and otherwise a serial
+stand-in whose `COMM_WORLD` has size 1 (see the
+[API reference](../api.md)):
+
+```python
+MPI = xp.mpi.get_mpi()
+comm = MPI.COMM_WORLD
+```
+
 ### 3. Check that MPI is CUDA-aware
 
 Only an MPI library built with CUDA support can send and receive CuPy arrays

src/cunumpy/LLM_GUIDE.md

@@ -142,6 +142,10 @@ Only transfers through cunumpy are counted (not raw `cupy.asarray`, `.get()`,
 MPI, accumulation and versions:
 
 ```python
+MPI = (
+    xp.mpi.get_mpi()
+)  # mpi4py.MPI under mpirun/srun, else a serial stand-in (no MPI_Init)
+xp.mpi.launched_under_mpi()  # from the launcher env, without importing mpi4py
 xp.mpi.mpi_is_cuda_aware(comm)  # collective, once at startup; remembered
 with xp.mpi.mpi_buffer(a) as buf:
     comm.Send(buf, ...)  # host array, CUDA-aware device

src/cunumpy/__init__.py

@@ -41,7 +41,19 @@
     **dict.fromkeys(kernels.__all__, "kernels"),
     **dict.fromkeys(rng.__all__, "rng"),
     **dict.fromkeys(algorithms.__all__, "algorithms"),
-    **dict.fromkeys(mpi.__all__, "mpi"),
+    # the names of cunumpy.mpi that were at the top level (not the later ones)
+    **dict.fromkeys(
+        (
+            "get_mpi_cuda_aware",
+            "local_rank",
+            "mpi_buffer",
+            "mpi_is_cuda_aware",
+            "require_cuda_aware_mpi",
+            "set_mpi_cuda_aware",
+            "synchronize_for_mpi",
+        ),
+        "mpi",
+    ),
     **dict.fromkeys(profiling.__all__, "profiling"),
     **dict.fromkeys(memory.__all__, "memory"),
     "petsc_vec": "petsc",

src/cunumpy/_mpi.py

@@ -3,53 +3,21 @@
 from __future__ import annotations
 
 import logging
-import os
 from collections.abc import Generator
 from contextlib import contextmanager
 from typing import Any
 
 import array_api_compat
 import array_api_compat.numpy as np
 
+from ._mpi_serial import _LOCAL_RANK_VARIABLES, local_rank  # noqa: F401 - re-exported
 from ._transfers import _ACTIVE as _COUNTERS
 from ._transfers import _describe, _record
 from .xp import array_backend, cupy_available, to_numpy
 
 _logger = logging.getLogger(__name__)
 
 
-# Node-local rank of the process, as exported by common MPI launchers. They are
-# set before ``MPI_Init``, so the device can be chosen before MPI starts.
-_LOCAL_RANK_VARIABLES = (
-    "OMPI_COMM_WORLD_LOCAL_RANK",  # Open MPI
-    "MV2_COMM_WORLD_LOCAL_RANK",  # MVAPICH2
-    "MPI_LOCALRANKID",  # Intel MPI, MPICH (Hydra)
-    "PMI_LOCAL_RANK",  # MPICH / PMI
-    "PALS_LOCAL_RANKID",  # Cray PALS
-    "SLURM_LOCALID",  # Slurm (srun)
-    "LOCAL_RANK",  # torchrun and others
-)
-
-
-def local_rank() -> int:
-    """Rank of this process within its node, from the MPI launcher's environment.
-
-    Reads the node-local rank that common launchers export (Open MPI, MVAPICH2,
-    Intel MPI/MPICH, PMI, Cray PALS, Slurm, ``LOCAL_RANK``). These variables are
-    set before ``MPI_Init``, so this works before MPI is initialized, and
-    without importing ``mpi4py``. Returns 0 if none is set (e.g. a serial run).
-    """
-    for variable in _LOCAL_RANK_VARIABLES:
-        value = os.environ.get(variable)
-        if value is None:
-            continue
-        try:
-            return int(value)
-        except ValueError:
-            continue
-    return 0
-
-
 def synchronize_for_mpi(*arrays: Any) -> None:
     """Wait for pending device work before MPI reads or writes `arrays`.