Numba：基于CUDA加速的高性能Python

LibraryPKU 2017-11-17

展开全文

[注意：这篇文章最初于2013年9月19日发布，后于2017年9月19日更新。]

Python是一种高效率的动态编程语言，广泛应用于科学，工程和数据分析等领域。导致python如此流行的原因有很多，主要有其简洁且易于理解的语法和标准的数据结构，广泛的功能齐备的标准库，优秀的文档，库和工具构成的良好的生态系统，专业支持的可用性以及大而开放的社区。也许最重要的是，像Python这样的动态类型解释语言的高效率。Python是灵活的，这使其成为可用于快速原型设计的一种很好的语言，同时也能用于构建完整的系统。

但Python最大的优势也可能是其最大的弱点：其灵活性和无类型的高级语法可能导致数据和计算密集型程序的性能不佳。因此，关心效率的Python程序员通常会用C重写最内层的循环，并从Python调用编译好的C函数。有许多旨在使此优化更容易的项目（如Cython），但它们通常需要学习新的语法。理想情况下，Python程序员希望使现有的Python代码更快，而不需要使用另一种编程语言，而且自然，很多人希望使用加速器从而使代码具有更高的性能。

Numba：高效的高性能计算

这篇文章将向您介绍基于Anaconda的Python编译器Numba，它可以编译Python代码，以便在支持CUDA在GPU或多核CPU上执行。由于Python通常不是编译语言，您可能会想知道为什么需要Python编译器。答案当然是运行本地已编译代码比运行动态解释代码快许多倍。Numba通过允许您指定Python函数的类型签名，从而在运行时进行编译（这叫“即时编译”或JIT编译）。Numba动态编译代码的能力意味着可以不失Python的灵活性。这是提供高效率编程和高性能计算的理想组合的巨大一步。

使用Numba可以编写标准的Python函数，并在支持CUDA的GPU上运行它们。Numba专为面向数组的计算任务而设计，就像广泛使用的NumPy库一样。面向数组的计算任务中的数据并行性对于像GPU这样的加速器来说是恰好适合的。Numba可以读取NumPy数组类型，并使用它们来生成有效的编译代码，用于在GPU或多核CPU上执行。所需的编程工作可以像添加函数装饰器来指示Numba为GPU编译一样简单。例如，以下代码中的@vectorize装饰器在运行时生成标量函数的一个编译的，向量化版本Add函数，以便可以在GPU上并行处理数据数组。

要在CPU上编译和运行相同的函数，我们只需将目标更改为“cpu”，从而在CPU上提供编译，向量化的C代码级别的性能。这种灵活性可以帮助您生成更多可重用的代码，并可让您在没有GPU的机器上进行开发。

Python的GPU加速库

CUDA并行计算平台的优势之一是其可用的GPU加速库的广泛性。Numba团队的另一个项目叫做pyculib，为CUDA cuBLAS（密集线性代数），cuFFT（快速傅里叶变换）和cuRAND（随机数生成）库提供了一个Python接口。许多应用程序将能够从使用这些库中得到明显加速，而无需编写任何特定于GPU的代码。例如，以下代码使用“XORWOW”伪随机数生成器在GPU上生成百万个均匀分布的随机数。

CUDA Python的大规模并行

Anaconda（以前的Continuum Analytics）认识到，在某些计算上实现大幅度加速需要一个更易于理解的编程接口，与库和自动循环向量化相比需要对并行度有更细节上的控制。因此，Numba具有另一个重要的功能，它们构成了被称为“CUDA Python”的功能。Numba公开了CUDA编程模型，就像CUDA C/C ++一样，但使用纯python语法，以便程序员可以创建自定义调优的并行内核，而不会丢失Python的舒适和优点。Numba的CUDA JIT（可通过装饰器或函数调用获得）在运行时编译CUDA Python函数，专门针对您使用的类型，其CUDA Python API提供对数据传输和CUDA流的显式控制以及其他功能。

以下代码示例使用简单的Mandelbrot集内核来示范。注意，mandel_kernel函数使用Numba提供的cuda.threadIdx, cuda.blockIdx, cuda.blockDim和cuda.gridDim结构来计算当前线程的全局X和Y像素索引。与其他CUDA语言一样，我们通过在函数名称和参数列表（mandel_kernel[griddim, blockdim](-2.0, 1.0, -1.0, 1.0, d_image, 20)）之间的括号中插入一个“执行配置”（CUDA--用于运行内核的线程数和线程块）来启动内核。您还可以看到使用to_host和to_device API函数将数据复制到GPU或从GPU复制数据。

您可以在Github上获得Mandelbrot示例的完整Jupyter Notebook。

在具有NVIDIA Tesla P100 GPU和Intel Xeon E5-2698 v3 CPU的服务器上，该CUDA Python Mandelbrot代码运行速度比纯Python版本快近1700倍。1700x似乎是一个不切实际的加速，但请记住，我们进行比较的是编译并行GPU加速的Python代码与CPU上的解释型单线程Python代码。

Numba入门

Numba为Python开发人员提供了一个轻松入门GPU加速计算的途径，为越来越复杂的CUDA代码提供了最新的语法和行话。您可以从简单的函数装饰器开始来自动编译您的函数，或使用由pyculib公开的功能强大的CUDA库。当您对并行编程概念的理解提高后，当您需要对并行线程进行易于理解且灵活的控制时，CUDA可以提供帮助，而不需要您在一开始就使用。

Numba是一个BSD许可的开源项目，它本身在很大程度上依赖于LLVM编译器的功能。Numba的GPU后端使用基于LLVM的NVIDIA编译器SDK。该pyculib围绕CUDA包装的库也是BSD许可的开源项目。

要开始使用Numba，第一步是下载并安装Anaconda Python发行版，这是一个“完全免费的企业级Python发行版，用于大规模数据处理，预测分析和科学计算”，包括许多流行的包（Numpy， Scipy，Matplotlib，iPython等）和“conda”，一个强大的包管理器。安装Anaconda后，通过输入conda install numba cudatoolkit pyculib安装所需的CUDA包。然后在ContinuumIO github资源库中查看CUDA的Numba教程。我也建议您在Anaconda博客上查看关于Numba的帖子。