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

问题描述：

我想写一些函数在numba，我可以交换使用不同的目标（cpu，cuda，并行）。我遇到的万阿英，蒋达清是一个新的数组的分配是CUDA设备代码，例如，不同：阵列分配的CPU和GPU功能在NUMBA

cuda.local.array(shape, dtype)

对比做了CPU的功能类似，即

np.empty(shape, dtype)

是否有聪明的方式如何处理这个，而不必编写单独的功能？

难道你不能在你的函数中测试类型吗？ – 0TTT0

问题是我不能在我的函数中有任何声明不起作用，因为numba编译代码并且吓坏了。否则，我会做一个简单的if/else或这样的。处理这种情况的自然方式是C中的预处理器指令，但是没有这样的东西可用于python –

答

我发现问题的一个肮脏的解决方法。这是我能做到的唯一方法。使用@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

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

相关推荐