Speed up Python numerical computation 660,000 times with Numba
Python 在現代科學運算上扮演很重要的角色，尤其是 Machine Learning 領域，其核心是一連串的矩陣運算和最佳化理論；但受限於其 Interpreter GIL 和 動態類型語言等等，在中大型運算上的效能一直為人所詬病，所以 Numpy 應運而生，大部分使用 C 語言撰寫使 Numpy 能夠較少的受上述兩點影響，許多常見且具有較高效能的 Python package 都是奠基於 Numpy。
但很多應用顯然不滿足於此，Numpy 還是太慢了，如果我們可以將某部分高運算需求的模塊編譯，就能夠大幅加速應用。本文討論的 Numba 是一種 Just-In-Time Compilation，介於如 C 的 Ahead-of-Time Compilation 和如 Python 的 Interpreter 之間；保留 Python 開發靈活性的同時，不需要深入最佳化 CPython Interpreter 也能夠大幅加速 Python；此外，還是一種異構計算，可以同時使用 CPU 與 GPU 運算，對於特定的應用，可以完全解放硬體性能。
本文提出一個範例實驗，設計一套演算法將點狀 (scatter) 資料轉為圖片(矩陣, grid)資料，分別使用 Python, Numpy, Numba, Numba CUDA 實現，並比較效能。
Sample Data
以下的程式碼產生五萬個點，其資料型式以三個向量組成，分別為 x 座標、y 座標以及其上的值；值的部分為 x, y 組成的 cos 函數：
import numpy as np
from matplotlib import pyplot as pltx = np.random.randint(1, 10, 50000) + np.random.random(50000)
y = np.random.randint(1, 10, 50000) + np.random.random(50000)
v = np.cos(x - y)
fig = plt.figure(figsize=(16, 12))
ax = plt.axes()
ax.scatter(x, y, c=v, cmap='Greens')
我們亦得到下圖：
這個點資料將做為演算法的輸入，轉換成矩陣型態，也就是圖片資料；圖片資料是在一個矩陣之上，每一個 pixel 儲存一個值，若演算法正確，我們應該能夠看到相同於上圖的矩陣(圖片)。除了正確性外，此處我們也要兼顧效能；點資料時常用於高精度的形式，如 Time of Flight 的 LiDAR，一份點資料很輕易就內含上百萬個點；但事實上我們很可能僅需要取很小一部分子集，如一張 320×320 的圖片(102400個點)就能得到足夠的資訊應用到其他演算法上，此外，轉換成圖片能夠讓我們輕易的使用傳統的電腦視覺演算法，或是 Object Detection 以及 Segmentation 等等。
以下演算法按照不同場景和應用可能有些細節還需要調校，但核心邏輯是將 grid 中的每一格在 scatter data 位置計算出來，並以 L2 求最近座標的點之值填入對應的格子。
以下實驗為了符合實際場景，我們改採用一百萬個點，並投射到 (320, 320) 的矩陣上。
Python
def to_image(image_np, x, y, value,
width, height, dist_threshold):
x_min, y_min = min(x), min(y)
x_max, y_max = max(x), max(y)
width_interval = (x_max - x_min) / width
height_interval = (y_max - y_min) / height
MAX_DISTANCE = 10 num_points = len(x)
for row in range(width):
x_point = width_interval * row + x_min
for col in range(height):
y_point = height_interval * col + y_min min_dist = MAX_DISTANCE
min_dist_idx = 0
for i in range(num_points):
dist = ((x[i] - x_point) ** 2) + \
((y[i] - y_point) ** 2)
if min_dist > dist:
min_dist_idx = i
min_dist = dist if min_dist > dist_threshold:
image_np[row, col] = 0
else:
image_np[row, col] = value[min_dist_idx]
使用原生 Python 的問題很明顯，因為沒有 vectorization 運算，所以 24行後for i in range(num_points) 以下的運算耗時非常久，這一段在 Numpy 可以輕易透過 np.argmax 和向量相加加速。
Numpy
def to_image_np(image_np, x, y, value,
width, height, dist_threshold):
width_interval = (x.max() - x.min()) / width
height_interval = (y.max() - y.min()) / height
image_np = np.empty((width, height), dtype=np.float32)
shift_x = x - x.min()
shift_y = y - y.min() for i in range(width):
x_dist_square = np.power(shift_x - (width_interval * i), 2)
for j in range(height):
y_dist_square = np.power(shift_y - \
(height_interval * j), 2) dist_square = x_dist_square + y_dist_square
nearest_points_idx = np.argmin(dist_square)
if dist_square[nearest_points_idx] > dist_threshold:
image_np[i, j] = 0
else:
image_np[i, j] = value[nearest_points_idx]
透過 line profiler ，我們大致可以觀察到 numpy code 的效能瓶頸：
主要集中在第 49 行計算 y_dist_square ，佔執行時間的將近九成；但 Per Hit 使用的時間並不比 x_dist_square 長，主要是因為運算次數達 102400 次。
而單次執行時間最久的為 47 和 49 的 np.power 與向量減去常數，這兩者都無法透過演算法再加速，僅能尋求底層的最佳化了。
我認為這樣的場景就是非常適合引入 Numba ：
1. 大型矩陣參與運算
2. 效能瓶頸不在於單一(或不嚴謹的說，atomic)複雜耗時的指令
3. for loop 之間獨立處理元素
以上條件讓使用者將 Numpy 改寫為 Numba 非常簡單，幾乎不需處理平行化的衍生議題，就能透過廣開執行緒的方式大幅加速應用。
Numba
事實上 njit 可以使用的更簡單，譬如僅加上 @njit ，但加上型別宣告可以更快速啟動 JIT 以及編譯最佳化，對於速度會有所幫助。
若型別錯誤，編譯警告會直接顯示輸入的型別是什麼，加上和 Numpy 還有部分兼容，甚至還可以使用 np.argmin ，可以說是相當親近 Python 使用者。
import numba
from numba import njit, prange
@njit((),
numba.types.Array(numba.float64, 1, 'C'),
numba.types.Array(numba.float64, 1, 'C'),
numba.types.Array(numba.float64, 1, 'C'),
numba.int64, numba.int64, numba.float64),
parallel=True, fastmath=True)
def to_image_signature(image_np, x, y, value,
width, height, dist_threshold):
width_interval = (x.max() - x.min()) / width
height_interval = (y.max() - y.min()) / height
shift_x = x - x.min()
shift_y = y - y.min() image_np = np.empty((width, height), dtype=np.float32)
for i in prange(width):
x_dist_square = np.power(shift_x - (width_interval * i), 2)
for j in prange(height):
y_dist_square = np.power(shift_y - \
(height_interval * j), 2) dist_square = x_dist_square + y_dist_square
nearest_points_idx = np.argmin(dist_square)
if dist_square[nearest_points_idx] > dist_threshold:
image_np[i, j] = 0
else:
image_np[i, j] = value[nearest_points_idx]
Numba CUDA
在使用 numba.cuda 前，要稍微介紹 CUDA 中 grid, block 和 thread 的概念。
這張圖的重點為，Thread ⊆ Block ⊆ Grid ⊆ GPU，GPU 中執行指令(kernel)的最小單位為 Thread。另一個概念是 Warp，像是 Thread 的小組別 (Block 像是宿舍房間)，一個 Warp 會含有 32 個在同一個 Block 的 Thread，多數時候；而當該 Block 的 Thread 數量小於 32 時，它們會被歸類在一個 Warp，所以多數時候也是 Thread ⊆ Warp ⊆ Block。
from numba import cuda@cuda.jit
def to_image_numba_cuda(image_np, x, y, value,
width, height, x_min, y_min,
x_max, y_max, dist_threshold):
width_interval = (x_max - x_min) / width
height_interval = (y_max - y_min) / height MAX_DISTANCE = 10
start_x, start_y = cuda.grid(2)
grid_x, grid_y = cuda.gridsize(2)
num_points = x.shape[0]
for row in range(start_x, width, grid_x):
x_point = width_interval * row + x_min
for col in range(start_y, height, grid_y):
y_point = height_interval * col + y_min min_dist_idx = 0
min_dist = MAX_DISTANCE
for i in range(num_points):
dist = ((x[i] - x_point) ** 2) + \
((y[i] - y_point) ** 2)
if min_dist > dist:
min_dist_idx = i
min_dist = dist if min_dist > dist_threshold:
image_np[row, col] = 0
else:
image_np[row, col] = value[min_dist_idx]
其中 cuda.grid 回傳整個 grid 中 thread的絕對位置，如 start_x = 0, 1, 2, … 即為第一個 block 的第一個 thread、第二個 thread、第三個 thread等；cuda.gridsize 則是 整個 grid 當中 thread的數量。
Performance
以下數字由 AMD Ryzen 3900X 12-Core Processor 及 RTX 2080 Ti on CUDA 11.2 得到；
在處理百萬個數值，並搜尋近似的 320×320 圖片的問題下，原生 Python 慢到幾乎不可能使用(註)；而 Numpy 則需約半小時，雖有顯著進步，但對於可能的即時運算場景還是捉襟見肘。
僅使用 CPU 的 njit 以簡易的平行化方案，有超出 10 倍於 Numpy 的加速；而 Numba CUDA 包含將資料 host copy to device , initialize cuda array , compilation 等等，僅需不到半秒鐘，運行時間太短連使用 nvprof 都有困難。
最終，相較於 Python ，Numba CUDA 能夠得到近 660,000 倍的速度提升。
Conclusion
Python 的易用性是其最大的優勢也是詛咒，雖然 Guido van Rossum 放話要為 CPython 解決效能問題，但面對越來越龐大精細的資料和複雜的演算法終究是杯水車薪；Numba 雖僅能治標，但藥效強勁，Python 使用者都應該要對其有基本的認識。Numba 值得討論的和坑還很多，如 CUDA stream , CUDA BlockDim , ThreadDim 的配置、使用 nvprof 和 nvvp 做效能調校，以及深入其原理的最佳化技巧等等，但篇幅龐大且艱深，就先交給明天吧。
除了 Numba ，如 Pypy , Cython 或是近日開源的 Triton 都值得使用 Python 且想為科學計算加速的使用者研究，一併推薦大家。
Reference
CUDA 的 Threading：Block 和 Grid 的設定與 Warp
註: CPython 實在是太久了，在將近十個小時候忍不住停下來看了狀況；此 244000s 是按照 line profiler 以比例大致計算。
