mpi並列化楽々化計画 in Fortran: mpi_allgathervを使いこなす

Fortran Advent Calendar 2019です。

mpi_allgathervを使う方法について、オブジェクト指向を使って述べます。
モジュールにしたので気軽に呼べるようになりました。

バージョン

gcc version 9.2.0 (Homebrew GCC 9.2.0_2)

はじめに

MPIでループを並列化したい！ということはよくあると思います。
例えば

program main
    implicit none
    integer::N
    real(8),allocatable::vec_a(:)
    integer::i

    N = 11
    allocate(vec_a(N))
    do i=1,N
        vec_a(i) = i
    end do
    write(*,*) vec_a
end program 

このdoループを並列計算したい、ということを考えます。

ループ数が並列数で割り切れる場合

並列数をnprocsとして、Nが並列数で割り切れる場合、

program main
    implicit none
    include "mpif.h"
    integer::N
    real(8),allocatable::vec_a(:)
    real(8),allocatable::vec_a_temp(:)
    integer::i,ista,iend,ierr,myrank,nprocs,nbun,count


    call mpi_init(ierr)


    call mpi_comm_rank(mpi_comm_world,myrank,ierr)
    call mpi_comm_size(mpi_comm_world,nprocs,ierr)

    N = 10
    nbun = N/nprocs
    ista = myrank*nbun + 1
    iend = ista + nbun -1
    allocate(vec_a(N))
    allocate(vec_a_temp(nbun))

    count = 0
    do i=ista,iend
        count = count + 1
        vec_a_temp(count) = i
    end do
    call mpi_allgather(vec_a_temp,nbun,mpi_double_precision,vec_a,nbun,mpi_double_precision,MPI_COMM_WORLD,ierr)   
    if (myrank == 0) then
        write(*,*) vec_a
    end if

    call mpi_finalize(ierr)

end program 

でできます。ここで、mpi_allgatherというMPIのサブルーチンを呼んでいます。
これは、長さN/nprocsのvec_a_tempという配列をnprocs分積み上げて、長さNのvec_aという配列を作るサブルーチンです。
ここで、それぞれのMPIプロセスでループの数が同じ場合には配列の長さN/nprocsは各プロセスで同じです。しかし、並列数で割り切れない場合には、この方法は使えません。

ループ数が並列数で割り切れない場合

ループ数が並列数で割り切れない場合、mpi_allgathervというサブルーチンを使います。しかし、いろいろ引数があって面倒です。
参考URL:
http://nkl.cc.u-tokyo.ac.jp/11e/03-MPI/reportS1.pdf
そこで、オブジェクト指向を使って煩雑な手続きを見えないようにしてしまいましょう。
まず、classとして、gatherv typeを作ります。

    type gatherv
        integer::nprocs
        integer::myrank
        integer::N
        integer,allocatable::rcounts(:)
        integer,allocatable::displs(:)
        integer::count

        contains
        procedure:: allgatherv_dble =>  allgatherv_dble
        procedure:: allgatherv_complex =>  allgatherv_complex
        procedure:: get_start
        procedure:: get_end
    end type gatherv

ここに、MPI並列で必要そうな情報を入れておきます。

次に、このtypeを使うためのコンストラクタ

    type(gatherv) function init_gatherv(N) result(result)
        implicit none
        integer,intent(in)::N
        integer::ierr,nprocs,myrank,i
        integer::count,m

        call mpi_comm_rank(mpi_comm_world,myrank,ierr)
        call mpi_comm_size(mpi_comm_world,nprocs,ierr)
        result%nprocs = nprocs
        result%myrank = myrank
        result%N = N
        allocate(result%rcounts(1:nprocs))
        allocate(result%displs(1:nprocs+1))

        count = 0
        m = mod(N,nprocs) !N*nprocs + m
        count = (N-m)/nprocs
        if (mod(myrank-1+nprocs,nprocs)+1 <= m) then 
            count = count + 1
        end if
        result%count = count
        call mpi_allgather(count,1,mpi_integer,result%rcounts,1,mpi_integer,mpi_comm_world,ierr)

        result%displs(1) = 0
        do i=1,nprocs
            result%displs(i+1) = result%displs(i) + result%rcounts(i)
        end do

        return
    end function init_gatherv

を作ります。MPI並列数をnprocs、自分自身のプロセス番号をmyrankとしています。
また、mpi_allgathervで必要なrcountsとdisplsを作成しています。
ここでは、まずループ数Nと並列数nprocsのあまりmを計算し、割り切れる数(N-m)/nprocsを求めます。全てのプロセスが少なくともこの数だけループを回すことになります。残りはあまりのmですが、残りのm回ループを各プロセスに割り振っています。myrank=0は何かと他の処理をする可能性があるためループが一番少なくなるようにしました。
例えば、m=5でnprocsが10の場合、myrankは0,1,2,3,...,9までですが、残りループ数m=5はmyrank=1,2,3,4,5の5つのプロセスが担当することになります。つまり、N=105,nprocs=10なら、あまりは5で、myrank=0は10回、myrank=1から5までが11回、残りは10回ループを回します。
これで、それぞれのプロセスにおけるベクトルの長さは、10,11,11,11,11,11,10,10,10,10となります。

