Edited at

[Julia] 行列を高速に更新する


訂正

[2019/03/22] 本文中の LinearIndices(CartesianIndices())p に関する記述は一旦保留とします。

おそらく例2について、意図した結果になりません。追って修正しますが、当該部分の記述とグラフは読み飛ばして頂くようにお願いします。


はじめに

Julia は表現力の豊かな言語で、同じ処理を行うにも、いくつもの書き方があります。

以下のような、行列を更新する計算について、実行時間を調べてみました。

$$W \leftarrow U + \alpha V$$

$U, V, W$ は、同じ寸法の密行列で、要素の型は浮動小数点型 Float64 とします。$U$ と $V$ は壊しません (書き換えません)。$\alpha$ は Float64型のスカラーです。

この計算を行うのに、以下のような方法が考えられます。(なお並列化は考えません)


  • 行列の計算

  • 行列の要素同士の計算

  • 二重ループ

  • 単純ループ


    • イテレータ eachindex() を使う

    • イテレータ CartesianIndices() を使う

    • イテレータ LinearIndices(CartesianIndices()) を使う



この他に、


  • 対応する BLAS や LAPACK の関数 (があれば。それら)を呼ぶ

という方法もありますが、割愛します ($U$ の値を保存しておいて、BLASの axpy 関数 $y \leftarrow ax + y$ に相当する関数を使うとよさそうです)。

以下に計測結果を紹介しますが、冒頭に実行時間の傾向をまとめておきます。


  1. 行列の計算は遅いです。

  2. 二重ループは、第1軸を最も内側で回しましょう。

  3. イテレータ eachindex(u)LinearIndices(CartersianIndices(u)) は簡潔に書けて、実行速度も落ちません。

ちなみに、1. と 2. および 3. の前半 eachindex(u)については、Julia Documentation や、過去の記事でも、解説されています。

ですが、3. の後半 LinearIndices(CartersianIndices(u)) は、日本語の紹介記事をまだ見たことがありません。


実行環境

Google Colab に興味があるので、Google Colab で実行した結果を、以下に示します。Qiita の記事 Google Colab で Julia を使ってみた の記事の通りに、Julia 1.0 の実行環境を準備しました。GPU は使っていません。

Julia の環境は、以下のようでした。

julia> versioninfo()

Julia Version 1.0.0
Commit 5d4eaca0c9 (2018-08-08 20:58 UTC)
Platform Info:
OS: Linux (x86_64-pc-linux-gnu)
CPU: Intel(R) Xeon(R) CPU @ 2.30GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-6.0.0 (ORCJIT, haswell)

この環境を用いて、以下のプログラムを実行した結果を、次のgist に置いておきます。

* https://gist.github.com/hsugawa8651/557a9237653962defeb8f55d9922aedb

参考に、手元の MacBookPro で同じプログラムを実行した結果を示します。

* https://gist.github.com/hsugawa8651/77f0cec335c2ed4b7d11c2db22179552

こちらの実行環境は、以下の通りです。Julia公式ページで配布の Julia 1.1 バイナリを用いました。各例題の末尾にグラフのみ示します。

julia> versioninfo()

Julia Version 1.1.0
Commit 80516ca202 (2019-01-21 21:24 UTC)
Platform Info:
OS: macOS (x86_64-apple-darwin14.5.0)
CPU: Intel(R) Core(TM) i5-7267U CPU @ 3.10GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-6.0.1 (ORCJIT, skylake)


例題1

では、具体的に、上の例題を Julia で書いてみます。


行列の計算

Julia は、行列の計算を、そのまま書き表すことができます。

w=u+v*alpha

これを、以下のような、関数として定義しました。


ex1_with_matrix.jl

