Confronto tra gli acceleratori Python (Cython, Numba, f2py) e Numpy einsum

Question

Sto confrontando gli acceleratori Python (Numba, Cython, f2py) con semplici cicli For e Numins einsum per un problema particolare (vedi sotto). Finora Numpy è il più veloce per questo problema (fattore 6x più veloce), ma volevo un feedback se ci fossero ulteriori ottimizzazioni dovrei provare, o se sto facendo qualcosa di sbagliato. Questo semplice codice si basa su un codice più grande che ha un numero di queste chiamate einsum, ma non esplicito per cicli. Sto controllando se qualcuno di questi acceleratori può fare di meglio.

Tempi eseguiti con Python 2.7.9 su Mac OS X Yosemite, con gcc-5.3.0 installato (--with-fortran --without-multilib) da Homebrew. Ha fatto anche chiamate% timeit; questi tempi di chiamata singola sono abbastanza precisi.

In [1]: %run -i test_numba.py 
test_numpy: 0.0805640220642 
Matches Numpy output: True 

test_dumb: 1.43043899536 
Matches Numpy output: True 

test_numba: 0.464295864105 
Matches Numpy output: True 

test_cython: 0.627640008926 
Matches Numpy output: True 

test_f2py: 5.01890516281 
Matches Numpy output: True 

test_f2py_order: 2.31424307823 
Matches Numpy output: True 

test_f2py_reorder: 0.507861852646 
Matches Numpy output: True 

Il codice principale:

import numpy as np 
import numba 
import time 
import test_f2py as tf2py 
import pyximport 
pyximport.install(setup_args={'include_dirs':np.get_include()}) 
import test_cython as tcyth 

def test_dumb(f,b): 
    fnew = np.empty((f.shape[1],f.shape[2])) 
    for i in range(f.shape[0]): 
     for l in range(f.shape[3]): 
      fnew += f[i,:,:,l] * b[i,l] 
    return fnew 


def test_dumber(f,b): 
    fnew = np.empty((f.shape[1],f.shape[2])) 
    for i in range(f.shape[0]): 
     for j in range(f.shape[1]): 
      for k in range(f.shape[2]): 
       for l in range(f.shape[3]): 
        fnew[j,k] += f[i,j,k,l] * b[i,l] 
    return fnew 

@numba.jit(nopython=True) 
def test_numba(f,b): 
    fnew = np.zeros((f.shape[1],f.shape[2])) #NOTE: can't be empty, gives errors 
    for i in range(f.shape[0]): 
     for j in range(f.shape[1]): 
      for k in range(f.shape[2]): 
       for l in range(f.shape[3]): 
        fnew[j,k] += f[i,j,k,l] * b[i,l] 
    return fnew 

def test_numpy(f,b): 
    return np.einsum('i...k,ik->...',f,b) 

def test_f2py(f,b): 
    return tf2py.test_f2py(f,b) 

def test_f2py_order(f,b): 
    return tf2py.test_f2py(f,b) 

def test_f2py_reorder(f,b): 
    return tf2py.test_f2py_reorder(f,b) 

def test_cython(f,b): 
    return tcyth.test_cython(f,b) 

if __name__ == '__main__': 

    #goal is to create: fnew = sum f*b over dim 0 and 3. 
    f = np.random.rand(32,33,2000,64) 
    b = np.random.rand(32,64) 

    f1 = np.asfortranarray(f) 
    b1 = np.asfortranarray(b) 

    f2 = np.asfortranarray(np.transpose(f,[1,2,0,3])) 

    funcs = [test_dumb,test_numba, test_cython, \ 
      test_f2py,test_f2py_order,test_f2py_reorder] 

    tstart = time.time()  
    fnew_numpy= test_numpy(f,b) 
    tstop = time.time() 
    print test_numpy.__name__+': '+str(tstop-tstart) 
    print 'Matches Numpy output: '+str(np.allclose(fnew_numpy,fnew_numpy)) 
    print '' 

    for func in funcs: 
     tstart = time.time() 
     if func.__name__ == 'test_f2py_order': 
      fnew = func(f1,b1) 
     elif func.__name__ == 'test_f2py_reorder': 
      fnew = func(f2,b1) 
     else: 
      fnew = func(f,b) 
     tstop = time.time() 
     print func.__name__+': '+str(tstop-tstart) 
     print 'Matches Numpy output: '+str(np.allclose(fnew,fnew_numpy)) 
     print '' 

Il file f2py (compilato con f2py -c -m test_f2py test_f2py.F90):

!file: test_f2py 
subroutine test_f2py(f,b,fnew,n1,n2,n3,n4) 

integer :: n1,n2,n3,n4 
real(8), dimension(n1,n2,n3,n4) :: f 
real(8), dimension(n1,n4) :: b 
real(8), dimension(n2,n3) :: fnew 
!f2py intent(in) f 
!f2py intent(in) b 
!f2py intent(out) fnew 
!f2py intent(in) n1 
!f2py intent(in) n2 
!f2py intent(in) n3 
!f2py intent(in) n4 

integer :: i1,i2,i3,i4 

do i1=1,n1 
    do i2=1,n2 
     do i3=1,n3 
      do i4=1,n4 
       fnew(i2,i3) = fnew(i2,i3) + f(i1,i2,i3,i4)*b(i1,i4) 
      enddo 
     enddo 
    enddo 
enddo 

end subroutine test_f2py 

subroutine test_f2py_reorder(f,b,fnew,n1,n2,n3,n4) 

