Python 用 Numba 调用 GPU 算力的 3 步快速上手

使用 Numba 调用 GPU 只需三步即可实现高性能计算加速。以下是完整实现方案：

1. 环境准备

安装必需库：

pip install numba cudatoolkit

验证 GPU 支持：

from numba import cuda
print(cuda.gpus)  # 输出检测到的 GPU 设备

2. 编写 GPU 核函数

用 @cuda.jit 装饰器标记函数，自动编译为 GPU 指令：

from numba import cuda
import numpy as np

# GPU 核函数：向量加法
@cuda.jit
def vector_add(a, b, result):
    idx = cuda.grid(1)  # 获取线程索引
    if idx < a.size:
        result[idx] = a[idx] + b[idx]  # 并行计算

3. 执行计算

# 生成数据（10^7 规模）
n = 10_000_000
a = np.ones(n, dtype=np.float32)
b = np.ones(n, dtype=np.float32)
result = np.empty_like(a)

# 配置并行参数
threads_per_block = 256
blocks_per_grid = (n + threads_per_block - 1) // threads_per_block

# 执行 GPU 计算
vector_add[blocks_per_grid, threads_per_block](a, b, result)

# 验证结果
print("结果校验:", np.allclose(result, a + b))  # 输出 True

⚡ 性能对比

关键说明

• 线程配置公式：blocks = (n + threads - 1) // threads
• 数据自动传输：Numba 自动处理 CPU-GPU 内存传输
• 适用场景：大规模并行计算（如矩阵运算、物理仿真）

完整代码可直接复制运行，无需额外配置。