目次

Numba (Python)

Pythonを高速に動くようにするモジュール。やってることは、LLVMでの機械語へのコンパイル。

「既存のPythonコードにあまり手を加えず、速くなればいいな程度でとりあえず軽く使う」こともできるし、「きちんとコードを書き換えて高速化する」こともできる。

他手段との比較

高速化の手段としては、他にPyPy、Cythonなどがある。Cythonは詳しく知らないが、PyPyと比較してはざっくり以下のようなイメージ。

速度比較は、以下の参考文献のyniji氏の記事が詳しい。

PyPyの方が、労力が少なく効果がそこそこでコスパが良い感はある。

Numbaは、後述のNoPythonモードによって高速化できないコードをエラーにできるので、着実に高速化させやすい。

また、未対応の外部モジュールを使っていると、PyPyなら丸ごと実行できなくなる一方、Numbaでは関数を切り分けて、使っていない部分だけコンパイルさせることができる。 2)

参考文献

環境

Objectモード NoPythonモード

高速化する上で重要な概念。

基本的にNumbaは関数のデコレータに「@jit」を付けるだけで実行時コンパイラが動くが、まだ発展途上であり、どんな型でもちゃんとコンパイルできるわけではない。 未対応の関数やデータ型が含まれる場合、「Objectモード」でコンパイルされる。Objectモードだと、高速化の効果は少ない。

基本的に何が出来て何が出来ないかは、下記を見ればよい。

全てがNumbaが対応する機能のみで記述されていれば、「NoPythonモード」となり、高速化の効果が大きくなる。

デコレータとして @njit を使うと、NoPythonモードを強制し、未対応の型があると警告やエラーを出してくれる。 一方、@jit だとObjectモードでも許容する。

from numba import jit, njit

# Numba非対応の機能
from collections import defaultdict

@jit
def func1():
    d = defaultdict(int)
    return d

print(func1())

# =>
# NumbaWarning: 
# Compilation is falling back to object mode WITH looplifting enabled
# because Function "func1" failed type inference due to:
# Untyped global name 'defaultdict': cannot determine Numba type of <class 'type'>
# ...
# NumbaWarning: Function "func1" was compiled in object mode without forceobj=True.
# 
# defaultdict(<class 'int'>, {})
# 
# 警告は出るが実行はされる


@njit
def func2():
    d = defaultdict(int)
    return d

print(func2())

# =>
# numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
# Untyped global name 'defaultdict': cannot determine Numba type of <class 'type'>
#
# エラー終了

既存コードに手を加えたくない場合は @jit、高速化目的なら @njit を使うとよい。

ただ、@jit の「NoPythonモードでのコンパイルに失敗した場合、自動的にObjectモードでコンパイル」という挙動はDeprecatedであり、そのうち @jitでも明示的にobjectモードを許容しない限りはエラーになるかもしれない。

以下、基本的にNoPythonモードを満たすように記述することを目標とする。

使える型

通常のPythonから大きく制限される。 数値型とnumpy.ndarrayしか使えないくらいに思っておいてよい。

一応SetやDictも使える。(ただし高速化の恩恵は若干少なくなる)

引数や戻り値のような、Numba関数の中と外の橋渡し的な役割をする変数は、特に留意が必要となる。

引数・戻り値に使える主な型

その他もあるかも知れない。

なるべく★を付けた3つのみを使うようにした方がよさそう。

(*1)はdeprecatedで今後使えなくなる可能性が高い。
(*2)は、対応はしているが、まだ十分高速に動くようコンパイル出来ない場合があるとリファレンスに書かれている。

(※)を付けたものは、下記でもう少し詳しく記述している。

UniTuple, Tuple

タプルは、複数の要素をまとめることが出来る。戻り値として複数の値を返す場合に重宝する。

要素の型が全て統一されているかどうかで大きく以下の2つに区別される。

混在しているTupleは、変数によるindexアクセスなどいくつかの機能が使えない(型を特定できないのでそりゃそう)。 また、イテレートするときに特殊な書き方が必要となる。

引数や戻り値のために一瞬使うだけなら大した影響はないだろうが、内部で使う場合はなるべくUniTupleな構造にした方がよい。

型指定の書き方は、以下のようにする。

@njit('UniTuple(i8, 5)()')    ... 64bit整数5個のUniTupleを返す

