Skip to content

API: bcast, gather and allgather of Python objects, as mpi4py's - #243

Merged
pastewka merged 1 commit into
mainfrom
enh/mpi4py-style-objects
Oct 22, 2026
Merged

pastewka merged 1 commit into
mainfrom
enh/mpi4py-style-objects

Conversation

@pastewka

Copy link
Copy Markdown
Contributor

PR 1b of the mpi4py-optional series; builds on #242 (PR 1a).

Objects over the wire. Communicator.bcast(obj, root=0), gather(obj, root=0) and allgather(obj) take any Python object with mpi4py's lower-case semantics: objects travel pickled as uint8 arrays; gather returns a list on the root and None elsewhere; allgather returns a list on every rank. Arrays are the exception and go as buffers: bcast of a NumPy/CuPy array arrives as an array of the same kind (a CuPy array stays on its device), allgather of an array or NumPy scalar still comes back stacked. All pickling lives in Python; C++ only ever moves bytes.

C++ primitives: Communicator::bcast(buffer, count, type, root, loc) (in place) and gatherv(send, n, recv, counts, displs, root, type, loc), with the same direct-or-staged device handling as the other buffer exchanges; stubs for no-MPI builds. Bound as _muGrid.Communicator.bcast(a, root=0) (in place on an array every rank pre-allocated) and gatherv(a, root=0, axis=0) (concatenation on the root, sharing the header/plan code with allgatherv). The old scalar bcast overloads of the bare communicator are now bcast_scalar.

Reductions: allreduce/iallreduce also take mpi4py's MPI.SUM/MAX/MIN, compared by identity when mpi4py.MPI is loaded (muGrid does not import it); a Python number comes back as a Python number (a NumPy scalar still as a 0-d array, since it carries an array interface).

This replaces the deprecated column-concatenating gather from #241 (now allgatherv). Docs and CHANGELOG updated. Companion: muSpectre/muTimer#1 lets muTimer gather its summaries through this gather without mpi4py; verified together on 2 ranks with mpi4py not importable.

Verified locally: ctest 35/35 (MPI) and 16/16 (no-MPI); Python suite serially and on 2/4 ranks with mpi4py and on 4 ranks with mpi4py blocked; the no-MPI build's communicator/lifecycle tests; flake8/isort.

🤖 Generated with Claude Code

The communicator moved numbers; a dict of settings or a tuple of timers
had to go through mpi4py. bcast(obj, root=0), gather(obj, root=0) and
allgather(obj) now take any Python object with mpi4py's lower-case
semantics: objects travel pickled, as uint8 arrays, over two new buffer
primitives of the C++ communicator -- bcast in place and gatherv to a
root -- and gather returns a list on the root and None elsewhere. Arrays
are the exception and go as buffers, so that a CuPy array arrives on its
device and a large one is not copied through pickle. This replaces the
deprecated column-concatenating gather of #241 (now allgatherv).

allreduce and iallreduce also take mpi4py's MPI.SUM, MAX and MIN,
compared by identity when mpi4py happens to be loaded, and a Python
number comes back as a Python number, as from mpi4py.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Base automatically changed from enh/mpi-without-mpi4py to main October 22, 2026 05:36
@pastewka
pastewka merged commit 2fe85cf into main Oct 22, 2026
42 of 43 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant