如何编写可同时在CPU与CUDA设备模式下运行的Numba函数？

核心思路：复用函数逻辑，避免代码重复 #

你不需要编写两个完全相同的函数体，只需把核心计算逻辑抽离出来，再分别用njit和cuda.jit(device=True)装饰器生成对应的CPU和CUDA设备版本：

from numba import cuda, njit

# 抽离核心计算逻辑，无需装饰器
def _core_calc(a, b):
    return a + b

# 生成CPU可用的JIT函数
func_cpu = njit("i4(i4, i4)")(_core_calc)

# 生成CUDA设备可用的函数
func_gpu = cuda.jit("i4(i4, i4)", device=True)(_core_calc)

这样既保证了代码的复用性，又符合Numba对不同执行环境的要求。

补充说明 #

• 若你的函数是元素级运算（如逐元素加减乘除），也可以使用numba.vectorize并指定多目标，同时生成CPU和CUDA版本：

  from numba import vectorize
  
  @vectorize(["i4(i4, i4)"], targets=['cpu', 'cuda'])
  def vec_func(a, b):
      return a + b
  

  这种方式下，vec_func可以直接在CPU上调用，也可以在CUDA内核中作为设备函数使用，同时支持数组的向量化运算。
• 注意：CUDA设备函数只能在CUDA内核函数内部调用，无法直接在CPU环境下执行，因此无法做到同一个函数对象同时适配两种执行场景，复用逻辑生成两个目标函数是最合理的折中方案。

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