@njit('Tuple(i8, string, f4)()')   ... (64bit整数,文字列,32bit小数)のTupleを返す
reflected list と numba.typed.List

どちらも、Pythonのlist機能に似せたようなデータ構造。一応、以下の違いがある。

reflected listは deprecated。ネスト等で複雑になると限界があるのでtyped listに置きかえていく方針らしい。
そのため、(今のところ問題なく使えるものの)ver.0.45以降ではその旨の警告が出る。

また、typed listの方は「実験的機能」とされていて、バグがあったり、高速化の恩恵が少なくなる可能性が言及されている。

両者細かい注意点はあるが、現状、過渡期なこともありどちらも中途半端な状態。 今後の進化に期待しつつ、今のところ配列に関しては numpy.ndarray でいい気がする。

違いとしては、以下のようなものがある。 いずれも、要素の型は統一されている必要があり、統一できないような組合せで初期化・追加したりするとエラーとなる。

また、Pythonのデータ構造にはListの他にSet, Dictがあるが、

Pythonのdictは使えないため、辞書構造を引数にしたければ numba.typed.Dict で生成して使うしかない。 ただこれも「実験的機能」であり、実際、動作はちょっと遅い。

typed.Listの初期化方法例

コンパイル済み関数

ほとんどの処理が同じで一部だけ異なる場合、関数を外部注入したいことはある。

Numbaでは、関数オブジェクトは引数としてのみ使え、戻り値には出来ない。また、(おそらく)関数を引数とした関数はキャッシュや事前コンパイルは出来ない。

まず、事前コンパイルには型指定が必須だが、その記述方法が不明。 一応、後述の「推論させた型指定」の方法を使えば、記述方法がわからなくても型指定することは出来るが……

@njit('i8(i8)')
def double(a):    # 引数として渡したい関数
    return 2 * a

@njit
def fumidai(func, a):    # 踏み台関数
    return func(a)

fumidai(double, 2)    # 実際に呼び出す

@njit(fumidai.nopython_signatures[0])
def hontai(func, a):    # fumidaiに記録された型定義をもって、hontaiを型指定
    return func(a)      # →通る

hontai(double, 3)       # => 6

しかし、

実体のあるオブジェクトでなく参照として渡されるので、関数を置きかえたり、ファイルとして残すことはできないらしい。

@njit をキャッシュを使わない本来の実行時コンパイルとして使っている場合は、これらの制約は大きく影響しない。
せいぜい、異なる関数を引数とするたびに別々の関数としてコンパイルされるので、与える関数の種類は少ないに越したことはない、というくらい。

キャッシュしたい場合は、関数は引数としても使えないと考えておいた方がよいだろう。

関数内部での型

list, set は、通常のPython表記 [0, 1], {2, 3} のように書くと reflected list, reflected set となる。

dict は {2: 10, 3: 15} のように書くと typed dict となる。

いずれも、初期化時・追加時に型が混在しているとエラー。

特にdictは、イテレータのように渡されるのか、最初のkey-valueの組合せで型が確定する。例えば、int→floatはキャストできるが、float→intはできないため、以下のようになる。

d = {2: 2.0,   3: 3  }  はOK   (key, value) = (int64, float64)

d = {2: 2,     3: 3.0}  はエラー

ただ、使ってみた感触としてはこれらのパフォーマンスはあまり優れているとは言えない。 実行時間が数倍遅くなることもある(それでもNumbaを使わないよりは十分速いが)。なるべくならNumpy配列を使った方がよい。

また、Set, Dictに関しては、リスト内包表記で生成することは出来ない。(イテレート元とすることは出来る)

@njit
def function():
    a = [0, 1, 2]         # int64型のreflected list
    b = {0, 1, 1.5}       # float64型のreflected set
    c = {1: 1.5, 2: 2.5}  # (int64, float64)型のtyped dict
    d = {1: 1, 2: 2.5}    # エラー
    
    g = [v + 1 for v in a]             # [1, 2, 3]
    h = [v + 1 for v in b]             # [1.0, 2.0, 3.0]
    i = [k + v for k, v in c.items()]  # [2.5, 4.5]
    j = {v + 1 for v in a}             # エラー
    k = {v: v+1 for v in a}            # エラー

使える関数

ビルトイン関数

