Optimieren des Zugriffs auf numpy Arrays für numba

Question

Ich bin vor kurzem über numba gestolpert und dachte darüber nach, einige hausgemachte C-Erweiterungen mit eleganteren Autojitted Python-Code zu ersetzen. Leider war ich nicht glücklich, als ich einen ersten, schnellen Benchmark ausprobierte. Es scheint, wie numba ist nicht viel besser als normaler Python hier tun, obwohl ich fast C-ähnliche Leistung erwartet hätte:

from numba import jit, autojit, uint, double 
import numpy as np 
import imp 
import logging 
logging.getLogger('numba.codegen.debug').setLevel(logging.INFO) 

def sum_accum(accmap, a): 
    res = np.zeros(np.max(accmap) + 1, dtype=a.dtype) 
    for i in xrange(len(accmap)): 
     res[accmap[i]] += a[i] 
    return res 

autonumba_sum_accum = autojit(sum_accum) 
numba_sum_accum = jit(double[:](int_[:], double[:]), 
         locals=dict(i=uint))(sum_accum) 

accmap = np.repeat(np.arange(1000), 2) 
np.random.shuffle(accmap) 
accmap = np.repeat(accmap, 10) 
a = np.random.randn(accmap.size) 

ref = sum_accum(accmap, a) 
assert np.all(ref == numba_sum_accum(accmap, a)) 
assert np.all(ref == autonumba_sum_accum(accmap, a)) 

%timeit sum_accum(accmap, a) 
%timeit autonumba_sum_accum(accmap, a) 
%timeit numba_sum_accum(accmap, a) 

accumarray = imp.load_source('accumarray', '/path/to/accumarray.py') 
assert np.all(ref == accumarray.accum(accmap, a)) 

%timeit accumarray.accum(accmap, a) 

Das auf meinem Rechner gibt:

10 loops, best of 3: 52 ms per loop 
10 loops, best of 3: 42.2 ms per loop 
10 loops, best of 3: 43.5 ms per loop 
1000 loops, best of 3: 321 us per loop 

Ich bin mit dem letzte numba-Version von pypi, 0.11.0. Irgendwelche Vorschläge, wie man den Code repariert, damit er mit Numba ziemlich schnell läuft?

Answer 1

@autojit 
def numbaMax(arr): 
    MAX = arr[0] 
    for i in arr: 
     if i > MAX: 
      MAX = i 
    return MAX 

@autojit 
def autonumba_sum_accum2(accmap, a): 
    res = np.zeros(numbaMax(accmap) + 1) 
    for i in xrange(len(accmap)): 
     res[accmap[i]] += a[i] 
    return res 

10 loops, best of 3: 26.5 ms per loop <- original 
100 loops, best of 3: 15.1 ms per loop <- with numba but the slow numpy max 
10000 loops, best of 3: 47.9 µs per loop <- with numbamax 

Answer 2

Ich habe es selbst herausgefunden. Numba konnte den Typ des Ergebnisses np.max(accmap) nicht ermitteln, auch wenn der Typ von accmap auf int festgelegt wurde. Dies verlangsamt irgendwie alles nieder, aber die Lösung ist einfach:

@autojit(locals=dict(reslen=uint)) 
def sum_accum(accmap, a): 
    reslen = np.max(accmap) + 1 
    res = np.zeros(reslen, dtype=a.dtype) 
    for i in range(len(accmap)): 
     res[accmap[i]] += a[i] 
    return res 

Das Ergebnis ist sehr beeindruckend, etwa 2/3 der C-Version:

10000 loops, best of 3: 192 us per loop