Как получить текущий индекс рабочей сетки в Cuda Python

Question

Привет, я пытаюсь понять каждый шаг ядра CUDA.Хорошо будет получить все индексы сетки, которые заняты данными.Мой код добавить 2 вектора и написан на Python Numba.

n = 10 
x = np.arange(n).astype(np.float32)
y = x + 1   

установка количества потоков и блоков в сетке

threads_per_block = 8
blocks_per_grid = 2

Ядро

def kernel_manual_add(x, y, out):
    threads_number = cuda.blockDim.x
    block_number = cuda.gridDim.x

    thread_index = cuda.threadIdx.x
    block_index = cuda.blockIdx.x

    grid_index = thread_index + block_index * threads_number  
    threads_range = threads_number * block_number

    for i in range(grid_index, x.shape[0], threads_range):
        out[i] = x[i] + y[i]

Инициализация ядра:

kernel_manual_add[blocks_per_grid, threads_per_block](x, y, out)

Когда япопробуйте распечатать grid_index я получаю все входные индексы 2 * 8.

Как получить сеточные индексы (10 из них), которые используются для вычисления данных?

Answer 1

Канонический способ написания вашего ядра будет выглядеть примерно так

@cuda.jit
def kernel_manual_add(x, y, out):

    i = cuda.grid(1)
    if i < x.shape[0]:
        out[i] = x[i] + y[i]

Вы должны запустить как минимум столько потоков, сколько есть элементов во входных массивах. Здесь нет никакой магии, вам нужно вычислить размеры сетки и блока вручную перед вызовом ядра. См. здесь и здесь для предложений.