编写设备函数

CUDA 设备函数只能从设备内部调用(通过内核或其他设备函数)。定义一个设备函数:

from numba import cuda

@cuda.jit(device=True)
def a_device_function(a, b):
    return a + b

与内核函数不同,设备函数可以像普通函数一样返回值。