Juliaで数値計算 その4:コードサンプル〜MPI並列計算編

Juliaで数値計算　その1：コードサンプル〜基本的計算編〜
https://qiita.com/cometscome_phys/items/31d0b811345a3e12fcef
Juliaで数値計算　その2：コードサンプル〜わかりやすい書き方編〜
https://qiita.com/cometscome_phys/items/dfd63f1f7b51bdda4ca4
Juliaで数値計算　その3：コードサンプル〜2次元プロット・可視化編〜
https://qiita.com/cometscome_phys/items/19b041b2f3364705f428
に引き続いて、FortranやCからやってきた人のためにJuliaの書き方を紹介する。
この記事はその4である。

昨今はノートPCでもCPUのコア数が複数ある場合も多く、数値計算をやる場合にはやはり並列計算がしたい、という需要があると思う。
Juliaでの並列計算はその1にも書いたが、今回は、よりFortranっぽく、MPIを使って並列化を行うこととする。MPIを使う利点は、既存のFortranやCのコードでのMPI並列と類似しているため、移植および理解がしやすい、という点である。
よく使いそうなものは適宜追記していく予定である。

バージョン

Julia: 1.1

インストール

MPI.jlのドキュメントを見つつインストールをする。
https://github.com/JuliaParallel/MPI.jl
このパッケージを入れる前に、MPIのインストールをしておかなければならないことに注意。例えば、Macの場合には、homebrewを用いて、

brew install --build-from-source openmpi

でopenmpiをインストールすることができる。

基本的には、]を押してパッケージモードとしてから、

add MPI

でインストールされる。

MPIについて

MPI並列計算は、異なる部屋での作業に似ている。並列計算の数だけ部屋があり、その中で何かの作業をしている。別々の作業をしている場合もあれば、異なる作業をしてい場合もある。もし、別の部屋で計算した結果が欲しかった場合、その部屋への電話(MPI 通信)をかけて、持ってきてもらう。他の部屋に電話をかけない限りそれぞれの作業は完全に独立(メモリが独立)である。最後に結果が欲しかった場合は、全員に連絡して、データをかき集めてもらう。

FortranにおいてのMPIの使用については
http://park.itc.u-tokyo.ac.jp/kato-yusuke-lab/nagai/note_170705_MPI.pdf
を参考にするとよい。

Forループの並列化:和

以下のようなForループを考える。

function test_0()    
    N=100
    wa = 0    
    for i=1:N
        wa += i^2
    end
    println("sum = $wa")    
    return wa
end

test_0()

このコードでは、i=1からi=Nまでの二乗の和を計算する。
実行結果は、

sum = 338350

となる。
このコードをMPI並列化してみよう。具体的なコードは、こちらである。

sample.jl

using MPI

function test()

    comm = MPI.COMM_WORLD
    nprocs = MPI.Comm_size(comm)
    myrank = MPI.Comm_rank(comm)

    N=100
    ista,iend,nbun = start_and_end(N,comm)

    wa = 0

    for i=ista:iend
        wa += i^2
    end
    println("$myrank: pertial sum = $wa")
    wa = MPI.Allreduce(wa,MPI.SUM,comm)
    println("$myrank: sum = $wa")


    return wa
end

function start_and_end(N,comm)
    nprocs = MPI.Comm_size(comm)
    myrank = MPI.Comm_rank(comm)
    if N % nprocs != 0
        println("error! N%procs should be 0.")
    end
    nbun = div(N,nprocs)
    ista = myrank*nbun+1
    iend = ista + nbun-1
    return ista,iend,nbun
end

MPI.Init() # MPI初期化
test()
MPI.Finalize() #MPI終了

ここで、start_and_endはForループの分割の方法を与える関数であり、今回はループの長さNがプロセス数procsで割り切れる時のみ動くようにしている。
このコードでは、ループの変数iを各MPIプロセスごとに割り振りをしており、Forループを分割している。そして、それぞれのプロセスでの計算が終わった後、最後にそれぞれのプロセスの和をとってきている。MPI 並列によるForループの並列計算はこれだけである。
MPI.Allreduceは各プロセスで計算した結果を集めて何かの処理をしてそれぞれのプロセスに再分配する関数である。今回はMPI.SUMという、集めた値の和をとり、それぞれのプロセスに再分配している。
このコードを並列で実行するには、

mpirun -np 4 julia sample.jl

とすればよい。ここでの4はプロセス数の数である。
出力結果は、

0: pertial sum = 5525
1: pertial sum = 37400
3: pertial sum = 194900
2: pertial sum = 100525
1: sum = 338350
3: sum = 338350
0: sum = 338350
2: sum = 338350

となる。

2次元

具体的な物理の計算においては、二次元での運動量空間での積分をすることがよくある。例えば、以下のコードである。

function test2d_0()
    Nx = 100
    Ny = 100
    wa = 0
    dkx = 2π/(Nx-1)
    dky = 2π/(Ny-1)
    for ikx=1:Nx
        for iky=1:Ny            
            kx = (ikx-1)*dkx -π
            ky = (iky-1)*dky -π
            wa += cos(kx)+cos(ky)
        end
    end
    wa /= Nx*Ny
    println("sum = $wa")

    return wa