通常の演算処理で使うものはほぼ使えるが、たまに一部のオプションが未実装だったり。

基本はリファレンス参照。注意すべきいくつかの関数は以下。

標準モジュール

よく使うもので関係ありそうなのは以下。

Numpy関数

公式にまとまってる。

それなりに対応しているが、関数自体は存在してもオプション引数は未対応のものも多い。

個人的に axis オプションはよく使うが、未対応なのはもどかしい。

速度

無理にNumpy関数を使わずとも、単純なforループで代替可能な処理はforループで書いた方が速いこともある。

NumPy関数の多くは非破壊的に、結果は新しい配列を生成して返すため、生成コストが発生する。
引数 out を渡すとそこに格納してくれる関数もあるが、Numbaでは out 引数は基本的に使えない。

破壊的な処理で問題ない場合は特に、forループで更新していく方が速い。

たとえば、2つの配列で同じindex同士の最大値を求める場合、np.maximum(aaa, bbb) よりforループで aaa[i] = max(aaa[i], bbb[i]) とした方が速い。

テスト例

グローバル変数

Numba関数内からグローバル変数にアクセスしても、それがNumbaに対応した型なら使える。

ただし、AOT や jit(cache=True) などでコンパイル結果をキャッシュする場合、グローバル変数はコンパイル当時のもので固定される

glb = 5

@njit('i8()', cache=True)
def global_test():
    return glb

print(global_test())  # => 5

glb = 6

print(global_test())  # => 5

グローバル変数がリストなどのオブジェクトの場合、エラーとなる(この辺の詳細な条件は要確認)。

from numba import njit

glb = [1, 2, 3]

@njit('i8()', cache=True)
def global_test():  # => コンパイルエラー
    return glb[-1]

一方で、「大枠のNumba関数内で、リストもクロージャー関数も定義し、クロージャー関数からリストにアクセス」する場合はOK。コンパイル後の変更もちゃんと反映される。

from numba import njit

@njit('void()', cache=True)
def solve():
    glb = [[1, 2, 3]]  # 変数定義自体は関数より前が必須

    def global_test():
        return glb[-1]

    print(global_test())  # => [1, 2, 3]  NestedなListだろうとアクセスできる

    glb.append([4, 5, 6])
    print(global_test())  # => [4, 5, 6]  変更も反映

    glb = [[7, 8, 9]]
    print(global_test())  # => [7, 8, 9]  代入も反映(ただし型は一致しないと代入の方でエラー)

solve()

型指定

@njit() の引数には、コンパイルする関数の引数の型と戻り値の型(Signature)を指定することができる。

事前コンパイルには必須となる。

型指定が無いと、必然的に実行時コンパイルとなり、実際に呼び出されたときの引数に応じて型推論してからコンパイルされる。

型を指定すると事前コンパイルが可能になるほか、以下の恩恵を期待できる。

記述方法

「戻り値の型(引数1の型, 引数2の型, …)」のように記述する。戻り値の型を省略した場合は推論される。

numbaで定義された型指定用クラスを使う方法と、文字列を使う方法があるが、文字列を使った方がimportなどの手間が省ける。

文字列とデータ型の対応は以下の通り。

型指定用のクラスは int64 など数値がbit数を表すのに対し、文字列指定は i8 などbyte数を表す。微妙な相違に注意。

また、型名[:] でその型の配列であることを示すが、この配列は np.ndarray を指す。

# 戻り値: 無し,   引数: int64, int64
@njit('void(i8, i8)')
def func(a, b):
    print(a + b)


# 戻り値: int8,   引数: float32の配列, complex128の二次元配列
@njit('i1(f4[:], c16[:, :])')


# 戻り値が複数ある場合(同じ型)
@njit('UniTuple(i8[:], 2)()')
def func():
    x = np.zeros(10, np.int64)
    y = np.ones(20, np.int64)
    return x, y


# 戻り値が複数ある場合(異なる型)
@njit('Tuple(i8, f8)()')
def func():
    return 1, 1.0

推論させた型指定

型をどう記述すればいいかわかんない、という場合、一旦JITを走らせて推論させた型でもって、AOT用の型指定とさせることも(一応)できる。

以下のコード、fumidai.nopython_signatures で型を参照できるので、それを hontai の型指定に用いる。