このfunctionはmodule内で

    interface gatherv
        module procedure::init_gatherv
    end interface gatherv

としてgathervとして呼べるようにしておきます。

次に、ループを分割した時のそれぞれのプロセスでのiの最初と最後を指定する関数を作っておきます。

    integer function get_start(self) result(ista)
        implicit none
        class(gatherv)::self
        ista = self%displs(self%myrank+1) + 1
        return
    end function get_start

    integer function get_end(self) result(iend)
        implicit none
        class(gatherv)::self
        iend = self%displs(self%myrank+1+1)
        return
    end function get_end  

そして、mpi_allgathervを呼ぶサブルーチンを

    subroutine allgatherv_dble(self,v_ip,v) 
        implicit none
        class(gatherv)::self
        real(8),intent(in)::v_ip(:)
        real(8),intent(out)::v(:)
        integer::ierr
        call mpi_allgatherv(v_ip(1:self%count),self%count,mpi_double_precision,v, &
            self%rcounts,self%displs,mpi_double_precision,mpi_comm_world,ierr)
        return
    end subroutine allgatherv_dble

    subroutine allgatherv_complex(self,v_ip,v) 
        implicit none
        class(gatherv)::self
        complex(8),intent(in)::v_ip(:)
        complex(8),intent(out)::v(:)
        integer::ierr
        call mpi_allgatherv(v_ip(1:self%count),self%count,mpi_double_complex,v, &
            self%rcounts,self%displs,mpi_double_complex,mpi_comm_world,ierr)
    end subroutine allgatherv_complex  

と作ります。

全体のmoduleは以下の通りです

module mpicalls
    implicit none
    include "mpif.h"
    private
    public::gatherv!,init_gatherv


    type gatherv
        integer::nprocs
        integer::myrank
        integer::N
        integer,allocatable::rcounts(:)
        integer,allocatable::displs(:)
        integer::count

        contains
        procedure:: allgatherv_dble =>  allgatherv_dble
        procedure:: allgatherv_complex =>  allgatherv_complex
        procedure:: get_start
        procedure:: get_end
    end type gatherv

    interface gatherv
        module procedure::init_gatherv
    end interface gatherv

    contains

    integer function get_start(self) result(ista)
        implicit none
        class(gatherv)::self
        ista = self%displs(self%myrank+1) + 1
        return
    end function get_start

    integer function get_end(self) result(iend)
        implicit none
        class(gatherv)::self
        iend = self%displs(self%myrank+1+1)
        return
    end function get_end   

    type(gatherv) function init_gatherv(N) result(result)
        implicit none
        integer,intent(in)::N
        integer::ierr,nprocs,myrank,i
        integer::count,m

        call mpi_comm_rank(mpi_comm_world,myrank,ierr)
        call mpi_comm_size(mpi_comm_world,nprocs,ierr)
        result%nprocs = nprocs
        result%myrank = myrank
        result%N = N
        allocate(result%rcounts(1:nprocs))
        allocate(result%displs(1:nprocs+1))

        count = 0
        m = mod(N,nprocs) !N*nprocs + m
        count = (N-m)/nprocs
        if (mod(myrank-1+nprocs,nprocs)+1 <= m) then 
            count = count + 1
        end if
        result%count = count
        call mpi_allgather(count,1,mpi_integer,result%rcounts,1,mpi_integer,mpi_comm_world,ierr)

        result%displs(1) = 0
        do i=1,nprocs
            result%displs(i+1) = result%displs(i) + result%rcounts(i)
        end do

        return
    end function init_gatherv

    subroutine allgatherv_dble(self,v_ip,v) 
        implicit none
        class(gatherv)::self
        real(8),intent(in)::v_ip(:)
        real(8),intent(out)::v(:)
        integer::ierr
        call mpi_allgatherv(v_ip(1:self%count),self%count,mpi_double_precision,v, &
            self%rcounts,self%displs,mpi_double_precision,mpi_comm_world,ierr)
        return
    end subroutine allgatherv_dble

    subroutine allgatherv_complex(self,v_ip,v) 
        implicit none
        class(gatherv)::self
        complex(8),intent(in)::v_ip(:)
        complex(8),intent(out)::v(:)
        integer::ierr
        call mpi_allgatherv(v_ip(1:self%count),self%count,mpi_double_complex,v, &
            self%rcounts,self%displs,mpi_double_complex,mpi_comm_world,ierr)
    end subroutine allgatherv_complex  


