10个使用Numba CUDA进行编程的例子

以下是10个使用Numba CUDA进行编程的例子，这些例子涵盖了基本的向量加法、矩阵乘法以及其他一些常见操作：

向量加法

python 复制代码

from numba import cuda
import numpy as np
@cuda.jit
def vector_add(a, b, c):
    i = cuda.grid(1)
    if i < len(a):
        c[i] = a[i] + b[i]
a = np.array([1, 2, 3], dtype=np.float32)
b = np.array([4, 5, 6], dtype=np.float32)
c = np.zeros_like(a)
threads_per_block = 32
blocks_per_grid = (len(a) + (threads_per_block - 1)) // threads_per_block
vector_add[blocks_per_grid, threads_per_block](a, b, c)
print(c)

矩阵乘法

python 复制代码

@cuda.jit
def matmul(A, B, C):
    i, j = cuda.grid(2)
    if i < C.shape[0] and j < C.shape[1]:
        tmp = 0.
        for k in range(A.shape[1]):
            tmp += A[i, k] * B[k, j]
        C[i, j] = tmp
A = np.array([[1, 2], [3, 4]], dtype=np.float32)
B = np.array([[5, 6], [7, 8]], dtype=np.float32)
C = np.zeros((2, 2), dtype=np.float32)
threads_per_block = (16, 16)
blocks_per_grid_x = math.ceil(A.shape[0] / threads_per_block[0])
blocks_per_grid_y = math.ceil(B.shape[1] / threads_per_block[1])
blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y)
matmul[blocks_per_grid, threads_per_block](A, B, C)
print(C)

向量点乘

python 复制代码

@cuda.jit
def dot_product(a, b, result):
    i = cuda.grid(1)
    tmp = 0.0
    if i < a.shape[0]:
        tmp = a[i] * b[i]
    cuda.atomic.add(result, 0, tmp)
a = np.array([1, 2, 3], dtype=np.float32)
b = np.array([4, 5, 6], dtype=np.float32)
result = np.zeros(1, dtype=np.float32)
threads_per_block = 32
blocks_per_grid = (len(a) + (threads_per_block - 1)) // threads_per_block
dot_product[blocks_per_grid, threads_per_block](a, b, result)
print(result[0])

向量乘法

python 复制代码

@cuda.jit
def vector_multiply(a, b, c):
    i = cuda.grid(1)
    if i < len(a):
        c[i] = a[i] * b[i]
# 同向量加法示例中的初始化和调用方式

向量除法

python 复制代码

@cuda.jit
def vector_divide(a, b, c):
    i = cuda.grid(1)
    if i < len(a):
        c[i] = a[i] / b[i]
# 同向量加法示例中的初始化和调用方式

向量求和

python 复制代码

@cuda.jit
def vector_sum(a, result):
    i = cuda.grid(1)
    tmp = 0.0
    if i < a.shape[0]:
        tmp = a[i]
    cuda.atomic.add(result, 0, tmp)
# 同向量点乘示例中的初始化和调用方式

向量求最大值

python 复制代码

@cuda.jit
def vector_max(a, result):
    i = cuda.grid(1)
    if i < a.shape[0]:
        cuda.atomic.max(result, 0, a[i])
# 同向量点乘示例中的初始化和调用方式

向量求最小值

python 复制代码

@cuda.jit
def vector_min(a, result):
    i = cuda.grid(1)
    if i < a.shape[0]:
        cuda.atomic.min(result, 0, a[i])
# 同向量点乘示例中的初始化和调用方式

向量排序（冒泡排序）

python 复制代码

@cuda.jit
def bubble_sort(arr):
    i = cuda.grid(1)
    n = arr.shape[0]
    if i < n:
        for j in range(n - 1):
            if arr[i] > arr[j]:
                arr[i], arr[j] = arr[j], arr[i]
# 同向量加法示例中的初始化