def distribute_vec_from_head(b):
from mfem.common.mpi_dtype import get_mpi_datatype
comm = MPI.COMM_WORLD
num_proc = MPI.COMM_WORLD.size
myid = MPI.COMM_WORLD.rank
if myid == 0:
partitioning = get_partition(b)
MPItype = get_mpi_datatype(b)
dtype = b.dtype
for i in range(num_proc - 1):
dest = i + 1
b0 = b[partitioning[dest]:partitioning[dest + 1]]
comm.send(b0.dtype, dest=dest, tag=dest)
comm.send(b0.shape, dest=dest, tag=dest)
comm.Send([b0, MPItype], dest=dest, tag=dest)
b0 = b[partitioning[0]:partitioning[1]]
else:
dtype = comm.recv(source=0, tag=myid)
shape = comm.recv(source=0, tag=myid)
b0 = np.zeros(shape, dtype=dtype)
MPItype = get_mpi_datatype(b0)
comm.Recv([b0, MPItype], source=0, tag=myid)
return b0
def scatter_vector(vector, mpi_data_type=None, rcounts=None, root=0):
# scatter data
#
# rcounts indicats the amount of data which each process
# receives
#
# for example: rcounts = fespace.GetTrueVSize()
if mpi_data_type is None:
from mfem.common.mpi_dtype import get_mpi_datatype
if MPI.COMM_WORLD.rank == root:
mpi_data_type = get_mpi_datatype(vector)
flag = mpi_type_flag(mpi_data_type)
else:
flag = None
flag = MPI.COMM_WORLD.bcast(flag)
mpi_data_type = mpi_type_from_flag(flag)
rcountss = MPI.COMM_WORLD.gather(rcounts, root=root)
# dprint1(rcountss)
disps = list(np.hstack((0, np.cumsum(rcountss)))[:-1])
recvdata = np.empty([rcounts], dtype=mpi_type_to_dtype(mpi_data_type))
if vector is not None:
senddata = [vector, rcountss, disps, mpi_data_type]
else:
senddata = None
MPI.COMM_WORLD.Scatterv(senddata, recvdata, root=root)
# MPI.COMM_WORLD.Barrier()
return recvdata
def gather_vector(data, mpi_data_type=None, root=0):
'''
gather vector to root node.
B: Vector to be collected
'''
from mfem.common.mpi_dtype import get_mpi_datatype
if mpi_data_type is None:
mpi_data_type = get_mpi_datatype(data)
myid = MPI.COMM_WORLD.rank
rcounts = data.shape[0]
rcounts = MPI.COMM_WORLD.gather(rcounts, root=root)
cm = np.hstack((0, np.cumsum(rcounts)))
disps = list(cm[:-1])
recvdata = None
senddata = [data, data.shape[0]]
if myid == root:
length = cm[-1]
recvbuf = np.empty([length], dtype=data.dtype)
recvdata = [recvbuf, rcounts, disps, mpi_data_type]
else:
recvdata = [None, rcounts, disps, mpi_data_type]
recvbuf = None
MPI.COMM_WORLD.Barrier()
MPI.COMM_WORLD.Gatherv(senddata, recvdata, root=root)
if myid == root:
# print 'collected'
MPI.COMM_WORLD.Barrier()
return np.array(recvbuf)
MPI.COMM_WORLD.Barrier()
return None
def gather_vector(data,
mpi_data_type=None,
parent=False,
world=MPI.COMM_WORLD,
root=0):
'''
gather vector to root
B: Vector to be collected
for intra-communication, leave parent False. data is gatherd
to root (myid = 0)
for inter-communication:
root group should call with parent = True
root group should call with data to tell the data type, like np.array(2)
world should be specified
'''
from mfem.common.mpi_dtype import get_mpi_datatype
myid = world.rank
if mpi_data_type is None:
mpi_data_type = get_mpi_datatype(data)
if world.Is_intra():
if myid == root:
parent = True
rcounts = data.shape[0]
senddata = [data, data.shape[0]]
elif parent:
root = MPI.ROOT if myid == root else MPI.PROC_NULL
rcounts = 0
senddata = [np.array(()), 0]
else:
rcounts = data.shape[0]
senddata = [data, data.shape[0]]
if myid == root:
parent = True
rcounts = world.allgather(rcounts)
cm = np.hstack((0, np.cumsum(rcounts)))
disps = list(cm[:-1])
if parent:
length = cm[-1]
recvbuf = np.empty([length], dtype=data.dtype)
recvdata = [recvbuf, rcounts, disps, mpi_data_type]
else:
recvdata = [None, rcounts, disps, mpi_data_type]
recvbuf = None
world.Barrier()
world.Gatherv(senddata, recvdata, root=root)
if parent:
# print 'collected'
world.Barrier()
return np.array(recvbuf)
else:
world.Barrier()
return None
def alltoall_vector(data, datatype):
'''
data = list of np.array.
each array goes to different MPI proc.
length of data must be equal to MPI size
'''
from mfem.common.mpi_dtype import get_mpi_datatype
num_proc = MPI.COMM_WORLD.size
senddata = np.hstack(data).astype(datatype, copy=False)
senddtype = get_mpi_datatype(senddata)
senddtypes = MPI.COMM_WORLD.allgather(mpi_type_flag(senddtype))
if len(np.unique(senddtypes)) != 1:
assert False, "data type is not unique"
sendsize = np.array([len(x) for x in data], dtype=int)
senddisp = list(np.hstack((0, np.cumsum(sendsize)))[:-1])
# (step 1) communicate the size of data
if len(sendsize) != num_proc:
assert False, "senddata size does not match with mpi size"
recvsize = np.empty(num_proc, dtype=int)
disp = list(range(num_proc))
counts = [1] * num_proc
dtype = get_mpi_datatype(recvsize)
s1 = [sendsize, counts, disp, dtype]
r1 = [recvsize, counts, disp, dtype]
MPI.COMM_WORLD.Alltoallv(s1, r1)
# (step 2) communicate the data
recvsize = list(recvsize)
recvdisp = list(np.hstack((0, np.cumsum(recvsize))))
recvdata = np.empty(np.sum(recvsize), dtype=senddata.dtype)
s1 = [senddata, sendsize, senddisp, senddtype]
r1 = [recvdata, recvsize, recvdisp[:-1], senddtype]
MPI.COMM_WORLD.Alltoallv(s1, r1)
data = [recvdata[recvdisp[i]:recvdisp[i + 1]] for i in range(num_proc)]
return data
def gather_allvector(data):
from mfem.common.mpi_dtype import get_mpi_datatype
from mpi4py import MPI
mpi_data_type = get_mpi_datatype(data)
myid = MPI.COMM_WORLD.rank
rcounts = data.shape[0]
rcounts = np.array(MPI.COMM_WORLD.allgather(rcounts))
for x in data.shape[1:]:
rcounts = rcounts * x
cm = np.hstack((0, np.cumsum(rcounts)))
disps = list(cm[:-1])
length = cm[-1]
recvbuf = np.empty([length], dtype=data.dtype)
recvdata = [recvbuf, rcounts, disps, mpi_data_type]
senddata = [data.flatten(), data.flatten().shape[0]]
MPI.COMM_WORLD.Allgatherv(senddata, recvdata)
return recvbuf.reshape(-1, *data.shape[1:])
def distribute_global_coo(A):
from mfem.common.mpi_dtype import get_mpi_datatype
comm = MPI.COMM_WORLD
num_proc = MPI.COMM_WORLD.size
myid = MPI.COMM_WORLD.rank
partitioning = get_partition(A)
row = A.row
col = A.col
data = A.data
ids = np.arange(num_proc)[::-1]
dtype = data.dtype
MPItype = get_mpi_datatype(data)
row2 = []
col2 = []
data2 = []
for i in range(num_proc):
ids = np.roll(ids, 1)
pair = ids[myid]
idx = np.logical_and(row >= partitioning[pair],
row < partitioning[pair + 1])
r0 = row[idx].astype(np.int32)
c0 = col[idx].astype(np.int32)
d0 = data[idx]
if pair < myid: # send first
comm.send(r0.shape, dest=pair, tag=i)
comm.Send([r0, MPI.INT], dest=pair, tag=i)
comm.Send([c0, MPI.INT], dest=pair, tag=i)
comm.Send([d0, MPItype], dest=pair, tag=i)
shape = comm.recv(source=pair, tag=i)
r = np.zeros(shape, dtype=np.int32)
c = np.zeros(shape, dtype=np.int32)
d = np.zeros(shape, dtype=dtype)
comm.Recv([r, MPI.INT], source=pair, tag=i)
comm.Recv([c, MPI.INT], source=pair, tag=i)
comm.Recv([d, MPItype], source=pair, tag=i)
elif pair > myid: # recv first
shape = comm.recv(source=pair, tag=i)
r = np.zeros(shape, dtype=np.int32)
c = np.zeros(shape, dtype=np.int32)
d = np.zeros(shape, dtype=dtype)
comm.Recv([r, MPI.INT], source=pair, tag=i)
comm.Recv([c, MPI.INT], source=pair, tag=i)
comm.Recv([d, MPItype], source=pair, tag=i)
comm.send(r0.shape, dest=pair, tag=i)
comm.Send([r0, MPI.INT], dest=pair, tag=i)
comm.Send([c0, MPI.INT], dest=pair, tag=i)
comm.Send([d0, MPItype], dest=pair, tag=i)
else:
r = r0
c = c0
d = d0
row2.append(r)
col2.append(c)
data2.append(d)
from scipy.sparse import coo_matrix
r = np.hstack(row2) - partitioning[myid]
c = np.hstack(col2)
d = np.hstack(data2)
rsize = partitioning[myid + 1] - partitioning[myid]
A = coo_matrix((d, (r, c)), shape=(rsize, A.shape[1]), dtype=d.dtype)
return A
def alltoall_vectorv(data, datatype):
'''
mulitidimension, arbitrary size version
data = list of list of np.array.
each list element goes to different MPI proc.
length of data must be equal to MPI size
'''
from mfem.common.mpi_dtype import get_mpi_datatype
num_proc = MPI.COMM_WORLD.size
datashape = [np.array([y.shape for y in x]).flatten() for x in data]
datadim = [
np.hstack([len(y.shape)
for y in x]) if len(x) > 0 else np.array([], dtype=int)
for x in data
]
senddata = [
np.hstack([y.flatten()
for y in x]) if len(x) > 0 else np.array([], dtype=datatype)
for x in data
]
sendsize = np.array([len(x) for x in senddata], dtype=int)
senddisp = list(np.hstack((0, np.cumsum(sendsize)))[:-1])
senddata = np.hstack(senddata).astype(datatype, copy=False)
senddtype = get_mpi_datatype(senddata)
senddtypes = MPI.COMM_WORLD.allgather(mpi_type_flag(senddtype))
if len(np.unique(senddtypes)) != 1:
assert False, "data type is not unique"
# (step 1) communicate the size of data
datashape = alltoall_vector(datashape, int)
datadim = alltoall_vector(datadim, int)
# (step 2) communicate the data
recvshapes = []
for dims, shapes in zip(datadim, datashape):
disps = np.hstack([0, np.cumsum(dims)])
recvshapes.append(
[shapes[disps[i]:disps[i + 1]] for i in range(len(disps) - 1)])
recvsize = [
np.sum([np.prod(s) for s in ss]).astype(int) for ss in recvshapes
]
recvdisp = list(np.hstack((0, np.cumsum(recvsize))))
recvdata = np.empty(np.sum(recvsize), dtype=senddata.dtype)
s1 = [senddata, sendsize, senddisp, senddtype]
r1 = [recvdata, recvsize, recvdisp[:-1], senddtype]
MPI.COMM_WORLD.Alltoallv(s1, r1)
data0 = [recvdata[recvdisp[i]:recvdisp[i + 1]] for i in range(num_proc)]
data = []
for d, shapes in zip(data0, recvshapes):
disps = np.hstack((0, np.cumsum([np.prod(s) for s in shapes])))
data.append([
d[disps[i]:disps[i + 1]].reshape(shapes[i])
for i in range(len(shapes))
])
return data
a_size = 1
orgdata = [(rank + 1) * np.arange(x + 1, dtype=int) for x in range(size)]
senddata = [(rank + 1) * np.arange(x + 1, dtype=int) for x in range(size)]
print("process %s sending %s " % (rank, senddata))
sendsize = np.array([len(x.flatten()) for x in senddata], dtype=int)
senddisp = list(np.hstack((0, np.cumsum(sendsize)))[:-1])
if len(sendsize) != size:
assert False, "senddata size does not match with mpi size"
recvsize = np.empty(size, dtype=int)
disp = list(range(size))
counts = [1] * size
dtype = get_mpi_datatype(sendsize)
s1 = [sendsize, counts, disp, dtype]
r1 = [recvsize, counts, disp, dtype]
comm.Alltoallv(s1, r1)
print("process %s receiving %s " % (rank, recvsize))
recvsize = list(recvsize)
recvdisp = list(np.hstack((0, np.cumsum(recvsize)))[:-1])
recvdata = np.empty(np.sum(recvsize), dtype=int)
senddata = np.hstack(senddata).flatten()
dtype = get_mpi_datatype(senddata[0])
s1 = [senddata, sendsize, senddisp, dtype]
r1 = [recvdata, recvsize, recvdisp, dtype]
