NumericExtensions.jl / NumericFuns.jl

JuliaLang Advent Calendar の10日目の記事です。

Julia 言語を使う動機として、開発速度と実行速度を共に向上させたい、 というものがやはり大きくなるのではないでしょうか。 NumericFuns.jlNumericExtensions.jl は、 Julia の実行速度、特に配列を用いた数値計算の速度を更に向上させる事のできるパッケージです。

予備知識 -- 時間測定

実行速度を向上させる、というためには、実際に実行速度(実行時間)を計測する必要があります。 Julia では@time を使うことで簡単に実行時間を計測できます。

# @time は与えた式を評価して、それと同時に標準出力に計算時間と必要メモリを表示
julia> @time sleep(0.5)
elapsed time: 0.501986776 seconds (472 bytes allocated)

なお、Julia には他にもコードカバレッジ測定ツールやプロファイラなどが 標準で 用意されています。詳しくはJuliaTokyo #2 での 発表スライド をご参照ください*1

なお、Julia では関数単位のJIT コンパイルがあるために、それぞれの関数を予め(小さなデータでもいいので)実行しておくことでより高速化を図れます。 逆に言えばパフォーマンス測定のためには予め空回ししておく必要があります。 以下では省スペース化のため省略していますが、時間測定前に一度は同じ計算を行っています。

何のパッケージ?

主に配列演算を高速化するパッケージです。 例えば2つのベクトルの差の自乗を取りたいときは、一番単純には

julia> N = 1000000
1000000

julia> xs, ys = rand(N), rand(N);

julia> @time sum(abs2(xs-ys))
elapsed time: 0.010056548 seconds (16000312 bytes allocated)
166877.0598863598

となります。これはこれで何をやっているのかわかりやすいのですが、 xs-ysabs2 とで2回走査してそのたびに新しいベクトルを作り、 更に最後に sum でもう1回操作するという流れになっています。 NumericExtensions.jl を使うと

julia> using NumericExtensions

# mapreduce
# foldl 部分で走査が1回減るが xs-ys がまだ残っている
julia> @time foldl(Add(), 0.0, Abs2Fun(), xs - ys)
elapsed time: 0.003482581 seconds (8000264 bytes allocated)
166877.05988635868

# 2つの配列の差に関して走る
julia> @time foldl_fdiff(Add(), 0.0, Abs2Fun(), xs, ys)
elapsed time: 0.001559327 seconds (176 bytes allocated)
166877.05988635868

# 縮約操作を和に特殊化
julia> @time sumfdiff(Abs2Fun(), xs, ys)
elapsed time: 0.001133061 seconds (176 bytes allocated)
166877.05988635944

# まさにやりたかったことを行う、そのものずばりな関数
julia> @time sumsqdiff(xs, ys)
elapsed time: 0.001081251 seconds (176 bytes allocated)
166877.05988635944

と言った具合に、様々なレベルで特殊化を行った関数が用意されているので、 より高速に配列操作を行うことができます。

関数オブジェクトによる高速化

数直線上原点からスタートして、単位時間あたりに ある決められた確率 p で+1 方向に進み、 または確率 1-p で-1 方向に進む粒子の運動(ランダムウォーク)の数値シミュレーションを行います。 同時にたくさんの粒子についてシミュレーションを行うことで、 粒子の存在位置の確率分布を(経験的に)求めることができます。

まず、次の関数は確率 px+1 を、1-px-1 を返す関数です。

julia> update(x::Integer, p::Real) = ifelse(rand()<p, x+1, x-1)
update (generic function with 1 method)

これを使って M 個の粒子を N ステップだけランダムウォークさせる シミュレーションは次のようにかけます *2 *3

julia> N,M = 100_000, 1000;

julia> xs = zeros(Int, M);

julia> for i in 1:N; xs = map(x->update(x,0.5), xs); end

julia> xs
1000-element Array{Int64,1}:
  -14
 -778
 -218
  374
 -570
  224
 -356
 -296
 -598
 -250
   96
    ⋮
 -316
 -208
  114
  390
   10
   -2
  228
 -144
   26
  154
   32

# 平均
julia> mean(xs)
-0.274

# 分散(不偏分散)
# なお、キーワード引数 corrected に false を渡すことで
# 不偏分散ではなく標本分散を計算できる
julia> var(xs)
93933.3422662663

問題がかなり簡単だという事情もありますが、 単純でわかりやすい部品の組立てによって簡単に開発することができます*4。 さて、こんな感じにもっと大きなプログラムを書いて、 解きたい問題のサイズも大きくなっていくと、実行時間も伸びてくるので、 高速化をしたくなります。 その時には、どこがホットスポットなのか、つまりどこが一番時間がかかっているかを プロファイラなどを使って測定するのですが、 とりあえずこの更新部分に最も時間がかかっていたとします。 そこでこの部分の高速化を行ってみましょう。

最初期のバージョンでは

