关于Numba你可能不了解的七个方面

更多深度文章，请关注：https://yq.aliyun.com/cloud

我最喜欢的事情之一是与人们谈论GPU计算和Python。 Python的生产力和互动性与GPU的高性能结合是科学和工程中许多问题的杀手。 有几种使用GPU加速Python的方法，但我最熟悉的是Numba，它是Python函数的即时编译器。 Numba在标准的Python翻译器中运行，因此您可以直接以Python语法编写CUDA内核，并在GPU上执行它们。 NVIDIA开发者博客最近推出了一篇对Numba的介绍，我建议阅读那篇文章，这对GPU上的Numba有一个简要的了解。

当我和人们谈论Numba时，我发现他们很快就在Python中编写了CUDA内核的基础知识。 但是我们经常没有时间使用Numba为GPU程序员提供的一些更先进的功能。 在这篇文章中，我想深入了解一下，并展示了在GPU上使用Numba的几个往往被忽视的方面。我会快速讲述一些主题，但是我会提供链接以供阅读。

1.Numba是100%开源的

你看到这是列表中的第一个项可能会感到惊讶，但我经常碰到一些人，他们并没有意识到Numba，特别是其CUDA支持，是完全开源的。 困惑是可以理解的，因为Numba从2012年的半专有性开始到现在的状态已经走了很长的路程。 当Numba项目开始时，实际上有两个不同的代码库：Numba，一个用于CPU的开源Python编译器，以及NumbaPro（后来更名为“Accelerate”），这是GPU的专有Python编译器。 在接下来的几年中，我们将来自NumbaPro的GPU支持的组件合并到开源的Numba项目中，最终于2017年中期发布Pyculib。

Pyculib是围绕标准CUDA算法的Python包装器的新名称。 它包括Python包装器，用于：

• CUBLAS
• CUFFT
• cuSPARSE
• cuRAND
• 从ModernGPU和的排序算法

这些包装器曾经是Anaconda Accelerate的一部分，Numba用户主要感兴趣的原因是它们兼容了CPU上的标准NumPy阵列以及Numba分配的GPU阵列。 因此，将标准操作（如FFT）与使用Numba编写的自定义CUDA内核组合起来非常容易，如下代码片段所示：

import pyculib.fft
import numba.cuda
import numpy as np
 
@numba.cuda.jit
def apply_mask(frame, mask):
    i, j = numba.cuda.grid(2)
    frame[i, j] *= mask[i, j]
 
# … skipping some array setup here: frame is a 720x1280 numpy array
 
out = np.empty_like(mask, dtype=np.complex64)
gpu_temp = numba.cuda.to_device(out)  # make GPU array
gpu_mask = numba.cuda.to_device(mask)  # make GPU array
 
pyculib.fft.fft(frame.astype(np.complex64), gpu_temp)  # implied host->device
apply_mask[blocks, tpb](gpu_temp, gpu_mask)  # all on device
pyculib.fft.ifft(gpu_temp, out)  # implied device->host

你可以在其文档中了解有关Pyculib功能的更多信息。 现在Pyculib是开源的，我们正在积极尝试扩展Pyculib，以包括其他CUDA库，如cuSOLVER和nvGRAPH。

2.Numba + Jupyter =快速CUDA原型开发

将Numba看作“用Python语法编写CUDA”是很容易的，但是Numba与Python数据科学生态系统中的其他工具的结合可以改变GPU计算的经验。 我们特别喜欢用Jupyter Notebook（和JupyterLab，下一代笔记本）使用Numba。 图1所示的Jupyter Notebook提供了一个基于浏览器的文档创建环境，允许将Markdown文本，可执行代码和图形和图像的图形输出相结合。Jupyter已经变得非常受欢迎，用于教学，记录科学分析和交互式原型。 事实上，这篇博文中的所有例子都是在Jupyter笔记本中创建的，你可以在这里找到。

图1.Jupyter Notebook的截图显示了一些在这篇文章中使用的CUDA Python代码。

为什么Numba和Jupyter如此适合GPU计算的实验？有几个原因：

