Numba的`guvectorize`针对CUDA，如何将变量指定为输入和输出？

Question

我想使用Numba's guvectorize method在我的CUDA卡上运行代码。我首先定义了一个CPU方法

from numba import guvectorize
import numpy as np

@guvectorize(['float32[:,:], float32[:,:]',
              'float64[:,:], float64[:,:]'],
             '(n,m)->(n,m)', nopython=True, target='cpu')
def update_a_cpu(A, Anew):
    n, m = A.shape
    for j in range(1, n-1):
        for i in range(1, m-1):
            Anew[j, i] = 0.25 * (A[j, i+1] + A[j, i-1] + A[j-1, i] + A[j+1, i])

给出测试矩阵的预期输出

>>> A = np.arange(16, dtype=np.float32).reshape(4,4)  # single precision for GTX card
>>> Anew = np.zeros((4,4), dtype=np.float32)

>>> res_cpu = update_a_cpu(A, Anew)

>>> print(res_cpu)
[[  0.   0.   0.   0.]
 [  0.   5.   6.   0.]
 [  0.   9.  10.   0.]
 [  0.   0.   0.   0.]]

实际上，在定位CPU时，Anew已发生变异，因此无需将输出分配给res_cpu

>>> res_cpu is Anew
True

将目标更改为'cuda'会以未记录的方式for Generalized CUDA ufuncs大幅改变guvectorize行为。这是修改后的ufunc定义

@guvectorize(['float32[:,:], float32[:,:]',
              'float64[:,:], float64[:,:]'],
             '(n,m)->(n,m)', nopython=True, target='cuda')
def update_a_cuda(A, Anew):
    n, m = A.shape
    for j in range(1, n-1):
        for i in range(1, m-1):
            Anew[j, i] = 0.25 * (A[j, i+1] + A[j, i-1] + A[j-1, i] + A[j+1, i])

现在该函数不接受第二个输入矩阵

>>> res_cuda = update_a_cuda(A, Anew)
... 
TypeError: invalid number of input argument

而是创建一个空矩阵以将值放入

>>> res_cuda = update_a_cuda(A)
>>> print(res_cuda)
array([[  1.55011636e-41,   1.55011636e-41,   1.55011636e-41,   1.55011636e-41],
       [  1.55011636e-41,   5.00000000e+00,   6.00000000e+00,   1.55011636e-41],
       [  1.55011636e-41,   9.00000000e+00,   1.00000000e+01,   1.55011636e-41],
       [  1.55011636e-41,   1.55011636e-41,   1.55011636e-41,   1.55011636e-41]], dtype=float32)

我希望通用的ufunc更新输入矩阵的适当值，而不是填充空矩阵。在定位CUDA设备时，有没有办法将变量指定为输入和输出？