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上並行處理資料陣列。　　

　　要在CPU上編譯和執行相同的函式，我們只需將目標更改為“CPU”，它將在編譯水平上帶來效能，在CPU上向量化C程式碼。這種靈活性可以幫助你生成更可重用的程式碼，並允許你在沒有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)
 

　　關於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的帖子。