• 作为一个即时编译器，Numba即时编译你的CUDA代码，因此可以通过重新执行Jupyter代码单元立即获得更改。不需要保存外部文件，不需要构建步骤。您的CUDA内核可以直接嵌入到笔记本本身中，并按照Shift-Enter快速更新。
• 如果将NumPy数组传递给CUDA函数，Numba将分配GPU内存并自动处理主机到设备和设备到主机副本。这可能不是最有效的使用GPU的方法，但是在原型设计时非常方便。那些NumPy数组可以随时更改为Numba GPU设备阵列。
• Jupyter在其“魔术”命令集中包含了一个基准测试工具。通过用带有％timeit的行前缀，Jupyter会自动运行该命令多次，以准确测量运行时间。 （不要忘记在运行内核后同步设备以获得准确的时间！）

Jupyter Notebook可以通过SSH进行隧道传输，从而可以在台式机或笔记本电脑上使用Web浏览器编辑笔记本，但在远程Linux服务器上执行代码。 （例如，我们用DGX-1做了这个）。在我的笔记本电脑上，我运行如下命令：

ssh -L 9999:localhost:9999 me@gpu-server.example.com

它记录我登录到我们的GPU服务器，并将端口9999转发回我的笔记本电脑。 然后我可以使用此命令在远程系统上启动Jupyter（这假设你在服务器上安装了Jupyter）：

jupyter notebook --port 9999 --no-browser

Jupyter将开始并打印一个URL粘贴到浏览器中以访问笔记本接口。 SSH转发端口将加密数据，并在远程服务器和本地计算机之间进行路由。 现在，你可以从你的网络浏览器中方便地运行Tesla P100上的算法实验！

3.Numba可以同时为CPU和GPU编译

在编写应用程序时，很方便的是在无需复制功能内容的情况下，让辅助功能在CPU和GPU上都可以工作。这样，你可以确定这两个地方的实现是相同的。此外，可以更容易地在CPU上对CUDA设备功能进行单元测试，以验证逻辑，而不必总是编写专门的CUDA内核包装器，只需在GPU上运行设备功能即可进行测试。

在CUDA C ++中，使用函数定义上的__host__和__device__关键字的组合可以从CPU（主机）或GPU（设备）调用。 例如，我可以在CUDA C ++中写这个：

__host__ __device__ float clamp(float x, float xmin, float xmax) {
    if (x < xmin){
        return xmin;
    } else if (x > xmin) {
        return xmax;
    } else {
        return x;
    }
}

那么我可以直接在主机和其他CUDA C ++函数中使用clamp（）函数。

使用Numba，我可以使用正常的CPU编译器装饰器在Python中编写相同的函数：

@numba.jit
def clamp(x, xmin, xmax):
    if x < xmin:
        return xmin
    elif x > xmax:
        return xmax
    else:
        return x

但是我可以直接从CUDA内核使用这个函数，而无需重新声明它，就像这样：

@numba.cuda.jit
def clamp_array(x, xmin, xmax, out):
    # Assuming 1D array
    start = numba.cuda.grid(1)
    stride = numba.cuda.gridsize(1)
    for i in range(start, x.shape[0], stride):
        out[i] = clamp(x[i], xmin, xmax)  # call "CPU" function here

当我从CUDA内核中调用它时，Numba编译器会自动编译一个CUDA版本的clamp（）。 请注意，Numba GPU编译器比CPU编译器要严格得多，因此某些功能可能无法重新编译GPU。 这里有一些提示。

• GPU支持NumPy数组，但数组函数和数组分配并不支持。
• 使用Python数学模块中的数学函数，而不是numpy模块。
• 不要在@jit装饰器中使用显式类型签名。 通常在CPU上使用64位数据类型，而在GPU上，32位类型更为常见。 Numba将自动重新编译正确的数据类型，无论何处需要它们。
• 你可以将共享内存数组作为参数传递到设备函数中，这样可以更轻松地编写可从CPU和GPU调用的实用程序函数。

4.Numba通过@vectorize轻松实现阵列处理

在Python中编写完整的CUDA内核的能力非常强大，但对于元素方面的数组函数来说，这可能是乏味的。你必须决定适用于数组维度的线程和块索引策略，选择合适的CUDA启动配置等。 值得庆幸的是，Numba提供了一种简单的创建这些特殊数组函数（NumPy中称为“通用函数”或“ufuncs”）的方法，这几乎不需要CUDA知识！

除了用于编译常规函数的正常的@jit装饰器，Numba还提供了一个@vectorize装饰器，用于从“内核函数”创建ufunc。 这个内核函数（不要与CUDA内核混淆）是一个标量函数，它描述了对所有输入的数组元素执行的操作。 例如，我可以实现一个高斯分布：

import numba
import math
import numpy as np
 
