Cython与C函数的结合

技术背景

在前面一篇博客中，我们介绍了使用Cython加速谐振势计算的方法。有了Cython对于计算过程更加灵活的配置（本质上是时间占用和空间占用的一种均衡），及其接近于C的性能，并且还最大程度上的保留了Python的编程语法特点，因此Cython确实是值得Python编程爱好者学习的一种加速手段。这里我们要介绍的是Cython与C语言相结合的一种方案，可以直接在pyx文件中加载C语言代码。

测试场景

我们测一个非常简单的场景————归约求和：

\[S=\sum_{i,j}A_{i,j} \]

当然了，像这种基本运算，在Numpy中已经优化的非常极致了。所以，这里我们并不是要展现Cython在性能上的优势，而是Cython对于C语言和Python语言两者的兼容性。首先我们用C语言实现一个归约求和的简单函数：

// array_sum.c
double reduce_sum(int arr_len, double* arr){
    double s=0.0;
    int i;
    for (i=0; i<arr_len; i++){
        s = s + *arr;
        arr++;
    }
    return s;
}

这里我们使用了一个指针数组，然后用for循环进行遍历计算。在Cython中，我们可以使用extern来直接加载C语言中的这个函数：

# test_pointer.pyx
import numpy as np
cimport numpy as np

cdef extern from "array_sum.c":
    double reduce_sum(int arr_len, double* arr)

cpdef rsum(int arr_len, np.ndarray[np.float64_t, ndim=2, mode="c"] arr):
    cdef:
        double* arr_ptr = <double *>arr.data
        double res = 0.0
    res = reduce_sum(arr_len, arr_ptr)
    return res

这里加载了C语言中的reduce_sum函数，然后以Cython中定义的rsum函数作为一个接口，将传入的numpy数组的内存地址作为指针传给C语言中写好的函数。然后需要对这个pyx文件进行编译构建：

$ cythonize -i test_pointer.pyx

编译完成后会在当前路径下生成*.c文件和*.so文件：

$ ll | grep test_pointer  
-rw-r--r-- 1 root root  374450 Jul 25 14:52 test_pointer.c
-rwxr-xr-x 1 root root  234848 Jul 25 14:52 test_pointer.cpython-37m-x86_64-linux-gnu.so*
-rw-r--r-- 1 root root     347 Jul 25 15:02 test_pointer.pyx

调用Cython函数

我们可以开启一个Ipython，或者直接在Python脚本文件中调用Cython函数：

In [1]: import numpy as np

In [2]: from test_pointer import rsum

In [3]: num=10000

In [4]: x=np.random.random((num,num))

In [5]: x.shape
Out[5]: (10000, 10000)

In [6]: %timeit s=np.sum(x)
38.3 ms ± 254 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [7]: %timeit rs=rsum(num*num,x)
51.7 ms ± 302 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [15]: np.sum(x)
Out[15]: 50003980.32921535

In [17]: rsum(num*num, x)
Out[17]: 50003980.32921728

经过测试，确实可以在Python中调用这个C语言实现的函数。当然，前面也提到过，Numpy对于这个简单的求和计算已经优化的非常好了，所以这里没有体现出性能上的优势，这里更多的是演示一个方法。