学习笔记 —— 基于C加速的Python高效计算 （Cython & pybind11）

TODO
Implementation Cmake --pybind11
Cython & pybind11 性能比较

引言

Python与C++，是目前比较流行的两种计算机。Python 易用但效率低，C++复杂但效率高。 所以有很多利用C来给python加速的工具。
这些工具一般分为两类： AOT(Ahead of Time) 和 JIT(Just in Time)
但是JIT有"cold start" 的问题，并且即便过了"cold start" ，其效率依旧不如经过人工精致优化的AOT，唯一的优点就是方便易用。 比如Numba，函数上面加个“@jit()”就完事儿了。

个人偏好更有水平的AOT方法，所以本文来介绍两个主流的AOT工具：Cython 和 pybind11

安装很简单：

conda install numpy
conda install cython
conda install pybind11

Cython

示例介绍

使用Cython一般需要三个文件：cython格式编写的库文件cythonfn.pyx, 用于编译cythonfn.pyx的setup.py, 以及调用cython库的python代码 main.py
考虑一种运算：

...
def calculate_z_serial_purepython(maxiter, zs, cs):
    """Calculate output list using Julia update rule"""
    output = [0] * len(zs)
    for i in range(len(zs)):
        n = 0
        z = zs[i]
        c = cs[i]
        while abs(z) < 2 and n < maxiter:
            z = z * z + c
            n += 1
        output[i] = n
    return output
    
def calc_pure_python(desired_width, max_iterations):
    """Create a list of complex coordinates (zs) and complex parameters (cs),
    build Julia set"""
    x_step = (x2 - x1) / desired_width
    y_step = (y1 - y2) / desired_width
    x = []
    y = []
    ycoord = y2
    while ycoord > y1:
        y.append(ycoord)
        ycoord += y_step
    xcoord = x1
    while xcoord < x2:
        x.append(xcoord)
        xcoord += x_step
    # build a list of coordinates and the initial condition for each cell.
    # Note that our initial condition is a constant and could easily be removed,
    # we use it to simulate a real-world scenario with several inputs to our
    # function
    zs = []
    cs = []
    for ycoord in y:
        for xcoord in x:
            zs.append(complex(xcoord, ycoord))
            cs.append(complex(c_real, c_imag))

    output = [0] * len(zs)
    start_time = time.time()
    output = calculate_z_serial_purepython(max_iterations, zs, cs)    ## attention!!!
    end_time = time.time()
    secs = end_time - start_time
    print(calculate_z_serial_purepython.__name__ + " took", secs, "seconds")
    ...
if __name__ == "__main__":
    secs = calc_pure_python(desired_width=1000, max_iterations=300)

经过分析不难发现（学习笔记 —— python代码耗时及内存占用测试方法 以及一些零碎的python小工具），进行循环计算的如下函数是耗时大户，也是我们最需要优化的对象。

由于*.pyx本身是基于python的，因此利用Cython我们可以看到python到C的演变过程，以及各阶段的优化力度。
现在优化前，我的耗时是这样的：

第一阶段优化

编写cythonfn.pyx文件， 把calculate_z_serial_purepython()函数原封不动copy过来：

def calculate_z_serial_purepython(maxiter, zs, cs):
    """Calculate output list using Julia update rule"""
    output = [0] * len(zs)
    for i in range(len(zs)):
        n = 0
        z = zs[i]
        c = cs[i]
        while abs(z) < 2 and n < maxiter:
            z = z * z + c
            n += 1
        output[i] = n
    return output

编写setup.py：

from distutils.core import setup
from Cython.Build import cythonize

setup(ext_modules=cythonize("cythonfn.pyx",
                            compiler_directives={'language_level': "3"}
                            ))

运行setup.py 编译cythonfn.pyx

你会发现多了两个文件：

现在我们用这个cython编译后的calculate_z_serial_purepython()替代之前的：

import cythonfn

