import numpy as np
from numba import cuda, int32
import time
BSP2 = 9
BLOCK_SIZE = 2**BSP2
#CUDA kernel to calculate prefix sum of each block of input array
@cuda.jit('void(int32[:], int32[:], int32[:], int32, int32)')
def prefix_sum_nzmask_block(a, b, s, nzm, length):
ab = cuda.shared.array(shape=(BLOCK_SIZE), dtype=int32)
tid = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x;
ab[cuda.threadIdx.x] = 0
if tid < length:
if nzm == 1:
ab[cuda.threadIdx.x] = int32(a[tid] != 0); #Load mask of input data into shared memory
else:
ab[cuda.threadIdx.x] = int32(a[tid]); #Load input data into shared memory
for j in range(0,BSP2):
i = 2**j
cuda.syncthreads()
if i <= cuda.threadIdx.x:
temp = ab[cuda.threadIdx.x]
temp += ab[cuda.threadIdx.x - i] #Perform scan on shared memory
cuda.syncthreads()
if i <= cuda.threadIdx.x:
ab[cuda.threadIdx.x] = temp
if tid < length:
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 pref_sum_update(b, s, length):
tid = (cuda.blockIdx.x + 1) * cuda.blockDim.x + cuda.threadIdx.x; #Skip first block
if tid<length:
b[tid] += s[cuda.blockIdx.x] #Accumulate last elements of all previous blocks
#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 pref_sum(a, asum, nzm):
block = BLOCK_SIZE
length = a.shape[0]
grid = int((length + block -1)/block)
#Create auxiliary array to hold the sum of each block
bs = cuda.device_array(shape=(grid), dtype=np.int32)
#Perform partial scan of each block. Store block sum in auxillary array named block_sum.
prefix_sum_nzmask_block[grid, block](a, asum, bs, nzm, length)
if grid > 1:
bssum = cuda.device_array(shape=(grid), dtype=np.int32)
pref_sum(bs, bssum, 0)
pref_sum_update[grid-1, block](asum, bssum, length)
def get_non_zeros(a):
#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.
pref_sum(ad, bd, int(1))
#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
block = BLOCK_SIZE
length = a.shape[0]
grid = int((length + block -1)/block)
map_non_zeros[grid, block](ad, bd, non_zeros, length)
#Return to host
return non_zeros.copy_to_host()
if __name__ == '__main__':
arr = np.zeros(500000000, dtype=np.int32)
for i in range(32,65000, 1024):
arr[i] = i
start_time=time.time()
nz = get_non_zeros(arr)
print(nz)
print("time elapsed=",time.time()-start_time,"s")
The above code takes approx. 2 seconds to run on Orin Nano in all the three power modes.
Also, all CPU cores are running on all power modes as per jtop utility.
Thanks,