Optimize For Gpu

uv add "numba-cuda[cu12]==0.30.*"    # For CUDA 12.x (pulls in numba and CUDA components)
uv add "numba-cuda[cu13]==0.30.*"    # For CUDA 13.x

from numba import cuda

# Verify GPU is available
print(cuda.is_available())   # True if CUDA works
cuda.detect()                # Prints GPU details

Grid (of blocks) → Blocks (of threads) → Threads

from numba import cuda

@cuda.jit
def my_kernel(input_array, output_array):
    i = cuda.grid(1)                    # Get this thread's global index
    if i < input_array.size:            # Bounds check — ALWAYS do this
        output_array[i] = input_array[i] * 2.0

import numpy as np
from numba import cuda

data = np.random.rand(1_000_000).astype(np.float32)
out = np.zeros_like(data)

# Transfer to GPU
d_data = cuda.to_device(data)
d_out = cuda.device_array_like(out)

# Calculate launch configuration
threads_per_block = 256
blocks_per_grid = (data.size + threads_per_block - 1) // threads_per_block

# Launch
my_kernel[blocks_per_grid, threads_per_block](d_data, d_out)

# Get results back
result = d_out.copy_to_host()

@cuda.jit
def kernel_2d(matrix, output):
    x, y = cuda.grid(2)
    if x < matrix.shape[0] and y < matrix.shape[1]:
        output[x, y] = matrix[x, y] * 2.0

threads = (16, 16)
blocks = (
    (matrix.shape[0] + threads[0] - 1) // threads[0],
    (matrix.shape[1] + threads[1] - 1) // threads[1],
)
kernel_2d[blocks, threads](d_matrix, d_output)

# Automatically computes grid dimensions for 1D
my_kernel.forall(len(data))(d_data, d_out)

@cuda.jit
def process_large(data, out):
    start = cuda.grid(1)
    stride = cuda.gridsize(1)
    for i in range(start, data.shape[0], stride):
        out[i] = data[i] * 2.0

# Host → Device
d_array = cuda.to_device(numpy_array)                    # Synchronous copy
d_array = cuda.to_device(numpy_array, stream=stream)     # Async copy

# Allocate on device (no copy)
d_array = cuda.device_array(shape=(1000,), dtype=np.float32)
d_array = cuda.device_array_like(numpy_array)

# Device → Host
host_array = d_array.copy_to_host()                      # New array
d_array.copy_to_host(existing_array)                     # Into pre-allocated
d_array.copy_to_host(stream=stream)                      # Async

# Allocate pinned host memory (page-locked — faster PCI-e transfers)
with cuda.pinned(host_array):
    d_array = cuda.to_device(host_array, stream=stream)
    # Transfer is faster because the OS can't page this memory out

# Or allocate directly
pinned = cuda.pinned_array(shape=(1000,), dtype=np.float32)

with cuda.defer_cleanup():
    # All GPU deallocation deferred here — avoids implicit synchronization
    # Use this in performance-critical sections
    run_many_kernels()
# Cleanup happens here

from numba import cuda, float32

@cuda.jit
def kernel_with_shared(data, output):
    # Allocate shared memory — visible to all threads in this block
    shared = cuda.shared.array(256, dtype=float32)

    tid = cuda.threadIdx.x
    i = cuda.grid(1)

    # Each thread loads one element into shared memory
    if i < data.size:
        shared[tid] = data[i]
    else:
        shared[tid] = 0.0

    # BARRIER: wait for ALL threads in block to finish loading
    cuda.syncthreads()

    # Now safe to read any element in shared[]
    if i < data.size:
        output[i] = shared[tid] + (shared[tid - 1] if tid > 0 else 0.0)

@cuda.jit
def kernel_dynamic_shared(data):
    # size=0 means "use dynamic shared memory"
    dyn = cuda.shared.array(0, dtype=float32)
    tid = cuda.threadIdx.x
    i = cuda.grid(1)
    dyn[tid] = data[i] if i < data.size else 0.0
    cuda.syncthreads()
    # ...

# Specify size at launch (4th parameter = bytes)
kernel_dynamic_shared[blocks, threads, stream, 1024](data)  # 1024 bytes of shared mem

@cuda.jit
def kernel_with_local(data):
    # Each thread gets its own private array
    local_buf = cuda.local.array(10, dtype=float32)
    i = cuda.grid(1)
    for j in range(10):
        idx = i * 10 + j
        local_buf[j] = data[idx] if idx < data.size else 0.0
    # Process local_buf...

@cuda.jit(device=True)
def compute_distance(x1, y1, x2, y2):
    return math.sqrt((x2 - x1)**2 + (y2 - y1)**2)

@cuda.jit
def kernel(points, distances):
    i = cuda.grid(1)
    if i < points.shape[0] - 1:
        distances[i] = compute_distance(
            points[i, 0], points[i, 1],
            points[i+1, 0], points[i+1, 1]
        )