def calc_pure_python(desired_width, max_iterations):
	...
    start_time = time.time()
    output = cythonfn.calculate_z_serial_purepython(max_iterations, zs, cs)   # cython
    end_time = time.time()
    secs = end_time - start_time
    print(calculate_z_serial_purepython.__name__ + " took", secs, "seconds")

效率快提高一倍了，厉害吧？但如标题所说，这才刚刚开始！
如果你图懒，觉得这就够了，建议你不要考虑AOT方法了，去找JIT的吧，比如Numba。

第二阶段优化

Cython Annotation tool

在第二阶段开始前，我们先介绍一个Cython自己的可视化工具，可以把Cython编译后的底层C代码展示给我们：

运行命令：

cython -a cythonfn.pyx

会生成一个 cythonfn.html的文件，打开：

现在我们通过点击原代码对应行左边的加号便可展开看到被编译后的底层C代码（我只知道你看不懂哈哈）。
我们需要关注的是颜色的深浅，颜色越深表示对Python virtual machine的调用越多，也就越慢； 纯白色表示就可以认为是C代码。

好了，现在我们的目标就是：让这个界面越白越好！

优化方法

Cython中可以通过 cdef 命令直接定义C层次的变量，避免Python的type识别：

def calculate_z_serial_purepython(maxiter, zs, cs):
    """Calculate output list using Julia update rule"""
    cdef unsigned int i,n
    cdef double complex z,c
    output = [0] * len(zs)
    for i in range(len(zs)):
        n = 0
        z = zs[i]
        c = cs[i]
        while abs(z) < 2 and n < maxiter:
            z = z * z + c
            n += 1
        output[i] = n
    return output

现在看是不是白了许多？尤其是最耗时的第14行，纯白色！这意味着完全脱离的Python到达了C的层次！

现在重新用 setup.py 编译后的结果如下：

比起最开始提高了一个数量级不止，有没有种恐怖如斯的感觉？

第三阶段优化

我们都知道python中常用numpy来进行科学计算，会加速很多。但numpy是一种封装好的类，而C喜欢直接对raw_data，如double[] 操作，因此np.array直接传递给C有点臃肿。

好在python使用了一种缓存协议(buffer protoco)，可以直接把缓存中的数据传给C。

import numpy as np
cimport numpy as np

def calculate_z_serial_purepython(int maxiter, double complex[:] zs, double complex[:] cs):
    """Calculate output list using Julia update rule"""
    cdef unsigned int i,n
    cdef double complex z,c
    cdef int[:] output = np.empty(len(zs), dtype=np.int32)
    # output = [0] * len(zs)
    for i in range(len(zs)):
        n = 0
        z = zs[i]
        c = cs[i]
        # while abs(z) < 2 and n < maxiter:
        while  n < maxiter and z.real*z.real+z.imag*z.imag < 4:
            z = z * z + c
            n += 1
        output[i] = n
    return output

因为用了numpy，所以setup.py也得做相应改动：

from distutils.core import setup
from Cython.Build import cythonize
import numpy

setup(ext_modules=cythonize("cythonfn.pyx",
                            compiler_directives={'language_level': "3"}
                            ), include_dirs=[numpy.get_include()])           

现在相比之前是不是又白了很多？


又在第二阶段的基础上快了一倍。
没有很大提高是因为最耗时的第18行已经优化到极致了（不能更白了），针对这个函数只是优化了数据存取而已。

好了，经过三个阶段的优化，我们把原本的5.49s耗时缩减到了 0.17s，足足32倍!

比对下 JIT的Numba

函数上方加个“@jit()”

from numba import jit
@jit()
def calculate_z_serial_purepython_numba(maxiter, zs, cs, output):
    """Calculate output list using Julia update rule"""
    for i in range(len(zs)):
        n = 0
        z = zs[i]
        c = cs[i]
        while abs(z) < 2 and n < maxiter:
            z = z * z + c
            n += 1
        output[i] = n

看看效果：

只能说和第一阶段的cython等同。
差的那几百毫秒怕是“cold start”造成的。

总结