SQRT_TWOPI = np.float32(math.sqrt(2 * math.pi))
 
@numba.vectorize(['float32(float32, float32, float32)'], target='cuda')
def gaussian(x, x0, sigma):
    return math.exp(-((x - x0) / sigma)**2 / 2) / SQRT_TWOPI / sigma

与正常函数编译器不同，我需要给ufunc编译器一个参数的类型签名列表。 现在我可以用NumPy数组来调用这个函数并返回数组结果：

x = np.linspace(-3, 3, 10000, dtype=np.float32)
g = gaussian(x, 0, 1)  # 1D result
 
x2d = x.reshape((100,100))
g2d = gaussian(x2d, 0, 1) # 2D result

我不必使用特殊的内核或者选择启动配置来启动调用约定。 Numba自动处理所有的CUDA详细信息，并将输入的数组从CPU复制到GPU，结果返回给CPU。 （或者，我可以通过GPU设备内存，并避免CUDA内存复制。）

请注意，在第一个调用中，x是一个1D数组，x0和sigma是标量。 标量被Numba隐含地视为1D数组，以通过称为广播的过程匹配另一个输入参数。 广播是NumPy的一个非常强大的概念，可用于组合不同但兼容的维度的数组。 Numba自动处理所有的并行化和循环，无论你的函数输入的尺寸如何。

要了解有关ufuncs和Numba的更多信息，请查看有关广播的NumPy文档和ufuncs上的Numba文档。

5.Numba配有一个CUDA模拟器

调试CUDA应用程序是棘手的，并且Python增加了一层复杂性。 使用Python和C中的函数调用堆栈，以及在CPU和GPU上运行的代码，都没有一个适合所有调试的解决方案。 因此，Numba开发人员一直在寻找新的方法来促进CUDA Python应用程序的调试。

几年前，我们引入了一个Numba功能，我们称之为CUDA模拟器。 模拟器的目的是直接在Python解释器中运行CUDA内核，以便使用标准的Python工具进行调试更容易。 有几个注意事项适用：

• 该模拟器旨在完全在Python解释器中重现并行内核执行的逻辑行为，但不能模拟GPU硬件特性。
• 模拟器并不是一个应用程序的高效CPU代码路径。 内核运行速度非常慢，只能用于测试目的。
• 在模拟器中运行的功能可以包含通常不允许在GPU上的代码。 这允许该函数执行诸如调用PDB（Python调试器）或执行其他日志记录。
• 模拟器可能不会重现设备上存在的竞争条件。

要调用CUDA模拟器，你必须在启动Python应用程序之前将NUMBA_ENABLE_CUDASIM环境变量设置为1。 这将迫使所有内核通过解释器代码路径。 你可以在CUDA模拟器的Numba文档中找到更多信息。

当然，这不是Numba中唯一可用的CUDA调试选项。 Numba还允许使用标准的Python打印函数/语句从GPU（常量字符串和标量）进行有限的打印。 另外，你可以使用nvprof（CUDA命令行剖析器），NVIDIA Visual Profiler和cuda-memcheck来运行Numba应用程序。 传递debug = True到@ numba.cuda.jit装饰器将允许cuda-memcheck显示检测到的内存错误的Python源代码行号。

6.你可以通过网络发送Numba功能

用于Python的分布式计算系统（如Dask和Spark Python API）通过分散许多工作人员的数据并将代码带到数据所在的位置来实现高性能。 这需要将代码序列化并通过网络进行传输的能力。 在Python中，分布式框架通常使用cloudpickle库（Python pickle模块的增强版），将对象（包括函数）转换为字节流。 这些字节可以从用户输入的功能的客户端发送到远程工作进程，并将它们转回到可执行功能中。

Numba编译的CPU和GPU功能（但不是ufuncs，由于一些技术问题）是专门设计来支持pickling。 当Numba编译的GPU功能被pickle时，NVVM IR和PTX都保存在序列化的字节流中。 一旦将此数据传输到远程工作人员，该功能将在内存中重新创建。 如果工作人员的GPU的CUDA架构与客户端匹配，则将使用PTX版本的功能。 如果CUDA架构不匹配，则CUDA内核将从NVVM IR重新编译，以确保最佳性能。 图2显示了这个过程。 最终的结果是您可以在移动开普勒GPU上测试和调试GPU代码，然后无缝地将其发送到Pascal GPU的Dask群集。

图2.Numba如何通过网络将GPU功能传输给群集工作人员。

