在 Python 函数中，如何判断它是在 GPU 中使用 Numba 执行还是在主机/CPU 上定期调用？

Question

我有一个函数，有时会使用 Numba 作为设备函数来调用以在 GPU 上执行，有时我会直接从主机上的常规 Python 中调用：

def process():
     # perform computation

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

有时我直接从 Python 调用

process()

，有时我使用 Numba 将

process_cuda

包装器作为内核调用。

我的问题，如何从

process

函数中判断它是直接从 Python 调用还是作为 Numba 设备函数执行？

Answer 1

我找到了一个临时方法：

import inspect
from numba import cuda

caller_name = ''

def process():
  global caller_name
  caller_frame = inspect.currentframe().f_back
  caller_name = caller_frame.f_code.co_name
  a  =  2 + 2

if __name__ == '__main__':
  process_cuda = cuda.jit(device=True)(process)
  print(caller_name)

  process()
  print(caller_name)

CUDA调用时，

caller_name

的长度似乎始终为0。