Numba：使用CUDA加速的高性能Python編譯器

Python是一種高效的動態編程語言，廣泛應用於科學、工程和數據分析應用領域。使python如此受歡迎的因素有很多，包括其乾淨的、表達性的語法和標準的數據結構，綜合的「內置電池」標準庫，優秀的文檔，庫和工具的廣泛生態系統，專業支持的可用性，以及大而開放的社區。不過，也許最重要的原因是，像Python這樣的動態類型化的解釋語言能夠提高生產率。Python足夠敏捷與靈活，使它成為快速原型開發的一種偉大語言，同時也是構建完整系統的語言。

但是Python的最大優點也可能是它最大的缺點:它的靈活性和無類型的高級語法會導致數據和計算密集型程序的性能不佳。出於這個原因，關心效率的Python程序員經常在C中重寫他們的最內層的循環，並從Python調用編譯的C函數。有許多項目旨在簡化這種優化，例如Cython，但這往往需要學習一種新的語法。理想情況下，Python程序員希望在不使用另一種編程語言的情況下使其現有的Python代碼更快，當然，許多人也希望使用加速器來獲得更高的性能。

Numba:高性能計算的高生產率

在這篇文章中，筆者將向你介紹一個來自Anaconda的Python編譯器Numba，它可以在CUDA-capable GPU或多核cpu上編譯Python代碼。Python通常不是一種編譯語言，你可能想知道為什麼要使用Python編譯器。答案當然是：運行本地編譯的代碼要比運行動態的、解譯的代碼快很多倍。Numba允許你為Python函數指定類型簽名，從而在運行時啟用編譯(這就是「Just-in-Time」，即時，也可以說JIT編譯)。Numba動態編譯代碼的能力意味著你不會因此而拋棄Python的靈活性。這是向提供高生產率編程和高性能計算的完美結合邁出的一大步。

使用Numba可以編寫標準的Python函數，並在CUDA-capable GPU上運行它們。Numba是為面向數組的計算任務而設計的，很像大家常用的NumPy庫。在面向數組的計算任務中，數據並行性對於像GPU這樣的加速器是很自然的。Numba了解NumPy數組類型，並使用它們生成高效的編譯代碼，用於在GPU或多核CPU上執行。所需的編程工作可以很簡單，就像添加一個函數修飾器來指示Numba為GPU編譯一樣。例如，在下面的代碼中，@ vectorize decorator會生成一個編譯的、矢量化的標量函數在運行時添加的版本，這樣它就可以用於在GPU上並行處理數據數組。

import numpy as np

from numba import vectorize

@vectorize(["float32(float32, float32)"], target="cuda")

def Add(a, b):

return a + b

# Initialize arrays

N = 100000

A = np.ones(N, dtype=np.float32)

B = np.ones(A.shape, dtype=A.dtype)

C = np.empty_like(A, dtype=A.dtype)

# Add arrays on GPU

C = Add(A, B)

要在CPU上編譯和運行相同的函數，我們只需將目標更改為「CPU」，它將在編譯水平上帶來性能，在CPU上向量化C代碼。這種靈活性可以幫助你生成更可重用的代碼，並允許你在沒有GPU的機器上開發。

關於Python 的GPU-Accelerated庫

CUDA並行計算平台的優勢之一是其可用的GPU加速庫的闊度。Numba團隊的另一個項目叫做pyculib，它提供了一個Python介面，用於CUDA cuBLAS(dense linear algebra，稠密線性代數)，cuFFT(Fast Fourier Transform，快速傅里葉變換)，和cuRAND(random number generation，隨機數生成)庫。許多應用程序都能夠通過使用這些庫獲得顯著的加速效果，而不需要編寫任何特定於GPU的代碼。例如，下面的代碼使用「XORWOW」偽隨機數生成器在GPU上生成100萬個均勻分布的隨機數。

import numpy as np

from pyculib import rand as curand

prng = curand.PRNG(rndtype=curand.PRNG.XORWOW)

rand = np.empty(100000)

prng.uniform(rand)

print rand[:10]

CUDA Python的高並行性

