Pythonアクセラレータ（Cython、Numba、f2py）とNumpy einsumの比較

Question

私はPythonアクセラレータ（Numba、Cython、f2py）を単純なForループと比較しています。これまでのところ、Numpyはこの問題の最も速い（6倍速い）が、私は試してみるべき追加の最適化がある場合や、何か間違っている場合にはフィードバックを求めていました。この単純なコードは、これらのeinsum呼び出しの数は多いがループは明示的ではない大きなコードに基づいています。私は、これらのアクセラレータのいずれかがうまくいくかどうかをチェックしています。

Mac OS XでPython 2.7.9で実行されたタイミング、Homebrewのgcc-5.3.0（--with-fortran --without-multilib）がインストールされています。また、％timeit呼び出しも行いました。これらの単一のコールタイミングはかなり正確です。 （f2py -c -m test_f2py test_f2py.F90でコンパイル）

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 '' 

f2pyファイル：

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 

メインコード

!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 

そしてCython .pyxファイル（メインルーチンでpyximportでコンパイル）：

#/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

文字列パラメータの解析が完了すると、einsumは、nditerのコンパイル済みバージョンを使用して、すべての軸に渡ってsum-of-productsを計算します。ソースコードはnumpy githubで簡単に見つかります。

私はeinsumを修正しましたが、パッチを書く作業の一部として働いていました。その一環として、私はcythonスクリプトを書いて製品の合計を出しました。私はeinsum速度で実行するように私のコードを取得しようとしませんでした

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

の位置は：このコードを見ることができます。私はそれがどう働いているのか理解しようとしていただけです。

Answer 2

通常、これらのアクセラレータは、Pythonループや多くの中間結果でコードを高速化するために使用されますが、einsumは既にかなりよく最適化されています（see source）。彼らは容易にeinsumを打つことを期待するべきではありませんが、あなたはそれをパフォーマンスに近づけるかもしれません。

Numbaでは、コンパイル時間をベンチマークから除外することが重要です。これは、ジット関数を2回実行することで（同じタイプの入力で）簡単に実行できます。例えば。 IPythonで私が取得：

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 

を多くの改良を行うことができますCythonコードの場合：配列の境界とラップアラウンドのため

1. 無効にチェック、compiler directivesを参照してください。
2. アレイが連続していることを指定します。
3. typed memoryviewsを使用してください。以下のよう

何か：最新のUbuntu 15日

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 

。10（x86）これは私にeinsumと同じ速度を与えます。しかし、Anacondaディストリビューションを搭載した同じPC上のWindows（x86）では、このCythonコードは約einsumの半分の速度です。私はこれがgccのバージョン（5.2.1対4.7.0）とSSE命令を挿入する能力（einsumはSSE2組み込み関数でコード化されている）と関係していると思います。おそらく異なるコンパイラオプションを提供すると助けになるかもしれませんが、わかりません。

私はほとんどすべてのFortranを知らないので、私はそれについてコメントすることはできません。

あなたの目標はeinsumを打ち負かすことなので、明らかに次のステップは並列度を上げることです。 cython.parallelでいくつかのスレッドを生成するのはかなり簡単です。それでもシステムのメモリ帯域幅が飽和しない場合は、AVX2やFused Multiply-Addのような最新のCPU命令を明示的に含めることができます。

もう1つ試みることができるのは、fの並べ替えと再編成で、np.dotで操作してください。あなたのNumpyが良いBLASライブラリを持っているならば、あなたは思うことができるあらゆる最適化を可能にするはずですが、一般性を失って、恐らくf配列の非常に高価なコピーを犠牲にしています。