Numba:用CUDA加速的高性能Python编译器-python cuda

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教程。