Skip to content

Commit 06edce1

Browse files
committed
Work around 32-bit limit on MPI buffer size
1 parent 4e71a53 commit 06edce1

1 file changed

Lines changed: 10 additions & 1 deletion

File tree

src/pshmem/shmem.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -543,7 +543,16 @@ def set(self, data, offset=None, fromrank=0):
543543
nodedata = np.zeros(datashape, dtype=self._dtype)
544544

545545
# Broadcast the data buffer
546-
self._rankcomm.Bcast([nodedata, self._mpitype], root=fromnode)
546+
nodedata = nodedata.reshape(-1, copy=False)
547+
maxlen = 2**30
548+
start = 0
549+
while start < nodedata.size:
550+
stop = min(start + maxlen, nodedata.size)
551+
self._rankcomm.Bcast(
552+
[nodedata[start:stop], self._mpitype], root=fromnode
553+
)
554+
start = stop
555+
nodedata = nodedata.reshape(datashape, copy=False)
547556

548557
# Now one process on every node has a copy of the data, and
549558
# can copy it into the shared memory buffer.

0 commit comments

Comments
 (0)