Numbaでアレイのスライスを取得する際の注意点

前置き

pythonでの処理を高速化するために、numbaを普段から利用しています。
今回は多次元アレイのスライスを取得する際に、どの次元に対してスライスを取るかでパフォーマンスが変わってしまう問題を取り上げます。

1. 環境

$ sw_vers                
ProductName:    macOS
ProductVersion: 11.4
BuildVersion:   20F71
$ system_profiler SPHardwareDataType
    Hardware Overview:

      Model Name: MacBook Air
      Model Identifier: MacBookAir10,1
      Chip: Apple M1
      Total Number of Cores: 8 (4 performance and 4 efficiency)
      Memory: 16 GB
...
$ python -V               
Python 3.9.6
$ pip show numba | grep Version
Version: 0.53.1
$ pip show numpy | grep Version
Version: 1.21.0

2. テスト内容

各次元の長さが等しい三次元アレイを特定の次元方向にスライスし、スライス処理の時間に差があるかを確認します。以下がテストコードになります。

import time

import numpy as np
import numba
import tqdm
import matplotlib.pyplot as plt

@numba.njit
def dim3_0(arr: np.ndarray):
    size, _, _ = arr.shape
    arr_r = np.zeros((size, size))
    for i in range(size):
        arr_r += arr[i,:,:]
    return arr_r

@numba.njit
def dim3_1(arr: np.ndarray):
    size, _, _ = arr.shape
    arr_r = np.zeros((size, size))
    for i in range(size):
        arr_r += arr[:,i,:]
    return arr_r

@numba.njit
def dim3_2(arr: np.ndarray):
    size, _, _ = arr.shape
    arr_r = np.zeros((size, size))
    for i in range(size):
        arr_r += arr[:,:,i]
    return arr_r


if __name__ == '__main__':
    size = 80
    arr_test = np.random.rand(size, size, size)
    dim3_0(arr_test)
    dim3_1(arr_test)
    dim3_2(arr_test)
    fnc_lst = [dim3_0, dim3_1, dim3_2]

    size_init = 10
    size_step = 10
    nsteps = 50
    num_iter = 10
    time_arr = np.empty((nsteps, len(fnc_lst), num_iter), dtype=np.float64)
    for i in tqdm.tqdm(range(nsteps)):
        size = size_init + size_step * i
        arr = np.random.rand(size, size, size)
        for j, fnc in enumerate(fnc_lst):
            for k in range(num_iter):
                st = time.perf_counter()
                fnc(arr)
                time_arr[i, j, k] = time.perf_counter() - st

    x = np.arange(size_init, size_init + size_step * nsteps, size_step)
    plt.loglog(x, np.median(time_arr[:, 0], axis=1), label='dim3_0')
    plt.loglog(x, np.median(time_arr[:, 1], axis=1), label='dim3_1')
    plt.loglog(x, np.median(time_arr[:, 2], axis=1), label='dim3_2')
    plt.xlabel('arr size')
    plt.ylabel('execution time [s]')
    plt.legend()
    plt.show()

3.結果

横軸：各次元のアレイの長さ, 縦軸：各関数の実行時間です。

4.考察

テストコードで定義した3つの関数(dim3_0,1,2)は3次元アレイを1,2,3の各次元方向にスライスし、for文でその方向に対して加算を行う。
加算処理の回数は共通であること、また、アレイは乱数で初期化していることから、加算に要する時間はどの関数でも同じであると考えられます。したがって、関数の処理時時間の差はスライスの処理に起因するものと考えられます。
図からはほとんどの場合で1次元目のスライスを取得するのが最も早いことがわかります。
また、アレイが長く(大きく)なるにしたがって、処理時間の差は広がる傾向にあります。

5.結論

numba.jit,njitでデコレートした関数の中で、アレイのスライスを取る際は、一次元目に対してスライスを取るようにするのが無難。