cuda.atomic.add(array, index, value)       # +=   (int32, float32, float64)
cuda.atomic.sub(array, index, value)       # -=   (int32, float32, float64)
cuda.atomic.max(array, index, value)       # max  (int/uint 32/64, float 32/64)
cuda.atomic.min(array, index, value)       # min  (same types)
cuda.atomic.nanmax(array, index, value)    # max ignoring NaN
cuda.atomic.nanmin(array, index, value)    # min ignoring NaN
cuda.atomic.and_(array, index, value)      # &=   (int/uint 32/64)
cuda.atomic.or_(array, index, value)       # |=   (int/uint 32/64)
cuda.atomic.xor(array, index, value)       # ^=   (int/uint 32/64)
cuda.atomic.exch(array, index, value)      # exchange
cuda.atomic.cas(array, index, old, value)  # compare-and-swap

@cuda.jit
def histogram(data, bins, min_value, max_value):
    i = cuda.grid(1)
    if i < data.size:
        value = data[i]
        n_bins = bins.size
        if min_value <= value <= max_value:
            bin_idx = int((value - min_value) * n_bins / (max_value - min_value))
            if bin_idx == n_bins:  # Include the rightmost edge.
                bin_idx = n_bins - 1
            cuda.atomic.add(bins, bin_idx, 1)

from numba import vectorize, float32, float64
import math

@vectorize([float32(float32, float32),
            float64(float64, float64)],
           target='cuda')
def gpu_hypot(x, y):
    return math.sqrt(x**2 + y**2)

# Usage — just call it like a NumPy ufunc
result = gpu_hypot(array_x, array_y)

# Pass device arrays to avoid transfers
d_x = cuda.to_device(x)
d_y = cuda.to_device(y)
d_result = gpu_hypot(d_x, d_y)

from numba import guvectorize, float32

@guvectorize([(float32[:,:], float32[:,:], float32[:,:])],
             '(m,n),(n,p)->(m,p)', target='cuda')
def gpu_matmul(A, B, C):
    for i in range(A.shape[0]):
        for j in range(B.shape[1]):
            total = 0.0
            for k in range(A.shape[1]):
                total += A[i, k] * B[k, j]
            C[i, j] = total

from numba import cuda

# Define reduction operation
sum_reduce = cuda.reduce(lambda a, b: a + b)

# Use it
result = sum_reduce(array)                    # Full reduction
result = sum_reduce(array, init=0)            # With initial value
sum_reduce(array, res=device_result)          # Write to device array (no D→H copy)
sum_reduce(array, stream=stream)              # Async

@cuda.reduce
def max_reduce(a, b):
    return a if a > b else b

maximum = max_reduce(data_array)

stream = cuda.stream()

# Async transfer → kernel → transfer back
d_data = cuda.to_device(host_data, stream=stream)
my_kernel[blocks, threads, stream](d_data, d_out)
result = d_out.copy_to_host(stream=stream)
stream.synchronize()  # Wait for everything on this stream

# Context manager that auto-synchronizes
with stream.auto_synchronize():
    d_data = cuda.to_device(host_data, stream=stream)
    my_kernel[blocks, threads, stream](d_data, d_out)
    result = d_out.copy_to_host(stream=stream)
# Synchronizes here automatically

stream1 = cuda.stream()
stream2 = cuda.stream()

# Chunk 1: transfer on stream1
d_chunk1 = cuda.to_device(data[:half], stream=stream1)
# Chunk 2: transfer on stream2 (overlaps with stream1 transfer)
d_chunk2 = cuda.to_device(data[half:], stream=stream2)

# Process chunk1 on stream1
kernel[blocks, threads, stream1](d_chunk1, d_out1)
# Process chunk2 on stream2 (overlaps with stream1 compute)
kernel[blocks, threads, stream2](d_chunk2, d_out2)

cuda.synchronize()  # Wait for all streams

from numba import cuda
from numba.cuda.random import (
    create_xoroshiro128p_states,
    xoroshiro128p_uniform_float32,
    xoroshiro128p_uniform_float64,
    xoroshiro128p_normal_float32,
    xoroshiro128p_normal_float64,
)

# Create RNG states — one per thread
n_threads = 256 * 128
rng_states = create_xoroshiro128p_states(n_threads, seed=42)

@cuda.jit
def monte_carlo_pi(rng_states, iterations, out):
    gid = cuda.grid(1)
    if gid < out.size:
        inside = 0
        for _ in range(iterations):
            x = xoroshiro128p_uniform_float32(rng_states, gid)
            y = xoroshiro128p_uniform_float32(rng_states, gid)
            if x**2 + y**2 <= 1.0:
                inside += 1
        out[gid] = inside / iterations * 4.0

monte_carlo_pi[128, 256](rng_states, 10000, d_out)

@cuda.jit
def iterative_kernel(M):
    col = cuda.grid(1)
    g = cuda.cg.this_grid()  # Get grid group

    for row in range(1, M.shape[0]):
        if col < M.shape[1]:
            M[row, col] = M[row - 1, col] + 1
        g.sync()  # Global barrier — all blocks wait here

