Встроенная функция Cython с массивом numpy в качестве параметра
Рассмотрим такой код:
import numpy as np
cimport numpy as np
cdef inline inc(np.ndarray[np.int32_t] arr, int i):
arr[i]+= 1
def test1(np.ndarray[np.int32_t] arr):
cdef int i
for i in xrange(len(arr)):
inc(arr, i)
def test2(np.ndarray[np.int32_t] arr):
cdef int i
for i in xrange(len(arr)):
arr[i] += 1
Я использовал ipython для измерения скорости теста1 и test2:
In [7]: timeit ttt.test1(arr)
100 loops, best of 3: 6.13 ms per loop
In [8]: timeit ttt.test2(arr)
100000 loops, best of 3: 9.79 us per loop
Есть ли способ оптимизировать test1? Почему cython не встраивает эту функцию, как сказано?
UPDATE:
На самом деле мне нужен многомерный код:
# 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
Сроки для него:
In [7]: timeit ttt.test1(arr)
1 loops, best of 3: 647 ms per loop
In [8]: timeit ttt.test2(arr)
100 loops, best of 3: 2.07 ms per loop
Явная вставка дает 300-кратное ускорение. И моя реальная функция довольно большая, поэтому она делает ее намного более удобной для обслуживания.
UPDATE2:
# cython: infer_types=True
# cython: boundscheck=False
# cython: wraparound=False
import numpy as np
cimport numpy as np
cdef inline inc(np.ndarray[np.float32_t, ndim=2] arr, int i, int j):
arr[i, j]+= 1
def test1(np.ndarray[np.float32_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.float32_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
cdef class FastPassingFloat2DArray(object):
cdef float* data
cdef int stride0, stride1
def __init__(self, np.ndarray[np.float32_t, ndim=2] arr):
self.data = <float*>arr.data
self.stride0 = arr.strides[0]/arr.dtype.itemsize
self.stride1 = arr.strides[1]/arr.dtype.itemsize
def __getitem__(self, tuple tp):
cdef int i, j
cdef float *pr, r
i, j = tp
pr = (self.data + self.stride0*i + self.stride1*j)
r = pr[0]
return r
def __setitem__(self, tuple tp, float value):
cdef int i, j
cdef float *pr, r
i, j = tp
pr = (self.data + self.stride0*i + self.stride1*j)
pr[0] = value
cdef inline inc2(FastPassingFloat2DArray arr, int i, int j):
arr[i, j]+= 1
def test3(np.ndarray[np.float32_t, ndim=2] arr):
cdef int i,j
cdef FastPassingFloat2DArray tmparr = FastPassingFloat2DArray(arr)
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
inc2(tmparr, i,j)
Сроки:
In [4]: timeit ttt.test1(arr)
1 loops, best of 3: 623 ms per loop
In [5]: timeit ttt.test2(arr)
100 loops, best of 3: 2.29 ms per loop
In [6]: timeit ttt.test3(arr)
1 loops, best of 3: 201 ms per loop
Ответы
Ответ 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)