Anaconda(原名Continuum Analytics)認識到，在某些計算上實現大的速度需要一個更具表現力的編程介面，它比庫和自動循環矢量化更詳細地控制並行性。因此，Numba有另一組重要的特性，構成了其非正式名稱「CUDA Python」。Numba公開了CUDA編程模型，正如CUDA C/ C++，但是使用純python語法，這樣程序員就可以創建自定義、調優的並行內核，而不會放棄python帶來的便捷和優勢。Numba的CUDA JIT(通過decorator或函數調用可用)在運行時編譯CUDA Python函數，專門針對你所使用的類型，它的CUDA Python API提供了對數據傳輸和CUDA流的顯式控制，以及其他特性。

下面的代碼示例演示了一個簡單的Mandelbrot設置內核。請注意，mandel_kernel函數使用Numba提供的cuda.threadIdx，cuda.blockIdx，cuda.blockDim和cuda.gridDim架構來計算當前線程的全局X和Y像素索引。與其他CUDA語言一樣，我們通過插入在括弧內一個「執行配置」(CUDA-speak用於線程數和線程塊啟動內核)，在函數名和參數列表之間中: mandel_kernel[griddim, blockdim](-2.0, 1.0, -1.0, 1.0, d_image, 20)。你還可以看到使用to_host和to_device API函數來從GPU中複製數據。

Mandelbrot的例子將在Github上持續更新。

@cuda.jit(device=True)

def mandel(x, y, max_iters):

"""

Given the real and imaginary parts of a complex number,

determine if it is a candidate for membership in the Mandelbrot

set given a fixed number of iterations.

"""

c = complex(x, y)

z = 0.0j

for i in range(max_iters):

z = z*z + c

if (z.real*z.real + z.imag*z.imag) >= 4:

return i

return max_iters

@cuda.jit

def mandel_kernel(min_x, max_x, min_y, max_y, image, iters):

height = image.shape[0]

width = image.shape[1]

pixel_size_x = (max_x - min_x) / width

pixel_size_y = (max_y - min_y) / height

startX = cuda.blockDim.x * cuda.blockIdx.x + cuda.threadIdx.x

startY = cuda.blockDim.y * cuda.blockIdx.y + cuda.threadIdx.y

gridX = cuda.gridDim.x * cuda.blockDim.x;

gridY = cuda.gridDim.y * cuda.blockDim.y;

for x in range(startX, width, gridX):

real = min_x + x * pixel_size_x

for y in range(startY, height, gridY):

imag = min_y + y * pixel_size_y

image[y, x] = mandel(real, imag, iters)

gimage = np.zeros((1024, 1536), dtype = np.uint8)

blockdim = (32, 8)

griddim = (32,16)

start = timer()

d_image = cuda.to_device(gimage)

mandel_kernel[griddim, blockdim](-2.0, 1.0, -1.0, 1.0, d_image, 20)

d_image.to_host()

dt = timer() - start

print "Mandelbrot created on GPU in %f s" % dt

imshow(gimage)

在一台帶有NVIDIA Tesla P100 GPU和Intel Xeon E5-2698 v3 CPU的伺服器上，這個CUDA Python Mandelbrot代碼運行的速度比純Python版本快1700倍。1700倍似乎有些不切實際，但請記住，我們正在比較編譯的、並行的、GPU加速的Python代碼來解釋CPU上的單線程Python代碼。

今天開始使用Numba吧

Numba為Python開發人員提供了一種簡單的進入GPU加速計算的方法，並提供了一種使用越來越複雜的CUDA代碼的方法，其中至少有新語法和術語。你可以從簡單的函數decorator開始實現自動編譯函數，或者使用pyculib的強大的CUDA庫。當你提高對並行編程概念的理解時，當你需要對於並行線程的有表現力且靈活的控制時，CUDA可以在不需要你第一天就完全了解的情況下使用。

Numba是一個BSD認證的開源項目，它本身嚴重依賴於LLVM編譯器的功能。Numba的GPU後端使用了基於LLVM的NVIDIA編譯器SDK。CUDA庫周圍的pyculib包裝器也是開源且經過BSD認證的。

要開始使用Numba，第一步是下載並安裝Anaconda Python發行版，這是一個「完全免費的、用於大規模數據處理、預測分析和科學計算的Python發行版」，其中包括許多流行的軟體包(Numpy、Scipy、Matplotlib、iPython等)和「conda」，這是一個強大的包管理器。一旦您安裝了Anaconda，通過鍵入conda安裝numba cudatoolkit pyculib，安裝所需的CUDA包。然後在ContinuumIO github存儲庫中查看CUDA的Numba教程。筆者建議你在Anaconda的博客上查看Numba的帖子。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！