阵列分配的CPU和GPU功能在NUMBA

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