end module mpicalls

使い方

使い方は、以下の通りです。

program main
    use mpicalls
    implicit none
    include "mpif.h"
    integer::N
    type(gatherv)::g
    real(8),allocatable::vec_a(:)
    real(8),allocatable::vec_a_temp(:)
    integer::i,ista,iend,ierr,myrank,nprocs,nbun,count

    call mpi_init(ierr)
    N = 111

    allocate(vec_a(N))

    g = gatherv(N)
    allocate(vec_a_temp(g%count))

    count = 0
    do i=g%get_start(),g%get_end()
        count = count + 1
        vec_a_temp(count) = i
    end do
    call g%allgatherv_dble(vec_a_temp,vec_a)

    if (g%myrank == 0) then
        do i=1,N
            write(*,*) vec_a(i)
        end do
    end if


    call mpi_finalize(ierr)

end program

初期化として、g = gatherv(N)を呼び、do i=g%get_start(),g%get_end()でループを分割し、call g%allgatherv_dble(vec_a_temp,vec_a)で呼べば、それでOKです。

全体のコード

全体のコードを置いておきます。

module mpicalls
    implicit none
    include "mpif.h"
    private
    public::gatherv!,init_gatherv


    type gatherv
        integer::nprocs
        integer::myrank
        integer::N
        integer,allocatable::rcounts(:)
        integer,allocatable::displs(:)
        integer::count

        contains
        procedure:: allgatherv_dble =>  allgatherv_dble
        procedure:: allgatherv_complex =>  allgatherv_complex
        procedure:: get_start
        procedure:: get_end
    end type gatherv

    interface gatherv
        module procedure::init_gatherv
    end interface gatherv

    contains

    integer function get_start(self) result(ista)
        implicit none
        class(gatherv)::self
        ista = self%displs(self%myrank+1) + 1
        return
    end function get_start

    integer function get_end(self) result(iend)
        implicit none
        class(gatherv)::self
        iend = self%displs(self%myrank+1+1)
        return
    end function get_end   

    type(gatherv) function init_gatherv(N) result(result)
        implicit none
        integer,intent(in)::N
        integer::ierr,nprocs,myrank,i
        integer::count,m

        call mpi_comm_rank(mpi_comm_world,myrank,ierr)
        call mpi_comm_size(mpi_comm_world,nprocs,ierr)
        result%nprocs = nprocs
        result%myrank = myrank
        result%N = N
        allocate(result%rcounts(1:nprocs))
        allocate(result%displs(1:nprocs+1))

        count = 0
        m = mod(N,nprocs) !N*nprocs + m
        count = (N-m)/nprocs
        if (mod(myrank-1+nprocs,nprocs)+1 <= m) then 
            count = count + 1
        end if
        result%count = count
        call mpi_allgather(count,1,mpi_integer,result%rcounts,1,mpi_integer,mpi_comm_world,ierr)

        result%displs(1) = 0
        do i=1,nprocs
            result%displs(i+1) = result%displs(i) + result%rcounts(i)
        end do

        return
    end function init_gatherv

    subroutine allgatherv_dble(self,v_ip,v) 
        implicit none
        class(gatherv)::self
        real(8),intent(in)::v_ip(:)
        real(8),intent(out)::v(:)
        integer::ierr
        call mpi_allgatherv(v_ip(1:self%count),self%count,mpi_double_precision,v, &
            self%rcounts,self%displs,mpi_double_precision,mpi_comm_world,ierr)
        return
    end subroutine allgatherv_dble

    subroutine allgatherv_complex(self,v_ip,v) 
        implicit none
        class(gatherv)::self
        complex(8),intent(in)::v_ip(:)
        complex(8),intent(out)::v(:)
        integer::ierr
        call mpi_allgatherv(v_ip(1:self%count),self%count,mpi_double_complex,v, &
            self%rcounts,self%displs,mpi_double_complex,mpi_comm_world,ierr)
    end subroutine allgatherv_complex  


end module mpicalls

program main
    use mpicalls
    implicit none
    include "mpif.h"
    integer::N
    type(gatherv)::g
    real(8),allocatable::vec_a(:)
    real(8),allocatable::vec_a_temp(:)
    integer::i,ista,iend,ierr,myrank,nprocs,nbun,count

    call mpi_init(ierr)
    N = 111

    allocate(vec_a(N))

    g = gatherv(N)
    allocate(vec_a_temp(g%count))

    count = 0
    do i=g%get_start(),g%get_end()
        count = count + 1
        vec_a_temp(count) = i
    end do
    call g%allgatherv_dble(vec_a_temp,vec_a)

    if (g%myrank == 0) then
        do i=1,N
            write(*,*) vec_a(i)
        end do
    end if


    call mpi_finalize(ierr)

end program