Cython 缓冲协议

python

Cython 缓冲协议

缓冲协议

协议（protocol）是一种接口规范，比如在 Python 中， Cats 类如果实现了 __len__() 和 __getitem__() 这两个方法，它就满足了序列协议，就可以被看作是一个序列，我们就可以用 cats[0] 或者 len(cats) 这样的方法来处理 Cats 对象。 Cats 具体是哪个类的子类无关紧要，只需要它满足了对应的协议就可以得到预期场景的行为，协议通过鸭子类型实现了多态。

面对不同的场合和需求， Python 规定了各式各样的协议。比如迭代器协议，规定了只要实现 __iter__() 和 __next__ 方法的对象就可以被迭代（比如用在 for 循环中）；容器协议，规定了只要实现 __contains__(self, item) 方法的对象，就支持检查该对象是否包含某个元素（使用 in 关键字）；数值协议，规定了一个对象如果想要进行 + 、 - 、 * 等数学运算，就需要实现 __add__ 、 __sub__ 、 __mul__ 等方法。

其中一个与高性能计算相关的 Python 底层协议，叫作缓冲协议（Buffer Protocol），它允许对象以一种标准的方式共享底层内存数组（或称为缓冲）数据。在 Python 中，一个对象要实现缓冲协议，它需要定义 __getbuffer__ 和 __releasebuffer__ 方法暴露其内部数据。在之前的博客 Python的C扩展 中提到过， Python 中一切皆对象，一切都是 PyObject 指针。比如一个浮点数，除了对应C语言层面的8个字节的 double ，还存储了包括引用计数、类型指针等各式各样的东西。而缓冲协议提供了一种方式，允许对象直接暴露其内部底层内存数据（通常是连续内存数组）给其他对象。这样一方面，通过共享内存而不是拷贝数据，我们可以大幅缩减内存使用，对于处理大型数据集非常有效；另一方面，缓冲不是 PyObject 指针而是简单的底层 C 结构，这也可以显著提高计算效率。

Python内置了一个叫作 memoryview 的类，它是建立在缓冲协议之上的一个高级接口。它允许 Python 代码以数组的形式安全地操作其他对象的底层数据，而无需实际复制任何数据。通过 memoryview ，我们可以对支持缓冲协议的对象（如 bytes , bytearray , array.array , numpy 数组等）的数据进行切片、访问和修改，而不会引起数据的复制，对于大型数据集的处理非常有效。

1
2
3
4
5
6
7
8
import numpy as np

array_1 = np.arange(5, dtype=np.float64)
m1 = memoryview(array_1)
array_2 = np.zeros(5, dtype=np.float64)
m2 = memoryview(array_2)
m1[2:] = m2[2:]
print(array_1)
1
[0. 1. 0. 0. 0.]

memoryview 对象不支持类似和 numpy 数组一样的广播操作，但是可以通过索引的方式来取值，或者通过切片的方式进行数值的修改。

在 Cython 中使用数组

在 Cython 中高效地操作数组，实际上也是利用了 numpy 数组的缓冲协议。大体上， Cython 提供了两种方式允许我们直接处理 numpy 数组的内存缓冲区，在 C 层面直接访问底层的内存，绕过 Python 的额外开销。一种是从早期版本就有的实现，叫作缓冲区语法，通过声明 numpy 数组的类型和维数， Cython 知道如何操作对应数组的底层内存。另一种是后来提供的通用接口，类似于 memoryview ——类型化内存视图，统一并简化了对所有类型数组的处理和访问。

缓冲区语法

缓冲区语法非常简单，就是在方括号内指定类型和维数来声明一个 numpy 数组，比如我们要用缓冲区语法声明一个二维的 double 数组，可以使用如下代码

1
cdef np.ndarray[double, ndim=2] arr

