LoginSignup
20
21

More than 3 years have passed since last update.

@oki_mebarun

諸々のテクニックでJuliaコードを高速化してみたら５０倍以上速くなった件

Posted at

はじめに

　Juliaでのプログラミングは、豊富な表現方法と数学との親和性のおかげで非常に素直に書けて楽しいのですが、”メモリ”を意識しないと相当効率が悪化してしまうことがあるようです。
　本記事では、前回書いたイジングモデルのコード　(Juliaで非正方格子のイジングモデルをシミュレーションしてみた)を例に、各種の高速化テクニックを使って高速化を試みます。
　

高速化の指針

　こちらのブログと

　黒木玄さん(@genkuroki)からのヒント

　をもとに、下記の指針でいきたいと思います。

  1. 無駄な計算を省く
    • アルゴリズムを見直して、不要な計算を省きます。
  2. 計算結果をキャッシュする
    • 何度も同じ計算をするのであれば、敢えてメモリに確保してしまいます。
  3. ヒープアロケーションを極力避ける
    • 何故かは知りませんが、現在のJulia(1.4)は変数に対してスタックを使わず何でもヒープに確保しようとしてしまうようです。これは、メモリアロケーションでの確保コストとガベージコレクションでの回収コストの２重の意味で損なことになってしまいます。
    • なので、特にループ内でのメモリアロケーションは避けるのが無難です。変数の確保をループ外に移すのが最善ですが、
    • StaticArrays.jl パッケージを使って、下記のように記述するとヒープではなくスタックにメモリが確保されるようです。ただし、サイズの大きな変数には向かないようです（あるサイズを超えるとヒープに確保されてしまう）。
    @SVector [a,b,c]
    @SVector [i^2 for i in 1:10]
  • また、配列のスライス(部分取り出し）をする際は、@viewマクロを使った方が良いようです。
  > a = fill(1,100000)
　> @time sum(a[1:90000])
　> @time sum(@view a[1:90000])
 　0.000371 seconds (3 allocations: 703.219 KiB)
   0.000058 seconds (2 allocations: 64 bytes)

viewマクロを挟んだだけで6倍近く速くなりました。

4.型を意識する。

  • 変数がAny型だと、型推論の過程（どの関数を当てはめるか）で無駄が多いようです。型を決めてしまっていい変数であれば予め固定型にしておくのが無難です。また、型の情報は伝搬するので、計算の初期段階で決めてあげれば十分なようです。

5.マクロを活用する

  • 高速化に寄与する便利マクロがいくつか用意されています。

    • @inbounds 配列のインデックスが範囲を超えないよと宣言して、実行時チェックを省く。
    • @inline 関数をインライン化する。
    • @simd CPUのSIMD命令(複数データへの同一命令の並列化)が適用できるなら適用する。

ソースコード

　では、具体的にソースコードを見ながら見直しの過程を追っていきます。全体はGitHubにアップロードしてあります。

下記の実行環境で動作確認済みです。

  • Julia 1.4.2
  • パッケージ：Plots, StaticArrays

1. 無駄な計算を省く

　前回のコードでは、MCMCの計算完了後にエネルギーを計算していましたが、MCMCの過程で⊿Eを毎回計算しているので、それを使えば再計算は不要でした。また、変数Sを確保する際に、Int8の型指定をするように変えました。
Before:

IsingMCMC02a.jl
    function MCMC( T, N, trial)
        # initialize
        sim = []
        S = ones(N)
        for i in 1:N
            if rand() < 0.5
                flipx!(S,i) # random flip at first
            end
        end
        # Gibbs sampling position
        gpos = [rand(1:N) for i in 1:trial]
        # random value in [0,1]
        rn = [rand() for i in 1:trial]

        # MCMC trial
        for t in 1:trial
            k = gpos[t]
            de = dE(S, k, N)
            if rn[t] < PrE(de, T) # MH criteria 
                flipx!(S, k) # change k
            end
            push!(sim, copy(S))
        end

        sim
    end

After:

IsingMCMC02a.jl
    function MCMC( T, N, trial)
        # initialize
        simE = zeros(Int, div(trial,1000)+1)
        S = ones(Int8, N)
        for i in 1:N
            if rand() < 0.5
                flipx!(S,i) # random flip at first
            end
        end

        # MCMC trial
        Ec = E(S,N) # current energy
        simE[1] = Ec
        local k, de
        @inbounds for t in 1:trial
            k  = rand(1:N) # Gibbs sampling position
            de = dE(S, k, N)
            if rand() < PrE(de, T) # MH criteria   
                flipx!(S, k) # change k
                Ec += de     # change E
            end
            (t % 1000 == 0) && (simE[div(t,1000)+1] = Ec)
        end

        (simE, S)
    end

2. 計算結果をキャッシュする

　前回のコードではgetJ, getBといった関数でイジングモデルの係数を求め、getNeighbors関数で近傍点を求めていました。しかし、どうせ同じ値なので辞書(Dict)を使ってキャッシュ化することにしました。

Before:

IsingSimMain02a.jl
function getJ(i, j, d=d)
    J0 = 1
    (j in getNeighbors(i,d)) ? J0 : 0
end

function getB(i, d=d)
    B0 = 0
    B0
end
IsingMCMC02a.jl
    function dE(S, k, N)
        dEt = 0
        J = Main.getNeighbors(k)
        dEt +=   S[k]*sum( [(getJ(k,j) + getJ(j,k)) * S[j] for j in J] )
        dEt += 2*S[k]*getB(k)
        dEt
    end

After:

IsingSimMain03a.jl
function genInitialJB(d=d)
    Jd = [Dict{Int,Int8}() for i in 1:d.N]
    Bd = zeros(Int8, d.N)
    (Jd, Bd)
end

function setJB!( Jd=Jd, Bd=Bd, d=d)
    J0 = 1
    B0 = 0
    for i in 1:d.N
        for j in getNeighbors(i,d)
            @inbounds Jd[i][j] = J0
        end
        @inbounds Bd[i] = B0
    end
end
IsingMCMC03a.jl
    function dE(S, k, N)
        local dEt = 0
        for (j,Jv) in Jd[k]
            @inbounds dEt += S[k]*Jv*S[j]*2
        end
        @inbounds dEt += S[k]*Bd[k]*2
        dEt
    end

3. ヒープアロケーションを極力避ける

@SVectorを使って、スタックを活用するように変えました。マクロなので、改行しないと上手く解釈できないようです。
Before:

HexaLatticeModel02a.jl
    function getNeighbors( k, d=d)
        (x, y) = getXY(k, d)
        [getK(x-2w,      y, d), getK(x+2w,      y, d),
         getK(x- w, y-1.5h, d), getK(x+ w, y-1.5h, d),
         getK(x- w, y+1.5h, d), getK(x+ w, y+1.5h, d)]
    end

    # polygon for each cell
    function genPoly(k, d=d)
        (x, y) = getXY(k, d)
        (x .+ [  -w, 0,   w,    w,  0,   -w,  -w],
         y .+ [ h/2, h, h/2, -h/2, -h, -h/2, h/2])
    end

After:

HexaLatticeModel03a.jl
    function getNeighbors( k, d=d)
        (x, y) = getXY(k, d)
        @SVector [
         getK(x-2w,      y, d), getK(x+2w,      y, d),
         getK(x- w, y-1.5h, d), getK(x+ w, y-1.5h, d),
         getK(x- w, y+1.5h, d), getK(x+ w, y+1.5h, d)]
    end

    # polygon for each cell
    function genPoly(k, d=d)
        (x, y) = getXY(k, d)
        (
         @SVector [   x-w,   x,   x+w,    x+w,  x,   x-w,  x-w]
            ,
         @SVector [ y+h/2, y+h, y+h/2, y-h/2, y-h, y-h/2, y+h/2]
        )
    end

また、vcatなどの動的なメモリ操作も遅いので事前に確保する方式に変えました。
Before:

IsingSimMain02a.jl
function drawSim(sim, f, d=d, N=N)
    snap = sim[f]

    p=Plots.plot([],[],legend=false, size=(500,500))

    p0xs = [];p0ys = []
    p1xs = [];p1ys = []
    for i in 1:N
        # bits
        (xs,ys) = genPoly(i, d)
        if (snap[i] == 1) # up-spin
            p1xs = vcat(p1xs, NaN, xs)
            p1ys = vcat(p1ys, NaN, ys)
        else # down-spin
            p0xs = vcat(p0xs, NaN, xs)
            p0ys = vcat(p0ys, NaN, ys)
        end
    end    
    poly0 = Shape(p0xs, p0ys) # creating polygon
    p = Plots.plot!(poly0, fillcolor = plot_color(:blue)  ,legend=false)
    poly1 = Shape(p1xs, p1ys)
    p = Plots.plot!(poly1, fillcolor = plot_color(:yellow),legend=false)
end

After:

IsingSimMain03a.jl
function drawSim(Sf, d=d, N=N)
    snap = Sf

    p=Plots.plot([],[],legend=false, size=(500,500))

    npoly = length(genPoly(1,d)[1])
    nbuff = (npoly+1)*N
    p0xs = fill(NaN, nbuff);p0ys = fill(NaN, nbuff)
    p1xs = fill(NaN, nbuff);p1ys = fill(NaN, nbuff)
    i1 = 1; i0 = 1
    @inbounds @simd for i in 1:N
        # bits
        (xs,ys) = genPoly(i, d)
        if (snap[i] == 1) # up-spin
            p1xs[i1:(i1+npoly-1)] .= xs
            p1ys[i1:(i1+npoly-1)] .= ys
            i1 += npoly+1
        else # down-spin
            p0xs[i0:(i0+npoly-1)] .= xs
            p0ys[i0:(i0+npoly-1)] .= ys
            i0 += npoly+1
        end
    end    
    poly0 = Shape(p0xs, p0ys) # creating polygon
    p = Plots.plot!(poly0, fillcolor = plot_color(:blue)  ,legend=false)
    poly1 = Shape(p1xs, p1ys)
    p = Plots.plot!(poly1, fillcolor = plot_color(:yellow),legend=false)
end

高速化の結果

それでは、高速化の結果を見てみましょう。
128x128メッシュのイジングモデルを10^5回MCMCで計算します。　
１回目の計算の後で最適化されている可能性があるので、２回目の結果で比較します。

Before:

julia> include("IsingSimMain02a.jl")
main (generic function with 1 method)

julia> main(2,1)
start MCMC... 12.268228 seconds (56.39 M allocations: 13.304 GiB, 16.37% gc time)
start drawEnergy... 24.559120 seconds (507.09 M allocations: 9.805 GiB, 4.76% gc time)
start drawSim...  7.025103 seconds (8.28 M allocations: 9.286 GiB, 16.56% gc time)
done.
julia> main(2,1)
start MCMC... 37.620995 seconds (55.74 M allocations: 13.272 GiB, 42.98% gc time)
start drawEnergy... 21.377200 seconds (502.59 M allocations: 9.589 GiB, 5.10% gc time)
start drawSim...  3.542199 seconds (2.24 M allocations: 8.697 GiB, 27.94% gc time)

After:

julia> include("IsingSimMain03a.jl")
main (generic function with 1 method)

julia> main(2,1)
start MCMC...  0.718519 seconds (6.43 M allocations: 170.372 MiB, 3.66% gc time)
start drawEnergy...  2.740000 seconds (3.99 M allocations: 195.505 MiB, 1.15% gc time)
start drawSim...  2.708490 seconds (7.20 M allocations: 350.252 MiB, 2.64% gc time)
done.
julia> main(2,1)
start MCMC...  0.511273 seconds (6.20 M allocations: 158.562 MiB, 3.24% gc time)
start drawEnergy...  0.000636 seconds (2.24 k allocations: 151.398 KiB)
start drawSim...  0.062203 seconds (946.81 k allocations: 36.983 MiB, 8.14% gc time)
read: No error
send: No error
done.

なんということでしょう！
MCMCでは約74倍, drawEnergyでは約33612倍，drawSimでは約57倍高速になりました!(#^.^#)
メモリアロケーションとGCタイムがかなり削減されています。

まとめ

　Juliaは書きやすいのでついつい素直に書いてしまいがちですが、特に、メモリアロケーションと無駄な計算を省くことを意識すると57～33612倍も速くなってしまうという実例でした。
　とは言え、Juliaの魅力の一つは可読性にあると思うので、やりすぎて訳の分からないコードにしないよう気を付けたいですね。
　

参考リンク

下記のページを参考にさせて頂きました。
1. Optimizing DiffEq Code
2. genkuroki/03_IsingMCMC.diff
3. Juliaのコードを更に高速化する方法
4. 高速な計算にむけて
5. Juliaで非正方格子のイジングモデルをシミュレーションしてみた

20
21
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
20
21