Numbaは再帰関数もコンパイル結果のキャッシュやAOT (事前) コンパイルできる

## はじめに

Numba で再帰関数を高速化する場合、コンパイル結果のキャッシュやAOT(事前)コンパイルができず、プログラム実行の度にコンパイル時間がとられるという制限があるらしい。

ところが挙動を試していたらキャッシュは特に工夫なくともできていたし、AOTは少しの工夫で成功してしまったので紹介する。


## 変更履歴
(2023-12-24) キャッシュ効果の測定コードと測定結果を追記

## コンパイル結果のキャッシュ

### 方法
特に工夫せず `cache=True` するだけでできた。


 - `@jit` に型指定しない場合 (初回呼び出し時にコンパイル)

```Python
from numba import jit

# 再帰関数に @jit、cache=True
@jit(nopython=True, cache=True)
def rec_fib(n):
    if n == 0 or n == 1:
        return n
    return rec_fib(n - 1) + rec_fib(n - 2)

# 初回呼び出しでコンパイル or キャッシュ読込み
print(rec_fib(6))
```

- `@jit` に型指定する場合 (関数定義時にコンパイル)

```Python
from numba import jit, boolean, int64, float64, typeof

# 再帰関数に @jit、型指定 + cache=True
@jit(int64(int64), nopython=True, cache=True)
def rec_fib(n):
    if n == 0 or n == 1:
        return n
    return rec_fib(n - 1) + rec_fib(n - 2)
# 関数定義時にコンパイル or キャッシュ読込み

print(rec_fib(6))
```
オーバーヘッド時間の測定結果はかなり長くなってしまったので後述。

なお、Numbaバージョンが最新の場合 (0.57.0) でも古めの場合 (0.48.0) でも同様の挙動だった。 


## AOT(事前)コンパイル

AOTは再帰関数に対して普通に試みるとエラーになる。

```Python
from numba.pycc import CC

cc = CC("my_module_rec")

@cc.export("rec_fib", "int64(int64)")
def rec_fib(n):
    if n == 0 or n == 1:
        return n
    return rec_fib(n - 1) + rec_fib(n - 2)


print("start-compile: rec_fib")
cc.compile()    # エラー!
print("finish-compile: rec_fib")
```


### AOT成功コード例
以下のようにするとAOTに成功した。

```Python
from numba import jit
from numba.pycc import CC

cc = CC("my_module_rec")


@jit(nopython=True)    # (a)
def _rec_fib(n):
    if n == 0 or n == 1:
        return n
    return _rec_fib(n - 1) + _rec_fib(n - 2)

@cc.export("rec_fib_caller", "int64(int64)")    # (c)
def rec_fib_caller(n):    # (b)
    return _rec_fib(n)


print("start-compile: rec_fib_caller")
cc.compile()
print("finish-compile: rec_fib_caller")


print("\nコンパイル結果の動作確認")
import my_module_rec
print(f"{my_module_rec.rec_fib_caller(20) =}", "\n成功!!")
```
- 工夫する点は以下の通り
  - 再帰関数には `@jit` するが `cc.export` しない (a)
  - 再帰関数を呼び出すだけの別の関数を用意する (b)
  - 呼び出し用関数を `cc.export` してコンパイルする (c)

AOTなのに `@jit` とは一体?、と思うだろうがこれが鍵だった。
これを外したときのエラーメッセージをから推測するに、`cc.export` した関数のコンパイル過程で内部で呼び出された関数があった時、JITコンパイル時と同様、呼び出し対象の関数が `@jit` されていればコンパイル結果に含まれるのではないかと思う。

こちらもNumbaバージョンが最新の場合 (0.57.0) でも古めの場合 (0.48.0) でも同様の挙動だった。 


## キャッシュの効果の検証測定 

Numba のオーバーヘッドが複雑なので真面目に検証したらかなり長くなってしまいました。
ここからは完全に興味のある人向け。


### 型指定なしの場合

キャッシュによるオーバーヘッド削減効果の挙動を知るため、まず基本形として、非再帰のNumba化関数を3つ ( `f0`, `f1`, `f2` ) 含む以下のコードを複数回実行した。

コンパイル時間が分かりやすくなるよう、関数定義では不要な分岐を追加してでたらめな処理を増やし、一方で関数呼び出しでは関数の中での計算が一瞬で終わるような値を与えている。



