Juliaのfor文はなぜ速い? - りんごがでている

1年以上ぶりのブログ更新ですね。

atsuoishimotoさんのブログ記事で，私のブログ記事が言及されていました。 atsuoishimoto.hatenablog.com

書いてある内容はまさにその通りで，for文を回すこと自体が取り立てて遅いわけではなく，その中で何らかの処理をする際に動的型付け言語であるPythonの性質上避けて通れない処理があり，その処理をfor文の中で何度も繰り返すことで結果としてfor文全体の処理が遅くなるという内容でした。

そうすると，ひとつの疑問として，なぜ同じ動的型付け言語であるはずのJulia言語では同様の処理が高速なのかが気になるところでしょう。この記事ではそのあたりを少し解説しようと思います。

Juliaの型推論

問題となっているコードは，以下のように２つのベクトルの内積を計算するコードです。

function dot(a, b)
    s = zero(eltype(a))
    for i in 1:endof(a)
        s += a[i] * b[i]
    end
    return s
end

このコードは，特に型の指定などがないにも関わらず，Pythonで似たコードを書いた場合と比べて非常に高速に動作します。

Julia:

julia> a = ones(100000); b = ones(100000);

julia> @benchmark dot(a, b)
BenchmarkTools.Trial:
  memory estimate:  16 bytes
  allocs estimate:  1
  --------------
  minimum time:     121.605 μs (0.00% GC)
  median time:      121.666 μs (0.00% GC)
  mean time:        132.996 μs (0.00% GC)
  maximum time:     1.456 ms (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1

Python:

In [2]: %timeit dot(a, b)
31.4 ms ± 680 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

ここでPythonがこれほど遅くなる理由は，元のブログ記事で解説されているとおり，for文の中で行われている一つ一つの計算の裏で様々な付随する処理が行われているためです。例えば，乗算を行う*演算子では，そもそも被演算子の型が乗算をサポートしているかなどの確認が必要です。動的言語では変数の型が動的に決定されるため，こうした処理が付随するのは避けて通れないようにも思えますが，Juliaはこうした処理を省略することで高速化を達成しています。

Juliaはなぜこうした処理を省略できるのでしょうか? その答えはJuliaの型推論にあります。実は，Juliaの処理系は関数を実行する前に型推論を行います。Juliaの型推論器は，関数のコードを見て回って，型が決定できるところに自動的に型注釈を付けて回ります。先ほどのコードを例に少し説明してみましょう。このコード例では，dot(a, b)という呼び出しが起きた際に引数であるaとbの型が決定します。上の例では両方ともVector{Float64} (これはArray{Float64,1}の別名) というFloat64型 (倍精度の浮動小数点数の型) を要素とする1次元配列です。こうすると当然for文の中にあるa[i]やb[i]はFloat64型の値を取り出すことが分かります。すると，a[i] * b[i]という配列のi番目の要素を取り出して積を計算する処理は，Float64同士の積だと分かりますので，被演算子が乗算をサポートしているかどうかをいちいちfor文の内部で確認する必要はありません。こうした処理を勝手に行ってくれるのが型推論です。

型推論の結果を見てみましょう。@code_typedというマクロを使うと，関数の型推論の結果を表示することができます。

julia> @code_typed dot(a, b)
CodeInfo(:(begin
        s = (Base.sitofp)(Float64, 0)::Float64 # line 3:
        $(Expr(:inbounds, false))
        # 省略
        SSAValue(4) = (Base.arraysize)(a, 1)::Int64
        # 省略
        $(Expr(:inbounds, :pop))
        SSAValue(5) = (Base.select_value)((Base.slt_int)(SSAValue(4), 0)::Bool, 0, SSAValue(4))::Int64
        SSAValue(6) = (Base.select_value)((Base.sle_int)(1, SSAValue(5))::Bool, SSAValue(5), (Base.sub_int)(1, 1)::Int64)::Int64
        #temp# = 1
        17:
        unless (Base.not_int)((#temp# === (Base.add_int)(SSAValue(6), 1)::Int64)::Bool)::Bool goto 27
        SSAValue(7) = #temp#
        SSAValue(8) = (Base.add_int)(#temp#, 1)::Int64
        i = SSAValue(7)
        #temp# = SSAValue(8) # line 4:
        s = (Base.add_float)(s, (Base.mul_float)((Base.arrayref)(a, i)::Float64, (Base.arrayref)(b, i)::Float64)::Float64)::Float64
        25:
        goto 17
        27:  # line 6:
        return s
    end))=>Float64

下の方を見てみると，(Base.mul_float)((Base.arrayref)(a, i)::Float64, (Base.arrayref)(b, i)::Float64)::Float64と書かれた部分があるのがわかると思います。この部分がまさにa[i] * b[i]の計算にに対応する部分で，それぞれの被演算子と演算結果がFloat64になることがちゃんと推論できています。また，一番最後に=>Float64と書かれているように，dot(a, b)の計算結果がFloat64型であることも推論されています。計算結果の型も推論できることで，他の関数からdot関数を呼び出したときもその計算結果が推論でき，型の情報がどんどん伝播していきます。

型推論は，ソースコードの上でインタープリタを実行するように行います。このインタープリタは引数の値などを扱うインタープリタではなく，型を扱うインタープリタです。型推論の詳しいアルゴリズムは実装を読むか，Inference Convergence Algorithm in Juliaと Inference Convergence Algorithm in Julia - Revisitedなどの解説を参照して下さい。

コードの生成

ソースコードに現れる変数や関数の返り値の型が型推論で判明すると，いくつかの確認処理を省くだけでなく，Juliaはその型に特化した命令を生成することができます。関数の実行時にはコード生成を行うことで処理速度を上げています。その中間状態を覗く@code_llvm dot(a, b)というマクロがありますので，結果を見てみましょう。

julia> @code_llvm dot(a, b)

define double @julia_dot_62600(i8** dereferenceable(40), i8** dereferenceable(40)) #0 !dbg !5 {
  # 省略
idxend2:                                          ; preds = %idxend
  %18 = add i64 %"#temp#.08", 1
  %19 = getelementptr double, double* %10, i64 %13
  %20 = load double, double* %19, align 8
  %21 = getelementptr double, double* %12, i64 %13
  %22 = load double, double* %21, align 8
  %23 = fmul double %20, %22
  %24 = fadd double %s.09, %23
  %25 = icmp eq i64 %"#temp#.08", %4
  br i1 %25, label %L27.loopexit, label %if
}

この結果はLLVM IRという言語で書かれています。LLVM IRは機械語に近い低レベルな命令列で，C言語などのコンパイラであるclangやRustのバックエンドでも使われているものです。この結果を見ると%20 = load double, double* %19, align 8という命令でメモリのある部分からdoubleの値を読み出しているのがわかります。これはちょうどa[i](もしくはb[i])に相当する処理です。また，読み出した値の積を計算する%23 = fmul double %20, %22という命令があるのも分かります。fmulという命令は浮動小数点数の積を計算するのに特化した命令です ('fmul' Instruction)。ではaやbの要素がFloat64でなく整数型のInt64だったらどうなるでしょうか? そういう場合でもJuliaは型推論を行ってInt64型に特化した積の演算命令を生成します。このように型推論の結果を使ってある処理に特化した命令を使うことで，非常に高効率なコードを生成します。ループ1回あたりでは大した違いはありませんが，長く回すループだと積もり積もって大きな差になります。こうして生成された関数のコードは，次回以降に同じ型の引数がくると再利用されます。

副作用

Juliaは関数の実行時に型推論とコード生成を行うことを説明しましたが，これには副作用もあります。初めて関数を呼び出す際に型推論とコード生成・最適化にかかるコストが上乗せされます。このコストがペイするほどのものかどうかを確認すべきなのでしょうが，現在のJulia (v0.6.2)は特殊なことをしない限り必ず型推論とコード生成を行います。ですので，初回の関数呼び出しでワンテンポ遅れたり，それが巨大なパッケージになると積もり積もって数秒から数十秒の遅れになったりします。とても良い例えをTwitterで見つけたので引用しようと思います。

これは人間と車が一メートル競走でどっちが早いかを見てるのと同じことだとおもいます．
— ごまふあざらし (@MathSorcerer) 2018年1月7日

将来的にもっとJuliaの処理系が賢くなれば，自分が今1メートル走をしようとしてるのかフルマラソンをしようとしてるのか分かるようになり，こうした副作用の効果も軽減されるようになると思います。

Pythonで同じことができないのか

Numbaというプロジェクトは，Juliaに似たような型推論とコード生成を行っているようです。使い方は簡単で，@jitというアノテーションを関数につけるだけです。

import numba

@numba.jit
def dot(a, b):
    s = 0.0
    for i in range(len(a)):
        s += a[i] * b[i]
    return s

実際にこれを試してみると，Juliaと同じレベルの高速なコードになりました。

In [2]: %timeit dot(a, b)
133 µs ± 1.43 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

どの程度複雑なコードまで上手く行くのかは分かりません(例えば，ユーザ定義型でもちゃんと最適化できるのか)が，非常に面白いプロジェクトだと思います。今後，RやRubyでも似たようなアプローチが出てくるかもしれませんね。