声明之后，Cython再访问这个数组时，就会直接操作它对应底层的内存了。下面这个例子比较了是否采用缓冲区语法对程序性能的影响

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
%%cython
import numpy as np
cimport numpy as np

cpdef numpy_bench_py():
    py_arr = np.random.rand(1000)
    cdef int i
    for i in range(1000):
        py_arr[i] += 1
        
cpdef numpy_bench_c():
    cdef np.ndarray[np.float64_t, ndim=1] c_arr
    c_arr = np.random.rand(1000)
    cdef int i
    for i in range(1000):
        c_arr[i] += 1

这里类型命名的 _t 是 Cython 的命名规范代表底层C的类型。从测量结果可以看出来，类型化注解后程序速度快了25倍。

1
2
3
4
%timeit numpy_bench_py()
162 µs ± 1.38 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit numpy_bench_c()
6.7 µs ± 44.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

类型化内存视图

类型化内存视图和最前面介绍的 memoryview 比较像，比如我们要定义一个 int 内存视图和一个二维的 double 内存视图，可以声明

1
2
cdef int[:] a
cdef double[:, :] b

这里的 double[:] 并不代表声明了一个 double 数组，而是说是声明了一个 double 类型的内存视图。内存视图是一个对象，维护着一个指向特定内存区域的引用。该内存区域实际上并不归内存视图所有，但内存视图能够读取和修改其内容；换而言之，内存视图是一个有关底层数据的视图，而不是一个实际数组。这种语法也可以用于函数定义中，当 Cython 中的函数声明了一个内存视图类型的参数，任何暴露了缓冲区接口的对象（如 numpy数组）都会自动绑定到该内存视图。理解内存视图最关键的就是，内存视图并不拥有与之绑定的数据，而只是提供了一种访问和修改它们的途径。通过内存视图修改数据时，操作的是底层的内存区域，因此这种修改将在原始数组中反映出来（反之亦然）。

1
2
3
4
5
6
7
8
%%cython
import numpy as np
cimport numpy as np

cpdef double[:] miaomiao(double[:] arr):
    cdef double[:] b
    b = arr
    return b

我们执行这个函数后，得到是一个类似于 MemoryView 的对象：

1
2
3
arr = np.array([1.0, 2.0, 3.0], dtype = np.float64)
b = miaomiao(arr)
b
1
<MemoryView of 'ndarray' at 0x7fe461aea1e0>

我们可以通过索引的方式直接返回对应内存的数值，直接修改 b 的值， arr 的值也会随之改变。

一些注意事项

  • 对于所有的变量我们要都声明其类型，让 Cython 能够进行优化。否则相比于纯粹的 Python 代码，编译没有指定静态类型信息的 Cython 模块没有太多优势。

  • 禁用 boundscheck 边界检查和 wraparound 负索引可以在基本不影响程序的基础上显著提升性能。boundscheck 控制 Cython 是否生成检查数组索引是否越界的代码。当启用（默认情况下为启用）时，Cython会检查每次数组访问是否在数组的有效范围内，如果索引越界，则抛出一个错误； wraparound指令控制 Cython 是否允许使用负索引按照 Python 的语义进行数组访问。在 Python 中，负索引表示从数组的末尾开始反向索引，如 -1 是最后一个元素。启用 wraparound （默认情况下为启用）时， Cython 会对负索引进行特殊处理，以支持这种访问模式。

    在Cython中调整这些编译指令可以使用cython装饰器或者在.pyx文件顶部通过编译指令注释的方式：

    1
    2
    3
    4
    5
    6
    7
    # cython: boundscheck=False
    # cython: wraparound=False

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def some_function():
        ...

  • 类型化内存视图返回的是对应对象的内存视图，而不是一个新的数组。当我们把原先的 Python 函数重构为 Cython 实现以优化性能的时候，要注意到这里的差异以防止产生错误。

参考

Python高性能-第二版.

流畅的Python-第二版.

Python/C API参考手册.

Cython 3.0文档.