end

test2d_0()

このコードでは、変数kxとkyが-πからπまで動いている。このコードのMPI並列計算コードは、

function test2d()
    comm = MPI.COMM_WORLD
    myrank = MPI.Comm_rank(comm)

    Nx = 100
    Ny = 100
    N = Nx*Ny
    wa = 0
    dkx = 2π/(Nx-1)
    dky = 2π/(Ny-1)

    ista,iend,nbun = start_and_end(N,comm)
    for i=ista:iend
        ikx,iky=i2xy(i,Nx)
        kx = (ikx-1)*dkx -π
        ky = (iky-1)*dky -π
        wa += cos(kx)+cos(ky)
    end
    wa /= Nx*Ny
    println("$myrank: pertial sum = $wa")
    wa = MPI.Allreduce(wa,MPI.SUM,comm)    
    println("$myrank: sum = $wa")

    return wa
end

function i2xy(i,Nx)
    x = ((i-1) % Nx)+1
    y = div((i-x),Nx)+1
    return x,y
end

MPI.Init() # MPI初期化
test2d()
MPI.Finalize() #MPI終了

である。ここで、Forループを分割するために、二重だったループを一重に変更してある。
その対応関係をi2xyという関数に書いてある。
このコードでのwaを計算する部分を適当に置き換えれば、運動量空間の積分の計算ができる。

複素数かつ配列

上記のコードの複素数配列版は

function test2d_2x2()
    comm = MPI.COMM_WORLD
    myrank = MPI.Comm_rank(comm)

    Nx = 100
    Ny = 100
    N = Nx*Ny
    wa = zeros(ComplexF64,2,2)
    dkx = 2π/(Nx-1)
    dky = 2π/(Ny-1)

    ista,iend,nbun = start_and_end(N,comm)
    for i=ista:iend
        ikx,iky=i2xy(i,Nx)
        kx = (ikx-1)*dkx -π
        ky = (iky-1)*dky -π
        wa[1,1] += cos(kx)+cos(ky)
        wa[2,2] += cos(kx)+cos(ky)+1im
    end
    wa ./= Nx*Ny
    println("$myrank: pertial sum = $wa")
    watemp = similar(wa)
    MPI.Allreduce!(wa,watemp,4,MPI.SUM,comm)    
    wa = watemp
    println("$myrank: sum = $wa")

    return wa
end

MPI.Init() # MPI初期化
test2d_2x2()
MPI.Finalize() #MPI終了

となる。
注意しなければならないのは、MPI.Allreduceではエラーが出る、ということである。
最初なぜエラーが出るか全くわからなかったが、これは、MPI.jl v0.7.2のmpi-base.jl
の745行目のところでMPIで送る長さが1に固定されているからである。
これを回避するには、
MPI.Allreduce!(wa,watemp,4,MPI.SUM,comm)を使えばよく、こうするとwatempの長さだけMPIで通信してくれる。
多分これはバグではないかと思う。

並列要素ごとの作業：配列を合体

次に、配列を小さな配列に分割し、その配列の配列要素ごとに計算をして、最後に大きな配列にまとめる、というコードについて述べる。
まず、並列していないコードを

function test3_0()
    N = 100

    vec_a = zeros(Float64,N)
    ii = 0
    for i=1:N
        vec_a[i] += internalfunc(i)
    end
    println(" sum= $(sum(vec_a))")
    return vec_a

end

function internalfunc(i)
    value = 0
    for j=1:100000
        value += cos(i^2)*cos(j^2)
    end
    return value
end

test3_0()

とする。ここでは、配列vec_aの要素ごとに関数internalfuncを呼び出している。
これを並列化すると、

function test3()
    comm = MPI.COMM_WORLD
    myrank = MPI.Comm_rank(comm)
    N = 100

    vec_a = zeros(Float64,N)
    ista,iend,nbun = start_and_end(N,comm)

    vec_atemp = zeros(Float64,nbun)
    ii = 0
    for i=ista:iend
        ii += 1
        vec_atemp[ii] += internalfunc(i)
    end
    println("$myrank: pertial sum= $(sum(vec_atemp))")
    vec_a = MPI.Allgather(vec_atemp,comm)
    println("$myrank: sum= $(sum(vec_a))")
    return vec_a
end

MPI.Init() # MPI初期化
test3()
MPI.Finalize() #MPI終了

となる。ここで、vec_atempは分割された小さな配列である。そして、MPI.Allgatherはそのような配列を合体させて一つの大きな配列にする。

(2020/6/11)追記：Julia 1.4におけるMPI.jlのインストールについて

Mac OS 10.14.6でMPI.jlのインストールに問題がある場合についてです。
Julia 1.3ではMPI.jlがちゃんとインストールできたのにJulia 1.4ではできない、という状況になりました。
その場合には、.bashrcや.bash_profileに

export JULIA_MPI_PATH=/usr/local/Cellar/open-mpi/4.0.3

と入れると解決する場合があります（4.0.3はバージョンなので環境に合わせてください）。なお、ここでのMPIはbrewでopen-mpiを入れた場合を想定しています。