GPU Kernels in Python (Numba/CuPy Style)

What You’re Building (Theory)

Python GPU programming usually follows this shape:

1. Create device data (allocate/copy to GPU)
2. Launch a kernel with enough threads to cover the work
3. Copy results back to host for verification or further Python processing

Two common approaches:

• Numba CUDA: write kernels with a Python-like syntax, compiled to GPU
• CuPy: NumPy-like API that runs operations on GPU (and supports custom raw kernels)

Code Example 1 — Numba CUDA Vector Add (Beginner)

# vector_add_numba.py
import numpy as np
from numba import cuda

@cuda.jit
def vec_add(a, b, c):
    i = cuda.grid(1)  # global thread index
    if i < c.size:
        c[i] = a[i] + b[i]

def main():
    N = 1 << 20
    a = np.ones(N, dtype=np.float32)
    b = np.full(N, 2.0, dtype=np.float32)

    # Copy inputs to GPU (device arrays)
    d_a = cuda.to_device(a)
    d_b = cuda.to_device(b)
    d_c = cuda.device_array_like(a)

    threads_per_block = 256
    blocks_per_grid = (N + threads_per_block - 1) // threads_per_block

    # Launch kernel
    vec_add[blocks_per_grid, threads_per_block](d_a, d_b, d_c)

    # Copy result back
    c = d_c.copy_to_host()

    print("c[0] =", c[0], "c[N-1] =", c[-1])

if __name__ == "__main__":
    main()

Run:

python vector_add_numba.py

Code Example 2 — Correctness Guard Patterns

For GPU kernels, the if i < N: guard is the difference between:

• correct results
• and memory access errors

Common pattern:

@cuda.jit
def scale_kernel(x, y, alpha):
    i = cuda.grid(1)
    if i < y.size:
        y[i] = alpha * x[i]

Common Gotchas

• Forgetting bounds checks: grid size often rounds up; excess threads must be guarded.
• Accidentally mixing host/device arrays: kernels expect device arrays.
• Assuming prints from inside kernels work normally: debug output is limited and slow; validate with host-side checks.

Quick Checklist

• Use cuda.to_device() / device_array_like() for device allocations
• Launch with enough threads to cover N
• Always guard out-of-range threads
• Copy results back for verification