```Python
import sys, time

TIMES = []
TIMES.append(("start", time.perf_counter()))

import numpy as np
TIMES.append(("import numpy", time.perf_counter()))

import numba
from numba import jit
TIMES.append(("import numba", time.perf_counter()))

print("Python: ", sys.version)
print("Numpy: ", np.__version__)
print("Numba: ", numba.__version__)
print()
TIMES.append(("print_versions", time.perf_counter()))


@jit(nopython=True, cache=True)
def f0(n):
    if n >= 0:
        return n + 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[0])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f0: def", time.perf_counter()))

r00 = f0(100)
TIMES.append(("f0: 1st call", time.perf_counter()))
r01 = f0(101)
TIMES.append(("f0: 2nd call", time.perf_counter()))


@jit(nopython=True, cache=True)
def f1(n):
    if n >= 0:
        return n - 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[1])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f1: def", time.perf_counter()))

r10 = f1(10)
TIMES.append(("f1: 1st call", time.perf_counter()))
r11 = f1(11)
TIMES.append(("f1: 2nd call", time.perf_counter()))

@jit(nopython=True, cache=True)
def f2(n):
    if n >= 0:
        return n * 2
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[2])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f2: def", time.perf_counter()))

r20 = f2(20)
TIMES.append(("f2: 1st call", time.perf_counter()))
r21 = f2(21)
TIMES.append(("f2: 2nd call", time.perf_counter()))



ELAPSED = [(TIMES[i][0], TIMES[i][1] - TIMES[i - 1][1]) for i in range(1, len(TIMES))]
for result in ELAPSED:
    print(result[0], "\t{:.6f}".format(result[1]))

print("\n", r00, r01, r10, r11, r20, r21)
```

測定結果 (型指定なし、非再帰) 
- 単位:sec (Python 3.11.4, Numpy 1.24.1, Numba 0.57.0)

|   | プログラム実行1回目<br>(キャッシュ未作成) | プログラム実行2回目<br>(キャッシュ作成済み) |
|---|---|---|
| f0関数定義      | 0.024984            | 0.004102             | 
| f0呼び出し1回目 | **1.752147**            | **0.141097**             | 
| f0呼び出し2回目 | 0.000003            | 0.000003             | 
| f1関数定義      | 0.000335            | 0.000381             | 
| f1呼び出し1回目 | **0.286833**            | **0.005060**             | 
| f1呼び出し2回目 | 0.000003            | 0.000002             | 
| f2関数定義      | 0.000322            | 0.000357             | 
| f2呼び出し1回目 | **0.290225**            | **0.004915**             | 
| f2呼び出し2回目 | 0.000003            | 0.000001             | 



この結果からオーバーヘッドの出方を整理すると以下のようになる。
- キャッシュ未作成時、プログラム実行1回目の各関数の初回呼び出しで大きなオーバーヘッド発生 (コンパイル)
- キャッシュ未作成時、一度のプログラム実行内で複数の関数を扱うと、1つ目の関数と2つ目以降の関数でオーバーヘッドの大きさが異なる
- キャッシュ利用時は、1つ目の関数の初回呼び出しは100ミリ秒単位 (おそらくキャッシュファイルの読み込み)
- キャッシュ利用時は、2つ目以降の関数の初回呼び出しが数ミリ秒単位
- 一度のプログラム実行内で同一の関数を複数回呼び出すと、2回目以降の呼び出しはマイクロ秒単位


次に、再帰関数でもキャッシュが効くかを確かめるため、上記のコードから `f2` を再帰関数 `f2_rec` に置き換えた以下のコードを測定した。

```Python
import sys, time

TIMES = []
TIMES.append(("start", time.perf_counter()))

import numpy as np
TIMES.append(("import numpy", time.perf_counter()))

import numba
from numba import jit
TIMES.append(("import numba", time.perf_counter()))

print("Python: ", sys.version)
print("Numpy: ", np.__version__)
print("Numba: ", numba.__version__)
print()
TIMES.append(("print_versions", time.perf_counter()))


@jit(nopython=True, cache=True)
def f0(n):
    if n >= 0:
        return n + 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[0])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f0: def", time.perf_counter()))

r00 = f0(100)
TIMES.append(("f0: 1st call", time.perf_counter()))
r01 = f0(101)
TIMES.append(("f0: 2nd call", time.perf_counter()))


@jit(nopython=True, cache=True)
def f1(n):
    if n >= 0:
        return n - 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[1])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f1: def", time.perf_counter()))

r10 = f1(10)
TIMES.append(("f1: 1st call", time.perf_counter()))
r11 = f1(11)
TIMES.append(("f1: 2nd call", time.perf_counter()))

@jit(nopython=True, cache=True)
def f2_rec(n):
    if n == 0 or n == 1:
        return n
    if n > 1:
        return f2_rec(n - 1) + f2_rec(n - 2)
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[2])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f2_rec: def", time.perf_counter()))

r20 = f2_rec(20)
TIMES.append(("f2_rec: 1st call", time.perf_counter()))
r21 = f2_rec(21)
TIMES.append(("f2_rec: 2nd call", time.perf_counter()))



ELAPSED = [(TIMES[i][0], TIMES[i][1] - TIMES[i - 1][1]) for i in range(1, len(TIMES))]
for result in ELAPSED:
    print(result[0], "\t{:.6f}".format(result[1]))

print("\n", r00, r01, r10, r11, r20, r21)
```

