Numba編譯器的介紹與應用

1.介紹

Numba 是 python 的即時（Just-in-time）編譯器，即當你呼叫 python 函式時，你的全部或部分程式碼就會被轉換為“即時”執行的機器碼，它將以你的本地機器碼速度執行！它由 Anaconda 公司贊助，並得到了許多其他組織的支援。

在 Numba 的幫助下，你可以加速所有計算負載比較大的 python 函式（例如迴圈）。它還支援 numpy 庫！所以，你也可以在你的計算中使用 numpy，並加快整體計算，因為 python 中的迴圈非常慢。你還可以使用 python 標準庫中的 math 庫的許多函式，如 sqrt 等。有關所有相容函式的完整列表，請檢視 此處。

2.為什麼選擇 Numba？

那麼，當有像 cython 和 Pypy 之類的許多其他編譯器時，為什麼要選擇 numba？

原因很簡單，這樣你就不必離開寫 python 程式碼的舒適區。是的，就是這樣，你根本不需要為了獲得一些的加速來改變你的程式碼，這與你從類似的具有型別定義的 cython 程式碼獲得的加速相當。那不是很好嗎？

你只需要新增一個熟悉的 python 功能，即新增一個包裝器（一個裝飾器）到你的函式上。類的裝飾器也在開發中了。

所以，你只需要新增一個裝飾器就可以了。例如：

cd ~/pythia/data
from numba import jit
@jit
def function(x):
    # your loop or numerically intensive computations
    return x

這仍然看起來像一個原生 python 程式碼，不是嗎？

3.如何使用 Numba？

Numba 使用 LLVM 編譯器基礎結構 將原生 python 程式碼轉換成最佳化的機器碼。使用 numba 執行程式碼的速度可與 C/C++ 或 Fortran 中的類似程式碼相媲美。

以下是程式碼的編譯方式：

首先，Python 函式被傳入，最佳化並轉換為 numba 的中間表達，然後在型別推斷（type inference）之後，就像 numpy 的型別推斷（所以 python float 是一個 float64），它被轉換為 LLVM 可解釋程式碼。然後將此程式碼提供給 LLVM 的即時編譯器以生成機器碼。

你可以根據需要在執行時或匯入時 生成 機器碼，匯入需要在 CPU（預設）或 GPU 上進行。

4.使用 numba 的基本功能

（只需要加上 @jit ！）

為了獲得最佳效能，numba 實際上建議在你的 jit 裝飾器中加上 nopython=True 引數，加上後就不會使用 Python 直譯器了。或者你也可以使用 @njit。如果你加上 nopython=True的裝飾器失敗並報錯，你可以用簡單的 @jit 裝飾器來編譯你的部分程式碼，對於它能夠編譯的程式碼，將它們轉換為函式，並編譯成機器碼。然後將其餘部分程式碼提供給 python 直譯器。

所以，你只需要這樣做：

from numba import njit, jit
@njit      # or @jit(nopython=True)
def function(a, b):
    # your loop or numerically intensive computations
    return result

當使用 @jit 時，請確保你的程式碼有 numba 可以編譯的內容，比如包含庫（numpy）和它支援的函式的計算密集型迴圈。否則它將不會編譯任何東西，並且你的程式碼將比沒有使用 numba 時更慢，因為存在 numba 內部程式碼檢查的額外開銷。

還有更好的一點是，numba 會對首次作為機器碼使用後的函式進行快取。因此，在第一次使用之後它將更快，因為它不需要再次編譯這些程式碼，如果你使用的是和之前相同的引數型別。

如果你的程式碼是 可並行化 的，你也可以傳遞 parallel=True 作為引數，但它必須與 nopython=True 一起使用，目前這隻適用於CPU。

你還可以指定希望函式具有的函式簽名，但是這樣就不會對你提供的任何其他型別的引數進行編譯。例如：

 from numba import jit, int32
 @jit(int32(int32, int32))
 def function(a, b):
     # your loop or numerically intensive computations
     return result
 # or if you haven t imported type names
 # you can pass them as string
 @jit( int32(int32, int32) )
 def function(a, b):
    # your loop or numerically intensive computations
    return result

現在你的函式只能接收兩個 int32 型別的引數並返回一個 int32 型別的值。透過這種方式，你可以更好地控制你的函式。如果需要，你甚至可以傳遞多個函式簽名。

你還可以使用 numba 提供的其他裝飾器：

• @vectorize：允許將標量引數作為 numpy 的 ufuncs 使用，
• @guvectorize：生成 NumPy 廣義上的 ufuncs，
• @stencil：定義一個函式使其成為 stencil 型別操作的核函式
• @jitclass：用於 jit 類，
• @cfunc：宣告一個函式用於本地回撥（被C/C++等呼叫），
• @overload：註冊你自己的函式實現，以便在 nopython 模式下使用，例如：@overload（scipy.special.j0）。