方法	Time Usage
无优化	5.4992077350616455 seconds
cython第一阶段	3.16896915435791 seconds
cython第二阶段	0.37267589569091797 seconds
cython第三阶段	0.17330312728881836 seconds
Numba	3.286240577697754 seconds

pybind11

Links

Pybind11 github
Pybind11 doc

Introduction

Cython使用了用C编写的独立共享库，并将它们转换为Python模块。实际还是在Python上写库。
PyBind11相反，它提供C++ API，写的是C++代码库，然后转成Python模块。
但Pybind11只支持C++，使用C++ 11特性。
此外，Pybind11的另一个优点是可以轻松处理来自Python的底层数据和类型(如Numpy array)。
小知识：TensorFlow也用了Pybind11。

Pybind11的编译可以用g++或者Cmake，二者在VS上都要有很多配置，个人觉得很麻烦，所以建议摆脱win，到linux上去用。

Implementation

考虑一个AXPY形式 的Python 矩阵运算：
$$y\leftarrow \alpha x+y$$
AXPY 有两种变体：SAXPY 和 DAXPY。 分别对应 Single Precision (32-bit float) 和 Double Precision (64-bit double)
这意味着我们需要区分double和float类型才能正常使用pybind。

以SAXPY 为例，编写核函数：

void axpy(size_t n, float a, float *x, float *y)
{
  for (size_t i = 0; i < n; i++)
    y[i] = a * x[i] + y[i];
}

前面我们说Pybind11可以轻松处理Numpy array，这还是float* ， 没看到Numpy啊，难道这就行了，可以自动转换？当然不行，自动就需要多余的识别，自然就慢了。

别着急，我们再写个 wrapper to handle the Numpy arrays：（使用模板类型 py::array_t<T> 来存取Numpy数据）

小知识：Python supports a very general and convenient way for exchanging data between libraries through a buffer protocol. In this sense, the py::array_t<T> can be considered a buffer object that is restricted to Numpy arrays. The information of the buffer, such as the dimension and size, are stored in an information object called py::buffer_info.

#include <pybind11/numpy.h>
using namespace py = pybind11;
void axpy_wrapper(float a, py::array_t<float> py_x, py::array_t<float> py_y)
{
  // Request for buffer information
  py::buffer_info x_buffer = py_x.request();
  py::buffer_info y_buffer = py_y.request();

  // check dim
  if (x_buffer.ndim != 1 || y_buffer.ndim != 1) {
    throw std::runtime_error("Error: dimension of vector should be 1");
  }

  // check shape match
  if (x_buffer.shape[0] != y_buffer.shape[0]) {
    throw std::runtime_error("Error: size of X and Y do not match");
  }

  // extract raw pointer
  float *x = (float*)x_buffer.ptr;
  float *y = (float*)y_buffer.ptr;
  
  return axpy(x_buffer.shape[0], a, x, y);
}

For more information on format checking and performance optimization, refer to here and here

完整example.cpp代码如下:

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>     // Include header that supports Numpy arrays

namespace py = pybind11;

void axpy(size_t n, float a, float *x, float *y)
{
  for (size_t i = 0; i < n; i++)
    y[i] = a * x[i] + y[i];
}

void axpy_wrapper(float a, py::array_t<float> py_x, py::array_t<float> py_y)
{
  // Request for buffer information
  py::buffer_info x_buffer = py_x.request();
  py::buffer_info y_buffer = py_y.request();

  // check dim
  if (x_buffer.ndim != 1 || y_buffer.ndim != 1) {
    throw std::runtime_error("Error: dimension of vector should be 1");
  }

  // check shape match
  if (x_buffer.shape[0] != y_buffer.shape[0]) {
    throw std::runtime_error("Error: size of X and Y do not match");
  }

  // extract raw pointer
  float *x = (float*)x_buffer.ptr;
  float *y = (float*)y_buffer.ptr;

  return axpy(x_buffer.shape[0], a, x, y);
}

PYBIND11_MODULE(math_kernels, m)
{
  m.def("axpy", &axpy_wrapper);
}

