从 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 编译之后（并直接编写一个采用 (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_addlocal（相关部分）

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 中默认假定），这将解决此问题。