Numba 還有 Ahead of time（AOT）編譯，它生成不依賴於 Numba 的已編譯擴充套件模組。但：

• 它只允許常規函式（ufuncs 就不行），
• 你必須指定函式簽名。並且你只能指定一種簽名，如果需要指定多個簽名，需要使用不同的名字。

它還根據你的CPU架構系列生成通用程式碼。

5.@vectorize 裝飾器

透過使用 @vectorize 裝飾器，你可以對僅能對標量操作的函式進行轉換，例如，如果你使用的是僅適用於標量的 python 的 math 庫，則轉換後就可以用於陣列。這提供了類似於 numpy 陣列運算（ufuncs）的速度。例如：

from numba import jit, int32
@vectorize
def func(a, b):
    # Some operation on scalars
    return result

你還可以將 target 引數傳遞給此裝飾器，該裝飾器使 target 引數為 parallel 時用於並行化程式碼，為 cuda 時用於在 cudaGPU 上執行程式碼。

@vectorize(target="parallel")
def func(a, b):
    # Some operation on scalars
    return result

使 target=“parallel” 或 “cuda” 進行向量化通常比 numpy 實現的程式碼執行得更快，只要你的程式碼具有足夠的計算密度或者陣列足夠大。如果不是，那麼由於建立執行緒以及將元素分配到不同執行緒需要額外的開銷，因此可能耗時更長。所以運算量應該足夠大，才能獲得明顯的加速。

這個影片講述了一個用 Numba 加速用於計算流體動力學的Navier Stokes方程的例子：

6.在GPU上執行函式

你也可以像裝飾器一樣傳遞 @jit 來執行 cuda/GPU 上的函式。為此你必須從 numba 庫中匯入 cuda。但是要在 GPU 上執行程式碼並不像之前那麼容易。為了在 GPU 上的數百甚至數千個執行緒上執行函式，需要先做一些初始計算。實際上，你必須宣告並管理網格，塊和執行緒的層次結構。這並不那麼難。

要在GPU上執行函式，你必須定義一個叫做 核函式 或 裝置函式 的函式。首先讓我們來看 核函式。

關於核函式要記住一些要點：

• 核函式在被呼叫時要顯式宣告其執行緒層次結構，即塊的數量和每塊的執行緒數量。你可以編譯一次核函式，然後用不同的塊和網格大小多次呼叫它。
• 核函式沒有返回值。因此，要麼必須對原始陣列進行更改，要麼傳遞另一個陣列來儲存結果。為了計算標量，你必須傳遞單元素陣列。

# Defining a kernel function
from numba import cuda
@cuda.jit
def func(a, result):
    # Some cuda related computation, then
    # your computationally intensive code.
    # (Your answer is stored in  result )

因此，要啟動核函式，你必須傳入兩個引數：

• 每塊的執行緒數，
• 塊的數量。

例如：

threadsperblock = 32
blockspergrid = (array.size + (threadsperblock - 1)) // threadsperblock
func[blockspergrid, threadsperblock](array)

每個執行緒中的核函式必須知道它在哪個執行緒中，以便了解它負責陣列的哪些元素。Numba 只需呼叫一次即可輕鬆獲得這些元素的位置。

@cuda.jit
def func(a, result):
    pos = cuda.grid(1)  # For 1D array
    # x, y = cuda.grid(2) # For 2D array
    if pos < a.shape[0]:
        result[pos] = a[pos] * (some computation)

為了節省將 numpy 陣列複製到指定裝置，然後又將結果儲存到 numpy 陣列中所浪費的時間，Numba 提供了一些 函式 來宣告並將陣列送到指定裝置，如：numba.cuda.device_array，numba.cuda。device_array_like，numba.cuda.to_device 等函式來節省不必要的複製到 cpu 的時間（除非必要）。

另一方面，裝置函式 只能從裝置內部（透過核函式或其他裝置函式）呼叫。比較好的一點是，你可以從 裝置函式 中返

from numba import cuda
@cuda.jit(device=True)
def device_function(a, b):
    return a + b

你還應該在這裡檢視 Numba 的 cuda 庫支援的功能。

Numba 在其 cuda 庫中也有自己的原子操作，隨機數生成器，共享記憶體實現（以加快資料的訪問）等功能。

ctypes/cffi/cython 的互用性：

• cffi – 在 nopython 模式下支援呼叫 CFFI 函式。
• ctypes – 在 nopython 模式下支援呼叫 ctypes 包裝函式。
• Cython 匯出的函式是 可呼叫的。