Обработка и создание массива в функции устройства numba.cuda - PullRequest
Новая
       31

Обработка и создание массива в функции устройства numba.cuda

1 голос
/ 17 июня 2019

Я хочу передать функции устройства кусок массива, а затем создать там новый массив и вернуть их комбинацию.Но это кажется не общим способом решения таких проблем с помощью cuda, потому что numba не поддерживает списки и массивы внутри функции устройства.Таким образом, как я должен обойти это ограничение? 

import numpy as np
from numba import cuda
from numba import *


@cuda.jit(device=True)
def comp_array_2(y): # y is a 3d vector
    """Returns y + 1"""
    new_arr = y
    mod_arr = new_arr+1
    return mod_arr

@cuda.jit(device=True)
def comp_array_1(y): # y is a 3d vector
    """Returns y*(y + 1)"""
    new_arr = y
    mod_arr = comp_array_2(new_arr)
    return [mod_arr[0]*new_arr[0],
            mod_arr[1]*new_arr[1],
            mod_arr[2]*new_arr[2]]  

@cuda.jit
def euler_sol_cuda(init, sol): 
    start = cuda.grid(1)      
    stride = cuda.gridsize(1) 

    # Threads loop
    for i in range(start, init.shape[0], stride): 
         sol[i] = comp_array_1(init[i]) #



if __name__=="__main__":
    # CPU
    n = 128*1024 
    blocks_per_grid = 32
    threads_per_block = 64

    # To GPU
    d_init = cuda.to_device(np.random.rand(n,3))
    d_sol = cuda.device_array_like(d_init)

    euler_sol_cuda[blocks_per_grid, threads_per_block](d_init, d_sol)
    print(sol_d.copy_to_host())

Желаемый вывод - это массив размером (2 ^ 17, 3), который похож на последовательный код numpy: 

d_init = np.random.rand(5,3)
print(d_init)
d_init * (d_init + 1)

Вывод 

`[[0.44905369 0.28315775 0.06539836]
 [0.71547266 0.77073541 0.76203624]
 [0.82483726 0.14611916 0.8317247 ]
 [0.22490635 0.83327892 0.48090544]
 [0.79825294 0.07037226 0.5455043 ]]`

`array([[0.6507029 , 0.36333606, 0.06967531],
       [1.2273738 , 1.36476849, 1.34273547],
       [1.50519376, 0.16746996, 1.52349068],
       [0.27548921, 1.52763267, 0.71217548],
       [1.43546068, 0.07532451, 0.84307924]])
`
...