python - 从 Cython 代码生成 SIMD 指令

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 数组的索引不是 efficient
2. 你忘记了cdef sum=0中的int

考虑到这一点，我们得到:

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 数组中看到的那样，为此必须考虑可能的 stride!=1。

在这种情况下，cython 生成的 C 代码导致 same assembler作为the code我会写的。正如 crisb 所指出的，添加 -march=native 会导致向量化，但在这种情况下，两个函数的汇编器会再次略有不同。

但是，根据我的经验，编译器在优化由 cython 创建的循环时经常会遇到一些问题，并且/或更容易遗漏一个细节，从而阻止生成真正好的 C 代码。所以我的工作马循环策略是用纯 C 编写它们并使用 cython 来包装/访问它们 - 通常它会更快一些，因为人们也可以使用专用的编译器标志来截断代码而不影响整个 Cython-模块。