julia> @time for i in 1:N; xs = map(x->update(x, 0.5), xs); end
elapsed time: 5.398709802 seconds (936531816 bytes allocated, 10.51% gc time)

という実行時間になりました。 実行時間のおよそ1/10 がガベージコレクタにかかる時間なのですが、 map は非破壊な関数、つまり引数として与えた配列に触らずに新しい配列をつくる関数なので、 毎回メモリ確保をすることが問題なのではないかと考えるわけです。 そこで代わりに破壊的な関数、つまり引数を直接書き換える関数であるmap! を使ってみます。

julia> @time for i in 1:N; map!(x->update(x, 0.5), xs); end
elapsed time: 7.686995863 seconds (873712320 bytes allocated, 7.36% gc time)

使用メモリがちょっと減って、ガベコレ時間の割合も減っているのですが、 実行時間が5割増しぐらいになってしまっています。 使用メモリも大して減っていないし、ぶっちゃけバグなんじゃないかとも思いますが、 とにかく別の手段を考える必要があります。

そもそも最初の例が遅い*5のは、 map 関数に渡す関数が無名関数リテラルで表現されており、 map の呼び出しごとに新しい無名関数を作り出しているため、 JIT コンパイラが機能しないからです *6。 しかしそもそもなぜ無名関数を作っているかというと、 もともと2引数関数だったものを部分適用して1引数関数にしたいからです。 今回の例では、部分適用する引数である確率p は (大体のケースで)シミュレーション全体を通して固定なので、 予め部分適用した関数を定義しておけますが、 そういうことができない場合もあります。 こういう時には、あらかじめ部分適用した関数をラッピングした型(関数オブジェクト)を作ってしまい、 束縛する値を型のフィールドとして持っておくことで、 部分適用関数にJIT コンパイルをかけることができます。 関数オブジェクトやそれに関する高階関数を定義しているのがNumericExtensions.jlNumericFuns.jl です。 NumericFuns にはN 個の引数を取る関数オブジェクトを表す抽象型 Functor{N} があります*7N=1,2 についてはそれぞれ UnaryFunctor, BinaryFunctor という別名があります。 update を1引数関数オブジェクトにしてみましょう。 もともと引数を2つ持つので、残った1つ(p) に束縛される実引数をフィールドとして保存しておく必要があります。

julia> using NumericFuns, NumericExtensions

julia> type UpdateFun <: NumericFuns.UnaryFunctor
       p :: Float64
       end

関数適用はNumericFuns.evaluate で行うので、メソッドを定義しておきます。 また型推論のための関数である NumericFuns.result_type も必要です。 なおNumericExtensions は これらの関数を import および export しています。

julia> NumericExtensions.evaluate(uf::UpdateFun, x::Integer) = update(x, uf.p)
evaluate (generic function with 143 methods)

julia> NumericExtensions.result_type{T<:Integer}(::UpdateFun, ::Type{T}) = T
result_type (generic function with 100 methods)

julia> evaluate(UpdateFun(0.5), 0)
-1

ここまでくれば後はNumericExtensins.map に渡すだけです。

julia> @time for i in 1:N; xs = map(UpdateFun(0.5), xs); end
elapsed time: 1.695740634 seconds (822383688 bytes allocated, 29.14% gc time)

最初のケースと比べて3倍以上早くなりました。 毎回関数オブジェクトが作られますが、呼ばれる関数は既に名前のついた NumericExtensions.evaluate なので、JIT コンパイルが働いて性能が向上します。 重要なポイントとして、UpdateFunコンストラクタに渡す値が変わっても、 JIT コンパイル自体されていることに変わりはないという点です *8。 実際にやってみると、

julia> @time for i in 1:N; xs = map(UpdateFun(rand()), xs); end
elapsed time: 1.620748121 seconds (822383688 bytes allocated, 30.23% gc time)

あれ、rand() の時間を考えるとむしろ早くなっているような……

なお、UpdateFunimmutable にするとむしろ性能が悪化します。

# immutable UpdateFun2
julia> @time for i in 1:N; xs = map(UpdateFun2(0.5), xs); end
elapsed time: 12.022501026 seconds (4637697080 bytes allocated, 23.65% gc time)

要求メモリが格段に増えているのが原因でガベコレが発生しまくっているのですが、 これは素朴に考えると、immutable 型ではC/C++ でいう参照渡しから値渡しに切り替わって、 コピーのコストがかかるのではないかと思います*9

さて、NumericExtensions.jl 版と初期版とを比較すると、 必要メモリ量は殆ど変わらず、従ってガベコレの時間はほとんど変わらないので、相対的にガベコレが3倍ほど重くなっています。 そのため、破壊的な関数を使うことでメモリ使用量を減らしてみましょう。 map! ではなく map1! となっていることに注意してください。

julia> @time for i in 1:N; map1!(UpdateFun(0.5), xs); end
elapsed time: 1.592824116 seconds (7983688 bytes allocated)

