JuliaでMPI片方向通信を行う

Juliaで並列計算したくありませんか？
並列計算として有名なのはMPIですが、これをJuliaで使うことができます。パッケージはMPI.jlです。これをインストールすれば、JuliaでMPIが使えます。

例えば、

using MPI
function test()
    MPI.Init()
    comm = MPI.COMM_WORLD
    print("Hello world, I am rank $(MPI.Comm_rank(comm)) of $(MPI.Comm_size(comm))\n")
    myrank = MPI.Comm_rank(comm)
    MPI.Barrier(comm)
end
test()

というコードを作って、

mpirun -np 2 test.jl

のように実行すると、

Hello world, I am rank 0 of 2
Hello world, I am rank 1 of 2

のように並列に動作させることができます。

MPI片方向通信

MPIでの通信方法にはsendとrecv、あるいはisendとirecvが有名ですが、これ以外に片方向通信と呼ばれるものがあります。
こちらのPDFが詳しいです。
片方向通信を使うと、相手のプロセスの配列にデータを「置いてくる」ことが可能となります。受け取るために待機する必要がなく、便利です。

例

まず、2並列の場合を考えます。プロセスには番号（rank）がつけられており、0と1になります。
まず、rank0と1で、

    n = 10
    if myrank == 0
        a1 = ones(n)
    else
        a1 = 2*ones(n)
    end

のように、[1,1,1,1,1,1,1,1,1,1]と[2,2,2,2,2,2,2,2,2,2]という配列を持っているとします。
ここで、rank 0の最初の半分の1をrank 1の前半へ、rank 1の後半半分の2をrank 0の後半へ移すことを考えます。つまり、前半と後半の入れ替えです。これは、rank 0も1も[1,1,1,1,1,2,2,2,2,2]という配列を持つ形にする、ということになります。これをやってみましょう。

まず、書き込みOKの配列を確保します。

 a = zeros(Float64,n)
win = MPI.Win_create(a,comm)

これで、aという配列は別のプロセスからアクセスできるようになりました。
次に、書き込みを開始します。書き込みをしている間はMPI.Win_fence(0, win)というもので囲む必要がありまして、

MPI.Win_fence(0, win)

書き込み

MPI.Win_fence(0, win)

のような形になります。ここでは、

    MPI.Win_fence(0, win)

    if myrank == 0
        a[1:div(n,2)] .= a1[1:div(n,2)]
        target_rank = 1
        MPI.Put(a1[1:div(n,2)], target_rank,win)
    else   
        a[1+div(n,2):end] .= a1[1:div(n,2)]
        target_rank = 0
        MPI.Put(a1[1+div(n,2):n],target_rank, div(n,2),win)
    end

    MPI.Win_fence(0, win)

とします。これは、rank 0ではaの前半にa1の前半（つまり[1,1,1,1,1]）をコピーしてから、MPI.Putという関数でrank 1のaにa1の前半を書き込んでいます。一方、rank 1ではaの後半にa1の前半（つまり[1,1,1,1,1]）をコピーして、MPI.Putという関数でaの後半にa1の後半を書き込みしています。

ここで、

MPI.Put(a1[1:div(n,2)], target_rank,win)

はwinを作る際に定義したaという配列の先頭からa1[1:div(n,2)]を書き込みするという意味です。また、

MPI.Put(a1[1+div(n,2):n],target_rank, div(n,2),win)

は引数が一つ多く div(n,2)というものが入っていますが、これはaという配列の先頭から div(n,2)個をスキップしてその後にa1[1+div(n,2):n]を書き込む、という意味です。

全体コード

コード全体は

using MPI
function test()
    MPI.Init()
    comm = MPI.COMM_WORLD
    print("Hello world, I am rank $(MPI.Comm_rank(comm)) of $(MPI.Comm_size(comm))\n")
    myrank = MPI.Comm_rank(comm)
    MPI.Barrier(comm)

    n = 10
    a = zeros(Float64,n)

    win = MPI.Win_create(a,comm)
    
    if myrank == 0
        a1 = ones(n)
    else
        a1 = 2*ones(n)
    end

    MPI.Win_fence(0, win)

    if myrank == 0
        a[1:div(n,2)] .= a1[1:div(n,2)]
        target_rank = 1
        MPI.Put(a1[1:div(n,2)], target_rank,win)
    else   
        a[1+div(n,2):end] .= a1[1:div(n,2)]
        target_rank = 0
        MPI.Put(a1[1+div(n,2):n],target_rank, div(n,2),win)
    end

    MPI.Win_fence(0, win)

    if myrank == 0
        println("myrank = $myrank ",a)
    end
    MPI.Barrier(comm)
    if myrank == 1
        println("myrank = $myrank ",a)
    end

    MPI.free(win)

end

test()

となります。片側通信が終わったら必ずMPI.free(win)で終了させます。
このコードを2並列で実行すると、

Hello world, I am rank 0 of 2
Hello world, I am rank 1 of 2
myrank = 0 [1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 2.0]
myrank = 1 [1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 2.0]

となります。