測定結果 (型指定なし、再帰あり) 
- 単位:sec (Python 3.11.4, Numpy 1.24.1, Numba 0.57.0)

|   | プログラム実行1回目<br>(キャッシュなし) | プログラム実行2回目<br>(キャッシュ作成済み) |
| ------------------- | ---------------- | -------------------- | 
| f0関数定義          | 0.003905         | 0.003968             | 
| f0呼び出し1回目     | 1.511317         | 0.135252             | 
| f0呼び出し2回目     | 0.000003         | 0.000003             | 
| f1関数定義          | 0.000286         | 0.000358             | 
| f1呼び出し1回目     | 0.273137         | 0.004888             | 
| f1呼び出し2回目     | 0.000003         | 0.000002             | 
| f2_rec関数定義      | 0.000285         | 0.000318             | 
| f2_rec呼び出し1回目 | **0.334119**         | **0.004873**             | 
| f2_rec呼び出し2回目 | 0.000130         | 0.000130             | 

キャッシュ利用時のf2_recの初回呼び出しオーバーヘッドはミリ秒単位と非再帰の場合と同等。
つまり、再帰関数でも関係なくキャッシュできている。

なお、Numbaバージョンが古い場合 (Numba 0.48.0) でも測定したが同様の挙動だった。 







### 型指定ありの場合

型指定ありでも測定した。
まずは非再帰関数3つを含むコード。

<details>
<summary> 型指定あり非再帰関数の測定コード (クリックで展開/折りたたみ) </summary>
<div>

```Python
import sys, time

TIMES = []
TIMES.append(("start", time.perf_counter()))

import numpy as np
TIMES.append(("import numpy", time.perf_counter()))

import numba
from numba import jit, boolean, int64, float64, typeof
TIMES.append(("import numba", time.perf_counter()))

print("Python: ", sys.version)
print("Numpy: ", np.__version__)
print("Numba: ", numba.__version__)
print()
TIMES.append(("print_versions", time.perf_counter()))


@jit(int64(int64), nopython=True, cache=True)
def f0(n):
    if n >= 0:
        return n + 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[0])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f0: def", time.perf_counter()))

r00 = f0(100)
TIMES.append(("f0: 1st call", time.perf_counter()))
r01 = f0(101)
TIMES.append(("f0: 2nd call", time.perf_counter()))


@jit(int64(int64), nopython=True, cache=True)
def f1(n):
    if n >= 0:
        return n - 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[1])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f1: def", time.perf_counter()))

r10 = f1(10)
TIMES.append(("f1: 1st call", time.perf_counter()))
r11 = f1(11)
TIMES.append(("f1: 2nd call", time.perf_counter()))

@jit(int64(int64), nopython=True, cache=True)
def f2(n):
    if n >= 0:
        return n * 2
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[2])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f2: def", time.perf_counter()))

r20 = f2(20)
TIMES.append(("f2: 1st call", time.perf_counter()))
r21 = f2(21)
TIMES.append(("f2: 2nd call", time.perf_counter()))



ELAPSED = [(TIMES[i][0], TIMES[i][1] - TIMES[i - 1][1]) for i in range(1, len(TIMES))]
for result in ELAPSED:
    print(result[0], "\t{:.6f}".format(result[1]))

print("\n", r00, r01, r10, r11, r20, r21)
```
--- 折りたたみここまで ---
</div>
</details>

測定結果 (型指定あり、非再帰) 
- 単位:sec (Python 3.11.4, Numpy 1.24.1, Numba 0.57.0)

|   | プログラム実行1回目<br>(キャッシュなし) | プログラム実行2回目<br>(キャッシュ作成済み) |
| --------------- | ---------------- | -------------------- | 
| f0関数定義      | **1.932430**         | **0.149574**             | 
| f0呼び出し1回目 | 0.000006         | 0.000004             | 
| f0呼び出し2回目 | 0.000001         | 0.000001             | 
| f1関数定義      | **0.288116**         | **0.005622**             | 
| f1呼び出し1回目 | 0.000005         | 0.000003             | 
| f1呼び出し2回目 | 0.000001         | 0.000001             | 
| f2関数定義      | **0.290771**         | **0.005422**             | 
| f2呼び出し1回目 | 0.000005         | 0.000003             | 
| f2呼び出し2回目 | 0.000001         | 0.000001             | 

