Numba (Python)

• Numba: A High Performance Python Compiler
• Numba

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

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

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

• PyPy
  • 対応している機能が多い
    • 標準ライブラリの処理はほぼ書き換えなしでそれなりに高速化される
    • Numpy, Scipyなどの外部ライブラリも、一部はPyPy用が提供されている¹⁾
  • Pythonとは別の実行環境としてインストールが必要
    • 最新のPythonの機能がPyPyに取り込まれるまで、若干のタイムラグはある
  • 基本的に実行時コンパイル（JIT）
    • なので、実行時間に常にコンパイル時間が含まれる
  • コンパイル時間は短め
• Numba
  • 単純な関数ならともかく、通常は書き換え無しではそこまで高速化されない
    • 恩恵を受けようと思うと、素Pythonでは却って遅くなるような書き方をすることもある
    • ただしCythonのように素Pythonで動かなくなるようなコードにすることはない
  • Pythonのモジュールとして、pipなどでインストールできる
  • JITも出来るし、事前コンパイル（AOT）もできる
  • コンパイル時間はPyPyと比較すると長め
  • 効果は高い（勿論、書き方や得意分野による）

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

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

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

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

• 公式リファレンス
  • Reference Manual
    • Supported Python features
    • Supported NumPy features
• yniji氏のQiita記事
  • AtCoderで Python を高速化する Numpy + Numba を使う - Qiita
  • PythonをNumbaでAOT (事前) コンパイルして高速化する - Qiita
• numbaでざっくりPython高速化 - Qiita
• チューニング例
  • NumPyをつかおう | 株式会社アイ・プライド

• Python 3.8.2
• Numba 0.47.0

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

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

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

• リファレンス
  • Supported Python features

全てが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モードを許容しない限りはエラーになるかもしれない。

• Deprecation of object mode fall-back behaviour when using @jit

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

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

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

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

• ★数値（byte, int, float, complex）
• ★numpy.ndarray
• ★UniTuple, Tuple（※）
• (*1) 通常のlist, set（reflected list, reflected set）（※）
• (*2) Numba独自のList, Dict（typed.List, typed.Dict）（※）
• (*2) Unicode文字列（Python3における通常の文字列）
• コンパイル済み関数（※）

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

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

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

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

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

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

• UniTuple: 統一されている（homogeneous）
• Tuple: 複数の型が混在している（heterogeneous）

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

• Supported Python features — Numba documentation

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

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

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

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

• Supported Python features — Numba documentation

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

• reflected list: 見た目上の扱いはPythonのlistのまま、制約を加えてNumbaで解釈できるようにしたもの
• typed list: Numba独自のリスト構造クラス

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

• Deprecation Notices — Numba 0.50.1 documentation

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

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

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

• reflected list
  • 外部から与える際はネスト不可、Numba関数内部で生成する場合はネスト可
  • Numba関数内部で「a = [0, 1]」とすると、こちらになる
• typed.List
  • ネスト可
  • 一度に生成することはできなくて、numba.typed.List() で空のインスタンスを生成後、1つずつappendする³⁾

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

• setについては、reflected set はあるが、typed set は未実装
• dictについては逆に、typed dict はあるが、reflected dict は存在しない

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

typed.Listの初期化方法例

from numba import typed, types

a = typed.List()             # 型未定義のリスト
a.append(1)                  # 初挿入時、aの型が固定される（int型）
a.append(1.5)                # エラー

a = typed.List()
numba_compiled_function(a)   # エラー（Numba関数に渡すときは型を定義してから）

a = typed.List.empty_list(types.int64)   # 型を定義しての空リスト生成
numba_compiled_function(a)   # OK


# ver.0.49.0で追加された、コンストラクタにリストを渡したときの挙動

a = typed.List([1, 2, 3])    # ←この時点で、aはint型と固定される
a.append(1.5)                # floatを入れようとするとエラー

a = typed.List([1, 2, 3.5])  # 同様にエラー

a = typed.List([1.5, 2, 3])  # 先頭がfloatならaはfloat型で定義され、
                             # int→floatのキャストは出来るため通る

# 辞書
d = typed.Dict.empty(key_type=types.int64, value_type=types.float64)
d[5] = 3.8
d[9] = 1

• Supported Python features — Numba documentation

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

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

しかし、

