如何使用Numba CUDA JIT装饰器？

Question

我按照这个教程使用了Numba的CUDA JIT装饰器：https://www.youtube.com/watch?v=-lcWV4wkHsk&t=510s。

这是我的Python代码：

import numpy as np
from timeit import default_timer as timer
from numba import cuda, jit

# This function will run on a CPU
def fill_array_with_cpu(a):
      for k in range(100000000):
            a[k] += 1

# This function will run on a CPU with @jit
@jit
def fill_array_with_cpu_jit(a):
      for k in range(100000000):
            a[k] += 1         

# This function will run on a GPU
@jit(target_backend='cuda')
def fill_array_with_gpu(a):
      for k in range(100000000):
            a[k] += 1   

# Main
a = np.ones(100000000, dtype = np.float64)

for i in range(3):
      start = timer()
      fill_array_with_cpu(a)
      print("On a CPU:", timer() - start)

for i in range(3):
      start = timer()
      fill_array_with_cpu_jit(a)
      print("On a CPU with @jit:", timer() - start)

for i in range(3):
      start = timer()
      fill_array_with_gpu(a)
      print("On a GPU:", timer() - start)

这是运行后的输出：

On a CPU: 24.228116830999852
On a CPU: 24.90354355699992
On a CPU: 24.277727688999903
On a CPU with @jit: 0.2590671719999591
On a CPU with @jit: 0.09131158500008496
On a CPU with @jit: 0.09054700799993043
On a GPU: 0.13547917200003212
On a GPU: 0.0922475330000907
On a GPU: 0.08995077999998102

使用@jit装饰器可以大大提高处理速度。不过，我不太明白@jit(target_backend='cuda')这个装饰器是否真的能让函数在GPU上运行。处理时间和使用@jit的函数差不多。我猜@jit(target_backend='cuda')并没有使用GPU。实际上，我在一台没有NVIDIA GPU的机器上试过这段代码，结果和之前一样，没有任何警告或错误。

我该怎么才能让它在我的GPU上运行呢？我有一块GeForce GT 730M显卡。

Answer 1

没有什么叫做 target_backend='cuda' 的东西。代码中的所有功能都是在CPU上执行的（所以如果不算编译时间，时间是一样的）。我知道很久以前有过这样的选项，但现在没有了。视频中的基准测试实际上是不正确的，原因有很多，我觉得不应该相信它。

---

不仅视频中的基准测试现在不正确，而且在进行测试时也存在偏见。实际上，即使它存在并且能像我们希望的那样工作，它也不会高效，因为目标数组存储在主机内存中（通常是RAM）。因此，数组必须先传输到GPU设备内存，进行计算，然后再从设备传回主机内存。问题是这种数据传输是非常耗费资源的（而且速度不能快于主机内存）。此外，计算是非常便宜的，所以CPU的计算应该是受内存限制的，尽管一个核心可能不足以充分利用RAM的带宽。大多数平台上需要使用并行CPU实现来充分利用RAM。比较并行CPU实现和GPU实现也是更好的选择，因为后者本质上是并行的。提供的基准测试充其量是有偏见的。最终，GPU实现不能比并行CPU实现快，因为数据传输的速度不能超过并行CPU实现，因为两者都受限于主机RAM。实际上，GPU实现应该更慢，因为CPU-GPU之间的连接（通常是PCIe）往往无法达到主机RAM带宽那么高的吞吐量。

最后但同样重要的是，数组的类型是float64，而所有主流的客户端Nvidia GPU并不适合这个：它们在进行64位浮点（FP）计算时非常慢。实际上，它们慢到主流的CPU计算速度更快。例如，你的GT 730M GPU（非常老旧的低端Kepler GPU）在32位FP计算时可以达到552 GFlops，而在64位FP计算时只有23 GFlops。相比之下，同年发布的i5-4258U移动CPU可以达到92 GFlops。这是快4倍！如果你想在GPU上快速进行64位FP计算，那么你需要一款支持64位FP计算的服务器级Nvidia GPU（大多数都支持）。不过要注意，这种GPU要贵得多。

注意，第一次调用Numba函数时会包括编译时间。在基准测试中必须排除这个开销（可以通过提前编译函数、缓存它或者直接忽略第一次调用的时间来实现）。

---

简而言之，这是一个糟糕的教程，你的GPU肯定无法比你的CPU更快地计算这个特定操作。我建议你阅读Numba的文档，它要可靠得多，且更新得多。你还可以阅读CUDA编程手册以获取更多信息，以及这篇维基百科页面以了解你的GPU的信息。