如何使用mpi4py的Create_vector结构体发送矩阵列？

需求说明 #

要实现的是：用mpi4py将矩阵A的右列发送到另一个进程，并写入矩阵B的左列，全程不进行缓冲区数据拷贝。

初始矩阵示例：

[[1,2,3]   [[0,0,0]
 [4,5,6]    [0,0,0]
 [7,7,9]],  [0,0,0]]

期望发送后状态：

[[1,2,3]   [[3,0,0]
 [4,5,6]    [6,0,0]
 [7,7,9]],  [9,0,0]]

尝试的代码 #

测试脚本通过mpirun -n 2 python3 mpi_type_tester.py运行：

import numpy as np
from mpi4py import MPI

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

# Size of my send and receive matrix
height = 3
width  = 3

# Variables used to define the struct
count = 3
blocklength = 1
stride = 3

# Int seemingly used to define how many of the structs are being sent?
sending_int = 1

# Here I define the struct with Create_vector:
column_type = MPI.INT.Create_vector(count = count,blocklength = blocklength,stride = stride)
column_type.Commit()


if rank == 0:
    send_array = np.arange(width*height).reshape(height,width)
    send_array += 1

    comm.Send([send_array,sending_int, column_type], dest = 1, tag = 0) 

    print(send_array)

if rank == 1:
    rec_array = np.zeros(width*height, dtype = int).reshape(height, width)
    comm.Recv([rec_array,sending_int,column_type], source = 0, tag = 0)

    print(rec_array)

遇到的问题 #

尝试用MPI.INT.Create_vector定义自定义数据类型，但始终无法得到正确结果。调整count、blocklength、stride或sending_int参数时，发送结果混乱，希望能理解Create_vector的用法，或者得到相关学习指导。

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