りんごがでている

何か役に立つことを書きます

Juliaではfor文を使え。ベクトル化してはいけない。

タイトルはいささか誇張です。最後のセクションが一番言いたいことなので、そこまで読んでいただけたら嬉しいです。

きっかけは昨日こんなツイートをしたら割りとRTとFavがあったので、分かっている人には当たり前なのですが、解説を書いたらいいかなと思った次第です。

Pythonのfor文

PythonやRを使っている人で、ある程度重い計算をする人達には半ば常識になっていることとして、いわゆる「for文を使ってはいけない。ベクトル化*1しろ。」という助言があります。 これは、PythonやRのようなインタープリター方式の処理系をもつ言語では、極めてfor文が遅いため、C言語Fortranで実装されたベクトル化計算を使うほうが速いという意味です。

まずはこれを簡単な例で検証してみましょう。 以下のように、ベクトルaとbの内積を計算するdot(a, b)という関数を考えてみます。

def dot(a, b):
    s = 0
    for i in xrange(len(a)):
        s += a[i] * b[i]
    return s

これと同じ動作をする関数がNumPyには用意されており、こちらはPythonではなくより下のCPUに近いレベルでfor文を実行しています。 2つの大きさのベクトル(10要素と100,000要素)に対して、この2つの実装の速度の違いを検証してみました。

~/tmp $ ipython
Python 2.7.9 (default, Feb  4 2015, 03:28:46)
Type "copyright", "credits" or "license" for more information.

IPython 3.0.0 -- An enhanced Interactive Python.
?         -> Introduction and overview of IPython's features.
%quickref -> Quick reference.
help      -> Python's own help system.
object?   -> Details about 'object', use 'object??' for extra details.

In [1]: def dot(a, b):
   ...:     s = 0
   ...:     for i in xrange(len(a)):
   ...:         s += a[i] * b[i]
   ...:     return s
   ...:

In [2]: import numpy as np

In [3]: a = np.ones(10)

In [4]: b = np.ones(10)

In [5]: dot(a, b)
Out[5]: 10.0

In [6]: np.dot(a, b)
Out[6]: 10.0

In [7]: timeit dot(a, b)
The slowest run took 7.89 times longer than the fastest. This could mean that an intermediate result is being cached
100000 loops, best of 3: 4.32 µs per loop

In [8]: timeit np.dot(a, b)
The slowest run took 13.83 times longer than the fastest. This could mean that an intermediate result is being cached
1000000 loops, best of 3: 793 ns per loop

In [9]: a = np.ones(100000)

In [10]: b = np.ones(100000)

In [11]: timeit dot(a, b)
10 loops, best of 3: 35.3 ms per loop

In [12]: timeit np.dot(a, b)
10000 loops, best of 3: 51.9 µs per loop

Python実装とNumPy実装の速度差は歴然としています。 10要素の短いベクトルでも1桁の違いがあり、100,000要素では実に3桁もNumPyの方が速いことになります。

大きさ Python NumPy
10 4.32μs 793ns
100,000 35.3ms 51.9μs

この違いは、たとえ大きな配列の確保などのオーバーヘッドがあったとしてもそれを相殺することができるほど大きく、 したがって、Pythonでは「for文を使ってはいけない。ベクトル化しろ。」というスローガンは正しいと言えます。

Juliaのfor文

では本題のJuliaの場合はどうでしょうか。Juliaでは、for文はC言語などと同じ速度で実行されるため、for文でループしても遅くはなりません。

ここでは、純粋なJuliaのコードと、JuliaからC言語を呼び出すコードを比較してみましょう。

C言語での実装は以下の通りです。

#include <stddef.h>

double mydot(double* a, double* b, size_t len)
{
    double s = 0;
    size_t i;
    for (i = 0; i < len; i++)
        s += a[i] * b[i];
    return s;
}

これを動的ライブラリとしてコンパイルしておきます。

~/tmp $ clang -dynamiclib -O3 -o libdot.dylib dot.c

Julia側の実装と、先のCのコードを呼び出す関数は以下のようになります。

