Почему реализация этой простой добавляемой функции в Cython медленнее, чем в Numba?

Question

У меня есть 2 версии функции, которая добавляет строку в 2d массив; один на Cython, а другой в Numba.

Производительность версии Cython намного ниже, чем у версии Numba. Я хотел бы оптимизировать версию Cython, чтобы она работала как минимум так же, как и версия Numba.

Я синхронизирую код с этими timer.py модулями: время импорта

class Timer(object):
    def __init__(self, name='', output=print):
        self._name = name
        self._output = output

    def __enter__(self):
        self.start = time.time()
        return self

    def __exit__(self, a, b, c):
        self.end = time.time()
        self.time_taken = self.end - self.start
        self._output('%s Took %0.2fs seconds' % (self._name, self.time_taken))

Мой append_2d_cython.pyx модуль:

#!python
#cython: boundscheck=False
#cython: wraparound=False


import numpy as np
cimport numpy as cnp

cnp.import_array()  # needed to initialize numpy-API


cpdef empty_2d(int d1, int d2):
    cdef:
        cnp.npy_intp[2] dimdim

    dimdim[0] = d1
    dimdim[1] = d2

    return cnp.PyArray_SimpleNew(2, dimdim, cnp.NPY_INT32)


cpdef append_2d(int[:, :] arr, int[:] value):

    cdef int[:, :] result

    result = empty_2d(arr.shape[0]+1, arr.shape[1])

    result[:arr.shape[0], :] = arr

    result[arr.shape[0], :] = value

    return result

Мой append_2d_numba.py модуль:

import numba as nb
import numpy as np


@nb.jit(nopython=True)
def append_2d(arr, value):

    result = np.empty((arr.shape[0]+1, arr.shape[1]), dtype=arr.dtype)
    result[:-1] = arr
    result[-1] = value
    return result

Я сравниваю версии * 1017 для Numba и Cython с этим сценарием:

import pyximport
import numpy as np
pyximport.install(setup_args={'include_dirs': np.get_include()})

from timer import Timer
from append_2d_cython import append_2d as append_2d_cython
from append_2d_numba import append_2d as append_2d_numba

arr_2d = np.random.randint(0, 100, size=(5, 4), dtype=np.int32)
arr_1d = np.array([0, 1, 2, 3], np.int32)

num_tests = 100000

with Timer('append_2d_cython'):
    for _ in range(num_tests):
        r_cython = append_2d_cython(arr_2d, arr_1d)

# # JIT Compile it
append_2d_numba(arr_2d, arr_1d)

with Timer('append_2d_numba'):
    for _ in range(num_tests):
        r_numba = append_2d_numba(arr_2d, arr_1d)

Какие отпечатки:

make many with cython Took 0.36s seconds
make many with numba Took 0.12s seconds

Итак, для этого кода numba в 3 раза быстрее, чем Cython. Я хотел бы провести рефакторинг кода Cython так, чтобы он был по крайней мере таким же быстрым, как и код Numba. Как я могу это сделать?

Answer 1

Это исследование покажет, что большие издержки Cython являются причиной плохой производительности Cython. Кроме того, будет предложена альтернативная (несколько хакерская) альтернатива, чтобы избежать большинства из них - поэтому решение numba будет побито фактором 4.

Давайте начнем с установления базовой линии на моем компьютере (я называю ваши функции cy_append_2d и nb_append_2d и использую магию %timeit для измерения времени работы):

arr_2d = np.arange(5*4, dtype=np.int32).reshape((5,4))
arr_1d = np.array([0, 1, 2, 3], np.int32)

%timeit cy_append_2d(arr_2d, arr_1d)
# 8.27 µs ± 141 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit nb_append_2d(arr_2d, arr_1d)
# 2.84 µs ± 169 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Numba-версия примерно в три раза быстрее - похоже на время, которое вы наблюдаете.