毎回、型推論のためにダミー実行させるのも無駄なので、以下の例ではコンパイルさせたい場合に限り実行時オプションに compile を付けることで区別している。


if sys.argv[1] == 'compile':
    from numba import njit
    from numba.pycc import CC

    def func(a, b):          # コンパイルしたい関数
        return a + b
        
    fumidai = njit(func)     # 型指定無しでnjit化

    fumidai(1.5, 2)          # 使用を想定している引数型で呼び出す
                             # この時点でこの引数型に対するJITが走る
                             # float64(float64, int64) になってるはず

    # 同じ関数を、fumidaiで推論された型でもって型指定、事前コンパイル
    cc = CC('my_module')
    cc.export('hontai', fumidai.nopython_signatures[0])(func)
    cc.compile()

from my_module import hontai

print(hontai(3.5, 4))  # => 7.5
print(hontai(5, 6.5))  # => 11.0(6.5はintにキャストされる)

また、ここまで自動化しなくても、一度ダミー実行させた関数の noython_signatures をチェックすることで、型指定の方法のヒントを得ることができるかも。 ただし、nopython_signatures で表示される型と、型指定の記述方法は若干異なるっぽい。

クラス

numba.experimental.jitclass で、クラスのコンパイルも出来る。モジュール名の通り、実験的機能っぽい。

複雑なクラスになるとコンパイルに何十秒とかかってしまう。
もちろん何回も使うのであれば十分にお釣りが来るが、(今のところ)AOTやキャッシュが出来ない以上、使いどころは選ぶ必要がある。

型指定の例

numba.int64 64bit整数(numba.types.int64 も同じ)
numba.float64[:] 64bit浮動小数の1次numpy.ndarray
numba.int64[:,:,:] 64bit整数の3次numpy.ndarray
numba.types.List(numba.int64) 64bit整数のList
a = [0] のように初期化されるもの
numba.types.ListType(numba.int64) 64bit整数のtyped.List
a = numba.typed.List.empty_list(numba.int64) のように初期化されるもの

numba.typeof(○○) で型を調べられるので、わからない場合はそれを使う。

typeofを使えば一応は関数や独自クラスの型もわかるが、type(CPUDispatcher(<function <lambda> at 0x00000000123456789>)) のようになり、 結局、型と言うよりは「メモリのどこどこに記録されたオブジェクト」という感じで、 typeofで調べる際に用いたその「○○」しか入れられない。 同質の他のオブジェクトを渡すとエラーになる。


function_for_type_check = numba.njit()(lambda: 0)
function_truly_want_to_use = numba.njit()(lambda: 1)

spec = [
    ('func', numba.typeof(function_for_type_check))
]

@numba.experimental.jitclass(spec)
class A:
    def __init__(self):
        self.func = function_truly_want_to_use  # => Error

事前コンパイル(AOT)

numba.pycc モジュールを使うことで、関数をコンパイルした結果を保存することができる。

PCにCコンパイラがインストールされている必要がある。 未インストールの場合、Windowsならこの辺とかを参考に Build Tools for Visual Studio をインストールすればいけるはず。

型指定が必須となる。

使う際は、通常のモジュールと同様 import すればよい。 コンパイル後に出来るファイルは my_module.cp38-win_amd64.pyd などとプラットフォーム名が付く場合があるが、import 時の記述は import my_module でよい。

2023/10

pycc モジュールは、近く非推奨になるらしい。

ざっと見たところでは、distutils モジュールが Python 3.12 で削除される予定であるが、pyccはそれに依存しているとのこと。

で、pyccで生成されるコードはjitでコンパイルされるコードと非互換だし制限が多いしで、これを機に他の方法を作った方がよいと判断されたらしい。

代替手段が用意されるまではpyccは維持される、と書いてはいる(本当に信じていいかはわかんない。Python3.12 がメジャーになるまでのタイムリミットもあるし)ので、しばらくはAOTが必要なら使い続けてよいだろうが、Python本体の3.12へのアップデートは少し調べてからの方がいいし、どうしてもというわけでなければJITへの切り替えも考えた方がよいかも。

1)
要インストール。どこまでちゃんと動くのかは未確認
2)
PyPyでも、モジュール不使用の部分だけファイルを切り分けて、本流のPythonからsubprocessで実行するなど可能っちゃ可能
3)
ver.0.49で、コンストラクタにイテレータを渡せるようになった