如何在 GPU 上运行 numba.jit 修饰函数？

Question

我正在努力进入numba gpu processing。我有这个MWE：

import numpy as np
import numba

@numba.njit
def function():
    ar = np.zeros((3, 3))
    for i in range(3):
        ar[i] = (1, 2, 3)
    return ar

ar = function()
print(ar)

输出：

[[1. 2. 3.]
 [1. 2. 3.]
 [1. 2. 3.]]

现在我想在我的gpu. 我尝试使用以下decorators：

@numba.njit(target='cuda')
@numba.njit(target='gpu')
@numba.cuda.jit

这些都不起作用。以下是上面的错误消息decorators：

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 4, in <module>
    @numba.jit(target='cuda')
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 171, in jit
    targetoptions=options, **dispatcher_args)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 179, in _jit
    dispatcher = registry.dispatcher_registry[target]
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/registry.py", line 96, in __getitem__
    return super(TargetRegistry, self).__getitem__(item)
KeyError: 'cuda'

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 4, in <module>
    @numba.njit(target='gpu')
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 236, in njit
    return jit(*args, **kws)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 171, in jit
    targetoptions=options, **dispatcher_args)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 179, in _jit
    dispatcher = registry.dispatcher_registry[target]
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/registry.py", line 96, in __getitem__
    return super(TargetRegistry, self).__getitem__(item)
KeyError: 'gpu'

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 4, in <module>
    @numba.cuda.jit()
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/__init__.py", line 140, in __getattr__
    ) from None
AttributeError: module 'numba' has no attribute 'cuda'

numba.cuda甚至不被认为是module function。我已经numba 49.1安装cudatoolkit 9.0了。

我是否必须更改function才能使其正常工作？我有一个巨大的numba.njit函数需要在gpu.

我试图了解更多信息numba's cuda implementation： http://numba.pydata.org/numba-doc/0.16.0/modules/numba.cuda.html

先感谢您。

编辑：

正如 @talonmies 提议的那样，我cuda从numba模块中显式导入并外包了数组创建：

import numpy as np
import numba
from numba import cuda

@numba.njit(target='cuda')
def function(ar=None):
    for i in range(3):
        ar[i] = (1, 2, 3)
    return ar

ar = np.zeros((3, 3))
ar_result = function(ar=ar)
print(ar_result)

输出：

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 12, in <module>
    ar_result = function(ar=ar)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/dispatcher.py", line 40, in __call__
    return self.compiled(*args, **kws)
TypeError: __call__() got an unexpected keyword argument 'ar'

除了 之外，error上述所有情况都会发生这种情况。decorators@numba.njit

编辑_2：

当我尝试运行时：

import numpy as np
import numba
from numba import cuda

@numba.jit(target='cuda')
def function(ar):
    for i in range(3):
        ar[i] = (1,2,3)

ar = np.zeros((3, 3))
function(ar)
print(ar)

输出是：

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 11, in <module>
    function(ar)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/dispatcher.py", line 40, in __call__
    return self.compiled(*args, **kws)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/compiler.py", line 758, in __call__
    kernel = self.specialize(*args)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/compiler.py", line 769, in specialize
    kernel = self.compile(argtypes)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/compiler.py", line 785, in compile
    **self.targetoptions)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
TypeError: compile_kernel() got an unexpected keyword argument 'boundscheck'

Answer 1

三点：

1. 你必须从 numba 显式导入 cuda 模块才能使用它（这不是 numba 特有的，所有 python 库都是这样工作的）
2. nopython 模式 ( njit) 不支持 CUDA 目标
3. CUDA 代码的 Numba 不支持数组创建、返回值、关键字参数

我可以像这样解决所有问题：

...: import numpy as np 
...: import numba 
...: from numba import cuda 
...:  
...: @numba.jit(target='cuda') 
...: def function(ar): 
...:     for i in range(3): 
...:         ar[i] = (1,2,3)   
...:  
...: ar = np.zeros((3, 3)) 
...: function(ar) 
...: print(ar) 

[[1. 2. 3.]
 [1. 2. 3.]
 [1. 2. 3.]]