Newer
Older
from __future__ import print_function
from mpi4py import MPI
import numpy
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
# Simple message exchange
meta = {'rank': rank}
if rank == 0:
comm.send(meta, dest=1)
msg = comm.recv(source=1)
elif rank == 1:
msg = comm.recv(source=0)
comm.send(meta, dest=0)
print("Rank {0} received a message from rank {1}.".format(rank, msg['rank']))
# Simple message exchange using numpy arrays
n = 100000
data = numpy.zeros(n, int) + rank
buff = numpy.empty(n, int)
if rank == 0:
comm.Send(data, dest=1)
comm.Recv(buff, source=1)
elif rank == 1:
# One MPI tasks needs to start with send and the other one with
# receive, otherwise the program dead locks with large message
# sizes in most MPI implementations
comm.Recv(buff, source=0)
comm.Send(data, dest=0)
print("Rank {0} received an array filled with {1}s.".format(rank, buff[0]))