5.4。编写设备功能

原文： http://numba.pydata.org/numba-doc/latest/roc/device-functions.html

HSA 设备功能是将在设备上运行的功能，但只能从内核或其他设备功能调用。与内核函数不同，设备函数可以返回与普通函数类似的值。要定义设备功能，必须在roc.jit装饰器中将 kwarg device设置为True：

from numba import roc

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

使用设备功能的示例：

from numba import roc
import numpy as np

@roc.jit
def kernel(an_array):
    pos = roc.get_global_id(0)
    if pos < an_array.size:  # Check array boundaries
        an_array[pos] = a_device_function(1, pos) # call device function

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

n = 16
x = np.zeros(n)

kernel[1, n](x)

print(x)

---

我们一直在努力

apachecn/AiLearning

---