なぜNumbaを使うと高速になるのか調べてみた

この記事は、HUIT アドベントカレンダー 2021 の13日目の記事です。

処理系初心者が少し調べてみた程度なので、間違った点があればご指摘して頂けると幸いです。

はじめに

Pythonのforがなぜ遅いのかを調べた記事を以前に書きました。
これによると、主にPythonの仮想マシンでコードの実行をしていることが遅い原因です。

この記事を書いた後に、具体的にPythonをどのように高速化するのかが気になってきました。

Pythonの高速化は

PyPyやCythonを使って実行する
Numbaで関数やループの最適化をする
キャッシュや適切なアルゴリズムにする

などの方法が考えられます。

今回はNumbaがどうやって高速化しているのかが気になったので軽く調べてみました。

Numbaとは

NumbaとはPython向けのJITコンパイラで、実行時にPythonのコードを機械語にコンパイルしています。
主にnumpyの配列や、関数、ループなどを使ったコードの高速化をしてくれます。

関数にデコレータを付けるだけで、その関数を最適化してくれるのが特徴です。
objectモードとnopythonモードの2種類があり、型推論に失敗するとobjectモードでコンパイルされてしまいます。objectモードでは高速化は期待できず、かえって遅くなる可能性があります。

Numbaを使うとどれぐらいの性能向上が見られるかを調べてみます。

次のコードで時間を計ってみます。

from numba import njit

@njit
def main():
    s = 0
    for _ in range(100000000):
        s += 1

if __name__ == "__main__":
    main()

Numbaなし	Numbaあり
0m5.115s	0m0.855s

njitはnopythonモードでコンパイルするためのデコレーターです。
Numba使った方が確かに高速化されています。
では、どのようにして高速化をしたのかを調べていきます。

Numbaの仕組み

まずNumbaがどのようにPythonのコードを機械語に変換しているのでしょうか。
これは大雑把に Python code -> LLVM IR -> 機械語の順で変換しています。
公式で詳しく解説している記事を見つけたので、そちらを簡潔にまとめます。
詳細を知りたい方は公式を覗いてください。

Numbaはフロントエンドとバックエンドに分かれていて、以下のような構成となっています。

フロントエンド
Pythonのバイトコードの解析をして、Numba IRと呼ばれるコードの内部表現に変換する。
ここでnopythonモードでLLVM IRを生成するために、全ての変数の型推論を行う。
失敗するとobjectモードになる。

バックエンド
フロントエンドで生成したNumba IRをLLVM IRに変換する。
そして最適化の後に機械語が生成される。

nopythonモードとobjectモードの違い

多くの記事ではobjectモードになると大して早くならない、場合によっては遅くなるのでnopythonモードで使いましょうと書かれていますが、これらは何が違うのでしょうか。

どちらもLLVM IR->機械語のように変換されますが、LLVMの部分が根本的に異なります。これに応じて生成される機械語も異なってきます。

何が異なるのか、objectモードだと全ての変数をPyObjectとして扱い、かつLLVMからとPython C APIを使っているのに対して、nopythonモードだと使用しません。

実際に生成されるLLVM IRで違いを見てみましょう。
環境変数でNUMBA_DUMP_OPTIMIZEDを0以外の値にすると、最適化されたLLVM IRが出力されるようになります。
以下のコードをobjectモードでコンパイルする方法が分からなかったので、公式ドキュメントからLLVMを引用します。

使用コード

from numba import jit

@jit()
def add(a, b):
    return a + b

def main():
    add(1, 2)

if __name__ == "__main__":
    main()

nopythonモード

define i32 @"_ZN8__main__7add$241Exx"(i64* noalias nocapture %retptr, { i8*, i32, i8* }** noalias nocapture readnone %excinfo, i64 %arg.a, i64 %arg.b) local_unnamed_addr #0 {
entry:
  %.6 = add nsw i64 %arg.b, %arg.a
  store i64 %.6, i64* %retptr, align 8
  ret i32 0
}

objectモード

@PyExc_SystemError = external global i8
@".const.Numba_internal_error:_object_mode_function_called_without_an_environment" = internal constant [73 x i8] c"Numba internal error: object mode function called without an environment\00"
@".const.name_'a'_is_not_defined" = internal constant [24 x i8] c"name 'a' is not defined\00"
@PyExc_NameError = external global i8
@".const.name_'b'_is_not_defined" = internal constant [24 x i8] c"name 'b' is not defined\00"

