Как Numba `guvectorize` нацелен на CUDA, как мне указать переменную как входную и выходную?

Я хочу использовать Нумбу guvectorize метод для запуска кода на моей карте CUDA. Я сначала определил метод процессора

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

На самом деле, при ориентации на процессор, Anew мутирован на месте, поэтому не было необходимости назначать вывод res_cpu

>>> res_cpu is Anew
True

Изменение цели на 'cuda' резко меняет guvectorize поведение способом, не задокументированным для обобщенных функций CUDA. Вот модифицированное определение 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, есть ли способ указать переменную как входной, так и выходной?

0 ответов

Другие вопросы по тегам