如何用MPI派生数据类型非扁平化分散不等尺寸3D Fortran数组

核心思路 #

通过MPI_Type_create_subarray为每个进程定义对应子数组的数据类型，结合MPI_Type_create_resized调整类型的内存范围，直接使用MPI_Scatterv完成分散操作，无需创建扁平化临时数组。

关键步骤与代码实现 #

1. 参数定义与初始化 #

全局数组维度为(5,3,2)，进程0接收(3,3,2)子数组，进程1接收(2,3,2)子数组。提前为每个进程分配对应大小的局部数组。

2. 派生数据类型创建 #

针对每个进程的子数组，先创建子数组类型，再调整其内存范围为单个元素的大小，确保位移计算基于元素个数而非原类型的内存布局。

3. Scatterv分散执行 #

主进程使用调整后的派生类型，通过MPI_Scatterv将全局数组的对应子数组直接发送到各进程的局部数组中。

完整代码示例：

program scatter_unequal_subarray
  use mpi
  implicit none

  integer :: ierr, rank, size
  integer, parameter :: global_dims(3) = [5,3,2]
  integer :: local_dims(3), start_indices(3)
  integer :: counts(2) = [1,1], displs(2)
  integer :: type_subarray, send_types(2)
  real :: global_array(global_dims(1), global_dims(2), global_dims(3))
  real, allocatable :: local_array(:,:,:)

  call MPI_Init(ierr)
  call MPI_Comm_rank(MPI_COMM_WORLD, rank, ierr)
  call MPI_Comm_size(MPI_COMM_WORLD, size, ierr)

  ! 限制仅支持2个进程
  if (size /= 2) then
    if (rank == 0) print *, "该程序仅支持2个MPI进程运行"
    call MPI_Finalize(ierr)
    stop
  end if

  ! 设置各进程的局部数组维度与子数组起始索引
  select case(rank)
  case(0)
    local_dims = [3,3,2]
    start_indices = [1,1,1]
    ! 初始化全局数组（示例数据）
    global_array = reshape([(real(i), i=1,5*3*2)], global_dims)
  case(1)
    local_dims = [2,3,2]
    start_indices = [4,1,1]
  end select

  allocate(local_array(local_dims(1), local_dims(2), local_dims(3)))

  ! 主进程创建并配置派生数据类型
  if (rank == 0) then
    ! 为进程0创建子数组类型
    call MPI_Type_create_subarray(3, global_dims, [3,3,2], [1,1,1], &
                                  MPI_ORDER_FORTRAN, MPI_REAL, type_subarray, ierr)
    call MPI_Type_create_resized(type_subarray, 0, sizeof(real), send_types(1), ierr)
    call MPI_Type_commit(send_types(1), ierr)

    ! 为进程1创建子数组类型
    call MPI_Type_create_subarray(3, global_dims, [2,3,2], [4,1,1], &
                                  MPI_ORDER_FORTRAN, MPI_REAL, type_subarray, ierr)
    call MPI_Type_create_resized(type_subarray, 0, sizeof(real), send_types(2), ierr)
    call MPI_Type_commit(send_types(2), ierr)

    ! 计算位移：进程1的起始偏移为进程0子数组的元素总数
    displs = [0, 3*3*2]
  end if

  ! 执行Scatterv分散
  call MPI_Scatterv(global_array, counts, displs, send_types, &
                    local_array, product(local_dims), MPI_REAL, &
                    0, MPI_COMM_WORLD, ierr)

  ! 验证输出（可选）
  print *, "进程 ", rank, " 接收的局部数组："
  print *, local_array

  ! 清理资源
  if (rank == 0) then
    call MPI_Type_free(send_types(1), ierr)
    call MPI_Type_free(send_types(2), ierr)
    call MPI_Type_free(type_subarray, ierr)
  end if

  deallocate(local_array)
  call MPI_Finalize(ierr)
end program scatter_unequal_subarray

关键细节说明 #

• 列主序处理：Fortran采用列主序存储，MPI_Type_create_subarray需指定MPI_ORDER_FORTRAN，确保子数组的索引与内存布局匹配。
• 类型范围调整：MPI_Type_create_resized将派生类型的下界设为0，范围设为单个元素大小，使得displs可以直接用元素个数计算偏移，避免内存布局导致的位移错误。
• 类型管理：派生类型必须通过MPI_Type_commit提交后才能使用，使用完毕需调用MPI_Type_free释放资源。

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