numbaでざっくりPython高速化

[numba](http://numba.pydata.org/)というライブラリを使うと、Pythonのコードを比較的簡単に高速化できます。
うまくいけば、`from numba import jit`を書いて、高速化したい関数の前の行に`@jit`を書くだけで高速化できます。

仕組みとしては、numbaはPythonの仮想マシンコードを取得し、LLVM IRにコンパイルし、LLVMを使ってネイティブコードにするようです。
初回実行時は、コンパイル処理が走るので、若干遅くなりますが、重い処理だと、コンパイル時間を考えてもnumbaの方が速いこともあります。

# 利点と欠点
先に述べておきます。

## 利点
- 場合によっては、コード自体は改造せずに手軽に高速化できる
    - コードの改造があったとしても、軽微な改造で済むことも多い
- 別ファイルに分けてビルドする、みたいな手間なことが必要なく、手軽に`.py`ファイルの中で使える

## 欠点
- すべてのPython機能がサポートされているわけではないらしく、場合によっては`@jit`をつけるだけでは済まない
    - 型の扱いが厳格になるので、今まで「なぁなぁ」で動いていた部分がエラーになったりするようになります。また、numbaにうまく型推論してもらうよう工夫が必要になる場合があります。
- 当然だが、動かすのにnumbaが必要になる。環境によっては、numbaのインストールに苦労する
    - conda環境だと簡単なようです。pipで入る環境も結構あります。私の使っているArch Linuxでは、今現在、Python 3.8、LLVM 9.0になっており、どちらも現時点ではnumbaが対応していないため、ビルドするのは諦めてDockerでconda環境を使ってます
- コンパイルに時間がかかるので、闇雲につけたら逆効果になる

# 例

とてもうまくいく事例を紹介します。
なんかしらんけど、とっても遅い関数があります。

```py
import sys
sys.setrecursionlimit(100000)

def ack(m, n):
    if m == 0:
        return n + 1
    if n == 0:
        return ack(m - 1, 1)
    return ack(m - 1, ack(m, n - 1))
```

ちょっと時間測ってみますね。

```py
import time
from contextlib import contextmanager

@contextmanager
def timer():
    t = time.perf_counter()
    yield None
    print('Elapsed:', time.perf_counter() - t)

with timer():
    print(ack(3, 10))
```

> 8189
Elapsed: 10.270420542001375

10秒もかかっちゃいました。

数字を大きくするともっと時間かかりますが、本当に時間がかかるのでおすすめしません。
特に、3を4に増やすと、多分、死ぬまでかかっても終わらないので、全くおすすめしません。
この関数は[アッカーマン関数](https://ja.wikipedia.org/wiki/%E3%82%A2%E3%83%83%E3%82%AB%E3%83%BC%E3%83%9E%E3%83%B3%E9%96%A2%E6%95%B0)として知られています。

これをnumbaで速くしてみます。

```py
from numba import jit

@jit
def ack(m, n):
    if m == 0:
        return n + 1
    if n == 0:
        return ack(m - 1, 1)
    return ack(m - 1, ack(m, n - 1))

# 1回目
with timer():
    print(ack(3, 10))

# 2回目
with timer():
    print(ack(3, 10))

# 3回目
with timer():
    print(ack(3, 10))
```

> 8189
Elapsed: 0.7036043469997821
8189
Elapsed: 0.4371343919992796
8189
Elapsed: 0.4372558859977289

なんと、10秒かかっていたものが、初回0.7秒、2回目以降0.4秒にまで縮みました。
1行付け足すだけでこうなるなら、本当にお得ですね。

# 思ったより速くない場合は
Objectモードが使われている可能性があります。

numbaには、No PythonモードとObjectモードがあり、一旦No Pythonモードでコンパイルされますが、失敗したらObjectモードでコンパイルされます。
(ただし、この仕様は将来なくなり、デフォルトでNo Pythonモードのみ、Objectモードはオプションになるようです)

前者は、型が全部直接扱われるのに対し、後者は、PythonのオブジェクトをPython C APIから叩き、前者のほうがより高速です。
さらに、後者だと、ループを効率よくネイティブコードに書き換えられない場合がありますが、前者だとループが効率化されます。(loop-jitting)

No Pythonモードを強制するためには`@jit(nopython=True)`または`@njit`と書きますが、こうすると「型が分からない」といった類のエラーがよく発生するようになります。

基本的には、

- ひとつの変数には、ひとつの型のみを入れる
- `return`が複数ある場合、どの`return`でも同じ型を返すようにする
- No Pythonモードの関数で呼び出している関数もNo Pythonモードにする

などの方法で、型を明確にすることができます。

## 比較的簡単にNo Pythonモードに対応させるには
- 計算が重い部分だけを別関数に切り出してnumbaに対応させる
- 切り出した中でも、型をごちゃごちゃいじったり、オブジェクトにあれこれ代入したりしている部分は、外に出すことを考える
    - numba対応する部分の前処理あるいは後処理の形で、なんとかnumbaでやらずに済ませることを考える

などがおすすめです。
Pythonが得意なことはPythonでやる、Pythonが苦手なことはnumbaでやる、を意識します。

# 並列化してみる
`@jit(parallel=True)`とすると、forループで`range`の代わりに`prange`が使えます(`from numba import prange`が必要)。

`prange`で書いたループは、並列化されます。

# コンパイル結果のキャッシュ
`@jit(cache=True)`で、コンパイル結果をキャッシュファイルに書き出し、毎回コンパイルする手間を避けることができます。

# `fastmath`を使う
`@jit(fastmath=True)`で使えます。
gccやclangにもあるfastmathを有効にします。やや危なっかしい最適化でfloatの計算を早くするやつです。

# CUDAを使う
あまりお手軽ではありませんが、一応、CUDAも使えます。
CUDA使ったことある人なら、以下のコードでなんとなく分かるかと思います。

個人的には、これだったらcupyでいいかなぁ、と思いました。

```py
import numpy as np
from numba import cuda

@cuda.jit
def add(a, b, n):
    idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    if idx < n:
        a[idx] += b[idx]

N = 1000000
a_host = np.array(np.ones(N))
b_host = np.array(np.ones(N))
a_dev = cuda.to_device(a_host)
b_dev = cuda.to_device(b_host)

n_thread = 64
n_block = N // n_thread + 1

add[n_block, n_thread](a_dev, b_dev, N)

a_dev.copy_to_host(a_host)

print(a_host) # Expect: [2, 2, ..., 2]
```

もし`libNVVM`がない、などと怒られたら、CUDAをインストールしていない(`conda install cudatoolkit`などでインストールできます)か、環境変数の設定が必要です。

Google Colabなどでの設定例:

```py
import os
os.environ['NUMBAPRO_LIBDEVICE'] = "/usr/local/cuda-10.0/nvvm/libdevice"
os.environ['NUMBAPRO_NVVM'] = "/usr/local/cuda-10.0/nvvm/lib64/libnvvm.so"
```

# まとめ
numbaを使うと、ざっくりとPythonコードを高速化できます。
簡単な使い方を見ていきました。

参考になれば幸いです。

# 参考文献
[公式ドキュメント](https://numba.pydata.org/numba-doc/dev/index.html)は、必要なところだけ読めば、さほど長くなく、非常に参考になります。
特に[Performance Tips](https://numba.pydata.org/numba-doc/dev/user/performance-tips.html)は役立ちます。

CUDAを使う際のの環境変数設定については、
https://colab.research.google.com/github/cbernet/maldives/blob/master/numba/numba_cuda.ipynb
を参考にしました。