删除numba cuda中数组的零值

Question

我有一个很长的数组

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

我想在 numba cuda 中删除该数组的所有零值，因为实际数组非常大而 numpy 非常慢。

有谁知道如何做到这一点？

Answer 1

正如评论中已经提到的，您正在寻找的算法称为 Stream Compaction。由于上述算法在 numba 中不是开箱即用的，因此必须手动实现。基本上，过程如下：

创建输入值的二进制掩码以识别非零值。
计算mask的前缀和
对于每个非零值，将其复制到该值对应索引处前缀和输出指定的索引。

对于第 3 步的解释，请考虑以下示例：

假设输入数组的索引号 5 处有一个非零值。我们在前缀和数组的索引号 5 处选择值（假设该值为 3）。使用此值作为输出索引。这意味着输入的索引号 5 处的元素将转到最终输出中的索引号 3。

前缀和的计算是整个过程中的难点。我冒昧地将前缀和算法的 C++ 版本（改编自 GPU Gems 3 书）移植到 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

numba 0.42

Ubuntu 18.04

免责声明： 该代码仅用于演示目的，并未profiled/tested 严格地进行性能测量。

删除numba cuda中数组的零值

remove zero values of an array in numba cuda

python

cuda

numba