diff --git a/docs/source/guides/mpi.md b/docs/source/guides/mpi.md index 9d1b56f..5b6e6e5 100644 --- a/docs/source/guides/mpi.md +++ b/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 diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 7adcf62..4064732 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/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 diff --git a/src/cunumpy/_mpi.py b/src/cunumpy/_mpi.py index e382423..c923ddd 100644 --- a/src/cunumpy/_mpi.py +++ b/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() diff --git a/src/cunumpy/mpi.py b/src/cunumpy/mpi.py index 32b794f..6da8454 100644 --- a/src/cunumpy/mpi.py +++ b/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", diff --git a/testing/smoke_wheel.py b/testing/smoke_wheel.py new file mode 100644 index 0000000..8ba93c6 --- /dev/null +++ b/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() diff --git a/tests/unit/test_mpi_serial.py b/tests/unit/test_mpi_serial.py index 75a5ddd..f616e16 100644 --- a/tests/unit/test_mpi_serial.py +++ b/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.""" diff --git a/tests/unit/test_mpi_staging.py b/tests/unit/test_mpi_staging.py index 30dbc3a..dc1bd0f 100644 --- a/tests/unit/test_mpi_staging.py +++ b/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):