Однако мы должны знать, что мы измеряем не время, необходимое для копирования данных, а накладные расходы. Это не то, что Numba делает что-то причудливое - просто у него меньше накладных расходов (но все же довольно много - почти 3 мкс для создания numpy-массива и копирования 24 целых чисел!)

Если мы увеличим объем копируемых данных, мы увидим, что Cython и Numba имеют довольно схожие характеристики - ни один модный компилятор не может значительно улучшить копирование:

N=5000
arr_2d_large = np.arange(5*N, dtype=np.int32).reshape((5,N))
arr_1d_large = np.arange(N, dtype=np.int32)
%timeit cy_append_2d(arr_2d_large, arr_1d_large)
# 35.7 µs ± 597 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit nb_append_2d(arr_2d_large, arr_1d_large)
# 44.8 µs ± 1.36 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

здесь Cython немного быстрее, но для разных машин и разных размеров это может варьироваться, для наших целей мы можем считать их почти одинаково быстрыми.

Как указал @DavidW, создание cython-массивов из cython-ndarray из numpy-массивов приводит к некоторым накладным расходам. Рассмотрим эту фиктивную функцию:

%%cython
cpdef dummy(int[:, :] arr, int[:] value):
    pass

%timeit dummy(arr_2d, arr_1d)
# 3.24 µs ± 47.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Это означает, что из первоначальных 8 мкс 3 уже потрачено до запуска первой операции в функции - здесь вы можете увидеть затраты на создание представлений памяти.

Как правило, никто не заботится об этих издержках - потому что если вы вызовете функциональность numpy для таких маленьких блоков данных, производительность все равно не будет звездной.

---

Однако, если вы действительно занимаетесь этим видом микрооптимизации, вы можете использовать C-API Numpy напрямую, без помощи ndarray Cythons. Мы не можем ожидать, что результат будет таким же быстрым, как копирование 24 целых чисел - для этого создание нового буферного / numpy-массива просто дорого, однако наши шансы побить 8 мкс довольно высоки!

Вот прототип, который показывает, что может быть возможно:

%%cython
from libc.string cimport memcpy

# don't use Cythons wrapper, because it uses ndarray
# define only the needed stuff
cdef extern from "numpy/arrayobject.h":
    ctypedef int npy_intp            # it isn't actually int, but cython doesn't care anyway
    int _import_array() except -1
    char *PyArray_BYTES(object arr)
    npy_intp PyArray_DIM(object arr, int n)
    object PyArray_SimpleNew(int nd, npy_intp* dims, int typenum)
    cdef enum NPY_TYPES:
        NPY_INT32

# initialize Numpy's C-API when imported.
_import_array()  

def cy_fast_append_2d(upper, lower):
    # create resulting array:
    cdef npy_intp dims[2]
    dims[0] = PyArray_DIM(upper, 0)+1
    dims[1] = PyArray_DIM(upper, 1)
    cdef object res = PyArray_SimpleNew(2, &dims[0], NPY_INT32)
    # copy upper array, assume C-order/memory layout
    cdef char *dest = PyArray_BYTES(res)
    cdef char *source = PyArray_BYTES(upper)
    cdef int size = (dims[0]-1)*dims[1]*4 # int32=4 bytes
    memcpy(dest, source, size)
    # copy lower array, assume C-order/memory layout
    dest += size
    source = PyArray_BYTES(lower)
    size = dims[1]*4
    memcpy(dest, source, size)
    return res

Время теперь:

%timeit cy_fast_append_2d(arr_2d, arr_1d)
753 ns ± 3.13 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

, что означает, что Cython превосходит Numba в 4 раза.

Однако безопасность теряется - например, он работает только для массивов C-порядка, а не для массивов Fortran-порядка. Но моя цель состояла не в том, чтобы дать водонепроницаемое решение, а в том, чтобы выяснить, насколько быстрым может стать непосредственное использование C-API Numpy - это ваше решение, следует ли выбирать этот хакерский путь.