• キャッシュや事前コンパイルはできない
  • hontai を @njit(cache=True) や cc.export() でキャッシュしようとするとエラーが発生

  • TypeError: cannot pickle 'weakref' object

• 異なる関数を与えることはできない
  • hontai に double 以外の関数を定義して渡すと、たとえ引数型が同じでもエラーが発生

  • TypeError: No matching definition for argument type(s) type(CPUDispatcher(<function hontai at 0x......))

  • 結局決まった関数しか渡せないなら、引数にしなくてもグローバルで定義して呼び出すのと変わんない

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

@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}            # エラー

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

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

• pow(x, a, mod)
  • mod で返す機能は未実装
  • 繰り返し二乗法などで自力実装する
• List.sort(), sorted()
  • key は指定できない
• int(n, base)
  • 何進数かを指定する機能は未実装

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

• collections
  • NamedTuple以外は未実装
  • deque, Counter, defaultdict 等は使えない
• bisect
  • 未実装
  • numpy.searchsorted はあるのでそちらを使う
• itertools
  • 未実装
  • accumulate, combinations 等は使えない
  • accumulateに関しては、numpy.cumsum、numpy.cumprod は使える

• heapq
  • 意外と(?)、ほぼ全ての機能を使える

公式にまとまってる。

• Supported NumPy features

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

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

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

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

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

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

テスト例

import sys
from timeit import timeit

import numpy as np


def solve1(n, aaa, bbb):
    return np.maximum(aaa, bbb)


def solve2(n, aaa, bbb):
    for i in range(n):
        aaa[i] = max(aaa[i], bbb[i])
    return aaa


from numba import njit

solve1 = njit('(i8,i8[:],i8[:])', cache=True)(solve1)
solve2 = njit('(i8,i8[:],i8[:])', cache=True)(solve2)
print('compiled', file=sys.stderr)

n = 10 ** 7
aaa = np.random.randint(0, 1 << 30, size=n, dtype=np.int64)
bbb = np.random.randint(0, 1 << 30, size=n, dtype=np.int64)

print(timeit('solve1(n, aaa, bbb)', globals=globals(), number=100))
print(timeit('solve2(n, aaa, bbb)', globals=globals(), number=100))

# =>
# 4.919
# 1.482

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

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

• Notes on Caching — Numba 0.50.1 documentation

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）を指定することができる。

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

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

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

• 予期せぬ重複コンパイルを防止できる
  • JITは、引数の型が異なる呼び出しのたびに、別々の関数としてコンパイルされる
  • float型でコンパイル済みのJIT関数に、うっかりint型を与えると、必要ないのに新規コンパイルが走ってしまう
• 読んだだけで挙動を想定しやすいコードになる
• よりメモリ効率を改善できる
  • intやfloatは、（実行環境によるが、通常）64bit型としてコンパイルされる
  • 8bitで済むような変数はそう明記することで、メモリ効率を改善でき、若干の高速化にも繋がる

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

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

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

• Types and signatures — Numba documentation

型指定用のクラスは 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 の型指定に用いる。

• AtCoderで Python を高速化する Numpy + Numba を使う - Qiita（再掲）

毎回、型推論のためにダミー実行させるのも無駄なので、以下の例ではコンパイルさせたい場合に限り実行時オプションに 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 で、クラスのコンパイルも出来る。モジュール名の通り、実験的機能っぽい。

• Compiling Python classes with @jitclass

• 事前コンパイル(AOT)はできず、実行時コンパイル(JIT)のみ
  • AOT compilation of jitclass · Issue #6648 · numba/numba · GitHub
• cache=Trueなどのキャッシュもできないっぽい
  • cache=True for jitclasses · Issue #4830 · numba/numba · GitHub
• 事前にクラス変数の型を、Numba type instancesに用意された中から指定する必要がある
  • Types and signatures
  • そのため、関数や他のクラスなど、型指定できない（難しい）型はクラス変数に持ちにくい
  • どのように指定するかは上記ドキュメント参照
• nopythonモードでコンパイルを試みられる
  • うっかりobjectモードになってて高速化の恩恵が少なかった、などの心配は無い

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

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

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

• Compiling code ahead of time — Numba documentation

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

• Visual Studio ビルドツール 2019 （Build Tools for Visual Studio 2019） のインストール（Windows 上）

型指定が必須となる。

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