Add mpi exchange

docs/source/guides/mpi.md

@@ -141,6 +141,34 @@ pinned staging buffers instead of allocating them per call, pass an
 `mpi_is_cuda_aware()` call and no `set_mpi_cuda_aware()`), a device array
 raises instead of guessing.
 
+### Exchange several buffers: `exchange`
+
+For a group of point-to-point transfers, `exchange()` manages the buffer
+contexts and waits for all requests internally:
+
+```python
+xp.mpi.exchange(
+    comm,
+    sends=[(send_r, right, 0), (send_l, left, 1)],
+    receives=[(recv_l, left, 0), (recv_r, right, 1)],
+)
+# recv_l and recv_r are ready, including any copies back to the GPU.
+```
+
+Each entry is `(array, peer_rank, tag)`. The helper prepares every buffer,
+posts all receives followed by all sends, and waits before closing any
+context. `cuda_aware=` has the same meaning as for `mpi_buffer()`. Each
+device receive slice gets its own staging buffer when staging is needed.
+Receive buffers must not overlap each other or send buffers; buffers passed
+directly to MPI must be contiguous. Empty lists do nothing, while zero-length
+arrays still send or receive messages and need matching peer operations.
+
+Preparation errors propagate before any local requests are posted. An error
+while posting or waiting aborts the communicator with `comm.Abort(1)`, because
+outstanding MPI requests may still access the buffers. On success the helper
+returns `None`; use `mpi_buffer()` directly when you need request statuses or
+control over asynchronous progress.
+
 ## Reproducible random numbers
 
 Each rank needs its own random stream, and a run is reproducible only if every

src/cunumpy/LLM_GUIDE.md

@@ -171,6 +171,9 @@ with xp.mpi.mpi_buffer(a) as buf:
 with xp.mpi.mpi_buffer(a, send=False, recv=True) as buf:
     ...  # array, or pinned staging copy
 xp.mpi.set_mpi_cuda_aware(True | False | None), xp.mpi.get_mpi_cuda_aware()
+xp.mpi.exchange(
+    comm, sends=[(send, dest, tag)], receives=[(recv, source, tag)]
+)  # waits for all transfers and GPU copy-back; MPI posting/wait errors abort comm
 xp.algorithms.segment_sum(
     values, keys, n_segments
 )  # out[k] = sum(values[keys == k]); keys < 0 dropped

src/cunumpy/_mpi.py

@@ -4,8 +4,8 @@
 
 import logging
 import operator
-from collections.abc import Generator
-from contextlib import contextmanager
+from collections.abc import Generator, Iterable
+from contextlib import ExitStack, contextmanager
 from typing import Any
 
 import array_api_compat
