remove zero values of an array in numba cuda

I have a long array

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

I would like to remove all zero values of this array in numba cuda since the real array is very large and numpy is very slow.

Does anyone know how this can be done?


  • As already mentioned in the comments, the algorithm that you are looking for is called Stream Compaction. Since the said algorithm is not available out-of-the-box in numba, so it has to be implemented by hand. Basically, the process is as follows:

    1. Create binary mask of input values to identify non-zero values.
    2. Calculate prefix sum of the mask
    3. For each non-zero value, copy it to the index specified by the output of prefix sum at the corresponding index of the value.

    For step 3 explanation, consider the following example:

    Lets say there is a non-zero value at index number 5 of the input array. We pick the value at index number 5 of the prefix sum array (lets say that value was 3). Use this value as the output index. It means that element at index number 5 of the input will go to index number 3 in the final output.

    Calculation of prefix sum is the difficult part in this whole process. I took the liberty of porting the C++ version of prefix sum algorithm (adapted from the GPU Gems 3 book) to numba. If we implement prefix sum ourselves, we can merge step 1 and 2 into a single CUDA kernel.

    Following is a complete working example of numba based stream compaction with the provided use case.

    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)

    Tested with the following environment:

    Python 3.6.7 in virtual environment

    CUDA 10.0.130

    NVIDIA Driver 410.48

    numba 0.42

    Ubuntu 18.04

    Disclaimer: The code is for demo purpose only and has not been profiled/tested rigorously for performance measurements.