Сравнение ускорителей Python (Cython, Numba, f2py) с Numpy einsum

Question

Я сравниваю ускорители Python (Numba, Cython, f2py) с простыми циклами для циклов и einsum Numpy для конкретной проблемы (см. Ниже). Пока что Numpy является самым быстрым для этой проблемы (фактор 6x быстрее), но мне нужна была некоторая обратная связь, если есть дополнительные оптимизации, которые я должен попробовать, или если я делаю что-то неправильно. Этот простой код основан на более крупном коде, который содержит несколько этих вызовов einsum, но не содержит явных для циклов. Я проверяю, может ли любой из этих ускорителей работать лучше.

Сроки выполнения с Python 2.7.9 на Mac OS X Yosemite, с установленным gcc-5.3.0 (--with-fortran --without-multilib) от Homebrew. Также было вызвано% timeit; эти однократные тайминги достаточно точны.

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 

Основной код:

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 (скомпилирован с 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 

И в 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 для вычисления суммы продуктов по всем осям. Исходный код легко найти на numpy github.

A a back назад Я разработал рабочий стол einsum как часть написания патча. В рамках этого я написал сценарий cython, который делает сумму продукта. Вы можете увидеть этот код:

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

Я не пытался получить код для запуска на einsum скорости. Я просто пытался понять, как это работает.

Answer 2

Обычно эти ускорители используются для ускорения кода с петлями Python или большого количества промежуточных результатов, тогда как einsum уже довольно хорошо оптимизирован (see source). Вы не должны ожидать, что они легко победят einsum, но вы можете приблизиться к нему в производительности.

Для Numba важно исключить время компиляции из эталона. Это может быть достигнуто просто за счет одновременной работы джиттовой функции (с теми же типами входов). Например. с 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.

Что-то вроде:

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 

На уточненный Ubuntu 15.10 (x86) это дает мне ту же скорость, что и einsum. Однако в Windows (x86) на том же ПК с дистрибутивом Anaconda этот код Cython составляет примерно половину скорости einsum. Я думаю, что это может быть связано с версиями gcc (5.2.1 vs 4.7.0) и возможностью вставлять инструкции SSE (einsum кодируется с помощью встроенных SSE2). Может быть, поставка различных вариантов компилятора поможет, но я не уверен.

Я почти не знаю Fortran, поэтому я не могу прокомментировать это.

Поскольку ваша цель - победить einsum Я думаю, что следующий следующий шаг - смотреть на увеличивающийся параллелизм. Должно быть довольно легко создавать некоторые потоки с помощью cython.parallel. Если это еще не насыщает вашу пропускную способность системной памяти, вы можете попытаться явно включить в себя новейшие инструкции процессора, такие как AVX2 и Fused Multiply-Add.

Еще одна вещь, которую вы могли бы попробовать - изменить порядок и изменить форму f и выполнить свою операцию с помощью np.dot. Если ваш Numpy поставляется с хорошей библиотекой BLAS, это должно обеспечить практически любую оптимизацию, о которой вы можете думать, хотя ценой потери общности и, возможно, очень дорогостоящей копии массива f.