define i32 @"__main__.add$1.pyobject.pyobject"(i8** nocapture %retptr, { i8*, i32 }** nocapture readnone %excinfo, i8* readnone %env, i8* %arg.a, i8* %arg.b) {
entry:
  %.6 = icmp eq i8* %env, null
  br i1 %.6, label %entry.if, label %entry.endif, !prof !0

entry.if:                                         ; preds = %entry
  tail call void @PyErr_SetString(i8* @PyExc_SystemError, i8* getelementptr inbounds ([73 x i8]* @".const.Numba_internal_error:_object_mode_function_called_without_an_environment", i64 0, i64 0))
  ret i32 -1

entry.endif:                                      ; preds = %entry
  tail call void @Py_IncRef(i8* %arg.a)
  tail call void @Py_IncRef(i8* %arg.b)
  %.21 = icmp eq i8* %arg.a, null
  br i1 %.21, label %B0.if, label %B0.endif, !prof !0

B0.if:                                            ; preds = %entry.endif
  tail call void @PyErr_SetString(i8* @PyExc_NameError, i8* getelementptr inbounds ([24 x i8]* @".const.name_'a'_is_not_defined", i64 0, i64 0))
  tail call void @Py_DecRef(i8* null)
  tail call void @Py_DecRef(i8* %arg.b)
  ret i32 -1

B0.endif:                                         ; preds = %entry.endif
  %.30 = icmp eq i8* %arg.b, null
  br i1 %.30, label %B0.endif1, label %B0.endif1.1, !prof !0

B0.endif1:                                        ; preds = %B0.endif
  tail call void @PyErr_SetString(i8* @PyExc_NameError, i8* getelementptr inbounds ([24 x i8]* @".const.name_'b'_is_not_defined", i64 0, i64 0))
  tail call void @Py_DecRef(i8* %arg.a)
  tail call void @Py_DecRef(i8* null)
  ret i32 -1

B0.endif1.1:                                      ; preds = %B0.endif
  %.38 = tail call i8* @PyNumber_Add(i8* %arg.a, i8* %arg.b)
  %.39 = icmp eq i8* %.38, null
  br i1 %.39, label %B0.endif1.1.if, label %B0.endif1.1.endif, !prof !0

B0.endif1.1.if:                                   ; preds = %B0.endif1.1
  tail call void @Py_DecRef(i8* %arg.a)
  tail call void @Py_DecRef(i8* %arg.b)
  ret i32 -1

B0.endif1.1.endif:                                ; preds = %B0.endif1.1
  tail call void @Py_DecRef(i8* %arg.b)
  tail call void @Py_DecRef(i8* %arg.a)
  tail call void @Py_IncRef(i8* %.38)
  tail call void @Py_DecRef(i8* %.38)
  store i8* %.38, i8** %retptr, align 8
  ret i32 0
}

declare void @PyErr_SetString(i8*, i8*)

declare void @Py_IncRef(i8*)

declare void @Py_DecRef(i8*)

declare i8* @PyNumber_Add(i8*, i8*)

objectモードではPy_IncRefなどの仮想マシンで使われていた関数・マクロが呼ばれていることが分かります。
確かに39行目で変数aと変数bを足していますが、これもPython C API経由です。
nopythonモードは非常に簡潔で、変数a, bを足して保存する処理だけとなっています。

変数の型に着目するとobjectモードではarg.a,arg.bがi8*で、nopythonモードではi64になっています。
これは前者ではPyObjectへのポインターで、後者は数値型となっています。

高速化のポイント

なぜNumbaを使うと高速化できるのかといった話に戻ると、Python C APIを使用しない機械語を生成して実行することで高速化を実現していることが分かりました。

このAPIは仮想マシンで使われているPyNumber_AddやPy_DecRefなどの処理の関数を呼びだしているので、遅くなる原因となっています。

またコンパイルする過程でコードの最適化も行われているので、そちらも一つの要因になっていると考えています。

終わりに

どうしても中身がブラックボックスだと魔法で高速化してる感がありましたが、中身を知ると納得する内容でした。

他にNumbaでは並列処理、コンパイル結果のキャッシュ、CUDAに対応...などの高速化の方法があります。こちらもいずれ調べてみたいと思います。

ここまで読んで頂きありがとうございました。