# Query max grid size for cooperative launch
overload = iterative_kernel.overloads[signature]
max_blocks = overload.max_cooperative_grid_blocks(block_dim)

from numba import cuda, float32
import numpy as np

TPB = 16  # Tile/block size

@cuda.jit
def matmul_shared(A, B, C):
    sA = cuda.shared.array((TPB, TPB), dtype=float32)
    sB = cuda.shared.array((TPB, TPB), dtype=float32)

    x, y = cuda.grid(2)
    tx, ty = cuda.threadIdx.x, cuda.threadIdx.y

    tmp = float32(0.0)
    n_tiles = (A.shape[1] + TPB - 1) // TPB
    for tile in range(n_tiles):
        # Load tile into shared memory (with bounds check)
        col = tx + tile * TPB
        row = ty + tile * TPB
        sA[ty, tx] = A[y, col] if (y < A.shape[0] and col < A.shape[1]) else 0
        sB[ty, tx] = B[row, x] if (x < B.shape[1] and row < B.shape[0]) else 0
        cuda.syncthreads()

        # Compute partial dot product from this tile
        for k in range(TPB):
            tmp += sA[ty, k] * sB[k, tx]
        cuda.syncthreads()

    if y < C.shape[0] and x < C.shape[1]:
        C[y, x] = tmp

@cuda.jit
def block_inclusive_scan(data, output):
    shared = cuda.shared.array(256, dtype=float32)
    tid = cuda.threadIdx.x
    i = cuda.grid(1)

    shared[tid] = data[i] if i < data.size else 0
    cuda.syncthreads()

    # Hillis-Steele scan within one block. Two barriers prevent read/write races.
    offset = 1
    while offset < cuda.blockDim.x:
        addend = float32(0.0)
        if tid >= offset:
            addend = shared[tid - offset]
        cuda.syncthreads()
        if tid >= offset:
            shared[tid] += addend
        cuda.syncthreads()
        offset *= 2

    if i < data.size:
        output[i] = shared[tid]

@cuda.jit
def block_reduce_sum(data, partial_sums):
    shared = cuda.shared.array(256, dtype=float32)
    tid = cuda.threadIdx.x
    i = cuda.grid(1)

    shared[tid] = data[i] if i < data.size else 0.0
    cuda.syncthreads()

    # Tree reduction in shared memory
    s = cuda.blockDim.x // 2
    while s > 0:
        if tid < s:
            shared[tid] += shared[tid + s]
        s //= 2
        cuda.syncthreads()

    # Thread 0 of each block writes the block's sum
    if tid == 0:
        partial_sums[cuda.blockIdx.x] = shared[0]

@cuda.jit
def stencil_1d(data, output, radius):
    shared = cuda.shared.array(288, dtype=float32)  # blockDim + 2*radius
    tid = cuda.threadIdx.x
    i = cuda.grid(1)

    # Load center + halo into shared memory
    shared[tid + radius] = data[i] if i < data.size else 0
    if tid < radius:
        shared[tid] = data[i - radius] if i >= radius else 0
        shared[tid + cuda.blockDim.x + radius] = (
            data[i + cuda.blockDim.x] if i + cuda.blockDim.x < data.size else 0
        )
    cuda.syncthreads()

    if i < data.size:
        total = float32(0.0)
        for j in range(-radius, radius + 1):
            total += shared[tid + radius + j]
        output[i] = total / (2 * radius + 1)

export NUMBA_ENABLE_CUDASIM=1
python your_script.py

@cuda.jit
def debug_kernel(data, out):
    i = cuda.grid(1)
    if cuda.threadIdx.x == 0 and cuda.blockIdx.x == 0:
        # Only thread (0,0) hits the debugger
        from pdb import set_trace; set_trace()
    if i < data.size:
        out[i] = data[i] * 2

@cuda.jit(debug=True, opt=False)
def kernel_debug(data):
    # Enables CUDA exception checking — much slower but catches errors
    ...

import cupy as cp
from numba import cuda

@cuda.jit
def add_kernel(x, y, out):
    i = cuda.grid(1)
    if i < x.shape[0]:
        out[i] = x[i] + y[i]

# CuPy arrays work directly — zero copy
a = cp.arange(1000, dtype=cp.float32)
b = cp.ones(1000, dtype=cp.float32)
out = cp.zeros(1000, dtype=cp.float32)
add_kernel[4, 256](a, b, out)

import torch
from numba import cuda

t = torch.tensor([1, 2, 3], dtype=torch.float32, device="cuda")
d_array = cuda.as_cuda_array(t)  # Zero-copy Numba view of PyTorch tensor

cuda.is_cuda_array(obj)       # True if obj has __cuda_array_interface__
cuda.as_cuda_array(obj)       # Wrap as Numba device array (zero copy)