5.4。编写设备功能

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

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

  1. from numba import roc
  2. @roc.jit(device=True)
  3. def a_device_function(a, b):
  4. return a + b

使用设备功能的示例:

  1. from numba import roc
  2. import numpy as np
  3. @roc.jit
  4. def kernel(an_array):
  5. pos = roc.get_global_id(0)
  6. if pos < an_array.size: # Check array boundaries
  7. an_array[pos] = a_device_function(1, pos) # call device function
  8. @roc.jit(device = True)
  9. def a_device_function(a, b):
  10. return a + b
  11. n = 16
  12. x = np.zeros(n)
  13. kernel[1, n](x)
  14. print(x)