mpi4py Scatterv 函数沿什么轴拆分 numpy 数组？

Question

我有以下 MWE 使用 comm.Scatterv 和 comm.Gatherv 在给定数量的核心上分布 4D 阵列 (size)

import numpy as np
from mpi4py import MPI
import matplotlib.pyplot as plt

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

if rank == 0:
    test = np.random.rand(411,48,52,40) #Create array of random numbers
    outputData = np.zeros(np.shape(test))
    split = np.array_split(test,size,axis = 0) #Split input array by the number of available cores

    split_sizes = []

    for i in range(0,len(split),1):
        split_sizes = np.append(split_sizes, len(split[i]))

    displacements = np.insert(np.cumsum(split_sizes),0,0)[0:-1]

    plt.imshow(test[0,0,:,:])
    plt.show()

else:
#Create variables on other cores
    split_sizes = None
    displacements = None
    split = None
    test = None
    outputData = None

#Broadcast variables to other cores
test = comm.bcast(test, root = 0)
split = comm.bcast(split, root=0) 
split_sizes = comm.bcast(split_sizes, root = 0)
displacements = comm.bcast(displacements, root = 0)

output_chunk = np.zeros(np.shape(split[rank])) #Create array to receive subset of data on each core, where rank specifies the core
print("Rank %d with output_chunk shape %s" %(rank,output_chunk.shape))

comm.Scatterv([test,split_sizes, displacements,MPI.DOUBLE],output_chunk,root=0) #Scatter data from test across cores and receive in output_chunk

output = output_chunk

plt.imshow(output_chunk[0,0,:,:])
plt.show()

print("Output shape %s for rank %d" %(output.shape,rank))

comm.Barrier()

comm.Gatherv(output,[outputData,split_sizes,displacements,MPI.DOUBLE], root=0) #Gather output data together

if rank == 0:
    print("Final data shape %s" %(outputData.shape,))
    plt.imshow(outputData[0,0,:,:])
    plt.show()

这会创建一个 4D 随机数数组，原则上应该在重新组合之前将其划分为 size 个核心。我希望 Scatterv 根据向量 split_sizes 和 displacements 中的起始整数和位移沿轴 0（长度 411）划分。但是，在与 Gatherv (mpi4py.MPI.Exception: MPI_ERR_TRUNCATE: message truncated) 重新组合时出现错误，并且每个核心上的 output_chunk 图显示大部分输入数据已丢失，因此看起来拆分沿第一个轴没有发生。

我的问题是：为什么不沿着第一个轴发生分裂，我怎么知道分裂发生在哪个轴上，是否可以 change/specify 这发生在哪个轴上？

Answer 1

comm.Scatterv 和 comm.Gatherv 对 numpy 数组维度一无所知。他们只是将 sendbuf 视为一块内存。因此，在指定 sendcounts 和 displacements 时有必要考虑到这一点（有关详细信息，请参阅 http://materials.jeremybejarano.com/MPIwithPython/collectiveCom.html）。还假设数据在内存中以 C 风格（主要行）布局。

下面给出了二维矩阵的示例。此代码的关键部分是正确设置 split_sizes_input/split_sizes_output 和 displacements_input/displacements_output。该代码将第二维大小考虑在内，以指定内存块中的正确划分：

split_sizes_input = split_sizes*512

对于更高的维度，此行将更改为：

split_sizes_input = split_sizes*indirect_dimension_sizes

哪里

indirect_dimension_sizes = npts2*npts3*npts4*....*nptsN

split_sizes_output.

也是如此

代码创建一个二维数组，其中数字 1 到 512 在一个维度上递增。从图中很容易看出数据是否已正确拆分和重新组合。

import numpy as np
from mpi4py import MPI
import matplotlib.pyplot as plt

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

if rank == 0:
    test = np.arange(0,512,dtype='float64')
    test = np.tile(test,[256,1]) #Create 2D input array. Numbers 1 to 512 increment across dimension 2.
    outputData = np.zeros([256,512]) #Create output array of same size
    split = np.array_split(test,size,axis = 0) #Split input array by the number of available cores

    split_sizes = []

    for i in range(0,len(split),1):
        split_sizes = np.append(split_sizes, len(split[i]))

    split_sizes_input = split_sizes*512
    displacements_input = np.insert(np.cumsum(split_sizes_input),0,0)[0:-1]

    split_sizes_output = split_sizes*512
    displacements_output = np.insert(np.cumsum(split_sizes_output),0,0)[0:-1]


    print("Input data split into vectors of sizes %s" %split_sizes_input)
    print("Input data split with displacements of %s" %displacements_input)

    plt.imshow(test)
    plt.colorbar()
    plt.title('Input data')
    plt.show()

else:
#Create variables on other cores
    split_sizes_input = None
    displacements_input = None
    split_sizes_output = None
    displacements_output = None
    split = None
    test = None
    outputData = None

split = comm.bcast(split, root=0) #Broadcast split array to other cores
split_sizes = comm.bcast(split_sizes_input, root = 0)
displacements = comm.bcast(displacements_input, root = 0)
split_sizes_output = comm.bcast(split_sizes_output, root = 0)
displacements_output = comm.bcast(displacements_output, root = 0)

output_chunk = np.zeros(np.shape(split[rank])) #Create array to receive subset of data on each core, where rank specifies the core
print("Rank %d with output_chunk shape %s" %(rank,output_chunk.shape))
comm.Scatterv([test,split_sizes_input, displacements_input,MPI.DOUBLE],output_chunk,root=0)

output = np.zeros([len(output_chunk),512]) #Create output array on each core

for i in range(0,np.shape(output_chunk)[0],1):
    output[i,0:512] = output_chunk[i]

plt.imshow(output)
plt.title("Output shape %s for rank %d" %(output.shape,rank))
plt.colorbar()
plt.show()

print("Output shape %s for rank %d" %(output.shape,rank))

comm.Barrier()

comm.Gatherv(output,[outputData,split_sizes_output,displacements_output,MPI.DOUBLE], root=0) #Gather output data together



if rank == 0:
    outputData = outputData[0:len(test),:]
    print("Final data shape %s" %(outputData.shape,))
    plt.imshow(outputData)
    plt.colorbar()
    plt.show()
    print(outputData)

mpi4py Scatterv 函数沿什么轴拆分 numpy 数组？

Along what axis does mpi4py Scatterv function split a numpy array?

python

arrays

numpy

mpi4py