陣列分配的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