integer :: n1,n2,n3,n4 
real(8), dimension(n1,n2,n3,n4) :: f 
real(8), dimension(n3,n4) :: b 
real(8), dimension(n1,n2) :: fnew 
!f2py intent(in) f 
!f2py intent(in) b 
!f2py intent(out) fnew 
!f2py intent(in) n1 
!f2py intent(in) n2 
!f2py intent(in) n3 
!f2py intent(in) n4 

integer :: i1,i2,i3,i4 

do i3=1,n3 
    do i4=1,n4 
     do i1=1,n1 
      do i2=1,n2 
       fnew(i1,i2) = fnew(i1,i2) + f(i1,i2,i3,i4)*b(i3,i4) 
      enddo 
     enddo 
    enddo 
enddo 

end subroutine test_f2py_reorder 

e il file Cython .pyx (compilato con pyximport nella routine principale):

#/usr/bin python 
import numpy as np 
cimport numpy as np 

def test_cython(np.ndarray[np.float64_t,ndim=4] f, np.ndarray[np.float64_t,ndim=2] b): 
    # cdef np.ndarray[np.float64_t,ndim=4] f 
    # cdef np.ndarray[np.float64_t,ndim=2] b 
    cdef np.ndarray[np.float64_t,ndim=2] fnew = np.empty((f.shape[1],f.shape[2]),dtype=np.float64) 
    cdef int i,j,k,l 
    cdef int Ni = f.shape[0] 
    cdef int Nj = f.shape[1] 
    cdef int Nk = f.shape[2] 
    cdef int Nl = f.shape[3] 

    for i in range(Ni): 
     for j in range(Nj): 
      for k in range(Nk): 
       for l in range(Nl): 
        fnew[j,k] += f[i,j,k,l] * b[i,l] 
    return fnew 

Answer 1

Una volta eseguita l'analisi del parametro stringa, einsum utilizza una versione compilata di nditer per eseguire un calcolo di somma di prodotti su tutti gli assi. Il codice sorgente è facilmente reperibile sul github numpy.

Qualche istante fa ho elaborato un lavoro simile a einsum come parte della scrittura di una patch. Come parte di ciò ho scritto uno script cython che esegue la somma di prodotto. Si può vedere questo codice:

https://github.com/hpaulj/numpy-einsum

Io non ho cercato di ottenere il mio codice per l'esecuzione a einsum velocità. Stavo solo cercando di capire come funzionava.

Answer 2

Normalmente questi acceleratori vengono utilizzati per velocizzare il codice con loop Python o molti risultati intermedi, mentre lo einsum è già ottimizzato (see source). Non dovresti aspettarti che battano facilmente einsum, ma potresti avvicinarti ad esso in termini di prestazioni.

Per Numba è importante escludere il tempo di compilazione dal benchmark. Questo può essere ottenuto semplicemente eseguendo la funzione jitted due volte (con lo stesso tipo di input). Per esempio. con IPython ottengo:

f = np.random.rand(32,33,500,64) 
b = np.random.rand(32,64) 

%time _ = test_numba(f,b) # First invocation 
# Wall time: 466 ms 
%time _ = test_numba(f,b) 
# Wall time: 73 ms 
%timeit test_numba(f, b) 
# 10 loops, best of 3: 72.7 ms per loop 
%timeit test_numpy(f, b) 
# 10 loops, best of 3: 62.8 ms per loop 

Per il codice Cython una serie di miglioramenti può essere fatta:

1. disabilitare le limiti di matrice e avvolgente, vedere compiler directives.
2. Specificare che gli array sono contigui.
3. Utilizzare typed memoryviews.

Qualcosa di simile:

cimport cython 
import numpy as np 

@cython.boundscheck(False) 
@cython.wraparound(False) 
def test_cython(double[:,:,:,::1] f, double[:,::1] b): 
    cdef int i, j, k, l, Ni, Nj, Nk, Nl 
    Ni = f.shape[0] 
    Nj = f.shape[1] 
    Nk = f.shape[2] 
    Nl = f.shape[3] 

    fnew = np.empty((Nj, Nk)) 
    cdef double[:,::1] fnew_v = fnew 

    for i in range(Ni): 
     for j in range(Nj): 
      for k in range(Nk): 
       for l in range(Nl): 
        fnew_v[j,k] += f[i,j,k,l] * b[i,l] 
    return fnew 

Su un up-to-date Ubuntu 15.10 (x86) questo mi dà la stessa velocità di einsum. Tuttavia, su Windows (x86) sullo stesso PC con distribuzione Anaconda questo codice Cython è circa la metà della velocità di einsum. Penso che questo potrebbe avere a che fare con le versioni di gcc (5.2.1 contro 4.7.0) e la possibilità di inserire istruzioni SSE (einsum è codificato con intrinseco SSE2). Forse fornire opzioni di compilazione diverse potrebbe aiutare, ma non ne sono sicuro.

Quasi non conosco nessun Fortran, quindi non posso commentare in merito.

Dal momento che il tuo obiettivo è battere einsum penso che il prossimo passo ovvio sia guardare al parallelismo crescente. Dovrebbe essere abbastanza facile generare alcuni thread con cython.parallel. Se ciò non satura ancora la larghezza di banda della memoria del tuo sistema, allora potresti provare ad includere esplicitamente le ultime istruzioni della CPU come AVX2 e Fused Multiply-Add.

Un'altra cosa che si potrebbe provare è di riordinare e rimodellare f e fare le operazioni con np.dot. Se il tuo Numpy viene fornito con una buona libreria BLAS, questo dovrebbe consentire praticamente qualsiasi ottimizzazione a cui puoi pensare, anche se a costo di una perdita di generalità e magari una copia molto costosa dell'array f.