2

我想使用Numba 的guvectorize方法在我的 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

将目标更改为以未针对 Generalized CUDA ufuncs'cuda'记录的方式彻底改变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 设备为目标时，有没有办法将变量指定为输入和输出？

4

0 回答 0