function ex1_with_matrix(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
w=u+v*alpha
0
end

行列 u, v, w は、関数の呼び出し側で用意しておき、この関数を呼び出すことにします。nxny は行列の寸法 $nx \times ny$ です。

計算時間の計測のため、@assert はコメントアウトしました。

関数は 0 で終わることにしました。w=u+v*alpha で終わってしまうと、関数の戻り値を作るために、余計な計算やメモリ確保が起こり、それが実行時間に含まれる可能性があるからです。

行列 u, v, w を作成する処理も、関数として定義します。


make_uvw.jl

function make_uvw(nx,ny,rng = Random.MersenneTwister(1234))

u=rand(nx,ny)
v=rand(Float64, size(u))
w=similar(u);
return u,v,w
end

uv の要素は乱数で満たすことにします。w の中身は未定です。


BenchmarkTools - 実行時間を計測する

BenchmarkTools パッケージを用いて、プログラムの実行時間を計測します。

julia> using Pkg

julia> Pkg.add("BenchmarkTools")

関数の実行に必要なパッケージを読み込んでおきます。

using BenchmarkTools

using Random
using Statistics

早速 ex1_with_matrix の実行時間を計測してみます。

実行時間の計測には、@benchmark マクロを使います。@benchmark マクロの直後に計測したい式を書きます。 setup=() には、計測前に行う準備を書きます。

しばらく、行列の寸法を、nx1=200, ny1=100 とします。@benchmark マクロの中で、 変数 nx1ny1 の値を使うため、nx1ny1 の前に、$ を付けます。

julia> nx1=200; ny1=100;

julia> rng=Random.MersenneTwister(12345);

julia> # ex1_with_matrix
@benchmark ex1_with_matrix(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 312.66 KiB
allocs estimate: 4
--------------
minimum time: 37.794 μs (0.00% GC)
median time: 51.663 μs (0.00% GC)
mean time: 66.438 μs (18.64% GC)
maximum time: 20.115 ms (99.51% GC)
--------------
samples: 10000
evals/sample: 1

まず、BenchmarkTools.Trial: の次の 2行はメモリ確保に関する統計です。memory estimatealloc estimate: が 0 ではないことで、この計算の実行中に「メモリの確保」が発生したことに注目します。

その後、-------- とに囲まれているのが、実行時間の色々な統計値です。

実行時間を互いに比較するための目安として、中央値 median time が、よく用いられます。行列計算の実行時間(の中央値)は 51.663 μs でした。


行列の要素同士の計算

行列の演算の前に @. マクロをつけると、行列の各要素に対する計算を行うことを指示します。

@. w=u+v*alpha

正確には、@. マクロは、配列の要素毎に対する計算を指示する . 演算子を、付けられる演算子や関数の全てに付けることに相当します。この例題では問題ないですが、例えば、行列の積 a*bを含んだ計算に、うっかり @. を適用すると、対応する要素同士の積 a.*b の意味になってしまうので、要注意です (経験者は語る)。

関数として定義します。


ex1_with_dot_matrix.jl

function ex1_with_dot_matrix(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@. w=u+v*alpha
0
end

実行時間を計測します。

julia> # with_dot_matrix

@benchmark ex1_with_dot_matrix(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 10.588 μs (0.00% GC)
median time: 12.857 μs (0.00% GC)
mean time: 13.635 μs (0.00% GC)
maximum time: 503.366 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 4

今度は、メモリ確保は発生しません (allocs estimate: 0) 。

実行時間(の中央値)は、単純な行列計算の 51.663 μs から 12.857 μs へと大幅に (約 1/4 に) 短縮しました。


二重ループによる計算

この例題は、配列の各要素に対する計算としても書けます。

$$\begin{equation*}

w_{i,j} = u_{i,j} + \alpha v_{i,j}

\end{equation*}$$

ここで、$i, j$ は、行列 (2次元配列) の添字であり、全ての要素を巡ります。

この計算は、添字毎に独立しているので、どの順番で計算しても構いません。

添字 ij を、二重ループで (入れ子の for 文で)回せば、確実に全ての要素を巡ることができます。

この際、内側のループ変数を i にするか、それとも j にするかの、二つの選択があります。


  • 第2軸の添字 j を内側で回す二重ループ

配列の要素を a[i,j] と書くとき、最初の添字 i を第1軸 (first axis)の添字、2番目の添字 j を第2軸 (second axis)の添字、といいます

外側のループで第1軸の添字 i を回し、内側のループで第2軸の添字 j を回す書き方です。

for i in 1:nx

for j in 1:ny
w[i,j]=u[i,j]+v[i,j]*alpha
end
end


  • 第1軸の添字 i を内側で回す二重ループ

今度は、外側のループで第2軸の添字 j を回し、内側のループで第1軸の添字 i を回す書き方です。

for j in 1:ny

for i in 1:nx
w[i,j]=u[i,j]+v[i,j]*alpha
end
end

それぞれの場合を、関数として定義します。


  • 第2軸の添字を内側で回す二重ループ


ex1_double_loop_ij.jl

function ex1_double_loop_ij(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds for i in 1:nx
@simd for j in 1:ny
w[i,j]=u[i,j]+v[i,j]*alpha
end
end
0
end


  • 第1軸の添字を内側で回す二重ループ


ex1_double_loop_ji.jl

function ex1_double_loop_ji(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds for j in 1:ny
@simd for i in 1:nx
w[i,j]=u[i,j]+v[i,j]*alpha
end
end
0
end

上に定義した関数では、for 文を高速化するためのマクロを付けました。

- @inbounds を付けると、添字の範囲検査が省略されます。添字が配列の範囲を超えても例外を発生しません。

- @simd を 最も内側の for 文に付けると、CPUが提供する並列計算命令を(可能なら)適用します。

実行時間を計測します。

julia> # 第2軸を内側に

@benchmark ex1_double_loop_ij(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 51.452 μs (0.00% GC)
median time: 63.237 μs (0.00% GC)
mean time: 66.718 μs (0.00% GC)
maximum time: 2.000 ms (0.00% GC)
--------------
samples: 10000
evals/sample: 1

julia> # 第1軸を内側に
@benchmark ex1_double_loop_ji(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 10.536 μs (0.00% GC)
median time: 12.498 μs (0.00% GC)
mean time: 13.228 μs (0.00% GC)
maximum time: 474.144 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 4

どちらの場合もメモリ確保は起こりませんでした。

実行時間(の中央値)を比較すると、以下のように傾向をまとめられます。

- 第2軸を内側で回した二重ループの実行時間 (63.237 μs) は、行列計算の実行時間 (51.663 μs) よりも長かったです。

- 第1軸を内側で回した二重ループの実行時間 (12.498 μs) は、第2軸を内側で回した二重ループよりも短く (およそ 1/5倍)、行列要素に対する計算の実行時間 (12.857 μs) と同程度でした。

さて、第1軸を内側で回すほうが、第2軸を内側で回すより、速く実行されるのは、冒頭に挙げた Performance tips に解説されている通りです (Julia Documentation, Perfomance Tips > Access arrays in memory order, along columns )

配列の要素は、第1軸が先に回る順番で格納されますから、第1軸の添字を先に増やすと、すぐ隣の配列要素を順で読み書きします。ですので、データの先読みの機能 (キャッシュ)が有効に働く可能性が高くなります。

逆に、第2軸を内側で回すと、飛び飛びに (第1軸の寸法分を増やして)配列要素を読み書きすることになり、キャッシュが失敗する可能性が増えてしまいます。


eachindex で単純ループ

Julia の配列の添字は、連続する整数でも指定できます。「線形の添字 (linear index) 」といいます。

1から配列の長さ (配列全体の数) length(u) まで順番に線形な添字を増やすのは、第1軸を内側で回すのと同じです。

for ij in 1:length(u)

w[ij]=u[ij]+v[ij]*alpha
end

範囲 1:length(u) で書く代わりに、イテレータ eachindex(u) を使ってスッキリ書けます。

for ij in eachindex(u)

w[ij]=u[ij]+v[ij]*alpha
end

スッキリ書けただけでなく、実行速度が速くなります (と期待されます)。

軸をまたいで連続した配列の要素が読み書きされるので、キャッシュがさらに有効になりやすくなります。また、繰り返しブロックの最後で実行される繰返しの終了条件の判定が、単純ループにすると少なくなるからです。

@inbounds@simd のマクロを付けて、関数を定義します。


ex1_eachindex.jl

function ex1_eachindex(u,v,w,  nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds @simd for ij in eachindex(u)
w[ij]=u[ij]+v[ij]*alpha
end
0
end

実行時間を計測します。

julia> # eachindex

@benchmark ex1_eachindex(u,v,w, $nx1, $ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 8.945 μs (0.00% GC)
median time: 10.329 μs (0.00% GC)
mean time: 10.829 μs (0.00% GC)
maximum time: 48.689 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 5

eachindex を使った場合の実行時間 (10.329 μs)は、行列要素に対する計算 (12.857 μs) や第1軸を内側で回す二重ループ (12.498 μs) と同程度で、少し短かったです。


CartesianIndices で単純ループ

配列の要素は、CartesianIndex(i,j) を用いても参照できます。CartesianIndex(i,j) は、各軸の添字 i, j を指定します。

すなわち、a[CartesianIndex(i,j)]a[i,j] と同じ意味です。

関数 CartesianIndices(a) は、配列 a 全体を効率よく回るような CartesianIndex のイテレータを作ります。

eachindex と同じように、第1軸の添字が先に動きます。

ですから、例題は、以下のように書けます。

for ij in CartesianIndices(u)

w[ij]=u[ij]+v[ij]*alpha
end

計測対象の関数を定義します。


ex1_CartesianIndices.jl

function ex1_CartesianIndices(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds @simd for ij in CartesianIndices(u)
w[ij]=u[ij]+v[ij]*alpha
end
0
end

実行時間を計測します。

julia> # CartesianIndices

@benchmark ex1_CartesianIndices(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 9.372 μs (0.00% GC)
median time: 11.245 μs (0.00% GC)
mean time: 11.764 μs (0.00% GC)
maximum time: 41.984 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 5

CartesianIndicesを使った場合の実行時間 (11.245 μs)は、eachindex による繰返しよりも少し長く、行列要素による計算、第1軸を内側で回す二重ループよりも少し短かったです。


LinearIndices(CartesianIndices()) で単純ループ

関数 LinearIndices(I) は、CartesianIndex のイテレータ I を「線形な添字」のイテレータに変換します。

つまり、以下の書き方は、eachindex(a) を用いた繰返しと同じです。

for ij in LinearIndices(CartesianIndices(u))

w[ij]=u[ij]+v[ij]*alpha
end

関数を定義します。


ex1_LinearIndices_CartesianIndices.jl

function ex1_LinearIndices_CartesianIndices(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds @simd for ij in LinearIndices(CartesianIndices(u))
w[ij]=u[ij]+v[ij]*alpha
end
0
end

実行時間を計測します。

julia> # LinearIndices_CartesianIndices

@benchmark ex1_LinearIndices_CartesianIndices(u,v,w, $nx1, $ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 9.105 μs (0.00% GC)
median time: 10.437 μs (0.00% GC)
mean time: 11.001 μs (0.00% GC)
maximum time: 452.964 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 5

イテレータ LinearIndices(CartesianIndices(...)) を使った場合の実行時間の

中央値は (10.437 μs)は、eachindex による計算と同程度です。


全てを比較してみる。

BenchmarkTools パッケージでは、計測対象を BenchmarkGroup としてまとめておいて、実行時間を一気に計測できます。

各計測結果には、各計測対象を複数回実行した結果の実行時間の生データが含まれています。これを、boxplot (箱プロット)してみましょう。

行列の寸法が 200x100 の場合、100x200 の場合の二通りについて、グラフで示しました。行列の全体の寸法は変わりませんが、100x200 の方が第1軸の寸法が小さいので、第1軸を内側に回す二重ループが不利になると考えられます。

横軸は、実行時間 (単位 ns) を、対数で表示しました。左ほど実行結果が短いです。実行時間の中央値は、橙色の線で表しました。中央値が短い順に、関数の名前を整列しました。

比較として、MacBookProで実行した結果のグラフを示します。この記事の下書きを書くために Jupyter notebook、Google Chrome、Atom エディタを起動していて、それなりに高負荷な環境だと思います。

実行時間の短い方法の順位は、行列の寸法や実行環境で、少し入れ替わっていますが、傾向は変わりません。これまでの説明を復唱しますが、以下のようにまとめられます。


  • 単なる「行列の計算」は二重ループより遅く、メモリ確保が発生します。

  • 二重ループでは、第1軸を内側で回すほうが速いです

  • 単一ループ、すなわち、eachindex(...)LinearIndices(CartesianIndices(...)) は、第1軸を内側で回すのと同程度に速いです。

  • 行列の要素毎の計算も、第1軸を内側で回すのと同程度に速いです。小さいメモリ確保が発生します。

ですので、行列の全ての要素を巡回する場合は、

* 第1軸を内側にした二重ループ

* eachindex(a) を使ったイテレータ

のどちらかで書くのがオススメといえます。


例題2

もう一つの例題は、見かけは、前の例題と同じです。

$$\begin{align*}

w_{i,j} & \leftarrow u_{i+1,j} + \alpha v_{i,j-1}

\end{align*}$$

しかし、添字 ij は行列全体を巡りません。


  • 第1軸の添字 $i$ は、2 から、size(a,1)-1 すなわち、第1軸の寸法から 1 減じた値まで繰返します。

  • 第2軸の添字 $j$ は、2 から、size(a,2)-1 すなわち、第2軸の寸法から 1 減じた値まで繰返します。

寸法 4x5 の行列について模式的に書くと、以下のようになります。黒丸 $\bullet$ が読み書きされる要素です。

$$\begin{bmatrix}

\circ & \circ & \circ & \circ & \circ \\

\circ & \bullet & \bullet & \bullet & \circ \\

\circ & \bullet & \bullet & \bullet & \circ \\ \circ & \circ & \circ & \circ & \circ \end{bmatrix} =

%

\begin{bmatrix}

\circ & \circ & \circ & \circ & \circ \\

\circ & \circ & \circ & \circ & \circ \\

\circ & \bullet & \bullet & \bullet & \circ \\

\circ & \bullet & \bullet & \bullet & \circ \end{bmatrix}

+ \alpha

\begin{bmatrix}

\circ & \circ & \circ & \circ & \circ \\

\bullet & \bullet & \bullet & \circ & \circ \\

\bullet & \bullet & \bullet & \circ & \circ \\

\circ & \circ & \circ & \circ & \circ \end{bmatrix}$$

この例題は、離散化した微分方程式を行列方程式として表したときに、時間発展させる場合や、座標変数の境界条件を考慮すべき場合を、模式化したものです。

(この記事は、 @Bimaterial さんの 2次元波動方程式の解の導出と性質 記事への、私のコメント続きも意図しています )。

この計算を行うのには、以下のような方法が考えられます。


  • 部分行列の計算

  • 部分行列の要素同士の計算

  • 二重ループ

  • 単純ループ


    • イテレータ CartesianIndices() を使う

    • イテレータ LinearIndices(CartesianIndices())




部分行列の計算

この例題は、行列全体の計算ではなく、添字もズレていますから、単純な行列として書くことはできません。

代わりに、部分配列 (部分行列, サブ行列)を使ってみます。

これは、配列の一部を、部分配列として扱うものです。

例えば、

w_sub=view(w,2:nx-1, 2:ny-1)

とすると、配列 w の部分配列 w_sub が作れます。ここで、部分配列 w_sub は、配列 w において、第1軸の添字が 2 から nx-1 まで、第2軸の添字が 2 から ny-1 までの範囲です。

部分配列は、通常の配列と同じように計算できます。元の配列の複製は作りません。

この例題の核心の部分は、このように書けます。

w_sub=view(w, 2:nx-1, 2:ny-1)

u_sub=view(u, 3:nx, 2:ny-1)
v_sub=view(v, 2:nx-1, 1:ny-2)
w_sub=u_sub+v_sub*alpha

部分配列 u, v, w は、元の配列 u, v, w の異なる部分を指していますが、対応する添字同士で正しく計算されます。

実行時間の計測対象を、関数として定義します。


ex2_view.jl

function ex2_view(u,v,w, nx=size(u,1),ny=size(u,2), alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
w_sub=view(w,2:nx-1, 2:ny-1)
v_sub=view(v, 2:nx-1, 1:ny-2)
u_sub=view(u, 3:nx, 2:ny-1)
w_sub=u_sub+v_sub*alpha
0
end

例題1と同様に、しばらく、行列の寸法を、nx1=200, ny1=100 とします。

実行時間を計測します。

julia> nx1=200; ny1=100;

julia> # CartesianIndices
@benchmark ex2_view(u,v,w, $nx1, $ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 303.53 KiB
allocs estimate: 6
--------------
minimum time: 42.012 μs (0.00% GC)
median time: 52.443 μs (0.00% GC)
mean time: 103.169 μs (46.16% GC)
maximum time: 112.921 ms (99.88% GC)
--------------
samples: 10000
evals/sample: 1

例題1の行列計算と同様に、メモリ確保が発生してしました。

実行時間(の中央値)は 52.443 μs でした。例1の行列全体の計算 (51.663 μs) と同程度でした。


部分行列の要素同士の計算

部分配列の計算に @. を前につけると、部分配列の各要素の計算となります。

@. w_sub=u_sub+v_sub*alpha

関数を定義します。


ex2_dot_view.jl

function ex2_dot_view(u,v,w, nx=size(u,1),ny=size(u,2), alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
w_sub=view(w,2:nx-1, 2:ny-1)
v_sub=view(v, 2:nx-1, 1:ny-2)
u_sub=view(u, 3:nx, 2:ny-1)
@. w_sub=u_sub+v_sub*alpha
0
end

実行時間を計測します。

julia> # dot_view

    @benchmark ex2_dot_view(u,v,w, $nx1, $ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 192 bytes
allocs estimate: 3
--------------
minimum time: 12.832 μs (0.00% GC)
median time: 15.110 μs (0.00% GC)
mean time: 16.108 μs (0.00% GC)
maximum time: 38.263 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 3

例1の行列の全要素に対する計算では発生していなかった「メモリ確保」が、ほんの少し発生しました。

実行時間(の中央値)を比較すると、52.443 μs から 15.110 μs へと大幅に (約 1/3 に) 短縮しました。


二重ループによる計算

この例は、添字を使った二重ループで自然に書くことができます。

第2軸を内側で回す二重ループを関数として用意します。


ex2_double_loop_ij.jl

function ex2_double_loop_ij(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds for i in 2:nx-1
@simd for j in 2:ny-1
w[i,j]=u[i+1,j]+v[i,j-1]*alpha
end
end
0
end

今度は、第1軸を内側で回す二重ループを関数として定義します。


ex2_double_loop_ji.jl

function ex2_double_loop_ji(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds for j in 2:ny-1
@simd for i in 2:nx-1
w[i,j]=u[i+1,j]+v[i,j-1]*alpha
end
end
0
end

実行時間を計測します。

julia> # 第2軸を内側に

@benchmark ex2_double_loop_ij(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 54.765 μs (0.00% GC)
median time: 75.025 μs (0.00% GC)
mean time: 78.491 μs (0.00% GC)
maximum time: 949.207 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 1

julia> # 第1軸を内側に
@benchmark ex2_double_loop_ji(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 10.918 μs (0.00% GC)
median time: 13.104 μs (0.00% GC)
mean time: 13.779 μs (0.00% GC)
maximum time: 32.241 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 4

実行時間(の中央値)を比較すると、以下のようになります。

- 第2軸を内側で回した二重ループの実行時間 (75.025 μs) は、部分配列(部分行列)の実行時間 (52.443 μs) よりも長かったです。

- 第1軸を内側で回した二重ループの実行時間 (13.104 μs) は、第2軸を内側で回した二重ループよりも短く (およそ 1/5倍)、部分配列行列の要素の実行時間 (15.110 μs) と同程度で、少し短かったです。


CartesianIndices で単純ループ

例2では、添字は行列全体を巡るわけではないので eachindex(a) は使えません。

イテレータ CartesianIndices は、配列の一部を切り出すこともできます。

すなわち、 CartesianIndices((2:nx-1,2:ny-1)) は、


  • 第1軸を 1から nx-1 まで、

  • 第2軸を 1から ny-1 まで

巡回する CartesianIndex のイテレータです。第1軸の方が内側で回ります。

CartesianIndex(i,j) を関数 Tuple に適用すると、各軸の添字が得られます。

すなわち、計算の主な部分は、以下のように書けます。

for ij in CartesianIndices((2:nx-1,2:ny-1))

i,j=Tuple(ij)
w[i,j]=u[i+1,j]+v[i,j-1]*alpha
end

関数として定義します。


ex2_CartesianIndices.jl

function ex2_CartesianIndices(u,v,w, nx=size(u,1),ny=size(u,2),alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
@inbounds @simd for ij in CartesianIndices((2:nx-1,2:ny-1))
i,j=Tuple(ij)
w[i,j]=u[i+1,j]+v[i,j-1]*alpha
end
0
end

実行時間を計測します。

julia> # CartesianIndices

@benchmark ex2_CartesianIndices(u,v,w, $nx1,$ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 11.389 μs (0.00% GC)
median time: 13.550 μs (0.00% GC)
mean time: 14.200 μs (0.00% GC)
maximum time: 33.031 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 4

CartesianIndicesを使った場合の実行時間の中央値 (13.550 μs)は、

部分配列の要素に対する計算や第1軸を内側で回す二重ループと同程度で、少し短かったです。


LinearIndices(CartesianIndices()) で単純ループ

CartesianIndices()LinearIndices に適用すると、「線形な添字」が得られます。

u, v, w に対する線形な添字はズレていますから、各々の行列の要素を指定するイテレータを用意してしましょう。これは、同じ添字で取り出せます。

すなわち、核心の部分は、以下のように書けます。

ここで、

- w_Iw に対する添字のイテレータ

- u_Ju に対する添字のイテレータ

- v_Kv に対する添字のイテレータ

です。

w_I=LinearIndices(CartesianIndices( (2:nx-1, 2:ny-1)))

u_J=LinearIndices(CartesianIndices( (3:nx, 2:ny-1)))
v_K=LinearIndices(CartesianIndices( (2:nx-1, 1:ny-2)))

for m in 1:length(w_I)
i=w_I[m]
j=u_J[m]
k=v_K[m]
w[i]=u[j]+v[k]*alpha
end

関数として定義します。


ex2_LinearIndices_CartesianIndices.jl

function ex2_LinearIndices_CartesianIndices(u,v,w, nx=size(u,1),ny=size(u,2), alpha=0.3)

# @assert size(u)==(nx,ny)
# @assert size(v)==(nx,ny)
# @assert size(w)==(nx,ny)
w_I=LinearIndices(CartesianIndices( (2:nx-1, 2:ny-1)))
u_J=LinearIndices(CartesianIndices( (3:nx, 2:ny-1)))
v_K=LinearIndices(CartesianIndices( (2:nx-1, 1:ny-2)))
@inbounds @simd for m in 1:length(w_I)
i=w_I[m]
j=u_J[m]
k=v_K[m]
w[i]=u[j]+v[k]*alpha
end
0
end

実行時間を計測します。

julia> # LinearIndices_CartesianIndices

@benchmark ex2_LinearIndices_CartesianIndices(u,v,w, $nx1, $ny1) setup=( (u,v,w)=make_uvw($nx1,$ny1, $rng))
BenchmarkTools.Trial:
memory estimate: 0 bytes
allocs estimate: 0
--------------
minimum time: 8.906 μs (0.00% GC)
median time: 10.400 μs (0.00% GC)
mean time: 10.895 μs (0.00% GC)
maximum time: 47.163 μs (0.00% GC)
--------------
samples: 10000
evals/sample: 5

LinearIndices(CartesianIndices(...))を使った場合の実行時間 (10.400 μs)は、CartesianIndices(...)を使った場合の実行時間 (13.550 μs) と同程度で、少し短かったです。


全てを比較してみる

例題1 と同じように、例題2の全ての計測対象を比較してみましょう。

比較として、MacBookProで実行した結果のグラフを示します。

実行時間の短い方法の順位は、行列の寸法や実行環境で、少し入れ替わっています。が、傾向は変わっていません。以下のようにまとめられます。


  • 単なる「部分行列の計算」は、二重ループより遅く、メモリ確保が発生します。

  • 二重ループでは、第1軸を内側で回すほうが速いです。

  • 部分行列の要素毎の計算も、第1軸を内側で回すのと同程度に速いです。小さいメモリ確保が発生します。

  • 単一ループ、すなわち、LinearIndices(CartesianIndices(..)) は、第1軸を内側で回すのと同程度に速いです。

ですので、行列の全ての要素を巡回する場合は、まず


  • 第1軸を内側にした二重ループ

で書いておいて、計測しながら、LinearIndices(CartesianIndices()) に書き換えるのがオススメといえます。