Генерация SIMD-инструкций из кода Cython

Question

Мне нужно получить представление о производительности, которую можно получить, используя Cython в высокопроизводительном числовом коде. Одна вещь, которая меня интересует, это выяснить, может ли оптимизирующий компилятор C векторизовать код, сгенерированный Cython. Поэтому я решил написать следующий небольшой пример:

import numpy as np
cimport numpy as np
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef int f(np.ndarray[int, ndim = 1] f):
    cdef int array_length =  f.shape[0]
    cdef int sum = 0
    cdef int k
    for k in range(array_length):
        sum += f[k]
    return sum

Я знаю, что есть функции Numpy, которые делают эту работу, но я хотел бы иметь простой код, чтобы понять, что возможно с Cython. Оказывается, код, сгенерированный с помощью:

from distutils.core import setup
from Cython.Build import cythonize

setup(ext_modules = cythonize("sum.pyx"))

и позвонил с:

python setup.py build_ext --inplace

генерирует код C, который выглядит следующим образом для цикла:

for (__pyx_t_2 = 0; __pyx_t_2 < __pyx_t_1; __pyx_t_2 += 1) {
  __pyx_v_sum = __pyx_v_sum + (*(int *)((char *) 
    __pyx_pybuffernd_f.rcbuffer->pybuffer.buf +
    __pyx_t_2 * __pyx_pybuffernd_f.diminfo[0].strides)));
}

Основная проблема с этим кодом заключается в том, что компилятор не знает во время компиляции, что __pyx_pybuffernd_f.diminfo[0].strides таков, что элементы массива находятся близко друг к другу в памяти. Без этой информации компилятор не может эффективно векторизоваться.

Есть ли способ сделать такую ​​вещь из Cython?

Answer 1

У вас есть две проблемы в вашем коде (используйте опцию -a чтобы было видно)

1. Индексирование массива numpy не эффективно
2. Вы забыли int в cdef sum=0

Принимая это во внимание, мы получаем:

cpdef int f(np.ndarray[np.int_t] f):  ##HERE
    assert f.dtype == np.int
    cdef int array_length =  f.shape[0]
    cdef int sum = 0                  ##HERE
    cdef int k
    for k in range(array_length):
        sum += f[k]
    return sum

Для цикла используется следующий код:

int __pyx_t_5;
int __pyx_t_6;
Py_ssize_t __pyx_t_7;
....
__pyx_t_5 = __pyx_v_array_length;
for (__pyx_t_6 = 0; __pyx_t_6 < __pyx_t_5; __pyx_t_6+=1) {
   __pyx_v_k = __pyx_t_6;
   __pyx_t_7 = __pyx_v_k;
   __pyx_v_sum = (__pyx_v_sum + (*__Pyx_BufPtrStrided1d(__pyx_t_5numpy_int_t *, __pyx_pybuffernd_f.rcbuffer->pybuffer.buf, __pyx_t_7, __pyx_pybuffernd_f.diminfo[0].strides)));

}

Что не так уж и плохо, но не так просто для оптимизатора, как обычный код, написанный человеком. Как вы уже указали, __pyx_pybuffernd_f.diminfo[0].strides не известно во время компиляции, и это предотвращает векторизацию.

Тем не менее, вы получите лучшие результаты при использовании типизированных представлений памяти, а именно:

cpdef int mf(int[::1] f):
    cdef int array_length =  len(f)
...

что приводит к менее непрозрачному C-коду - тот, по крайней мере, мой компилятор, может лучше оптимизировать:

 __pyx_t_2 = __pyx_v_array_length;
  for (__pyx_t_3 = 0; __pyx_t_3 < __pyx_t_2; __pyx_t_3+=1) {
    __pyx_v_k = __pyx_t_3;
    __pyx_t_4 = __pyx_v_k;
    __pyx_v_sum = (__pyx_v_sum + (*((int *) ( /* dim=0 */ ((char *) (((int *) __pyx_v_f.data) + __pyx_t_4)) ))));
  }

Самая важная вещь здесь, это то, что мы проясняем для Cython, что память непрерывна, то есть int[::1] по сравнению с int[:] как это видно для NumPy-массивов, для которых возможно stride!=1 должны быть приняты во внимание.

В этом случае сгенерированный на Cthon C-код приводит к тому же ассемблеру, что и код, который я написал бы. Как указал Крисб, добавляя -march=native приведет к векторизации, но в этом случае ассемблер обеих функций снова будет немного отличаться.

Однако, по моему опыту, у компиляторов довольно часто возникают проблемы с оптимизацией циклов, созданных в Cython, и / или легче пропустить детали, которые препятствуют генерации действительно хорошего C-кода. Поэтому моя стратегия для рабочих контуров состоит в том, чтобы писать их на простом C и использовать Cython для их обертывания / доступа к ним - часто это происходит несколько быстрее, потому что можно также использовать выделенные флаги компилятора для этого кода, не затрагивая весь Cython. модуль.