使用g++编译：

$ g++ -O3 -Wall -shared -std=c++11 -fPIC $(python3 -m pybind11 --includes) example.cpp -o math_kernels$(python3-config --extension-suffix)

然后python中调用试试：

$ python
Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy as np
>>> import math_kernels
>>> x = np.random.random(5)
>>> y = np.random.random(5)
>>> a = 0.5
>>> z = a * x + y
>>> math_kernels.axpy(a, x, y)
>>> z
array([0.72109651, 0.80642238, 1.3581158 , 0.68677673, 0.47000476])
>>> y
array([0.67480229, 0.67799412, 0.97670182, 0.20975148, 0.35173144])
>>> z - y
array([0.04629421, 0.12842827, 0.38141398, 0.47702525, 0.11827332])
>>> 

哈哈哈，pybind11的结果完全是错的！！！
还是数据类型的问题呀！ Numpy arrays的默认数据类型是 np.float64，也就是 double， 咱们上面C库里写的可是float哦~
这下你知道即便是 float 和 double 甚至常常在C代码中都会被忽略数据类型冲突的两种类型，在pybind11里搞错后有多大影响了吧？

但是你有没有发现，数据类型不一样它不给你报错，然后给了个错的离谱的结果！ 这点是很坑的！因为pybind11会做自动类型转换，还转不对！所以大家养好习惯，在PYBIND11_MODULE中做一些声明，不让他自动转：

/* name of module-        ---a variable with type py::module_ to create the binding */
/* (这个才是python里   |       |                                                                                                */
/*  import的库名)     v       v                                                                                               */
PYBIND11_MODULE(math_kernels, m) {                            // Create a module using the PYBIND11_MODULE macro
    m.doc() = "pybind11 math module";    // optional module docstringm.def("sum", &sum, "sum two numbers in double"); // calls module::def() to create generate binding code that exposes sum()
    m.def("axpy", &axpy_wrapper, py::arg("a"), py::arg("x").noconvert(true), py::arg("y").noconvert(true)); // 第一个参数是python调用时用的函数名
}

现在我们编译再试：

$ python
Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy as np
>>> import math_kernels
>>> x = np.random.random(5).astype(np.float32)
>>> y = np.random.random(5)
>>> math_kernels.axpy(0.5, x, y)
Traceback (most recent call last):
  File "", line 1, in 
TypeError: axpy(): incompatible function arguments. The following argument types are supported:
    1. (a: float, x: numpy.ndarray[float32], y: numpy.ndarray[float32]) -> None

Invoked with: 0.5, array([0.583261  , 0.4329178 , 0.1882612 , 0.94220173, 0.56544894],
      dtype=float32), array([0.65019744, 0.12314781, 0.04364232, 0.54704202, 0.18679054])
>>> 

这次你数据类型不对就别想跑通了嘿嘿。

但是我们C里定义的是模板T呀，我们加个double的声明不就完了：

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>    // Include header that supports Numpy arrays

namespace py = pybind11;

template <typename T>
void axpy(size_t n, T a, T *x, T *y)
{
  for (size_t i = 0; i < n; i++)
    y[i] = a * x[i] + y[i];
}

template <typename T>
void axpy_wrapper(T a, py::array_t<T> py_x, py::array_t<T> py_y)
{
  // Request for buffer information
  py::buffer_info x_buffer = py_x.request();
  py::buffer_info y_buffer = py_y.request();

  // check dim
  if (x_buffer.ndim != 1 || y_buffer.ndim != 1) {
    throw std::runtime_error("Error: dimension of vector should be 1");
  }

  // check shape match
  if (x_buffer.shape[0] != y_buffer.shape[0]) {
    throw std::runtime_error("Error: size of X and Y not match");
  }

  // extract raw pointer
  T *x = (T*)x_buffer.ptr;
  T *y = (T*)y_buffer.ptr;
  
  return axpy<T>(x_buffer.shape[0], a, x, y);
}

