Numba のコードをちょっと読んだメモ
Numba は LLVM を使って Python のコードを JIT するライブラリ。
ちゃんと速い。https://gist.github.com/iwiwi/9228787711a353e115ffcdee21f1a882
@jit からコンパイル部分に到達するまで
@jit
→decorators._jit
→ dispatcher というものが返される- 適当に動かした例では、dispatcher は
registry.CPUDispatcher
のようだ - こいつにはほぼ全く実装がない
- 基底クラスは
dispatcher.Dispatcher
- こいつの
_compiler
とかいかにもコンパイルしそう
- こいつの
- その基底クラスは
_DispatcherBase
- その基底クラスは
_dispatcher.Dispatcher
で C extension 内 __call__
時に JIT してると思うけど、__call__
は_dispatcher.Dispatcher
のものが呼ばれているようだ_dispatcher.c
内の関数Dispatcher_call
に相当- (後で調べる①:どこで返り値の型を推定した? → 分かった:返り値の型はここでは推定する必要がなく引数の型だけで良い)
- (後で調べる②:なんで C extension に入ってるんだろう? → 速度のためだろうか・・・?)
_compile_for_args
は_DispatcherBase
で定義されており、self.compile
を呼ぶcompile
はDispatcher
で定義されており、キャッシュ等を調べたあとself._compiler.compile
を呼ぶ- 適当に動かした例では、
_compiler
は_FunctionCompiler
のようだ _FunctionCompiler.compile
はcompiler.py
内のcompile_extra
を呼ぶcompiler.compile_extra
はcompiler.Pipeline.complile_extra
を呼ぶ- (見逃してなければ Python の関数はまだ Python の関数オブジェクトそのまま来てる)
compiler.Pipeline.complile_extra
この辺が重要そう。- こいつは Python 関数オブジェクトをバイトコード化し、それをコンパイルする
- ここから呼び出される
_compile_core
を見るとコンパイルのプロセスが書かれている
コンパイルのパイプライン
以下は nopython モードでのコンパイルのパイプライン
- Python bytecode を取り出す
- Python bytecode を Numba IR に変換
- Numba IR の後処理をする
- 型を推定・アノテートする
- LLVM IR を生成・コンパイル
1. Python bytecode を取り出す
bytecode.py
内のget_code_object
で、関数.__code__
で取り出す
2. Python bytecode を Numba IR に変換
interpreter.py
内のInterpreter
クラスにop_HOGE_PIYO
関数が大量にある- Numba の IR は
ir.py
内に定義されるもので、Interpreter
はこれを生成する
3. Numba IR の後処理をする
postproc.py
内のPostProcessor
- ちゃんと読んでないけどあんまり長くないのであまり重要ではないと信じて進む
4. 型を推定・アノテートする
typeinfer.TypeInferer
が作られるTypeInferer
に色々な事前情報を与える(seed_ほげほげ
)build_constraint
で式等による伝搬ルールを列挙するpropagate
で constraint による情報を伝播し型を決めていく(収束までループするアルゴリズム)
5. LLVM IR を生成・コンパイル
Pipeine.stage_nopython_backend
→Pipeline._backend
→targetctx.codegen()
targetctx
は多分CPUContext
やCUDATargetContext
CPUContext.codegen
はJITCPUCodegen
に行くJITCPUCodegen
の基底クラスBaseCPUCodegen
で LLVM が llvmlite 経由で呼ばれている_engine
を見るとMCJIT
を使っていることが分かる
メモ
Python から LLVM 叩くならこいつ便利そう https://github.com/numba/llvmlite