オーバーヘッドが発生するタイミングは型指定なしでは初回呼び出し時だったのが、型指定ありでは関数定義時に移動した。
オーバーヘッドの大きさは同等だった。(少し数値が異なるが繰り返し測定でのバラツキの方が大きい)

次に、再帰関数でもキャッシュが効くかを確かめるため、上記のコードから `f2` を再帰関数 `f2_rec` に置き換えた。

<details>
<summary> 型指定あり再帰関数の測定コード (クリックで展開/折りたたみ) </summary>
<div>

```Python
import sys, time

TIMES = []
TIMES.append(("start", time.perf_counter()))

import numpy as np
TIMES.append(("import numpy", time.perf_counter()))

import numba
from numba import jit, boolean, int64, float64, typeof
TIMES.append(("import numba", time.perf_counter()))

print("Python: ", sys.version)
print("Numpy: ", np.__version__)
print("Numba: ", numba.__version__)
print()
TIMES.append(("print_versions", time.perf_counter()))


@jit(int64(int64), nopython=True, cache=True)
def f0(n):
    if n >= 0:
        return n + 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[0])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f0: def", time.perf_counter()))

r00 = f0(100)
TIMES.append(("f0: 1st call", time.perf_counter()))
r01 = f0(101)
TIMES.append(("f0: 2nd call", time.perf_counter()))


@jit(int64(int64), nopython=True, cache=True)
def f1(n):
    if n >= 0:
        return n - 1
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[1])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f1: def", time.perf_counter()))

r10 = f1(10)
TIMES.append(("f1: 1st call", time.perf_counter()))
r11 = f1(11)
TIMES.append(("f1: 2nd call", time.perf_counter()))

@jit(int64(int64), nopython=True, cache=True)
def f2_rec(n):
    if n == 0 or n == 1:
        return n
    if n > 1:
        return f2_rec(n - 1) + f2_rec(n - 2)
    # この後はコンパイル時間を増やすための無駄なコード
    arr = np.arange(0, 15, dtype=np.int64).reshape((3, 5))
    np.sort(arr[2])
    for i in range(arr.shape[0]): arr[i] += i
    return arr[0][0] + arr[1][1]

TIMES.append(("f2_rec: def", time.perf_counter()))

r20 = f2_rec(20)
TIMES.append(("f2_rec: 1st call", time.perf_counter()))
r21 = f2_rec(21)
TIMES.append(("f2_rec: 2nd call", time.perf_counter()))



ELAPSED = [(TIMES[i][0], TIMES[i][1] - TIMES[i - 1][1]) for i in range(1, len(TIMES))]
for result in ELAPSED:
    print(result[0], "\t{:.6f}".format(result[1]))

print("\n", r00, r01, r10, r11, r20, r21)
```
--- 折りたたみここまで ---
</div>
</details>

測定結果 (型指定あり、再帰あり) 
- 単位:sec (Python 3.11.4, Numpy 1.24.1, Numba 0.57.0)

|   | プログラム実行1回目<br>(キャッシュなし) | プログラム実行2回目<br>(キャッシュ作成済み) |
| ------------------- | ---------------- | -------------------- | 
| f0関数定義          | 1.843977         | 0.138856             | 
| f0呼び出し1回目     | 0.000006         | 0.000004             | 
| f0呼び出し2回目     | 0.000001         | 0.000001             | 
| f1関数定義          | 0.287179         | 0.005209             | 
| f1呼び出し1回目     | 0.000005         | 0.000002             | 
| f1呼び出し2回目     | 0.000001         | 0.000001             | 
| f2_rec関数定義      | **0.350976**         | **0.005049**             | 
| f2_rec呼び出し1回目 | 0.000088         | 0.000081             | 
| f2_rec呼び出し2回目 | 0.000132         | 0.000133             | 


キャッシュ利用時のf2_recの初回呼び出しオーバーヘッドはミリ秒単位と非再帰の場合と同等。
つまり、再帰関数でも関係なくキャッシュできている。

なお、Numbaバージョンが古い場合 (Numba 0.48.0) でも測定したが同様の挙動だった。 



### キャッシュ効果の測定まとめ

再帰関数でも非再帰関数でもキャッシュ利用により初回オーバーヘッドは同等に小さくなった。


## まとめ

実は再帰関数もキャッシュやAOTコンパイルができた。