如何正确使用Numba CUDA JIT装饰器让代码运行在GPU上？

我跟着教程使用Numba CUDA JIT装饰器编写了Python代码，但发现@jit(target_backend='cuda')并未让代码在GPU上运行——耗时和CPU+@jit版本接近，甚至在无NVIDIA GPU的机器上运行也没有警告或报错。我拥有GeForce GT 730M显卡，请问如何让代码真正在GPU上运行？

我的代码 #

import numpy as np
from timeit import default_timer as timer
from numba import cuda, jit

# 运行在CPU上的函数
def fill_array_with_cpu(a):
	for k in range(100000000):
		a[k] += 1

# 用@jit加速的CPU函数
@jit
def fill_array_with_cpu_jit(a):
	for k in range(100000000):
		a[k] += 1				

# 预期运行在GPU上的函数
@jit(target_backend='cuda')
def fill_array_with_gpu(a):
	for k in range(100000000):
		a[k] += 1	

# 主逻辑
a = np.ones(100000000, dtype = np.float64)

for i in range(3):
	start = timer()
	fill_array_with_cpu(a)
	print("On a CPU:", timer() - start)

for i in range(3):
	start = timer()
	fill_array_with_cpu_jit(a)
	print("On a CPU with @jit:", timer() - start)

for i in range(3):
	start = timer()
	fill_array_with_gpu(a)
	print("On a GPU:", timer() - start)

运行输出 #

On a CPU: 24.228116830999852
On a CPU: 24.90354355699992
On a CPU: 24.277727688999903
On a CPU with @jit: 0.2590671719999591
On a CPU with @jit: 0.09131158500008496
On a CPU with @jit: 0.09054700799993043
On a GPU: 0.13547917200003212
On a GPU: 0.0922475330000907
On a GPU: 0.08995077999998102

要让代码真正在GPU上运行，需要遵循Numba CUDA的并行编程范式，而不是直接用@jit(target_backend='cuda')装饰串行循环，具体修改步骤如下：

1. 使用Numba CUDA专用装饰器
   替换@jit(target_backend='cuda')为@cuda.jit，这个装饰器会明确将函数编译为CUDA内核，不会自动回退到CPU执行。

2. 重构为并行化GPU内核
   GPU的核心优势是并行计算，需要将单线程遍历数组的逻辑，改成由多个GPU线程并行处理元素的形式：

   @cuda.jit
   def fill_array_with_gpu(a):
       # 获取当前线程的全局索引
       idx = cuda.grid(1)
       # 避免索引超出数组范围
       if idx < a.size:
           a[idx] += 1
   

3. 处理数据的设备传输与内核启动
   CPU的numpy数组默认在主机内存，需要手动拷贝到GPU设备内存；同时要配置线程网格参数，指定GPU的线程块和线程数量：

   # 将主机数组拷贝到GPU设备
   d_a = cuda.to_device(a)
   # 配置线程网格：每个线程块256个线程，计算所需线程块数
   threads_per_block = 256
   blocks_per_grid = (d_a.size + threads_per_block - 1) // threads_per_block
   
   for i in range(3):
       start = timer()
       # 启动GPU内核
       fill_array_with_gpu[blocks_per_grid, threads_per_block](d_a)
       # 等待GPU计算完成（同步操作，确保计时准确）
       cuda.synchronize()
       print("On a GPU:", timer() - start)
   
   # 可选：将GPU上的结果拷贝回主机内存
   a = d_a.copy_to_host()
   

4. 验证GPU环境与设备识别
   添加代码确认程序是否正确识别到你的GeForce GT 730M：

   print("CUDA设备信息：")
   for i in range(cuda.gpus):
       gpu = cuda.gpus[i]
       print(f"设备{i}: {gpu.name}")
   
   # 检查CUDA环境是否正常
   from numba import cuda
   print(cuda.detect())
   

5. 确认显卡兼容性
   GeForce GT 730M属于Kepler架构（计算能力3.5），需确保你的Numba和CUDA Toolkit版本支持该架构。若版本不兼容，可能会导致内核无法在GPU上运行。

内容的提问来源于stack exchange，提问作者TVG