如何在Python函数内判断其是Numba GPU设备执行还是CPU调用？

可以通过Numba提供的cuda.is_currently_on_device()函数直接在process内部判断执行环境，这是最可靠的方式：

import numba
from numba import cuda

def process():
    # 检测当前执行上下文
    if cuda.is_currently_on_device():
        # 作为GPU设备函数执行的逻辑
        pass
    else:
        # CPU常规Python环境执行的逻辑
        pass

process_cuda = cuda.jit(device=True)(process)

关键说明 #

• cuda.is_currently_on_device()是Numba官方提供的上下文检测函数：
  • 当函数作为设备函数在GPU上执行时，返回True
  • 当直接在CPU的Python环境中调用时，返回False
• 注意：设备函数内无法直接使用部分Python特性（比如控制台打印），实际开发中需替换为GPU兼容的逻辑。

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