使用 numba 在 GPU 上运行 Python

Question

我正在尝试在我的 NVIDIA GPU 中运行 python 代码，谷歌搜索似乎告诉我 numbapro 是我正在寻找的模块。然而，根据this，numbapro 不再继续，而是已移至 numba 库。我尝试了 numba，它的

@jit

装饰器似乎确实大大加快了我的一些代码的速度。然而，当我阅读更多内容时，在我看来，

jit

只是在运行时编译代码，这样做时，它会进行一些大量的优化，从而提高速度。

jit

似乎并没有加速已经优化的numpy操作（例如

numpy.dot

等）这一事实进一步强化了这一点。

我是否感到困惑并且偏离了轨道？

jit

到底有什么作用？如果它不能让我的代码在 GPU 上运行，我还能怎么做？

Answer 1

您必须通过 ufunc 明确告诉 Numba 以 GPU 为目标：

http://numba.pydata.org/numba-doc/latest/cuda/ufunc.html

或者通过明确考虑 GPU 的方式对函数进行编程：

http://numba.pydata.org/numba-doc/latest/cuda/examples.html http://numba.pydata.org/numba-doc/latest/cuda/index.html

普通的

jit

函数不针对 GPU，通常不会加速对

np.dot

之类的调用。通常，Numba 擅长于避免创建中间临时 numpy 数组，或者您正在编写的代码一开始就很难以矢量化方式编写的情况。

Answer 2

用固定链接编辑了@JoshAdel 答案。由于编辑队列已满，无法提交编辑。

您必须通过 ufunc 明确告诉 Numba 以 GPU 为目标：

https://numba.readthedocs.io/en/stable/cuda/ufunc.html

或者通过明确考虑 GPU 的方式对函数进行编程：

https://numba.readthedocs.io/en/stable/cuda/examples.html https://numba.readthedocs.io/en/stable/cuda/

普通的

jit

函数不针对 GPU，通常不会加速对

np.dot

之类的调用。通常，Numba 擅长于避免创建中间临时 numpy 数组，或者您正在编写的代码一开始就很难以矢量化方式编写的情况。