使用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数组，或者你正在编写的代码很难以矢量化方式开始编写。