Как 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, есть ли способ указать переменную как входной, так и выходной?