MPI_Comm_split のテスト

MPIのプログラム書いててコミュニケータを分割するようなプログラムを書いたことなかったなぁと思ったのでメモ．

プログラム

split.f90

program main
    use mpi
    implicit none

    integer :: i, ierr, comm_rank, comm_size, color, key
    integer :: new_comm, new_rank, new_size

    ! program main
    !! init
    call MPI_INIT(ierr)

    !! rank, size for MPI_COMM_WORLD
    call MPI_COMM_RANK(MPI_COMM_WORLD, comm_rank, ierr)
    call MPI_COMM_SIZE(MPI_COMM_WORLD, comm_size, ierr)

    !! set colors
    key = 0
    color = comm_rank/4

    !! split MPI_COMM_WORLD
    call MPI_COMM_SPLIT(MPI_COMM_WORLD, color, key, new_comm, ierr)

    !! rank, size for `new_comm'
    call MPI_COMM_RANK(new_comm, new_rank, ierr)
    call MPI_COMM_SIZE(new_comm, new_size, ierr)

    if (comm_rank == 0) then
        write(*,'(A)') "old rank/size => new rank/size/color"
    end if

    call MPI_BARRIER(MPI_COMM_WORLD, ierr)

    do i=0,comm_size-1
        if (i == comm_rank) then
            write(*,'(I8,"/",I4," => ",I8,"/",I4,"/",I4)') comm_rank, comm_size, new_rank, new_size,color
        end if
        call MPI_BARRIER(MPI_COMM_WORLD, ierr)
    end do

    !! finalize
    call MPI_FINALIZE(ierr)
end program main

説明

MPI_COMM_WORLDを4プロセスごとに分割する．
やっていることは，各rankにcolorを割り当ててMPI_Comm_split関数を呼ぶだけ．

• 分割後のコミュニケータのrankはkeyの昇順に割り振られるが，同一値のkeyがある場合，それらは分割元のコミュニケータのrankの昇順で割り振られる
• 分割後はcolorを利用してプログラムの条件分岐を行えば良い

実行結果

[20:22:22,****@****:split]$ mpirun -np 8 ./a.out
old rank/size => new rank/size/color
       0/   8 =>        0/   4/   0
       1/   8 =>        1/   4/   0
       2/   8 =>        2/   4/   0
       3/   8 =>        3/   4/   0
       4/   8 =>        0/   4/   1
       5/   8 =>        1/   4/   1
       6/   8 =>        2/   4/   1
       7/   8 =>        3/   4/   1
[20:22:24,****@****:split]$ mpirun -np 6 ./a.out
old rank/size => new rank/size/color
       0/   6 =>        0/   4/   0
       1/   6 =>        1/   4/   0
       2/   6 =>        2/   4/   0
       3/   6 =>        3/   4/   0
       4/   6 =>        0/   2/   1
       5/   6 =>        1/   2/   1