Juliaで並列計算：DistributedArrays.jlを使う

Juliaで並列計算をするためにこれまで色々調べてきました。今回は配列をCPUコアに分散して持つことができるというDistributedArrays.jlを使ってみたいと思います。

最近、julia を googlecolab で使えるようにしたい。
という記事によって、非常に簡単にGoogleColabでJuliaを実行できるようになりましたので、ここではGoogle Colabを使いたいと思います。使い方はリンク先を参照してください。

DistributedArraysについて

DistributedArraysは配列をCPUコアで分割して持つことができます。一つのCPUコアが持つ領域がコア数に応じて変化するわけです。行列の場合はわかりやすいです。行列を分割して持っておくと、行列と何かの演算をする際に、複数のCPUコアが同時に実行することができます。例えば、行列の各要素の全部の和を計算したいときは、それぞれのCPUコアが自分の持っている領域の和を取ってから、最後に全体の和をとれば良いです。これによってCPUコアを増やせば増やすほど計算が速くなることがわかります。あるいは、配列を分散して持ってけば、配列の各要素ごとに計算したい場合には簡単に並列化ができます。物理学で言えば、空間を離散化して何らかの場を配列として持っていた時に、その配列を引数にして別の計算をしたい、ということがよくあります。このような時にも分散して持っていると並列化ができます。

インストール

インストールはいつものようにJuliaのREPLで]キーを押してパッケージモードにしてから

add DistributedArrays

で入れられます。並列計算をしますので、

add Distributed

も入れておいてください。
Google colabであれば、

using Pkg
Pkg.add("DistributedArrays")
Pkg.add("Distributed")
using Distributed
using DistributedArrays

でOKです。

バージョン

• Julia 1.6.0

使い方

Juliaでの並列計算はMPIとは少し異なります。詳しくは、Juliaでプロセス並列〜MPIとの違いをみてください。
Juliaでの並列計算ではワーカーというものを使います。ワーカーに作業をさせる、という感じです。ということで、Google Colabのノートブックにワーカーを追加します。

addprocs(4)
nworkers()

これで4つ追加されました。4並列で動くことになります。

次に、

@everywhere module Paralleltest
    using Distributed
    using DistributedArrays 

    function loops(a)
        for (ilocal,i) in enumerate(localindices(a)[1]),(jlocal,j) in enumerate(localindices(a)[2])
            localpart(a)[ilocal,jlocal] =1000*i+j
        end
    end

    function test()
        a = dzeros((10,10), workers(), [1,4])
        indices = [@fetchfrom p localindices(a) for p in workers()]
        println(indices)
        for id in workers()
            remotecall(D->println(localindices(D)),id,a)
        end

        for id in workers()
            remotecall(loops,id,a)
        end
        display(a)
    end
end
using .Paralleltest
Paralleltest.test()

が本体のコードです。
少し説明をします。

a = dzeros((10,10), workers(), [1,4])

は10x10行列を4分割して持つことを意味しています。要素はゼロになっています。2個目の引数workers()はワーカーのIDが収められています。ワーカーはIDで区別されます。今4並列ですので、これは4個の整数が入った配列です。3個目の引数は分割の仕方を意味しています。今、2次元配列を考えていますが、[1,4]は1次元目は1分割つまり分割せず、2次元目は4分割にしています。4並列なので4分割です。つまり、この4は要素数で割り切れる必要はありません。良い感じに分配してくれます。それぞれのワーカーがどのような範囲の配列を持っているかを知るには、

indices = [@fetchfrom p localindices(a) for p in workers()]
println(indices)

です。@fetchfrom pというのは、IDがpのワーカーからデータを取ってくる、という意味です。localindices(a)は指定されたワーカーが持つ配列を返す関数です。出力は

Tuple{UnitRange{Int64}, UnitRange{Int64}}[(1:10, 1:3), (1:10, 4:6), (1:10, 7:8), (1:10, 9:10)]

こんな感じになります。それぞれのワーカーが配列を分割してもっていることがわかりますね。
上のコードでは、他にも

for id in workers()
   remotecall(D->println(localindices(D)),id,a)
end

というのがあります。この出力結果は

      From worker 4:	(1:10, 7:8)
      From worker 5:	(1:10, 9:10)
      From worker 3:	(1:10, 4:6)
      From worker 2:	(1:10, 1:3)

となります。From workerみたいに書かれているものは、そのワーカーが出力したことを意味しています。remotecall(D->println(localindices(D)),id,a)は、remotecall(func,id,a)という形式になっていまして、idのワーカーにfuncという仕事をさせる、という関数です。aは引数です。ここでは、ラムダ式（無名関数）を使って、localindices(D)をprintするようにしました。

このremotecallはかなり有用な関数でして、上のコードでは

for id in workers()
            remotecall(loops,id,a)
end

としています。これは、idのワーカーにloop(a)を実行させるという意味です。loopという関数は

function loops(a)
        for (ilocal,i) in enumerate(localindices(a)[1]),(jlocal,j) in enumerate(localindices(a)[2])
            localpart(a)[ilocal,jlocal] =1000*i+j
        end
end

で定義されています。for (ilocal,i) in enumerate(localindices(a)[1]),(jlocal,j) in enumerate(localindices(a)[2])は、2変数のfor文ですね。

配列の書き換え

DistributedArraysで使われているDArray型は、自分の持っている配列の読み書きは可能ですが、他のワーカーが持っている配列は読み込むことのみしかできません。つまり、配列を書き換える際には、自分が持っている配列を書き換えるように指示する必要があります。DArray型aの自分が持っている配列を取り出すにはlocalpart(a)とします。つまり、

localpart(a)[ilocal,jlocal] =1000*i+j

とすると、自分が持っている配列を書き換えます。ここで少し注意してもらいたいのは、分割して持った時の全体の配列のインデックスと自分が持っている配列のインデックスは異なっている、ということです。例えば、3:5までの三つのインデックスを持っているワーカーがいたとします。この時、localpart(a)という配列のインデックスは1:3です。つまり、1から順番に格納されています。これは考えてみれば明らかで、各ワーカーがそれぞれの配列を確保しているわけですから、その確保した配列に順番に値を入れるなら1から順番に値が入っているわけです。
ともかく、localpartを使うことによって配列の書き換えができます。

Google Colab上での出力結果は

10×10 DArray{Float64, 2, Matrix{Float64}}:
  1001.0   1002.0   1003.0   1004.0  …   1007.0   1008.0   1009.0   1010.0
  2001.0   2002.0   2003.0   2004.0      2007.0   2008.0   2009.0   2010.0
  3001.0   3002.0   3003.0   3004.0      3007.0   3008.0   3009.0   3010.0
  4001.0   4002.0   4003.0   4004.0      4007.0   4008.0   4009.0   4010.0
  5001.0   5002.0   5003.0   5004.0      5007.0   5008.0   5009.0   5010.0
  6001.0   6002.0   6003.0   6004.0  …   6007.0   6008.0   6009.0   6010.0
  7001.0   7002.0   7003.0   7004.0      7007.0   7008.0   7009.0   7010.0
  8001.0   8002.0   8003.0   8004.0      8007.0   8008.0   8009.0   8010.0
  9001.0   9002.0   9003.0   9004.0      9007.0   9008.0   9009.0   9010.0
 10001.0  10002.0  10003.0  10004.0     10007.0  10008.0  10009.0  10010.0

となります。