Skip to content

Instantly share code, notes, and snippets.

@rcthomas
Created August 10, 2017 19:40
Show Gist options
  • Star 3 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save rcthomas/d1f0ad0f71a61d791b7e5d96ab259559 to your computer and use it in GitHub Desktop.
Save rcthomas/d1f0ad0f71a61d791b7e5d96ab259559 to your computer and use it in GitHub Desktop.
Use mpi4py Win_Allocate_shared
from mpi4py import MPI
import numpy as np
world_comm = MPI.COMM_WORLD
node_comm = world_comm.Split_type(MPI.COMM_TYPE_SHARED)
size = 1000
disp_unit = MPI.DOUBLE.Get_size()
win = MPI.Win.Allocate_shared(size * disp_unit if node_comm.rank == 0 else 0,
disp_unit, comm = node_comm)
buf, itemsize = win.Shared_query(0)
assert itemsize == MPI.DOUBLE.Get_size()
buf = np.array(buf, dtype='B', copy=False)
ary = np.ndarray(buffer=buf, dtype='d', shape=(size,))
if node_comm.rank == 1:
ary[:5] = np.arange(5)
node_comm.Barrier()
if node_comm.rank == 0:
print ary[:10]
@rcthomas
Copy link
Author

For instance

salloc -N 2 -p debug -t 10
srun -N 2 -n 4 -c 24 python example-allocate-shared.py

Output example has the first 5 entries of the array as read on the root rank of each node communicator as non-garbage as filled in by the other rank of the node.

[  0.00000000e+000   1.00000000e+000   2.00000000e+000   3.00000000e+000
4.00000000e+000   9.07681078e+223   1.31309109e+213   4.63461863e+228
5.83086691e+252   9.70158285e+189]
[  0.00000000e+000   1.00000000e+000   2.00000000e+000   3.00000000e+000
4.00000000e+000   9.07681078e+223   1.31309109e+213   4.63461863e+228
5.83086691e+252   9.70158285e+189]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment