With Numba's `guvectorize` targeted to CUDA, how do I specify a variable as both input and output?

Question

I want to use Numba's guvectorize method to run code on my CUDA card. I first defined a CPU method

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])

which gives the expected output for a test matrix

>>> 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.]]

Actually, when targeting the CPU, Anew is mutated in place so there was no need to assign the output to res_cpu

>>> res_cpu is Anew
True

Changing the target to 'cuda' drastically changes the guvectorize behavior in a manner not documented for Generalized CUDA ufuncs. Here is the modified ufunc definition

@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])

Now the function does not accept the second input matrix

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

and instead creates an empty matrix to put the value into

>>> 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)

I would like the generalized ufunc to update the appropriate values of an input matrix rather than populating an empty matrix. When targeting a CUDA device, is there a way to specify a variable as both input and output?

With Numba's `guvectorize` targeted to CUDA, how do I specify a variable as both input and output?

0 Answers0