如何通过mpi4py实现N×N矩阵按列拆分并分发至多进程？

你遇到的问题核心是numpy转置后的数组内存不连续，而MPI的Scatter是按连续内存块拆分数据的——哪怕你转置了矩阵，Scatter依然会按照原矩阵的行内存顺序分发，而非转置后的行（也就是原矩阵的列）。

下面给出两种直接可行的解决办法：

办法一：将转置数组转为连续内存 #

转置后用np.ascontiguousarray()把数组强制转为连续存储格式，这样Scatter就能正确按转置后的行（原矩阵的列）拆分。

修改代码中matrix_columns的赋值部分即可：

from mpi4py import MPI
import numpy as np

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

N = 4

if rank == 0:
    A = np.random.random((N,N))/N*2
    vector = np.random.random(N)
    print("Rank: ", rank)
    print("A: ", A)
    print("Vector: ", vector)
else:
    A=np.empty((N,N), dtype='float64')
    vector = np.empty(N, dtype='float64')

matrix_columns = np.empty((A.shape[0],A.shape[0]), dtype='float64')

# 修改这里：转置后转为连续内存
if rank == 0:
    matrix_columns = np.ascontiguousarray(np.transpose(A))
else:
    matrix_columns = np.empty((A.shape[0],A.shape[0]), dtype='float64')

if rank == 0:
    print("Columns >>>")
    print(matrix_columns)

received_columns = np.empty((matrix_columns.shape[0]//size,matrix_columns.shape[0]), dtype='float64')

comm.Scatter(matrix_columns, received_columns, root = 0)

print("My rank: ", rank,  "received columns: ", received_columns)

办法二：用Scatterv直接分发列（无需转置） #

如果不想做转置操作，可以直接计算每个进程要接收的列对应的内存偏移，用Scatterv精准分发，这样更节省内存和计算资源：

from mpi4py import MPI
import numpy as np

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

N = 4
cols_per_proc = N // size  # 每个进程分配的列数

if rank == 0:
    A = np.random.random((N,N))/N*2
    print("Rank: ", rank)
    print("A: ", A)
else:
    A = np.empty((N,N), dtype='float64')

# 定义每个进程接收的数据量和发送偏移量
sendcounts = [cols_per_proc * N] * size  # 每列N个元素，cols_per_proc列的总元素数
displs = [i * cols_per_proc * N for i in range(size)]  # 每个进程对应的内存起始位置

# 每个进程接收的数组：N行 × cols_per_proc列
received_columns = np.empty((N, cols_per_proc), dtype='float64')

# 用Scatterv分发列
comm.Scatterv([A, sendcounts, displs, MPI.DOUBLE], received_columns, root=0)

# 转置后匹配你预期的输出格式（每行对应原矩阵的一列）
received_columns = received_columns.T
print("My rank: ", rank,  "received columns: ", received_columns)

内容的提问来源于stack exchange，提问作者AmacOS