Question

我的数组很长

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

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

有人知道该怎么做吗？

Answer 1

正如评论中已经提到的那样，您正在寻找的算法称为流压缩。由于无法在numba中直接使用上述算法，因此必须手动实现。基本上，过程如下：

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

有关第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

免责声明：该代码仅用于演示目的，尚未针对性能测量进行严格的配置文件/测试。

删除numba cuda中数组的零值

1 个答案: