cuda加速实例

以下是一个简单的 CUDA 加速实例：

import numpy as np
from numba import cuda

@cuda.jit
def increment_by_one(x):
    thread_id = cuda.threadIdx.x
    block_id = cuda.blockIdx.x
    block_size = cuda.blockDim.x
    idx = thread_id + block_id * block_size
    if idx < x.size:
        x[idx] += 1

arr = np.ones(1024)
threads_per_block = 32
blocks_per_grid = (arr.size + (threads_per_block - 1)) // threads_per_block
increment_by_one[blocks_per_grid, threads_per_block](arr)

这个例子中，我们使用 Numba 库中的 `cuda.jit` 装饰器来定义一个可以在 CUDA 设备上运行的函数 `increment_by_one`。这个函数接受一个一维的 NumPy 数组 `x`，并将其每个元素加 1。在函数内部，我们使用 `cuda.threadIdx.x` 和 `cuda.blockIdx.x` 计算当前线程和当前块的索引，并使用 `cuda.blockDim.x` 计算当前块的大小。然后，我们使用这些索引来计算当前线程应该处理的元素的索引。最后，我们检查这个索引是否在数组的范围内，并在这种情况下将对应元素加 1。

在主程序中，我们首先创建一个全为 1 的长度为 1024 的 NumPy 数组 `arr`。然后，我们计算出每个块的大小和每个网格的大小，并使用 `increment_by_one` 函数将 `arr` 数组中的所有元素加 1。由于使用了 CUDA 加速，这个操作比在 CPU 上运行要快得多。