Attribution de tableau pour les fonctions CPU et GPU dans NUMBA

Question

J'essaye d'écrire certaines fonctions en numba que je peux utiliser de manière interchangeable pour différentes cibles (cpu, cuda, parallel). Le gros problème que je vais avoir est l'attribution d'un nouveau tableau est différent pour le code de l'appareil cuda, par exemple:

cuda.local.array(shape, dtype) 

vs faire quelque chose de similaire pour une fonction CPU, à savoir

np.empty(shape, dtype) 

Y at-il un manière intelligente comment faire face à cela sans avoir à écrire des fonctions séparées?

Answer 1

J'ai trouvé une solution de contournement sale pour le problème. C'est le seul moyen de le faire fonctionner. Utilisez le décorateur @myjit au lieu de @jit et @cuda.jit et attribuez tous les tableaux sous la forme 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