Numba

Pythonの関数に @jit のデコレーターを付けることで、実行時にJITコンパイルされる。

下記の例においては、Numba 0.27.0において、Pythonよりも200倍高速に動作し、NumPyの組み込みのsum()関数よりも30%高速に動作する^[3][4]。

from numba import jit

@jit
def sum1d(ary):
    total = 0.0
    for i in range(ary.shape[0]):
        total += ary[i]
    return total

Numba CUDAはNVIDIAが開発していて、通常はCUDAカーネルをC++で記述するが、Numba CUDAはそれをPythonで記述できる^[7]。なお、CUDAカーネルをC++で記述するので構わないが、CUDAをPythonで実行したい場合は、PyCUDA^[8]やCUDA Python^[9]などがある。

Numba CUDAのCUDAカーネル内では、Pythonの記法で記載するが、メモリ確保が出来ないため、NumPyのほとんどの機能は使えない^[10]。Numba CUDAはCuPyと併用することが出来て、CUDAカーネル外になるが、CuPyではNumPyの機能が使える。以下、CuPyを併用したときのB = A + A.Tのサンプルコード。

import cupy as cp
from numba import cuda

@cuda.jit
def add_kernel(A, B):
    y, x = cuda.grid(2)
    B[y, x] = A[y, x] + A[x, y]

A = cp.array([[1, 2], [3, 4]], dtype=cp.float32)
B = cp.zeros(A.shape, dtype=cp.float32)
add_kernel[1, A.shape](A, B)

print(B)
print(A + A.T)