Mpi4Py - 发送 numpy 子数组（非连续内存）而不复制答案

【问题标题】：Mpi4Py - sending numpy subarray (non-contiguous memory) without copyMpi4Py - 发送 numpy 子数组（非连续内存）而不复制
【发布时间】：2016-03-22 21:28:33
【问题描述】：

我正在重写通信瓶颈。我需要发送 3d numpy 数组的子数组。但是将子数组直接传递给MPI.Send() 失败：

ValueError: ndarray is not contiguous

MPI.Datatype.Create_vector 创建新的数据类型并没有帮助 - 由于同样的原因它失败了。

简化示例：

a = numpy.zeros([9,9])
sub = a[3:5, 3:5]

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

t = MPI.DOUBLE.Create_vector(2, 2, 9)
t.Commit()

if rank == 0:
    sub.flat[:] = range(1,9)
    comm.Send([sub, t], dest=1)
else:
    comm.Recv([sub, t], source=0)

在实际代码中，我使用异步发送/接收。目前我通过将子数组复制到具有连续内存布局的临时数组来解决它。

问题是缓冲区往往非常大，它会占用所有内存和可用交换空间。

我认为创建跨步数据类型是一种可行的方法，但由于我无法访问具有连续内存的原始数组，因此我无法创建跨步缓冲区。

也不能选择使用小写版本的发送/接收，因为正如我所说，我需要速度并且数据量很大。

目前我唯一的想法是创建 C 模块扩展，在其中我进行所有指针计算并返回 numpy 数组并访问包含我的子数组的连续内存段。

【问题讨论】：

标签： arrays numpy mpi mpi4py

【解决方案1】：

当我遇到同样的问题时，发现这个问题没有答案。您可以使用Create_subarray(sizes, subsizes, starts, int order=ORDER_C) 完成此操作：

import numpy as np
from mpi4py import MPI

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

a = np.zeros([9, 9])

if rank == 0:
    a[3:5, 3:5] = np.ones([2, 2])

t = MPI.DOUBLE.Create_subarray([9, 9], [2, 2], [3, 3])
t.Commit()

if rank == 0:
    comm.Send([a, t], dest=1)
else:
    comm.Recv([a, t], source=0)

额外提示：据我所知，mpi4py 的文档非常少。我找不到任何有关使用mpi4py 在线创建派生数据类型的帮助。如果您想查看可以使用 mpi4py 创建的数据类型列表，请在 Python 控制台中运行以下命令：

>>> from mpi4py import MPI
>>> help(MPI.DOUBLE)

【讨论】：

【解决方案2】：

另一种可能的解决方案（也允许任意子向量具有相同的 MPI.Datatype）

import numpy
from mpi4py import MPI

size = 9
subsize = 2
start = 1

a = numpy.zeros([size,size])

t = MPI.DOUBLE.Create_vector(subsize, subsize, size)
t.Commit()

# 2D index to 1D index 
begin = start * size + start
end = begin + t.Get_extent()[1] // a.itemsize

if MPI.COMM_WORLD.Get_rank() == 0:
    a[start : start + subsize, start : start + subsize] = 2
    MPI.COMM_WORLD.Send([a.reshape(size * size)[begin:end], t], dest=1)
else:
    MPI.COMM_WORLD.Recv([a.reshape(size * size)[begin:end], t], source=0)
    print(a)

输出：

$ mpirun -n 2 python src/python/test.py
[[0. 0. 0.]
 [0. 2. 2.]
 [0. 2. 2.]]

【讨论】：