iwiwi 備忘録

学んだことを殴り書きます。自分向けのメモです。

Numba のコードをちょっと読んだメモ

Numba は LLVM を使って Python のコードを JIT するライブラリ。

ちゃんと速い。https://gist.github.com/iwiwi/9228787711a353e115ffcdee21f1a882

@jit からコンパイル部分に到達するまで

  • @jit → decorators._jit → dispatcher というものが返される
  • 適当に動かした例では、dispatcher は registry.CPUDispatcher のようだ
  • こいつにはほぼ全く実装がない
  • 基底クラスは dispatcher.Dispatcher
  • その基底クラスは _DispatcherBase
  • その基底クラスは _dispatcher.Dispatcher で C extension 内
  • __call__ 時に JIT してると思うけど、__call__ は _dispatcher.Dispatcher のものが呼ばれているようだ
  • _dispatcher.c 内の関数 Dispatcher_call に相当
    1. 呼び出された Python 関数の引数の型を確認しコンパイルできるか等調べる
    2. コンパイルする時は関数 compile_and_invoke を呼ぶ
    3. compile_and_invoke は self._compile_for_args を呼ぶ
  • （後で調べる①：どこで返り値の型を推定した？ → 分かった：返り値の型はここでは推定する必要がなく引数の型だけで良い）
  • （後で調べる②：なんで 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 モードでのコンパイルのパイプライン

  1. Python bytecode を取り出す
  2. Python bytecode を Numba IR に変換
  3. Numba IR の後処理をする
  4. 型を推定・アノテートする
  5. 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