Встроенная функция Cython с массивом numpy в качестве параметра

Ответ 1

Проблема заключается в том, что назначение массива numpy (или, что то же самое, передача его как аргумент функции) - это не просто простое назначение, а "извлечение буфера", которое заполняет структуру и выводит информацию о шаге и указателе в локальные переменные, необходимые для быстрой индексации. Если вы выполняете итерацию по умеренному числу элементов, это O (1) накладные расходы легко амортизируются по циклу, но это, конечно, не относится к небольшим функциям.

Улучшение этого очень важно для многих людей, но это нетривиальное изменение. См., Например, обсуждение в http://groups.google.com/group/cython-users/browse_thread/thread/8fc8686315d7f3fe

Ответ 2

Прошло более 3 лет с момента публикации вопроса, и за это время был достигнут значительный прогресс. В этом коде (обновление 2 вопроса):

# cython: infer_types=True
# cython: boundscheck=False
# cython: wraparound=False
import numpy as np
cimport numpy as np

cdef inline inc(np.ndarray[np.int32_t, ndim=2] arr, int i, int j):
    arr[i, j]+= 1

def test1(np.ndarray[np.int32_t, ndim=2] arr):
    cdef int i,j    
    for i in xrange(arr.shape[0]):
        for j in xrange(arr.shape[1]):
            inc(arr, i, j)

def test2(np.ndarray[np.int32_t, ndim=2] arr):    
    cdef int i,j    
    for i in xrange(arr.shape[0]):
        for j in xrange(arr.shape[1]):
            arr[i,j] += 1

Я получаю следующие тайминги:

arr = np.zeros((1000,1000), dtype=np.int32)
%timeit test1(arr)
%timeit test2(arr)
   1 loops, best of 3: 354 ms per loop
1000 loops, best of 3: 1.02 ms per loop

Таким образом, проблема воспроизводится даже спустя более 3 лет. Cython теперь имеет типизированные просмотры памяти, AFAIK был введен в Cython 0.16, поэтому недоступен в момент публикации вопроса, При этом:

# cython: infer_types=True
# cython: boundscheck=False
# cython: wraparound=False
import numpy as np
cimport numpy as np

cdef inline inc(int[:, ::1] tmv, int i, int j):
    tmv[i, j]+= 1

def test3(np.ndarray[np.int32_t, ndim=2] arr):
    cdef int i,j
    cdef int[:, ::1] tmv = arr
    for i in xrange(tmv.shape[0]):
        for j in xrange(tmv.shape[1]):
            inc(tmv, i, j)

def test4(np.ndarray[np.int32_t, ndim=2] arr):    
    cdef int i,j
    cdef int[:, ::1] tmv = arr
    for i in xrange(tmv.shape[0]):
        for j in xrange(tmv.shape[1]):
            tmv[i,j] += 1

С этим я получаю:

arr = np.zeros((1000,1000), dtype=np.int32)
%timeit test3(arr)
%timeit test4(arr)
1000 loops, best of 3: 977 µs per loop
1000 loops, best of 3: 838 µs per loop

Мы почти там и уже быстрее, чем старомодный! Теперь функция inc() может быть объявлена ​​nogil, поэтому пусть объявляет ее так! Но oops:

Error compiling Cython file:
[...]

cdef inline inc(int[:, ::1] tmv, int i, int j) nogil:
    ^
[...]
Function with Python return type cannot be declared nogil

Ааа, я полностью пропустил, что возвращаемый тип void отсутствовал! Еще раз, но теперь с void:

cdef inline void inc(int[:, ::1] tmv, int i, int j) nogil:
    tmv[i, j]+= 1

И наконец я получаю:

%timeit test3(arr)
%timeit test4(arr)
1000 loops, best of 3: 843 µs per loop
1000 loops, best of 3: 853 µs per loop

Так же быстро, как ручная установка!

Теперь, просто для удовольствия, я пробовал Numba в этом коде:

import numpy as np
from numba import autojit, jit

@autojit
def inc(arr, i, j):
    arr[i, j] += 1

@autojit
def test5(arr):
    for i in xrange(arr.shape[0]):
        for j in xrange(arr.shape[1]):
            inc(arr, i, j)

Я получаю:

arr = np.zeros((1000,1000), dtype=np.int32)
%timeit test5(arr)
100 loops, best of 3: 4.03 ms per loop

Несмотря на то, что он 4.7x медленнее, чем Cython, скорее всего, потому что компилятор JIT не смог встроить inc(), я думаю, что он AWESOME!. Мне нужно было добавить @autojit и не пришлось испортить код с неуклюжими объявлениями типа; 88x ускорение почти ничего!

Я пробовал другие вещи с Numba, например

@jit('void(i4[:],i4,i4)')
def inc(arr, i, j):
    arr[i, j] += 1

или nopython=True, но не удалось его улучшить.

Улучшение вложения находится в списке разработчиков Numba, нам нужно только подать больше запросов, чтобы сделать его более приоритетным.;)

Ответ 3

Вы передаете массив в inc() как объект Python типа numpy.ndarray. Передача объектов Python является дорогостоящей из-за проблем, таких как подсчет ссылок, и, похоже, предотвращает inlining. Если вы передаете массив, то способ C, т.е. Как указатель, test1() становится еще быстрее, чем test2() на моей машине:

cimport numpy as np

cdef inline inc(int* arr, int i):
    arr[i] += 1

def test1(np.ndarray[np.int32_t] arr):
    cdef int i
    for i in xrange(len(arr)):
        inc(<int*>arr.data, i)