function mydot(a, b)
    s = 0.0
    for i in 1:length(a)
        s += a[i] * b[i]
    end
    s
end

function cmydot(a, b)
    ccall((:mydot, "libdot"), Float64, (Ptr{Float64}, Ptr{Float64}, Csize_t), a, b, length(a))
end

これらの時間を計測してみましょう。

push!(LOAD_PATH, ".")

let
    a = ones(10)
    b = ones(10)
    mydot(a, b)
    cmydot(a, b)
    println("mydot 10")
    @time for _ in 1:1000; mydot(a, b); end
    println("cmydot 10")
    @time for _ in 1:1000; cmydot(a, b); end
    a = ones(100_000)
    b = ones(100_000)
    println("mydot 100,000")
    @time for _ in 1:1000; mydot(a, b); end
    println("cmydot 100,000")
    @time for _ in 1:1000; cmydot(a, b); end
end

これを実行してみると、幾分C言語のほうが速いようですが、Pythonのときに見られたほどの大きな差は得られません。 しかもこの差は、Juliaの@inboundsマクロを使えば消えてしまいます。

~/tmp $ julia for.jl
mydot 10
elapsed time: 2.3714e-5 seconds (0 bytes allocated)
cmydot 10
elapsed time: 9.407e-6 seconds (0 bytes allocated)
mydot 100,000
elapsed time: 0.136645766 seconds (0 bytes allocated)
cmydot 100,000
elapsed time: 0.112500924 seconds (0 bytes allocated)

これまでのことから、Juliaのfor文は十分高速で、「for文を使っても良い」ということが分かります。

Juliaでのベクトル化は遅い

ブログタイトルのスローガンは「Juliaではfor文を使え。ベクトル化してはいけない。」でした。 「for文を使っても良い」のは大丈夫として、では何故ベクトル化してはいけないのでしょうか。

この根本的な原因は、今のJulia(v0.3.7)には破壊的な演算子が無いため、いつでも新しい配列を確保してしまう点にあります。

この効果を見るため、拙作のANMS.jlパッケージのコードを見てみましょう。

# reflect
@inbounds for j in 1:n
    xr[j] = c[j] + α * (c[j] - xh[j])
end

ANMS.jl#L101-103

ここで、xr, c, xhは同じ長さのベクトルで、αスカラーの係数です。 計算したいことは、二点cxhの座標から点xrの座標を計算しているだけなのですが、これをNumPyの感覚でxr = c + α * (c - xh)と書いてしまうと、各演算子の適用の度(+, *, -で都合3回)に新たな配列が確保されることになります。これでは、配列の確保とGCが頻発し、パフォーマンスが要求されるコードでは大きな損失になります。 ですので、ここではベクトル化された+*のような演算子は使えず、for文を書かざるを得なくなります。 もちろんNumPyでもこうした一次配列の確保は生じるのですが、Pythonではfor文があまりに遅いため、for文に書き換えると逆に遅くなってしまいます。

以上をもって、「Juliaではfor文を使え。ベクトル化してはいけない。」ということが実証されたことになります。

「早すぎる最適化は諸悪の根源である」

ただ、上に挙げたJuliaのベクトル化は、問題にならないことの方が多いです。 というのも、ほとんどのコードはプログラム全体から見るとボトルネックにはならず、したがってベクトル化計算をfor文に展開しても、得られる効果は殆ど無いに等しいものになります。 ですので、可読性や表記の簡便性から、Juliaにおいてもまずはベクトル化したコードを書き、ベンチマークによってそこがボトルネックになっていることが判明した時に初めてfor文を使って書き換えることが推奨されるのです。 Knuth先生の箴言を持って先のスローガンを修正するとすれば、「ボトルネックになっていることが確かめられたのなら、そこではfor文を使え。ベクトル化してはいけない。」となります。

参考

*1:ここでいうベクトル化とは、CPUレベルで言うSIMDのような並列化の話ではなく、配列やベクトルをひとまとまりの値とみなして計算することを指します。