我正在尝试用 numba 编写一些函数,我可以将它们交替用于不同的目标(cpu、cuda、并行)。我遇到的问题是新数组的分配对于 cuda 设备代码是不同的,例如:
cuda.local.array(shape, dtype)
与为 CPU 功能做类似的事情,即
np.empty(shape, dtype)
有没有一种聪明的方法来处理这个而不必编写单独的函数?
我找到了一个解决这个问题的肮脏的方法。这是我让它工作的唯一方法。使用@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
本文收集自互联网,转载请注明来源。
如有侵权,请联系[email protected] 删除。
我来说两句