这里有一个简短的例子，这里有一些启动本地Dask群集的代码，并使用dask.distributed futures API执行一个简单的CUDA内核：

@numba.cuda.jit
def gpu_cos(x, out):
    # Assuming 1D array
    start = numba.cuda.grid(1)
    stride = numba.cuda.gridsize(1)
    for i in range(start, x.shape[0], stride):
        out[i] = math.cos(x[i])
        
def do_cos(x):
    out = numba.cuda.device_array_like(x)
    gpu_cos[64, 64](x, out)
    return out.copy_to_host()
 
# check if works locally first
test_x = np.random.uniform(-10, 10, 1000).astype(np.float32)
result = do_cos(test_x)
 
# now try remote
from dask.distributed import Client
client = Client() # starts a local cluster
 
future = client.submit(do_cos, test_x)
gpu_result = future.result()

虽然此示例执行的工作量很小，但它显示了使用分布式系统的Numba的一般模式。 提交到集群的函数是一个常规的Python函数，它内部调用一个CUDA函数。 包装器函数提供了一个分配GPU内存并确定CUDA内核启动配置的地方，分布式框架无法实现。 当do_cos提交到群集时，cloudpickle还会检测到对gpu_cos函数的依赖性并将其序列化。 这确保do_cos具有在远程工作器上运行所需的一切。 通常在使用Dask时，我们倾向于更高级别的API来构建计算图，如dask.delayed，但是对于一些迭代算法，直接使用 future是最直接的方法。

Numba社区认为分布式GPU使用Numba计算仍然具有出色的优势。 在这种情况下，Numba和Dask绝对有一些改进，所以如果你尝试使用此功能，请与Google Group上的Numba社区联系，以便我们更多地了解你的需求并提供指导。

7.Numba开发人员正在使用GPU DataFrame

在2017年的GTC 大会上，与H2O，MapD，BlazingDB，Graphistry和Gunrock合作的Anaconda公司（Numba开发的主要赞助商）宣布成立GPU开放分析计划（简称“GOAI”）。我们都认识到需要在应用程序和库之间进行GPU数据交换，因为数据科学工作负载越来越需要多种工具的组合。GPU计算已经无处不在，所以我们不能再把GPU当作一个隐藏的实现细节。现在是更多应用程序和库公开允许直接在组件之间传递GPU内存的接口的时候了。想要对GOAI的深入了解，请查看GOAI项目中的NVIDIA Developer Blog文章。

团队成员自2017年3月起就一直在工作，最近还与Apache Arrow项目的Wes McKinney合作，共同创建了可以在应用程序和库之间共享的GPU DataFrame。 GPU DataFrame实现使用Arrow格式来表示GPU上的表格数据，我们希望将来将大部分实现直接转移到Arrow代码库中。作为该软件堆栈的一部分，Numba开发人员已经创建了PyGDF，这是一个用于使用Pandas API子集来操作GPU DataFrames的Python库。该库支持过滤，排序，列数学运算，缩减，加入，按组合运算，以及与其他进程零拷贝共享GPU DataFrames。为了实现这一点，PyGDF使用Numba to JIT编译CUDA内核以进行自定义分组，缩减和过滤操作。此外，PyGDF列可以传递给Numba CUDA函数，以执行不能表示为DataFrame操作的自定义转换。

到目前为止，GOAI已经取得了很大的进展，但为时尚早，我们还有更多的工作要做。 还有一些令人兴奋的未来！ 要了解GOAI活动的信息，请加入GOAI Google Group。

了解Numba的更多信息

我希望这篇文章向你展示了关于Numba的一些以前不了解的内容。如果你想了解更多有关这些高级Numba主题的信息，我建议以下资源：

• Anaconda Python Distribution：安装Numba的最简单的方法
  
• 来自2017年GTC大会的Numba GPU教程的材料
  
• Numba CUDA文档
  
• Gbub上的Numba问题跟踪器：用于错误报告和功能请求
  
• Numba博客简介
  

此外，如果你想提问或获取Numba的帮助，最好的地方是Numba Users Google Group。

本文由北邮@爱可可-爱生活老师推荐，阿里云云栖社区组织翻译。

文章原标题《Seven Things You Might Not Know about Numba》，作者：Stanley Seibert，Anaconda社区创新总监，曾是Mobi的首席数据科学家。译者：董昭男，审校：

文章为简译，更为详细的内容，请查看原文