使用メモリ量が大幅に減った結果、ガベコレが発生しなくなりました。 ただ、実行実時間はちょっとしか減っていないので、本質的には遅くなっているようです。 この辺りのトレードオフは、どのくらいのサイズの問題を解くかにもよるので、 実際に測定するまでは判断できないでしょう。

その他の拡張たち

NumericFuns.jl

標準ライブラリにない数学関数がいくつか追加されています。 何があるかはREADME を参照してください。

また、標準ライブラリにある数学関数も含めて、関数オブジェクト化されています。 型名は$(capitalize(funcname))Fun です。例えばsin => SinFun など。

NumericExtensions.jl

前述のmap など、主に配列操作の拡張がメインです。 標準ライブラリにあるものでも、関数オブジェクト NumericFuns.Functor{N} を受け付けるようになっていたり、 単純に実行速度が上がっていることがあります。 この場合、using NumericExtensions とするだけで高速化がなされることがあります *10。 非常に沢山あるのでドキュメント を見てください。

  • mapping
    • (いくつかの)配列から配列を作る操作 例えば `map(abs2, 1:3) ==> [1,4,9]
    • map(f, xs)
      • f として2引数、3引数関数を用いる map(f,xs,ys)map(f,xs,ys,zs) もあります
    • map!(f, dst, xs) ==> dst = map(f,xs)
    • map1!(f, xs) ==> x = map(f,xs)
    • mapdiff(f, xs, ys) ==> map(f, xs-ys)
    • add!(xs,ys) ==> xs += ys
      • 他にもsubtract!, multiply!divide! など
    • absdiff(xs,ys) ==> abs(xs-ys)
    • sqrdiff(xs,ys) ==> abs2(xs-ys)
    • fma(xs,ys,zs) ==> xs + ys .* zs
  • reduction
    • いわゆる畳み込み操作
      • 例えば sum(1:3) ==> 6
    • sum(xs) ==> reduce(+,xs)
    • sum(f, xs) ==> mapreduce(f,+,xs)
      • 複数引数パターンもあります
    • sumdiff(f, xs, ys) ==> mapreduce(f, +, xs-ys)
    • sumabs(xs) ==> sum(abs(xs))
    • sumsq(xs) ==> sum(abs2(xs))
  • scan
    • 畳み込みながら、途中経過を逐次残すことで配列を作る
      • 例えば cumsum(1:3) ==> [1,3,6]

Julia 0.4.0-dev がでたけれど

2日目の記事でも書きましたが、 Julia の次のバージョンである v0.4.0 では Base.callメソッドを書くことで、 関数呼び出し構文 a() を任意の型の値 a に適用できます。 普通の関数と同じ記法を使えるので、NumericFuns.Functor{N} よりも扱いやすいのですが、 残念なことに現状ではこうして作った関数オブジェクトは Base.map などに渡すことができません*11。 とは言え流石に標準ライブラリの高階関数は、将来的にはユーザ定義型を受け入れるようになるんじゃないかとは思います。 また、性能が上がっている拡張関数たちも、いつかは標準ライブラリに引き受けられるだろうので、 いつの日にかこれらの2パッケージは役目を終えるかもしれません。 しかし、どうやら今はまだその時ではないようなので、是非是非使ってみてください。

おまけ

twitter における、日本語話者な Julia ユーザのリストを作りました。 まだまだまったくもって不完全だと思うので、追加とか削除とかの希望を(twitter で直接)いただければと思います。

https://twitter.com/yomichi_137/lists/julialangjp

*1:ちなみにこの記事は当時割愛した「チューニングの実例」の補足となっています。

*2:ifelse(cond, T, F)cond が真ならT を、偽ならF を返す関数。 関数なのでTFともに必ず評価されるが、 その代わりif や三項条件演算子?: と比べて速い

*3:実際には毎ステップごとに平均や分散などの計算を行うことになるでしょうが、本記事の本質ではないのでおもいっきり簡便化しました。

*4:言うまでもなく、他の現代的な言語でもこのぐらい簡単だったりしますが

*5:正確に言えば、NumericExtensions.jl のような、より高速な比較対象があって初めて「遅い」といえるのですが

*6: 忘れがちですが、do 記法は関数の第一引数として無名関数リテラルを流し込むという仕様になっているため、便利だからと乱用すると確実に性能を落とします。何度も実行されるものはきちんと外に関数を書きましょう。

*7:C++ 的な意味でのFunctor

*8:例えば複雑な分岐がある関数において、めったに起きないコードパスを通るパラメータを渡してしまった場合などでは影響があったりするかもしれませんが、コンパイラの癖も含めて調べていないのでよくわかりません

*9:それにしても上がり過ぎな気もするけれど

*10:ただ、単純な高速化の場合、既に標準に取り入れられていることもあるので、その意味で過度な期待はしないこと

*11:引数の型がUnion(Function, DataType) なので。