从Cython代码生成SIMD指令

时间:2018-03-01 21:26:38

标签: python cython

我需要概述一下在高性能数字代码中使用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做这样的事情?

1 个答案:

答案 0 :(得分:6)

您的代码中存在两个问题(使用选项-a使其可见):

  1. numpy数组的索引不是efficient
  2. 您已忘记int
  3. 中的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在编译时是不可知的,这会阻止矢量化。

    然而,当使用typed memory views时,你会得到更好的结果,即:

    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-arrays,必须考虑可能的stride!=1

    在这种情况下,cython生成的C代码会产生same assembler the code我想写的http://localhost:9876/。正如crisb指出的那样,添加-march=native会导致向量化,但在这种情况下,两个函数的汇编程序会再次略有不同。

    但是,根据我的经验,编译器经常会出现一些问题来优化cython创建的循环和/或更容易错过一个阻止生成真正优秀C代码的细节。所以我的工作循环策略是用纯C编写它们并使用cython来包装/访问它们 - 通常它会更快一些,因为也可以使用专用的编译器标志来剪切这段代码而不影响整个Cython-模块。