@@ -366,6 +366,91 @@ def mpi_buffer(
             cp.cuda.get_current_stream().synchronize()
 
 
+def exchange(
+    comm: Any,
+    *,
+    sends: Iterable[tuple[Any, int, int]] = (),
+    receives: Iterable[tuple[Any, int, int]] = (),
+    cuda_aware: bool | None = None,
+) -> None:
+    """Exchange array buffers and wait for all transfers to finish.
+
+    Prepare buffers with :func:`mpi_buffer`, post all receives and sends,
+    then wait before closing any buffer context. On successful return,
+    received data is ready, including any copies back to the GPU.
+
+    Parameters
+    ----------
+    comm : MPI communicator
+        Communicator providing ``Irecv``, ``Isend`` and ``Abort``.
+    sends : iterable of (array, destination_rank, tag), optional
+        Outgoing buffers and their destinations.
+    receives : iterable of (array, source_rank, tag), optional
+        Incoming buffers and their sources.
+    cuda_aware : bool or None, optional
+        Passed to :func:`mpi_buffer`. None uses the recorded MPI capability.
+
+    See Also
+    --------
+    mpi_buffer : Manage a single buffer around caller-provided MPI operations.
+
+    Notes
+    -----
+    Receive buffers must not overlap each other or send buffers. Buffers
+    must be contiguous when passed directly to MPI. Each device buffer is
+    staged separately when MPI is not CUDA-aware. Empty iterables do no
+    communication; zero-length arrays still post MPI messages and require
+    matching operations on the peer. Self-transfers are supported when
+    matching sends and receives are supplied.
+
+    All local buffers are prepared before posting any requests. Preparation
+    errors propagate normally. Once posting starts, an exception during
+    posting or waiting calls ``comm.Abort(1)``: outstanding requests may
+    still use the buffers, so ordinary context cleanup is unsafe. This
+    helper does not provide recovery from MPI errors or return statuses.
+
+    Examples
+    --------
+    >>> comm = xp.mpi.get_mpi().COMM_WORLD
+    >>> send, recv = np.arange(3.0), np.empty(3)
+    >>> xp.mpi.exchange(
+    ...     comm,
+    ...     sends=[(send, (comm.rank + 1) % comm.size, 0)],
+    ...     receives=[(recv, (comm.rank - 1) % comm.size, 0)],
+    ... )
+    """
+    with ExitStack() as buffers:
+        prepared_receives = [
+            (
+                buffers.enter_context(
+                    mpi_buffer(array, send=False, recv=True, cuda_aware=cuda_aware)
+                ),
+                source,
+                tag,
+            )
+            for array, source, tag in receives
+        ]
+        prepared_sends = [
+            (
+                buffers.enter_context(mpi_buffer(array, cuda_aware=cuda_aware)),
+                destination,
+                tag,
+            )
+            for array, destination, tag in sends
+        ]
+        requests = []
+        try:
+            for buf, source, tag in prepared_receives:
+                requests.append(comm.Irecv(buf, source=source, tag=tag))
+            for buf, destination, tag in prepared_sends:
+                requests.append(comm.Isend(buf, dest=destination, tag=tag))
+            for request in requests:
+                request.Wait()
+        except BaseException:
+            comm.Abort(1)
+            raise
+
+
 def _device_buffers_in_use() -> bool:
     """Whether MPI calls of this process would carry device (CuPy) buffers."""
     return array_backend.backend == "cupy" and cupy_available()

src/cunumpy/mpi.py

@@ -36,6 +36,7 @@
 
 from cunumpy._mpi import (
     MPIStaging,
+    exchange,
     get_mpi_cuda_aware,
     mpi_buffer,
     mpi_is_cuda_aware,
@@ -66,6 +67,7 @@ def _count_serial_copy(kind: str, array: object) -> None:
     "SerialMPI",
     "SerialRequest",
     "SerialStatus",
+    "exchange",
     "get_mpi",
     "get_mpi_cuda_aware",
     "is_serial",

testing/smoke_wheel.py

@@ -0,0 +1,46 @@
+"""Run with an isolated Python from a fresh venv containing the release wheel."""
+
+import os
+import sys
+from importlib.resources import files
+from pathlib import Path
+
+os.environ["CUNUMPY_FAKE_CUPY"] = "1"
+os.environ["CUNUMPY_BACKEND"] = "cupy"
+
+import numpy as np
+
+import cunumpy as xp
+
+
+def main():
+    # Catch accidental imports from the checkout or another environment.
+    assert Path(xp.__file__).resolve().is_relative_to(Path(sys.prefix).resolve())
+    package = files("cunumpy")
+    for asset in (
+        "py.typed",
+        "__init__.pyi",
+        "LLM_GUIDE.md",
+        "cuda/include/cunumpy/array_view.cuh",
+    ):
+        assert package.joinpath(asset).is_file(), f"Missing wheel asset: {asset}"
+
+    expected = np.array(np.arange(12.0).reshape(3, 4), order="F")
+    device = xp.to_cupy(expected.copy(order="K"))
+    host = xp.to_numpy(device)
+    assert host.flags.f_contiguous
+    np.testing.assert_array_equal(host, expected)
+
+    def scale(array):
+        assert isinstance(array, np.ndarray) and array.flags.f_contiguous
+        array *= 2
+
+    xp.kernels.PyccelKernel(scale, outputs=(0,))(device)
+    np.testing.assert_array_equal(xp.to_numpy(device), 2 * expected)
+    with xp.use_backend("numpy", strict=True):
+        np.testing.assert_array_equal(xp.arange(3), np.arange(3))
+    print("Installed-wheel smoke test passed")
+
+
+if __name__ == "__main__":
+    main()

tests/unit/test_mpi_serial.py

@@ -31,6 +31,20 @@ def test_reexported_from_maybempi():
     assert xp.mpi.is_serial(MPI) and xp.mpi.is_serial(comm)
 
 
+@pytest.mark.parametrize("size", [0, 4])
+def test_exchange_host_self_transfer(size):
+    send = np.arange(size, dtype=float)
+    recv = np.zeros_like(send)
+    assert (
+        xp.mpi.exchange(comm, sends=[(send, 0, 17)], receives=[(recv, 0, 17)]) is None
+    )
+    np.testing.assert_array_equal(recv, send)
+
+
+def test_exchange_empty_does_not_touch_communicator():
+    xp.mpi.exchange(object())
+
+
 class _DeviceArray:
     """Duck-typed device array: `.get()` copies it to the host."""
 

tests/unit/test_mpi_staging.py

@@ -170,6 +170,111 @@ def test_host_arrays_are_passed_through_and_invalid_staging_is_rejected():
         MPIStaging(2, object)
 
 
+@pytest.mark.parametrize("cuda_aware", [False, True])
+def test_exchange_device_buffers_stay_open_until_all_requests_finish(
+    device, cuda_aware
+):
+    Array, log = device
+    sends = [Array([1.0, 2.0]), Array([3.0, 4.0])]
+    receives = [Array([0.0, 0.0]), Array([0.0, 0.0])]
+    posted = []
+
+    class Comm:
+        def Irecv(self, buf, source, tag):
+            posted.append(("recv", source, tag, buf))
+
+            def wait():
+                assert len(posted) == 4
+                assert not any(entry[0] == "set" for entry in log)
+                values = buf.values if cuda_aware else buf
+                values[:] = [source, tag]
+
+            return SimpleNamespace(Wait=wait)
+
+        def Isend(self, buf, dest, tag):
+            posted.append(("send", dest, tag, buf))
+
+            def wait():
+                assert not any(entry[0] == "set" for entry in log)
+                values = buf.values if cuda_aware else buf
+                np.testing.assert_array_equal(values, sends[dest - 1].values)
+
+            return SimpleNamespace(Wait=wait)
+
+        def Abort(self, code):
+            pytest.fail("successful exchange must not abort")
+
+    with xp.profiling.count_transfers() as counter:
+        xp.mpi.exchange(
+            Comm(),
+            sends=((array, i + 1, 10 + i) for i, array in enumerate(sends)),
+            receives=((array, i + 1, 20 + i) for i, array in enumerate(receives)),
+            cuda_aware=cuda_aware,
+        )
+    assert [entry[:3] for entry in posted] == [
+        ("recv", 1, 20),
+        ("recv", 2, 21),
+        ("send", 1, 10),
+        ("send", 2, 11),
+    ]
+    for i, array in enumerate(receives):
+        np.testing.assert_array_equal(array.values, [i + 1, 20 + i])
+    assert counter.to_host == counter.to_device == (0 if cuda_aware else 2)
+    if cuda_aware:
+        assert posted[0][3] is receives[0]
+        assert posted[2][3] is sends[0]
+
+
+def test_exchange_prepares_every_buffer_before_posting(device):
+    Array, log = device
+    recv = Array([7.0])
+    # A later invalid descriptor must fail before touching the communicator,
+    # and unwinding preparation must not copy an uninitialized receive buffer.
+    with pytest.raises(ValueError):
+        xp.mpi.exchange(
+            object(),
+            receives=[(recv, 1, 0)],
+            sends=[(Array([1.0]), 1)],
+            cuda_aware=False,
+        )
+    np.testing.assert_array_equal(recv.values, [7])
+    assert not any(entry[0] == "set" for entry in log)
+
+
+@pytest.mark.parametrize("failure", ["receive", "send", "wait"])
+def test_exchange_aborts_on_communication_failure(device, failure):
+    Array, log = device
+    events = []
+
+    class Comm:
+        def Irecv(self, buf, source, tag):
+            if failure == "receive":
+                raise RuntimeError("MPI failure")
+            return SimpleNamespace(Wait=self.wait)
+
+        def Isend(self, buf, dest, tag):
+            if failure == "send":
+                raise RuntimeError("MPI failure")
+            return SimpleNamespace(Wait=self.wait)
+
+        def wait(self):
+            raise RuntimeError("MPI failure")
+
+        def Abort(self, code):
+            assert not any(entry[0] == "set" for entry in log)
+            events.append(("abort", code))
+
+    with pytest.raises(RuntimeError, match="MPI failure"):
+        xp.mpi.exchange(
+            Comm(),
+            sends=[(Array([1.0]), 1, 0)],
+            receives=[(Array([0.0]), 1, 0)],
+            cuda_aware=False,
+        )
+    assert events == [("abort", 1)]
+    assert not any(entry[0] == "set" for entry in log)
+
+
 @pytest.mark.skipif(not xp.cupy_available(), reason="requires CUDA")
 @pytest.mark.parametrize("layout", ["C", "F", "strided"])
 def test_real_device_staging_reuses_host_storage_after_nondefault_producer(layout):