• January 14, 2025
• in zh
• 2 min read

使用 Numba 加速 Python —— GPU 程式設計入門

Python 因其簡潔性與強大的科學運算函式庫，成為開發者的首選。然而，對於計算密集型任務，Python 的執行效率可能成為瓶頸。這時，Numba —— 一款即時編譯器（Just-In-Time Compiler），能夠將數值運算為主的 Python 代碼在 CPU 和 GPU 上大幅提速。

在本文中，我們將探討如何使用 Numba 簡化基於 NVIDIA CUDA 平台的 GPU 程式設計，即便是對 C/C++ 不熟悉的開發者，也能輕鬆上手。

什麼是 Numba？

Numba 是一款 即時編譯（JIT）、類型專門化（Type-Specializing）、函式編譯器（Function Compiler），可將 Python 函式轉換為最佳化的機器碼。無論是 CPU 或 NVIDIA GPU，Numba 都能在最少代碼改動的情況下大幅提升效能。

Numba 的主要特點包括： - 函式編譯器：優化單獨的函式，而非整個程式。 - 類型專門化：根據參數類型生成高效實作。 - 即時編譯：函式執行時才編譯，以適應 Python 動態類型特性。 - 數值運算優化：專注於 int、float 和 complex 等數據類型。

為何選擇 GPU 程式設計？

GPU 具備大規模並行運算能力，能夠同時執行數千個線程，特別適用於矩陣運算、模擬計算、圖像處理等數據並行任務。NVIDIA 的 CUDA 平台釋放了 GPU 的潛能，而 Numba 提供了 Python 友好的介面，使開發者能夠利用 CUDA，而無需深入學習 C/C++。

開始使用 Numba

CPU 優化

在深入 GPU 運算前，先來看看 Numba 如何加速 Python 的 CPU 運算。透過 @jit 修飾器，Numba 可優化以下計算斜邊長度的函式：

from numba import jit
import math

@jit
def hypot(x, y):
    return math.sqrt(x**2 + y**2)

當函式首次被呼叫時，Numba 會將其編譯為機器碼，從而加快運行速度。

GPU 加速

Numba 提供對 CUDA 的支援，使 GPU 編程變得簡單。我們可以利用 @vectorize 修飾器，將 NumPy 通用函式（ufuncs）加速到 GPU。例如，向量化的標量加法可如下實現：

from numba import vectorize
import numpy as np

@vectorize(['int64(int64, int64)'], target='cuda')
def add(x, y):
    return x + y

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
print(add(a, b))  # 輸出: [5, 7, 9]

這行代碼會觸發 GPU 運算，包括記憶體分配、數據傳輸及核心（kernel）執行。

Numba 的進階功能

自訂 CUDA 核心函式

對於超越元素級別運算的場景，Numba 支援使用 @cuda.jit 修飾器編寫自訂 CUDA 核心（kernel），讓開發者可以更精細地控制執行緒行為，進一步優化複雜演算法。

共享記憶體與多維網格

在更高級的 GPU 程式設計中，Numba 支援 2D 和 3D 資料結構，以及 CUDA 共享記憶體機制，幫助開發者打造針對特定應用的高效 GPU 代碼。

CUDA 程式設計選項比較

Numba 不是唯一的 Python GPU 編程庫，下表對比了幾種常見選項：

框架	優勢	劣勢
CUDA C/C++	高效能，完整 CUDA API	需要 C/C++ 專業知識
pyCUDA	Python 介面，可直接使用 CUDA API	需要較多代碼改動
Numba	代碼改動少，Pythonic 語法	相較於 pyCUDA，效能稍遜一籌

GPU 編程的最佳實踐

雖然 GPU 能提供極大加速，但錯誤的使用方式可能導致效能低下。以下是幾個最佳實踐： - 使用大規模數據集：GPU 在高並行度的場景中表現最佳。 - 最大化計算密度：確保計算量足夠，以彌補記憶體存取的開銷。 - 優化數據傳輸：最小化 CPU 與 GPU 之間的數據移動，以減少傳輸延遲。

結論

Numba 讓 Python 開發者能夠輕鬆發揮 GPU 的強大運算能力，讓高效能計算變得更可及。無論是數據科學家、研究人員，還是開發者，Numba 都提供了一種簡單且高效的方式來加速 Python 應用。

準備好進一步探索了嗎？深入學習 Numba 和 CUDA，釋放 GPU 運算的潛力，提升你的計算工作負載！