PYBIND11_MODULE(math_kernels, m)
{
  m.def("axpy", &axpy_wrapper<float>,  py::arg("a"), py::arg("x").noconvert(true), py::arg("y").noconvert(true));
  m.def("axpy", &axpy_wrapper<double>, py::arg("a"), py::arg("x").noconvert(true), py::arg("y").noconvert(true));
}

再试试：

$ python
Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy as np
>>> import math_kernels
>>> x = np.random.random(5).astype(np.float32)
>>> y = np.random.random(5).astype(np.float32)
>>> a = 0.5
>>> z = a * x + y
>>> math_kernels.axpy(a, x, y)
>>> y - z
array([0., 0., 0., 0., 0.], dtype=float32)
>>> x = np.random.random(5).astype(np.float64)
>>> y = np.random.random(5).astype(np.float64)
>>> z = a * x + y
>>> math_kernels.axpy(a, x, y)
>>> y - z
array([0., 0., 0., 0., 0.])
>>>

OK，现在你只要保证输入数据类型同时是 float 或者同时是 double 就行啦。

现在你可以尝试一下用pybind11实现上面Cython中提到的**calculate_z_serial_purepython()**看看有多少提高啦！

Implementation Cmake

如果只是简单测试，g++当然足够了。但大多数情况下，我们需要Cmake来管理一整个项目，因此如果使用Cmake编译pybind11变得很重要！

参考文档

cytpes

还有一种方法与pybind11很像， ctypes，都是用g++编译后调用。
不过个人感觉pybind11更方便一点，效率暂且没做比较
下面是ctypes的一个简单示例：

定义C源文件 sum_src.c.

#include <stdio.h>
#include "sum.h"

double sum(double a, double b)
{
  printf("hello world from sum()\n");
  return a + b;
}
```h
和头文件**sum.h**:
```cpp
#ifndef __SUM_H__
#define __SUM_H__
double sum(double, double);
#endif

gcc编译为 libsum.so:

gcc -fPIC -shared -o libsum.so sum_src.c

然后python调用：

Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import ctypes
>>> _libsum = ctypes.CDLL('libsum.so')	# 加载共享库，定义函数签名，以便ctypes可以转换类型
>>> _libsum.sum.argtypes = [ctypes.c_double, ctypes.c_double]	# 指定参数类型
>>> _libsum.sum.restype = ctypes.c_double 	# 指定返回类型
>>> _libsum.sum(ctypes.c_double(3.1), ctypes.c_double(3.3)) # 函数调用
hello world from sum()
6.4
>>> 

python调的时候还需要再指定遍类型，是不是有点麻烦？另外，所有参数都必须转换为各自的基本数据类型， 也就是你用numpy的话还得转一下，而不是pybind那种可以直接调用。
ctypes的一个特性是可以定义回调函数以允许C程序调用Python函数。有关更多信息，请参阅 文档。

Cython & pybind11 性能比较

win 下(Pybind11未成功调用)：

总结

小矩阵下： cython>torch-cpu>numpy>torch-gpu>cupy

中大矩阵下： torch-gpu>cupy>cython>torch-cpu>numpy

Ubuntu下：

总结

小矩阵下： pybind11>cython>torch-cpu>numpy>torch-gpu>cupy

大矩阵下： torch-gpu>cupy>pybind11>cython>numpy>torch-cpu

容易理解，Cython还是在pyhon上编写的，依旧需要Python的一些相对于C++比较多余的解码机制；而pybind11直接在C++写库，并且一定要指定数据类型，没有多余的解码，自然会快一些。

以上分析来自代码：Assignments/Assignment 3/e2GaussSeidel.py 与 e2GaussSeidel_gpu.py

最后的总结

大矩阵能并行最好通过torch用GPU
小点的运算就用 pybind(Ubuntu) 或 cython(win) 提速

Reference

学习笔记 —— python代码耗时及内存占用测试方法 以及一些零碎的python小工具
Cython Doc
Cython github
Pybind11 github
Pybind11 doc