удалить нулевые значения массива в Numba Cuda - PullRequest
Новая
       6

удалить нулевые значения массива в Numba Cuda

0 голосов
/ 21 января 2019

у меня длинный массив 

arr = np.array([1,1,2,2,3,3,0,0,2,2])

Я бы хотел удалить все нулевые значения этого массива в numba cuda, поскольку реальный массив очень большой, а numpy очень медленный.

Кто-нибудь знает, как это можно сделать?

1 Ответ

0 голосов
/ 21 января 2019

Как уже упоминалось в комментариях, алгоритм, который вы ищете, называется Stream Compaction . Поскольку указанный алгоритм не доступен в numba, поэтому он должен быть реализован вручную. В основном, процесс выглядит следующим образом:

  1. Создание двоичной маски входных значений для определения ненулевых значений.
  2. Рассчитать префиксную сумму маски
  3. Для каждого ненулевого значения скопируйте его в индекс, указанный выводом суммы префикса в соответствующем индексе значения.

Для объяснения шага 3 рассмотрим следующий пример:

Допустим, в индексном номере 5 входного массива есть ненулевое значение. Мы выбираем значение по индексу № 5 массива префиксных сумм (допустим, это значение было 3). Используйте это значение в качестве выходного индекса. Это означает, что элемент с индексом № 5 на входе перейдет на индекс № 3 в конечном выводе.

Расчет суммы префикса - сложная часть всего этого процесса. Я позволил себе перенести версию алгоритма суммирования префиксов на C ++ (адаптированную из книги Gems 3 по GPU) на numba. Если мы сами реализуем префиксную сумму, мы можем объединить шаги 1 и 2 в одно ядро ​​CUDA.

Ниже приведен полный рабочий пример сжатия потока на основе numba с предоставленным вариантом использования. 

import numpy as np
from numba import cuda, int32

BLOCK_SIZE = 8

#CUDA kernel to calculate prefix sum of each block of input array
@cuda.jit('void(int32[:], int32[:], int32[:], int32)')
def prefix_sum_nzmask_block(a, b, s, length):
    ab = cuda.shared.array(shape=(BLOCK_SIZE), dtype=int32)

    tid = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x;

    if tid < length:
        ab[cuda.threadIdx.x] = int32(a[tid] != 0); #Load mask of input data into shared memory

    i = 1
    while i<=cuda.threadIdx.x:
        cuda.syncthreads() #Total number of cuda.synchthread calls = log2(BLOCK_SIZE).
        ab[cuda.threadIdx.x] += ab[cuda.threadIdx.x - i] #Perform scan on shared memory
        i *= 2

    b[tid] = ab[cuda.threadIdx.x]; #Write scanned blocks to global memory

    if(cuda.threadIdx.x == cuda.blockDim.x-1):  #Last thread of block
        s[cuda.blockIdx.x] = ab[cuda.threadIdx.x]; #Write last element of shared memory into global memory



#CUDA kernel to merge the prefix sums of individual blocks
@cuda.jit('void(int32[:], int32[:], int32)')
def prefix_sum_merge_blocks(b, s, length):
    tid = (cuda.blockIdx.x + 1) * cuda.blockDim.x + cuda.threadIdx.x; #Skip first block

    if tid<length:
        i = 0
        while i<=cuda.blockIdx.x:
            b[tid] += s[i] #Accumulate last elements of all previous blocks
            i += 1


#CUDA kernel to copy non-zero entries to the correct index of the output array
@cuda.jit('void(int32[:], int32[:], int32[:], int32)')
def map_non_zeros(a, prefix_sum, nz, length):
    tid = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x;

    if tid < length:
        input_value = a[tid]
        if input_value != 0:
            index = prefix_sum[tid] #The correct output index is the value at current index of prefix sum array
            nz[index-1] = input_value



#Apply stream compaction algorithm to get only the non-zero entries from the input array
def get_non_zeros(a):   
    length = a.shape[0]
    block = BLOCK_SIZE
    grid = int((length + block - 1)/block)

    #Create auxiliary array to hold the sum of each block
    block_sum = cuda.device_array(shape=(grid), dtype=np.int32)

    #Copy input array from host to device
    ad = cuda.to_device(a)

    #Create prefix sum output array
    bd = cuda.device_array_like(ad)

    #Perform partial scan of each block. Store block sum in auxillary array named block_sum.
    prefix_sum_nzmask_block[grid, block](ad, bd, block_sum, length)

    #Add block sum to the output
    prefix_sum_merge_blocks[grid, block](bd,block_sum,length);

    #The last element of prefix sum contains the total number of non-zero elements
    non_zero_count = int(bd[bd.shape[0]-1])

    #Create device output array to hold ONLY the non-zero entries
    non_zeros = cuda.device_array(shape=(non_zero_count), dtype=np.int32)

    #Copy ONLY the non-zero entries
    map_non_zeros[grid, block](a, bd, non_zeros, length)

    #Return to host
    return non_zeros.copy_to_host()


if __name__ == '__main__':

    arr = np.array([1,1,2,2,3,3,0,0,2,2], dtype=np.int32)

    nz = get_non_zeros(arr)

    print(nz)

Протестировано в следующих условиях:

Python 3.6.7 в виртуальной среде

CUDA 10.0.130

Драйвер NVIDIA 410.48

нумба 0,42

Ubuntu 18.04

Отказ от ответственности: Код предназначен только для демонстрационных целей и не был профилирован / строго проверен для измерения производительности.

...