Repository navigation
API: bcast, gather and allgather of Python objects, as mpi4py's - #243
Merged
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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)andallgather(obj)take any Python object with mpi4py's lower-case semantics: objects travel pickled asuint8arrays;gatherreturns a list on the root andNoneelsewhere;allgatherreturns a list on every rank. Arrays are the exception and go as buffers:bcastof a NumPy/CuPy array arrives as an array of the same kind (a CuPy array stays on its device),allgatherof 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) andgatherv(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) andgatherv(a, root=0, axis=0)(concatenation on the root, sharing the header/plan code withallgatherv). The old scalarbcastoverloads of the bare communicator are nowbcast_scalar.Reductions:
allreduce/iallreducealso take mpi4py'sMPI.SUM/MAX/MIN, compared by identity whenmpi4py.MPIis 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
gatherfrom #241 (nowallgatherv). Docs and CHANGELOG updated. Companion: muSpectre/muTimer#1 lets muTimer gather its summaries through thisgatherwithout 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