numba_cuda_elementwise_array_addition_kernel_quickstart.py

1from numba import cuda
2import numpy as np
3
4@cuda.jit
5def add_kernel(x, y, out):
6    # cuda.grid(1) returns the unique index of the current thread in the 1D grid
7    tx = cuda.grid(1)
8    
9    # Check if the index is within the bounds of the array
10    if tx < x.size:
11        out[tx] = x[tx] + y[tx]
12
13# Initialize data
14n = 64
15x = np.arange(n).astype(np.float32)
16y = 2 * x
17out = np.empty_like(x)
18
19# Configure the blocks and threads
20threads_per_block = 32
21blocks_per_grid = (n + (threads_per_block - 1)) // threads_per_block
22
23# Launch the kernel
24add_kernel[blocks_per_grid, threads_per_block](x, y, out)
25
26# Print the result
27print(out)