NUMBA 中 CPU 和 GPU 函数的数组分配

Question

我正在尝试在 numba 中编写一些函数，这些函数可以互换用于不同的目标（cpu、cuda、并行）。我遇到的问题是新数组的分配对于 cuda 设备代码是不同的，例如：

cuda.local.array(shape, dtype)

对比为 CPU 函数做类似的事情，即

np.empty(shape, dtype)

有没有不用编写单独的函数来处理这个问题的巧妙方法？

Answer 1

我找到了解决该问题的一种肮脏的解决方法。这是我让它发挥作用的唯一方法。使用 @myjit 装饰器代替 @jit 和 @cuda.jit 并将所有数组分配为 cuda.local.array.

def myjit(f):
'''
f : function
Decorator to assign the right jit for different targets
In case of non-cuda targets, all instances of `cuda.local.array`
are replaced by `np.empty`. This is a dirty fix, hopefully in the
near future numba will support numpy array allocation and this will
not be necessary anymore
'''
if target == 'cuda':
    return cuda.jit(f, device=True)
else:
    source = inspect.getsource(f).splitlines()
    assert '@myjit' in source[0]
    source = '\n'.join(source[1:]) + '\n'
    source = source.replace('cuda.local.array', 'np.empty')
    exec(source)
    fun = eval(f.__name__)
    newfun = jit(fun, nopython=True)
    # needs to be exported to globals
    globals()[f.__name__] = newfun
    return newfun

NUMBA 中 CPU 和 GPU 函数的数组分配

Array allocation for CPU and GPU functions in NUMBA

python

numpy

numba