使用numba在GPU上运行Python

Question

我试图在NVIDIA GPU中运行python代码，而谷歌搜索似乎告诉我numbapro是我正在寻找的模块。 然而，根据这个 ，numbapro不再继续，但已被移动到numba库。 我尝试了numba，它的@jit装饰器似乎确实可以大大加快我的某些代码的速度。 但是，当我阅读更多内容时，在我看来jit只是在运行时编译您的代码，并且这样做会进行一些繁重的优化并因此提高了速度。

jit似乎并不能加快已经优化的numpy操作（例如numpy.dot等）的速度， jit进一步加强了这一点。

我会感到困惑并偏离这里的轨道吗？ 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类的np.dot 。 通常，Numba擅长于此，您可以避免创建中间临时numpy数组，或者如果您正在编写的代码很难以矢量化的方式开始编写。