从 cython 数组循环的矢量化

Question

考虑以下在 Cython 内存视图上执行就地添加的示例：

#cython: boundscheck=False, wraparound=False, initializedcheck=False, nonecheck=False, cdivision=True
from libc.stdlib cimport malloc, free
from libc.stdio cimport printf
cimport numpy as np
import numpy as np


cdef extern from "time.h":
    int clock()


cdef void inplace_add(double[::1] a, double[::1] b):
    cdef int i
    for i in range(a.shape[0]):
        a[i] += b[i]


cdef void inplace_addlocal(double[::1] a, double[::1] b):
    cdef int i, n = a.shape[0]
    for i in range(n):
        a[i] += b[i]


def main(int N):
    cdef:
        int rep = 1000000, i
        double* pa = <double*>malloc(N * sizeof(double))
        double* pb = <double*>malloc(N * sizeof(double))
        double[::1] a = <double[:N]>pa
        double[::1] b = <double[:N]>pb
        int start
    start = clock()
    for i in range(N):
        a[i] = b[i] = 1. / (1 + i)
    for i in range(rep):
        inplace_add(a, b)
    printf("loop %i\n", clock() - start)
    print(np.asarray(a)[:4])
    start = clock()
    for i in range(N):
        a[i] = b[i] = 1. / (1 + i)
    for i in range(rep):
        inplace_addlocal(a, b)
    printf("loop_local %i\n", clock() - start)
    print(np.asarray(a)[:4])

使用这些 Cython 指令，看似等效的 inplace_add 和 inplace_addlocal 都编译为紧密的 C 循环。但是对于 N=128（我期望的近似大小）inplace_addlocal 比 inplace_add 快两倍（！），在使用 gcc -Ofast 编译之后（并直接编写一个 C 函数a (int, double*, double*) 或多或少与 addlocal 一样快，有或没有 #openmp simd）。将 -fopt-info 传递给 gcc 表明 inplace_addlocal 被向量化，但 inplace_add 没有。

这是 Cython 生成的 C 代码的问题（即，gcc 确实无法推断出向量化代码所需的任何保证），还是 gcc（即，缺少某些优化），或其他问题？

谢谢。

（交叉发布给 cython 用户）

Answer 1

生成的 C 代码的唯一区别是，在 inplace_addlocal 中，循环的结束变量是 int，而在 inplace_add 中，它是 Py_ssize_t。

由于您的循环计数器是 int，因此在 inplace_add 版本中，在执行比较时会因两种类型之间的强制转换而产生额外的开销。

inplace_add（相关部分）

Py_ssize_t __pyx_t_1;
int __pyx_t_2;
int __pyx_t_3;
int __pyx_t_4;

__pyx_t_1 = (__pyx_v_a.shape[0]);
for (__pyx_t_2 = 0; __pyx_t_2 < __pyx_t_1; __pyx_t_2+=1) {
  __pyx_v_i = __pyx_t_2;

inplace_add本地（相关部分）

int __pyx_t_1;
int __pyx_t_2;
int __pyx_t_3;
int __pyx_t_4;

__pyx_v_n = (__pyx_v_a.shape[0]);
__pyx_t_1 = __pyx_v_n;
for (__pyx_t_2 = 0; __pyx_t_2 < __pyx_t_1; __pyx_t_2+=1) {
  __pyx_v_i = __pyx_t_2;

这个 answer 提到最好使用 Py_ssize_t 作为索引（并且在 Cython 中必须默认假定），这将解决这个问题。

从 cython 数组循环的矢量化

vectorization of looping on an array from cython

gcc

cython

auto-vectorization