《Cython系列》7. Cython、numpy、以及类型化memoryview

楔子

Cython 的两个优秀的品质就是它的广度和成熟度,可以编译所有的 Python 代码,并且将 C 的速度代入了 Python,并且还能轻松的和 C、C++ 集成。而本篇文章的任务就是完善 Cython 的功能,并介绍 Cython 的阵列特性,比如:对 Numpy 数组的深入支持。

我们已经知道,Cython 可以很好的支持 list、dict、set、tuple 等内置容器,这些容器非常容易使用,可以包含指向任意 Python 对象的变量,并且对 Python 对象的查找、分配、检索都进行了高度的优化。尽管 list 类型和 dict 类型之间实现存储和检索的方式有很大差异,但是从实现角度来讲,所有的容器都有一个共同点:它们存储的都是对 Python 对象的一个引用。如果我们有一个包含100万个整型的列表,那么在 C 中,每一个元素都是一个 Python *。虽然整型在底层对应 PyLongObject,但是会转成 PyObject * 存在列表里面,因为 C 中的数组只能存放相同类型的元素(指针也是相同类型的指针),然后用的时候再转回对应的类型,而 PyObject 里面的 ob_type 成员便是指向对应类型对象的指针。所以 Python 中的列表可以存储任意类型的变量,因为它们都是一个 PyObject *。但我们说,由于每次运算的时候都要进行类型转化,所以效率低,尽管我们知道都是整型,但是 Python 不会做这种假设。

因此对于同质容器(只包含一种固定的类型),可以在存储开销和性能方面做得更好,比如 Numpy 中的数组。Numpy 中的数组就有一个准确的类型,这样的话在存储和计算的时候可以表现的更加优秀,我们举个栗子:

import numpy as np

# 创建一个整型数组
# 可以指定 dtype 为: int, int8, int16, int32, int64
# 或者是: uint, uint8, uint16, uint32, uint64
arr1 = np.zeros((3, 3), dtype="uint32")
print(arr1)
"""
[[0 0 0]
 [0 0 0]
 [0 0 0]]
"""
try:
    arr1[0, 0] = "xx"
except Exception as e:
    print(e)  # invalid literal for int() with base 10: 'xx'
# 我们看到报错了,因为已经规定了arr 是一个整型数组,那么在存储和计算都是按照整型来处理的
# 既然是整型数组,那么赋值一个字符串是不允许的

arr1[0, 0] = -123
print(arr1)
"""
[[4294967173          0          0]
 [         0          0          0]
 [         0          0          0]]
"""
# 因为是 uint32, 只能存储正整数, 所以结果是 uint32 的最大值 - 123
print((2 << 31) - 123)  # 4294967173


# 创建一个浮点数组, 可以指定 dtype 为: float, float16, float32, float64
arr2 = np.zeros((3, 3), dtype="float")
print(arr2)
"""
[[0. 0. 0.]
 [0. 0. 0.]
 [0. 0. 0.]]
"""


# 创建一个字符串类型, dtype可以为: U, str
# 如果是 U, 那么还可以加上一个数值, 比如: U3, 表示最多存储3个字符。
# 并且还可以通过 <U 或者 >U 的方式来指定是小端存储或者大端存储
arr3 = np.zeros((3, 3), dtype="U3")
print(arr3)
"""
[['' '' '']
 ['' '' '']
 ['' '' '']]
"""
arr3[0, 0] = "古明地觉"
print(arr3)
"""
[['古明地' '' '']
 ['' '' '']
 ['' '' '']]
"""
# 我们看到被截断了,并且截断是按照字符来的,不是按照字节


# 创建一个Python对象, 注意:没有 tuple、list、dict 等类型
# 特定的类型只有整型、浮点型、字符串,至于其它的类型统统用 object 表示
# 可以指定 dtype="O"
arr4 = np.zeros((3, 3), dtype="O")
print(arr4)
"""
[[0 0 0]
 [0 0 0]
 [0 0 0]]
"""
# 虽然打印的也是 0,但它是一个 object 类型
print(arr4.dtype)  # object
# 我们可以使用 empty 创建
print(np.empty((3, 3), dtype="O"))
"""
[[None None None]
 [None None None]
 [None None None]]
"""

而实现同质容器是通过缓冲区的方式,它允许我们将连续的简单数据使用单个数据类型表示,支持缓冲区协议的 Numpy 数组在 Python 中则是使用最广泛的数组类型,所以 Numpy 数组的高性能是有目共睹的。

有效地使用缓冲区通常是从 Cython 代码中获得 C 性能的关键,而幸运的是,Cython 使处理缓冲区变得非常的容易,它对新缓冲区协议和 Numpy 数组有着一流的支持。

缓冲区协议

下面来说一下缓冲区协议,缓冲区协议是一个 C 级协议,它定义了一个具有数据缓冲区和元数据的 C 级结构体,并用它来描述缓冲区的布局、数据类型和读写权限,并且还定义了支持协议的对象所必须实现的 API。而实现了该协议的 Python 对象之间可以向彼此暴露自身的原始字节数组,这在科学计算中非常有用,因为在科学计算中我们经常会使用诸如 Numpy 之类的包来存储和操作大型数组,因为要对数据做各种各样的变换,所以难免要涉及到数组的拷贝。而使用缓冲区协议,那么数组之间就可以不用拷贝了,而是共享同一份数据缓冲区,这些数组都可以看成是该缓冲区的一个视图,那么也意味着操作任何一个数组都会影响其它数组。

都有哪些类型实现了缓冲区协议呢?

Numpy 的 ndarray

Python 中最知名、使用最广泛的 Numpy 包中有一个 ndarray 对象,它是一个支持缓冲区协议的有效 Python 缓冲区。

Python2 中的 str

Python2 中的 str 也实现了该协议,但是 Python3 的 str 则没有。

Python3 中的 bytes 和 bytearray

既然 Python2 中的 str 实现了该协议,那么代表 Python3 的 bytes 也实现了,当然还有 bytearray。

标准库 array 中的 array

Python 标准库中有一个 array 模块,里面的 array 也实现了该协议,但是我们用的不多。

标准库 ctypes 中的 array

这个我们用的也不多。

其它的第三方数据类型

比如第三方库 PIL,用于处理图片的,将图片读取进来得到的对象的类型也实现了缓冲区协议。当然这个很好理解,因为它们读取进来可以直接转成 Numpy的 ndarray。

缓冲区协议最重要的特性就是它能以不同的方式来表示相同的底层数组,它允许 Numpy 数组、几个 Python 的内置类型、标准库的数组之间共享相同的数据,而无需再拷贝。当然 Cython 级别的数组也是可以的,并且使用 Cython,我们还可以轻松地扩展缓冲区协议去处理来自外部库的数据(后面说)。

我们举个栗子,看看共享不同类型的数据之间如何共享内存:

import array
import numpy as np

"""
'b'         signed integer    
'B'         unsigned integer  
'u'         Unicode character 
'h'         signed short    
'H'         unsigned short  
'i'         signed int   
'I'         unsigned int
'l'         signed long    
'L'         unsigned long  
'q'         signed long long    
'Q'         unsigned long long  
'f'         float    
'd'         double    
"""
arr = array.array("i", range(6))
print(arr)  # array('i', [0, 1, 2, 3, 4, 5])

# array(数组) 是 Python 标准库中提供的数据结构,它不是 Python 内置的
# 数组不同于列表,因为数组里面存储的都是连续的整数块,它的数据存储要比列表紧凑得多
# 因此一些操作也可以更快的执行
np_arr = np.asarray(arr)
print(np_arr)  # [0 1 2 3 4 5]

np_arr[0] = 123
print(arr)  # array('i', [123, 1, 2, 3, 4, 5])

Python 提供的数组使用的不是特别多,而 Numpy 中的数组使用的则是非常广泛,并且支持的操作非常丰富。Python 标准库提供的数组和 Numpy 提供的数组都实现了缓冲区协议,因此可以共享同一份数据缓冲区,它们在转化的时候是不用复制原始数据的。所以 np_arr 在将第一个元素修改之后,打印 arr 也发生了变化。

当然 np.asarray 等价于不拷贝的 np.array:

import array
import numpy as np

arr = array.array("i", range(6))

# np.array 内部有一个 copy 参数,默认是 True,也就是会将原始数组拷贝一份
np_arr1 = np.array(arr)
np_arr1[0] = 123
# 此时 arr 是没有变化的,因为操作的不是同一个数组
print(arr)  # array('i', [0, 1, 2, 3, 4, 5])

# 不进行拷贝,此时会共享缓冲区
np_arr2 = np.array(arr, copy=False)
np_arr2[0] = 123
# 因此结果变了
print(arr)  # array('i', [123, 1, 2, 3, 4, 5])

问题来了,如果我们将 array 换成 list 的话会怎么样呢?

import numpy as np

s = [1, 2, 3]
np_arr = np.asarray(s)
np_arr[0] = 123
print(s)  # [1, 2, 3]

因为列表不支持、或者说没有实现缓冲区协议,所以 Numpy 没办法与之共享数据缓冲区,因而会将数据拷贝一份。

可能有人觉得以现如今的硬件来说,根本不需要考虑内存占用方面的问题,但即便如此,共享内存也是非常有必要的。因为在科学计算中,大部分的经典算法都是采用编译型语言实现的,像我们熟知的 scipy 本质上就是基于 NetLib 实现的一些包装器,NetLib 才是负责提供大量高效算法的工具箱,而这些高效算法几乎都是采用 Fortran 和 C 编写的。Python 能够和这些编译库(NetLib)共享本地数据对于科学计算而言非常重要,这也是 Numpy 被开发的一个重要原因,更是缓冲区协议被提出、并添加到 Python 标准库的原因。

而这里我们提一下 PyPy,我们知道它是用 CPython 编写的 Python 解释器,它的速度要比 Python 快很多,但是对于使用 Python 进行科学计算的人来说却反而没什么吸引力。原因就是在科学计算时所使用的算法实际上都是采用 Fortran 和 C 等语言编写的、并被编译成库的形式,Python 只是负责对这些库进行包装、提供一个友好的接口,因此这意味着 Python 能够与之进行很好的交互。而 PyPy 还无法做到这一点,因此现在用的解释器基本都是 CPython,至于 PyPy 引入 JIT(即时编译)所带来的性能收益实际上用处不大。

Python 能成为有效的科学计算平台,主要得益于缓冲区协议的实现和 Numpy。

那缓冲区协议到底长什么样?

Python 的缓冲区协议本质上是一个结构体,它为多维数组定义了一个非常灵活的接口,我们看一下底层实现,源码位于 Include\cpython\object.h 中(Python 3.9),小于 Python 3.9 的话则位于 Include\object.h 中。

/* buffer interface */
typedef struct bufferinfo {
    void *buf;
    PyObject *obj;        /* owned reference */
    Py_ssize_t len;
    Py_ssize_t itemsize;  /* This is Py_ssize_t so it can be
                             pointed to by strides in simple case.*/
    int readonly;
    int ndim;
    char *format;
    Py_ssize_t *shape;
    Py_ssize_t *strides;
    Py_ssize_t *suboffsets;
    void *internal;
} Py_buffer;

以上就是缓冲区协议的底层定义,我们来解释一下里面的成员都代表什么含义,至于如何实现我们一会再说。

void *buf:

实现了缓冲区协议的对象的内部缓冲区(指针)。

PyObject *obj:

实现了缓冲区协议的对象(指针),比如 ndarray 对象、bytes 对象等等。

Py_ssize_t len:

不要被名字误导了,这里表示缓冲区的总大小。比如一个 shape 为 (3, 4, 5) 的数组,存储的元素是 8 字节的 int64,那么这个 len 就等于 3 *4 * 5 * 8。

Py_ssize_t itemsize:

缓冲区存储的元素都是同质的,每一个元素都占用相同的字节,而 itemsize 就表示每个元素所占的大小。比如缓冲区存储的元素是 int64,那么 itemsize 就是 8。

int readonly:

缓冲区是否是只读的,为 0 表示可读写,为 1 表示只读。

int ndim:

维度,比如 shape 为 (3, 4, 5) 的数组,那么底层缓冲区的维度就是 3。注意:如果 ndim 为 0,那么表示 buf 指向的缓冲区代表的只是一个标量,这种情况下,成员 shape、strides、suboffsets 都必须为 NULL。

而且维度最大不超过 64,但在 Numpy 里面支持的最大维度是 32。

char *format:

格式化字符,用于描述缓存区元素的类型,和 Python 标准库 struct 使用的 format 是一致的。比如:i 表示 C 的 int,L 表示 C 的 unsigned long 等等。

Py_ssize_t *shape:

这个很好理解,等同于 Numpy array 的 shape,只不过在 C 里面是一个数组。

Py_ssize_t *strides:

长度为 ndim 的数组,里面的值表示在某个维度下,从一个元素到下一个元素所需要跳跃的字节数。举个栗子,假设有一个 shape 为 (10, 20, 30) 的数组,里面的元素是 int64,那么 strides 就是 (4800, 240, 8)

因为有三个维度:对于第一个维度而言,每一个元素都是 (20, 30) 的二维数组,所以当前元素和下一个元素的地址差了 20 * 30 * 8 = 4800 个字节;对于第二个维度而言,每一个元素都是 (30,) 的一维数组,所以当前元素和下一个元素的地址差了 30 * 8 = 240 个字节;对于第三个维度而言,每一个元素都是一个标量,所以当前元素和下一个元素的地址差了 8 个字节。

根据 strides 我们可以引出一个概念:full strided array,直接解释的话比较费劲,我们用代码说明。

import numpy as np

arr1 = np.array(range(10), dtype="int64")
print(arr1.strides)  # (8,)

arr2 = arr1[:: 2]
print(arr2.strides)  # (16,)

显然 arr1 和 arr2 是共享缓冲区的,也就是说它们底层的 buf 指向的是同一块内存,但是显然它们的 strides 不同。因此 arr1 从一个元素到下一个元素需要跳跃 8 字节,但是 arr2 则是跳跃 16 个字节,否则就无法实现步长为 2 了。假设把步长从 2 改成 3,那么 arr2 的 strides 显然就变成了 (24,),所以此刻应该对 Numpy 数组有更深的认识了。使用切片,本质上是通过改变 strides 来实现跨步访问,但仍然共享同一个缓冲区,。

import numpy as np

arr1 = np.array(range(10), dtype="int64")
arr2 = arr1[:: 2]

print(arr2)  # [0 2 4 6 8]
arr1[0] = 111
print(arr2)  # [111   2   4   6   8]

arr1[:] = 0
print(arr2)  # [0 0 0 0 0]

既然共享同一个缓冲区,那么改变 arr1 是会影响 arr2 的。

回归正题,以 arr2 为例,由于它只有一个维度,所以 strides 的元素个数为 1,里面的 16 表示数组 arr2 从一个元素到下一个元素所跳跃的字节数。但是问题来了,arr2 里面的元素大小只有 8 字节,所以像这种元素大小和对应的 strides 不相等的数组,我们称之为 full strided 数组。多维数组也是一样,举个栗子:

import numpy as np

arr = np.ones((10, 20, 30), dtype="int64")
print(arr.strides)  # (4800, 240, 8)

arr2 = arr[:: 2]
# arr2 是 full strided,因为在第一个维度中,一个元素到下一个元素应该需要 4800 个字节
# 但是 arr2 的 strides 的第一个元素是 9600,因为不相等,所以是 full strided
print(arr2.strides)  # (9600, 240, 8)

arr3 = arr[:, :: 2]
# arr3 是 full strided,因为在第二个维度中,一个元素到下一个元素应该需要 240 个字节
# 但是 arr3 的 strides 的第二个元素是 480,因为不相等,所以是 full strided
print(arr3.strides)  # (4800, 480, 8)


arr4 = arr[:, :, :: 2]
# arr4 是 full strided,因为在第三个维度中,一个元素到下一个元素应该需要 8 个字节
# 但是 arr4 的 strides 的第三个元素是 16,因为不相等,所以是 full strided
print(arr4.strides)  # (4800, 240, 16)

说白了,只要任意维度出现了数组的跨步访问、且步长不为 1,那么这个数组就是 full strided 数组。之所以要说这个 full strided,是因为后面会用到。

下面我们来实现缓冲区协议,由于会涉及到 C 和 Python / C API,所以可能比较枯燥,而且需要你对 C 语言 和 CPython 底层有一定的了解。当然也可以跳过,因为 Cython 允许我们在不了解协议的情况下使用缓冲区,从而在不复制的情况下高效的访问底层数据、减少开销,有这一点就足够了。

实现缓冲区协议(准备工作)

我们说 Numpy 的数组已经实现了缓冲区协议,我们可以直接使用而不需要关注背后的细节,因为缓冲区协议本质上是一个 Python / C API 实体,如果我们想直接操作它那么必须要深入了解 Python / C API。而 Python 则是抽象了 C 级的混乱,给我们提供了一个简洁、美丽、干净的 Python 接口。

但有些时候我们不得不亲自实现,假设我们有一个 C 对象,需要对其进行包装然后让 Python 可用,那么或许你首先会想到用 Cython 来实现(后面会介绍)。当然这是绝对推荐的,因为这是非常正确的做法,Cython 的目的之一就是让我们能够更方便地调用现有的 C 库,此外除了 Cython 你还可以使用 f2py、swig 等等,这些工具的存在仍然可以让我们不用直接和缓冲区协议打交道。

但话虽如此,可实际了解一下缓冲区协议也是非常有必要的,而想要了解缓冲区协议的最好方式就是手动实现它。一旦了解如何实现缓冲区协议,那么你可以很轻松地将其它语言(比如 Julia)的数据结构,以灵活且优雅的方式暴露给 Python。

下面会涉及到 C 和 Python / C API,需要你对 C 语言 和 CPython 有一定的了解。

// 文件名:my_array.c
#include <stdio.h>
#include <stdlib.h>


// 定义一个一维的数组
typedef struct {
    int *arr;
    int length;
} MyArray;


// 初始化
void initial_MyArray(MyArray *array, int length) {
    // length >= 0
    array->length = length;
    if (length == 0) {
        array->arr = NULL;
    } else {
        array->arr = (int *) malloc(sizeof(int) * length);
        for (int i = 0; i < length; i++) {
            array->arr[i] = i;
        }
    }
}

// 释放内存
void dealloc_MyArray(MyArray *array) {
    if (array->arr != NULL) free(array->arr);
    array->arr = NULL;
}

// 将数组进行格式化方便打印
char *change_array_to_string(MyArray *array) {
    // 我们模仿 Python 中的列表,将 C 数组格式化成字符串,例如 "[1, 2, 3, 4, 5]"
    // 但是最多显示 10 个
    int max = array->length < 10 ? array->length : 10;
    char *output = (char *) malloc(100);

    // 长度为 0,直接返回 "[]"
    if (max == 0) {
        sprintf(&output[0], "[]");
        return output;
    }
    // 数据长度不为 0,那么先写入 "["
    int pos = sprintf(&output[0], "[");
    // 然后依次写入数字,由于最后一个数的后面不需要 ", " 所以先设置前 n - 1 个数
    for (int i = 0; i < max - 1; i++) {
        pos += sprintf(&output[pos], "%d, ", array->arr[i]);
    }
    // 然后再将最后一个元素设置进去
    pos += sprintf(&output[pos], "%d", array->arr[max - 1]);

    // 如果数组的长度比 10 大,多余的元素用 ", ..." 显示
    if (array->length > 10) pos += sprintf(&output[pos], ", ...");
    // 写入最后一个 "]"
    sprintf(&output[pos], "]");

    return output;
}

// 打印数组
void print_MyArray(MyArray *array) {
    char *s = change_array_to_string(array);
    printf("%s\n", s);
    free(s);
}

我们上面定义了一个结构体 MyArray,里面包含了一个数组以及该数组的长度信息,同时还定义了数组的初始化、销毁、格式化、打印等函数。当然上述这段 C 代码还是过于简单,以及错误检测我们也没有做,像初始化的时候传递的 length 小于 0 的话会造成段错误,不过这不是我们的重点,目前能用就行。

#include "my_array.c"

int main() {
    MyArray array;
    initial_MyArray(&array, 0);
    print_MyArray(&array);

    initial_MyArray(&array, 10);
    print_MyArray(&array);

    initial_MyArray(&array, 20);
    print_MyArray(&array);

    dealloc_MyArray(&array);
}

来编译测试一下:

D:\C\matsuri>gcc -std=c99 main.c -o a.exe

D:\C\matsuri>a.exe
[]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ...]

以上我们就实现了一个类似于 Python 的 range,而且还是 low 版的,但是使用的 C 而不是 Python。下面我们就来将 MyArray 包装一下,交给 Python 使用,这里会涉及到 Python / C API,需要事先对其有相应的了解。

#include <Python.h>
#include "my_array.c"

// Python 中的对象在 C 中都嵌套了 PyObject
typedef struct {
    PyObject_HEAD
    MyArray array;
} PyMyArray;

// 初始化函数 __init__
static int
PyMyArray_init(PyMyArray *self, PyObject *args, PyObject *kwargs) {
    if (self->array.arr != NULL) {
        dealloc_MyArray(&self->array);
    }
    int length = 0;
    static char *kwlist[] = {"length", NULL};
    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i", kwlist, &length)) {
        return -1;
    }
    // 因为给 Python 调用,所以这里额外对 length 进行一下检测
    if (length < 0) {
        PyErr_SetString(PyExc_ValueError, "argument 'length' can not be negative");
        return -1;
    }
    initial_MyArray(&self->array, length);
    return 0;
}

// 析构函数
static void
PyMyArray_dealloc(PyMyArray *self) {
    dealloc_MyArray(&self->array);
    Py_TYPE(self)->tp_free((PyObject *) self);
}

// __str__ 函数
static PyObject *
PyMyArray_str(PyMyArray *self) {
    char *s = change_array_to_string(&self->array);
    PyObject *ret = PyUnicode_FromString(s);
    free(s);
    return ret;
}

static PyTypeObject PyMyArrayType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "py_my_array.PyMyArray",         /* tp_name           */
    sizeof(PyMyArray),               /* tp_basicsize      */
    0,                               /* tp_itemsize       */
    (destructor) PyMyArray_dealloc,  /* tp_dealloc        */
    0,                               /* tp_print          */
    0,                               /* tp_getattr        */
    0,                               /* tp_setattr        */
    0,                               /* tp_reserved       */
    (reprfunc) PyMyArray_str,        /* tp_repr           */
    0,                               /* tp_as_number      */
    0,                               /* tp_as_sequence    */
    0,                               /* tp_as_mapping     */
    0,                               /* tp_hash           */
    0,                               /* tp_call           */
    (reprfunc) PyMyArray_str,        /* tp_str            */
    0,                               /* tp_getattro       */
    0,                               /* tp_setattro       */
    0,                               /* tp_as_buffer      */
    Py_TPFLAGS_DEFAULT,              /* tp_flags          */
    "PyMyArray object",              /* tp_doc            */
    0,                               /* tp_traverse       */
    0,                               /* tp_clear          */
    0,                               /* tp_richcompare    */
    0,                               /* tp_weaklistoffset */
    0,                               /* tp_iter           */
    0,                               /* tp_iternext       */
    0,                               /* tp_methods        */
    0,                               /* tp_members        */
    0,                               /* tp_getset         */
    0,                               /* tp_base           */
    0,                               /* tp_dict           */
    0,                               /* tp_descr_get      */
    0,                               /* tp_descr_set      */
    0,                               /* tp_dictoffset     */
    (initproc) PyMyArray_init,       /* tp_init           */
};


static PyModuleDef py_my_array_module = {
    PyModuleDef_HEAD_INIT,
    "py_my_array",
    "this is a module named py_my_array",
    -1,
    0,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC
PyInit_py_my_array(void) {
    PyObject *m;
    PyMyArrayType.tp_new = PyType_GenericNew;
    if (PyType_Ready(&PyMyArrayType) < 0) return NULL;
    m = PyModule_Create(&py_my_array_module);
    if (m == NULL) return NULL;

    Py_XINCREF(&PyMyArrayType);
    PyModule_AddObject(m, "PyMyArray", (PyObject *) &PyMyArrayType);
    return m;
}

这里我们算是用 Python / C API 编写了一个扩展模块,相信你此刻应该明白为什么要 Cython 了,因为 Python / C API 实在是太痛苦了。

然后我们来编译成扩展模块:

from distutils.core import setup, Extension

setup(ext_modules=[Extension("py_my_array", [r"D:\C\matsuri\py_my_array.c"])])

编译命令和 Cython 是一致的,然后我们来导入测试一下:

import py_my_array

print(py_my_array)  # <module 'py_my_array' from ...\\...\\py_my_array.cp38-win_amd64.pyd'>

arr = py_my_array.PyMyArray()
print(arr)  # []

arr = py_my_array.PyMyArray(5)
print(arr)  # [0, 1, 2, 3, 4]

arr = py_my_array.PyMyArray(20)
print(arr)  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ...]

arr = py_my_array.PyMyArray(-1)
"""
    arr = py_my_array.PyMyArray(-1)
ValueError: argument 'length' can not be negative
"""

下面我们来给 Numpy 使用,看看会有什么结果:

import numpy as np
import py_my_array

array = py_my_array.PyMyArray(5)
print(array)  # [0, 1, 2, 3, 4]

np_arr = np.asarray(array)
# 看到这里相信你已经有所察觉了,因为 Numpy 的数组在打印的时候,元素之间是没有逗号的
print(np_arr)  # [0, 1, 2, 3, 4]
# 而且 shape 为 空元组,很明显这也是不正常的,应该是 (5,) 才对
print(np_arr.shape)  # ()
# 然后我们来修改一下元素
try:
    np_arr[0] = 123
except IndexError as e:
    # 发现无法修改
    print(e)  # too many indices for array

# 我们再手动创建一个数组
np_arr = np.array([0, 1, 2, 3, 4])
# 此时一切正常
print(np_arr)  # [0 1 2 3 4]
print(np_arr.shape)  # (5,)
np_arr[0] = 123
print(np_arr)  # [123   1   2   3   4]

我们看到 Numpy 在处理 py_my_array.PyMyArray(5) 的时候,表现非常怪异,原因就在于 PyMyArray 实例化得到的不具备列表的性质,它只是在打印的时候以列表的形式进行打印的。而 Numpy 不知道该如何处理这个对象,于是它创建了一个标量,其值等于这个对象。至于其它对象也是一样的,举个例子:

import numpy as np

class A:
    def __str__(self):
        return "神乐七奈"

arr = np.asarray(A())
print(arr)  # 神乐七奈
print(type(arr))  # <class 'numpy.ndarray'>

显然这不是我们想要的,而解决这一点就是实现缓冲区协议。

使用 C 实现缓冲区协议

再来看一下结构体定义:

/* buffer interface */
typedef struct bufferinfo {
    void *buf;
    PyObject *obj;        /* owned reference */
    Py_ssize_t len;
    Py_ssize_t itemsize;  /* This is Py_ssize_t so it can be
                             pointed to by strides in simple case.*/
    int readonly;
    int ndim;
    char *format;
    Py_ssize_t *shape;
    Py_ssize_t *strides;
    Py_ssize_t *suboffsets;
    void *internal;
} Py_buffer;

buffer interface 所做的事情就是允许你定义一个函数,在函数中构造这样的一个指向特定数据的结构体。从结构体定义中我们可以看到,数组可以是多维的,并且支持以指定步长访问。当然对于我们的一维数组而言,我们只会使用一部分简单的功能。下面我们来创建扩展模块 py_my_array2,之前的代码直接搬过来,然后再加上一些缓冲区协议相关的钩子即可。

#include <Python.h>
#include "my_array.c"

// Python 中的对象在 C 中都嵌套了 PyObject
typedef struct {
    PyObject_HEAD
    MyArray array;
} PyMyArray;

// 初始化函数 __init__
static int
PyMyArray_init(PyMyArray *self, PyObject *args, PyObject *kwargs) {
    if (self->array.arr != NULL) {
        dealloc_MyArray(&self->array);
    }
    int length = 0;
    static char *kwlist[] = {"length", NULL};
    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i", kwlist, &length)) {
        return -1;
    }
    // 因为给 Python 调用,所以这里额外对 length 进行一下检测
    if (length < 0) {
        PyErr_SetString(PyExc_ValueError, "argument 'length' can not be negative");
        return -1;
    }
    initial_MyArray(&self->array, length);
    return 0;
}

// 析构函数
static void
PyMyArray_dealloc(PyMyArray *self) {
    dealloc_MyArray(&self->array);
    Py_TYPE(self)->tp_free((PyObject *) self);
}

// __str__ 函数
static PyObject *
PyMyArray_str(PyMyArray *self) {
    char *s = change_array_to_string(&self->array);
    PyObject *ret = PyUnicode_FromString(s);
    free(s);
    return ret;
}

// 定义 buffer interface 支持的函数
static int
PyMyArray_getbuffer(PyObject *obj, Py_buffer *view, int flags) {
    if (view == NULL) {
        PyErr_SetString(PyExc_ValueError, "NULL view in getbuffer");
        return -1;
    }
    PyMyArray* self = (PyMyArray*)obj;
    view->obj = (PyObject*)self;
    view->buf = (void*)self->array.arr;
    view->len = self->array.length * sizeof(int);
    view->readonly = 0;
    view->itemsize = sizeof(int);
    view->format = "i"; 
    view->ndim = 1;
    view->shape = &self->array.length;
    view->strides = &view->itemsize;
    view->suboffsets = NULL;
    view->internal = NULL;

    Py_INCREF(self); 
    return 0;
}

// 将上面的函数放入到 PyBufferProcs 结构体中
static PyBufferProcs PyMyArray_as_buffer = {
        (getbufferproc)PyMyArray_getbuffer,
        (releasebufferproc)0
};

static PyTypeObject PyMyArrayType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "py_my_array2.PyMyArray",        /* tp_name           */
    sizeof(PyMyArray),               /* tp_basicsize      */
    0,                               /* tp_itemsize       */
    (destructor) PyMyArray_dealloc,  /* tp_dealloc        */
    0,                               /* tp_print          */
    0,                               /* tp_getattr        */
    0,                               /* tp_setattr        */
    0,                               /* tp_reserved       */
    (reprfunc) PyMyArray_str,        /* tp_repr           */
    0,                               /* tp_as_number      */
    0,                               /* tp_as_sequence    */
    0,                               /* tp_as_mapping     */
    0,                               /* tp_hash           */
    0,                               /* tp_call           */
    (reprfunc) PyMyArray_str,        /* tp_str            */
    0,                               /* tp_getattro       */
    0,                               /* tp_setattro       */
    // 指定 tp_as_buffer
    &PyMyArray_as_buffer,            /* tp_as_buffer      */
    Py_TPFLAGS_DEFAULT,              /* tp_flags          */
    "PyMyArray object",              /* tp_doc            */
    0,                               /* tp_traverse       */
    0,                               /* tp_clear          */
    0,                               /* tp_richcompare    */
    0,                               /* tp_weaklistoffset */
    0,                               /* tp_iter           */
    0,                               /* tp_iternext       */
    0,                               /* tp_methods        */
    0,                               /* tp_members        */
    0,                               /* tp_getset         */
    0,                               /* tp_base           */
    0,                               /* tp_dict           */
    0,                               /* tp_descr_get      */
    0,                               /* tp_descr_set      */
    0,                               /* tp_dictoffset     */
    (initproc) PyMyArray_init,       /* tp_init           */
};


static PyModuleDef py_my_array_module = {
        PyModuleDef_HEAD_INIT,
        "py_my_array2",
        "this is a module named py_my_array2",
        -1,
        0,
        NULL,
        NULL,
        NULL,
        NULL
};

PyMODINIT_FUNC
PyInit_py_my_array2(void) {
    PyObject *m;
    PyMyArrayType.tp_new = PyType_GenericNew;
    if (PyType_Ready(&PyMyArrayType) < 0) return NULL;
    m = PyModule_Create(&py_my_array_module);
    if (m == NULL) return NULL;

    Py_XINCREF(&PyMyArrayType);
    PyModule_AddObject(m, "PyMyArray", (PyObject *) &PyMyArrayType);
    return m;
}

然后编译成扩展,方式不变,编译之后导入测试一下:

import numpy as np
import py_my_array2

array = py_my_array2.PyMyArray(5)
print(array)  # [0, 1, 2, 3, 4]

np_arr = np.asarray(array)
# 我们看到此刻是正常打印的
print(np_arr)  # [0 1 2 3 4]
print(np_arr.shape)  # (5,)

# 那么两者是不是共享内存的呢
np_arr[0] = 123
print(array)   # [123, 1, 2, 3, 4]
print(np_arr)  # [123   1   2   3   4]

显然此时就万事大吉了,因为实现了缓冲区协议,Numpy 知道了缓冲区数据,因此会在此基础之上建一个 view,因此 array 和 np_arr 是共享内存的。

因此核心就在于缓冲区协议的理解,它本质上就是一个结构体,内部的成员描述了缓冲区数据的所有信息。而我们只需要定义一个函数,然后根据数组初始化这些信息即可,最后构建 PyBufferProcs 实例作为 tp_as_buffer 成员的值。

我们目前实现的缓冲区协议非常简单,但是至少我们理解它是如何工作的了,至于想要实现更复杂的功能,我们会使用 Cython。

使用 Cython 实现缓冲区协议

下面我们来看看如何使用 Cython 实现缓冲区协议,看完之后你会发现使用 Cython 真是一个幸运的选择:

from cpython cimport Py_buffer
from cpython.mem cimport PyMem_Malloc, PyMem_Free

cdef class Matrix:
    cdef Py_ssize_t shape[2]  # 数组的形状
    cdef Py_ssize_t strides[2]  # 数组的 stride
    cdef float *array

    def __cinit__(self, row, col):
        self.shape[0] = <Py_ssize_t> row
        self.shape[1] = <Py_ssize_t> col
        self.strides[1] = sizeof(float)
        self.strides[0] = self.strides[1] * self.shape[1]

        self.array = <float *> PyMem_Malloc(self.shape[0] * self.shape[1] * sizeof(float))

    def set_item_by_index(self, int index, float value):
        """留一个接口,用来设置元素"""
        if index >= self.shape[0] * self.shape[1] or index < 0:
            raise ValueError("索引无效")

        self.array[index] = value

    def __getbuffer__(self, Py_buffer *buffer, int flags):
        """自定义缓冲区需要实现 __getbuffer__ 方法"""
        cdef int i;
        for i from 0 <= i < self.shape[0] * self.shape[1]:
            self.array[i] = float(i)

        buffer.obj = self
        buffer.buf = <void *> self.array  
        buffer.len = self.shape[0] * self.shape[1] * sizeof(float)  
        buffer.readonly = 0
        buffer.itemsize = sizeof(float)
        buffer.format = "f"
        buffer.ndim = 2
        buffer.shape = self.shape
        buffer.strides = self.strides
        buffer.suboffsets = NULL

    def dealloc(self):
        if self.array != NULL:
            PyMem_Free(<void *> self.array)

在 Cython 中我们只需要实现一个相应的魔法方法即可,真的是非常方法,当然我们为了验证是否共享内存,专门定义了一个方法。

import pyximport
pyximport.install(language_level=3)

import cython_test
import numpy as np

m = cython_test.Matrix(5, 4)
np_m = np.asarray(m)
print(m)  # <cython_test.Matrix object at 0x7f96ba55a3f0>
print(np_m)
"""
[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]
 [12. 13. 14. 15.]
 [16. 17. 18. 19.]]
"""

# 而且也是共享内存的
m.set_item_by_index(13, 666.666)
print(np_m)
"""
[[  0.      1.      2.      3.   ]
 [  4.      5.      6.      7.   ]
 [  8.      9.     10.     11.   ]
 [ 12.    666.666  14.     15.   ]
 [ 16.     17.     18.     19.   ]]
"""

结果上没有任何问题。

以上就是缓冲区协议,其实在日常工作中我们不需要直接面对它,但了解一下总是好的。然后我们来介绍一下 Python 的 memoryview,从而引出 Cython 的 memoryview,这样我们可以在不和缓冲区协议直接打交道的前提下使用缓冲区协议。

memoryview(内存视图)

一个内置的 Python 类型是 memoryview(内存视图),它存在的唯一目的就是在 Python 级别表示 C 级的缓冲区(避免数据的复制)。我们可以像 memoryview 传递一个实现了新缓冲区协议的对象(比如:bytes对象),来创建一个 memoryview 对象。

b = b"komeiji satori"
m = memoryview(b)

# 此时和m和b之间是共享内存的
print(m)  # <memory at 0x000001C3A54EFA00>

# 可以通过索引、切片的方式访问
print(m[0], m[-1])  # 107 105
# 通过索引获取得到的是数字,切片获取得到的还是一个 memoryview 对象
print(f"{m[0]:c}", f"{m[-1]:c}")  # k i
print(m[0: 2])  # <memory at 0x00000229D637F880>
print(m[0: 2][-1], f"{m[0: 2][-1]:c}")  # 111 o

我们可以通过切片的方式对 memoryview 进行任意截取,通过这种方式,灵活性就变得非常高。

那么能不能修改呢?答案是:因为 bytes 对象不可以修改,所以 memoryview 也不可以修改,也就是只读的。

b = b"komeiji satori"
m = memoryview(b)

print(m.readonly)  # True
try:
    m[0] = "K"  
except Exception as e:
    print(e)  # cannot modify read-only memory

bytes 对应的缓冲区是不可以修改的,所以如果想修改的话,我们应该使用 bytearray。

b = bytearray(b"my name is satori")
m = memoryview(b)

# 此时 m 和 b 共享内存
print(b)  # bytearray(b'my name is satori')
m[0] = ord("M")
print(b)  # bytearray(b'My name is satori')

b[1] = ord("Y")
print(chr(m[1]))  # Y

当然我们还可以传入一个 Numpy 的 ndarray,只要是实现了新缓冲区协议的,都可以传递到 memoryview 中。

import numpy as np

np_mv = memoryview(np.ones((10, 20, 30)))
# 查看维度
print(np_mv.ndim)  # 3
# 查看 shape
print(np_mv.shape)  # (10, 20, 30)
# strides 属性表示某个维度中, 一个元素和下一个元素之间差了多少个字节
print(np_mv.strides)  # (4800, 240, 8)
# 查看每个元素的大小
print(np_mv.itemsize)  # 8
# 查看整个数组占用的字节数, 等于 itemsize * 元素的个数
print(np_mv.nbytes)  # 48000
# 查看数组类型
print(np_mv.format)  # d

结构化数据也是支持的,首先我们来创建一个 Numpy 的 dtype,其中 a 和 b 的类型分别是 int8 和 字符串。

import numpy as np

dt = np.dtype([("a", "int8"), ("b", "U")])
print(dt)  # [('a', 'i1'), ('b', '<U')]
print(np.empty(5, dtype=dt))  # [(0, '') (0, '') (0, '') (0, '') (0, '')]

structured_mv = memoryview(np.empty(5, dtype=dt))
print(structured_mv.format)  # T{b:a:=0w:b:}
"""
这里的 format(格式化字符串) 来自标准库模块 struct 的规范, 对于结构化类型来说是相当神秘的 
但是我们将 memoryview 的格式化字符串的细节留给官方文档吧, 我们不需要与它们直接打交道
因为缓冲区和 memoryview 对象可以使用 简单的标量类型 以及 用户自定义的结构化类型
"""

那么问题来了,memoryview 对象和缓冲区如何转换到 Cython 中呢?考虑到 Cython 是专门用来连接 Python 和 C 的,所以它一定非常适合在 C 级别使用 memoryview 对象和缓冲区协议。

类型化 memoryview

Cython 有一个 C 级类型,即:类型化 memoryview,它在概念上和 Python 中的 memoryview 重叠、并且在其基础上展开。正如其名所示,类型化 memoryview 用于查看(共享)来自缓冲区对象的数据。并且类型化 memoryview 在 C 级别操作,所以它有着最小的 Python 开销,因此非常高效,而且比直接使用 C 级缓冲区更方便。此外类型化 memoryview 是为了和缓冲区一起工作而被设计的,因此它可以有效支持任何来自缓冲区的对象,从而允许在不复制的情况下共享缓冲区数据。

假设我们想在 Cython 中有效地处理一维数据的缓冲区,而不关心如何在 Python 级别创建数据,我们只是想以一种有效的方式访问它。

def summer(double[:] numbers):
    cdef double res = 0
    cdef double number
    # memoryview 对象可以像 Python 迭代器一样进行遍历
    for number in numbers:
        res += number
    return res

double[:] numbers 声明了 numbers 是一个类型化 memoryview 对象,而 double 指定了该 memoryview 对象的基本类型,[:] 则表明这是一个一维的 memoryview 对象。

当我们从 Python 中调用 summer 函数时,会传入一个 Python 对象,该对象会在调用函数的时候隐式的分配给参数 numbers。我们可以提供一个类型化 memoryview 对象,如果提供的不是 memoryview 对象,那么看该对象是否支持缓冲区协议,如果支持缓冲区协议,那么将会提供一个 C 级缓冲区来构建 memoryview;如果不支持该缓冲区协议(没有提供相应的缓冲区),那么会引发一个 ValueError。

import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test

# 必须传递支持缓存区协议的对象
print(cython_test.summer(np.array([1.2, 2.3, 3.4, 4.5])))  # 11.4

# 可以直接传入数组,也可以传入memoryview对象
print(cython_test.summer(memoryview(np.array([1.2, 2.3, 3.4, 4.5]))))  # 11.4

# 但是传递列表是不行的,因为它不支持缓冲区协议
print(cython_test.summer([1.2, 2.3, 3.4, 4.5]))
"""
    def summer(double[:] numbers):
  File "stringsource", line 658, in View.MemoryView.memoryview_cwrapper
  File "stringsource", line 349, in View.MemoryView.memoryview.__cinit__
TypeError: a bytes-like object is required, not 'list'
"""

不过当我们在编译类型化 memoryview 对象时,Cython 本质上还是将它当成通用的 Python 迭代器来看待的,因此我们可以做的更好。

C 级访问类型化 memoryview 数据

类型化 memoryview 对象是为 C 风格的访问而设计的,没有开销,因此也可以用另一种方式去遍历 number。

def summer(double[:] numbers):
    cdef double s = 0
    cdef int i, N
    # 调用 shape 拿到其长度
    N = numbers.shape[0]
    for i in range(N):
        s += numbers[i]
    return s
print(cython_test.summer(np.array([1.2, 2.3, 3.4, 4.5])))  # 11.4

这个版本会有更好的性能:对于百万元素数组来说大约是 1 毫秒,因为我们用了一个有类型的整数去作为索引,那而这种索引访问类型化 memoryview 时,Cython 会生成绕过 Python/C API 调用的代码,直接索引底层缓冲区,这是我们一个加速的来源,当然我们还可以做得更好。

用安全换取性能

每次我们访问 memoryview 对象时,Cython 都会检测索引是否越界,如果越界,那么 Cython 将引发一个 IndexError。而且 Cython 也允许我们像 Python 一样通过负数索引对 memoryview 对象进行访问。

对于我们上面的 summer 函数,我们在访问内部 memoryview 对象之前就已经获取了它的元素个数,因为我们在遍历的时候永远不会越界,所以我们可以指示 Cython 关闭这些检查以获取更高的性能。而关闭检查可以使用上下文的方式:

from cython cimport boundscheck, wraparound

def summer(double[:] numbers):
    cdef double s = 0
    cdef int i, N
    N = numbers.shape[0]
    # 关闭检查
    with boundscheck(False), wraparound(False):
        for i in range(N):
            s += numbers[i]
        return s

我们通过上下文管理的方式,在其内部对 memoryview 对象进行索引访问的时候,会关闭检测,当然这些修改仅在上下文管理器内部有效。得到的结果是性能会有小幅度提高(当然我们这里数据很少,看不出来),代码生成效率也更高。但效率提升的后果就是必须由我们来确保索引不会越界、并且不可以使用负数索引,否则的话可能会导致段错误(非常危险,不管程序崩溃,解释器就直接退出了)。因此如果没有百分之百的把握,不要关闭检查。

如果我们想关闭整个函数内部的检查,那么可以使用装饰器的形式。

from cython cimport boundscheck, wraparound

@boundscheck(False)
@wraparound(False)
def summer(double[:] numbers):
    cdef double s = 0
    cdef int i, N
    N = numbers.shape[0]
    for i in range(N):
        s += numbers[i]
    return s

如果想关闭全局的边界检测,那么可以使用注释的形式。

# cython: boundscheck=False
# cython: wraparound=False

所以关闭边界检测有多种方式,但是不同的方式对应不同的作用域。但是有了以上这些方法,我们的 summer 函数的性能,和 Numpy 中 sum 函数的性能便在一个数量级了。我们编译成扩展模块测试一下吧:

可以看到,虽然没有 Numpy 中的 sum 函数那么厉害,但至少在同一水平线上,反正都甩开内置的 sum 函数一条街。

那么到目前为止,我们都了解到了什么呢?首先我们知道如何在 Cython 中声明一个简单的类型化 memoryview,以及对它进行索引、访问内部缓冲区的数据。并且还可以通过使用 boundscheck 和 wraparound 关闭边界检查,来生成更加高效的代码,但前提是我们能确保不会出现索引越界,否则还是不要关闭检查。因为为了安全,这些都是值得的。

但是还没完,我们还有更多的细节要讲。

声明类型化 memoryview

当我们声明一个类型化 memoryview 时,我们可以控制很多的属性。

元素类型

类型化 memoryview 的元素类型可以任意,在 Cython 中凡是能拿来做变量类型声明的都可以,因此也可以是 ctypedef 起的别名。

维度

类型化 memoryview 最多可以有7个维度,我们之前声明一个一维的,使用的是 double[:] 这种形式,如果是 3 位,那么写成 double[:, :, :] 即可。当然类型不一定是 double,也可以是其它的。

C 和 Fortran 的连续性

通过指定数据打包约束的 C、Fortran 类型化内存视图是一个非常重要的特例,C 连续和 Fortran 连续都意味着缓冲区在内存中是连续的。但如果是多维度,C 连续的 memoryview 的最后一个维度是连续的,而 Fortran 连续的 memoryview 的第一个维度是连续的。如果可能的话,从性能的角度上来说,将数组声明为 C 或者 Fortran 连续是有利的,因为这使得 Cython 可以生成更快的代码。如果不是 C 连续,也不是 Fortran 连续,那么我们称之为 full strided。

我们通过 Numpy 来对比一下 C 连续和 Fortran 连续的区别。

import numpy as np

arr = np.arange(16)
print(arr.reshape((4 , 4), order="C"))
# 默认是 C 连续, 即 order="C", 最后一个维度是连续的
"""
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]
"""
print(arr.reshape((4, 4), order="F"))
# 如果 Fortran, 即 order="F", 那么第一个维度是连续的
"""
[[ 0  4  8 12]
 [ 1  5  9 13]
 [ 2  6 10 14]
 [ 3  7 11 15]]
"""
# 所以 C 连续的数组在转置之后就会变成 Fortran 连续

直接或间接访问

直接访问是默认的,它涵盖了几乎所有的情况,它指定对应的维度可以通过索引的方式直接访问底层数据。如果将某个维度指定为间接访问,则底层缓冲区将存储一个指向数组剩余部分的指针,该指针必须在访问时解引用(因此是间接访问)。而 Numpy 不支持间接访问,所以我们不使用这种访问规范,因此直接访问是默认的。事实上,官方一般设置为默认的都是最好的。

下面我们来举例说明:

import numpy as np

# 这是最灵活的声明方式,可以从任何一个元素为 int 的二维类型化 memoryview 对象中获取缓冲区
def func(int[:, :] ages):
    # 直接打印是一个 memoryview 对象,这里转成 ndarray, 这里说一句: Numpy 中数组的类型是 <class 'ndarray'>, 为了方便有时叫它 array
    print(np.array(ages))

然后我们测试一下:

import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test

# 这里类型要匹配,C 中的 int 对应 numpy 中的 int32,而 numpy 的整型默认是 int64
# 因为 memoryview 对类型要求很严格
arr = np.random.randint(1, 10, (3, 3), dtype="int32")

cython_test.func(arr)
"""
[[7 9 2]
 [9 6 9]
 [9 9 1]]
"""
# 也可以对 arr 进行切片
cython_test.func(arr[1: 3, 0: 2])
"""
[[9 6]
 [9 9]]
"""

当我们对数据进行索引的时候,Cython 会自动生成索引代码来考虑数据的跨步访问。但如果我们愿意用一些灵活性来换取速度的话,那么 C 连续或者 Fortran 连续的类型化内存视图可以更有效的建立索引。

声明一个 C 连续的类型化内存视图,需要对最后一个维度进行修改。前 n - 1 个维度不变,还是一个冒号,最后一个维度换成两个冒号并跟一个数字 1。举个栗子:之前的声明是 double [:, :, :],如果想要 C 连续,那么应该改成 double [:, :, :: 1] ,表示最后一个维度具有统一的步长。

import numpy as np

def func(int[:, :: 1] ages):
    print(np.array(ages))

然后我们讲一个 C 连续的数组赋值给这里的 ages,C 连续在 Numpy 中是所有数组创建函数的默认布局。

import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test

arr = np.arange(16, dtype="int32").reshape((4, 4))
cython_test.func(arr)
"""
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]
"""

但如果赋值一个 Fortran 连续,或者指定步长筛选之后的数组就会引发一个运行的 ValueError。

import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test

arr = np.arange(16, dtype="int32")
try:
    cython_test.func(arr.reshape((4, 4), order="F"))
except ValueError as e:
    print(e)  # ndarray is not C-contiguous

try:
    cython_test.func(arr.reshape((4, 4))[:, :: 2])
except ValueError as e:
    print(e)  # ndarray is not C-contiguous

如果我们声明一个 Fortran 连续的数组,那么很明显,我们需要将第一个维度指定为 :: 1

import numpy as np

def func(int[::1, :] ages):
    print(np.array(ages))
import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test

arr = np.arange(16, dtype="int32")
try:
    cython_test.func(arr.reshape((4, 4)))
except ValueError as e:
    print(e)  # ndarray is not Fortran contiguous

cython_test.func(arr.reshape((4, 4), order="F"))
"""
[[ 0  4  8 12]
 [ 1  5  9 13]
 [ 2  6 10 14]
 [ 3  7 11 15]]
"""

try:
    cython_test.func(arr.reshape((4, 4), order="F")[:: 2])
except ValueError as e:
    print(e)  # ndarray is not Fortran contiguous

一个多维数组要么是 C 连续,要么是 Fortran 连续,但不能同时是 C 连续和 Fortran 连续。但是一维数组除外,一维数组既可以是 C 连续、也可以是 Fortran 连续。

import numpy as np

def func(int[::1] ages):
    print(np.array(ages))

所以我们目前已经介绍了三种类型化内存视图,分别是:C 连续、Fortran 连续、fully strided。常见的情况下,所有数组都是 C 连续的,这是最常见的内存布局。特别是在需要和外部 C、C++ 库进行交互的时候,这种布局就显得尤为重要。并且在很多情况下,如果传递了非 C 连续的数组,比如:full strided 或者 Fortran 连续,将会引发一个 ValueError。

但如果你的应该是以 Fortran 为中心的,那么应该将数组声明为 Fortran 连续,这样会更好一些。

而 Numpy 提供了 ascontiguousarray 和 asfortranarray 这两个转换函数,可以接收一个数组并返回一个 C 连续或者 Fortran 连续的数组。当然,如果数组本身就是 C 连续或者 Fortran 连续,那么就不会转化了,这样更有效率。

import numpy as np

arr = np.arange(16).reshape((4, 4))
print(arr.flags["C_CONTIGUOUS"])  # True
print(arr.flags["F_CONTIGUOUS"])  # False

arr = np.asfortranarray(arr)
print(arr.flags["C_CONTIGUOUS"])  # False
print(arr.flags["F_CONTIGUOUS"])  # True

# 这两个方法底层都调用了 np.array, 所以下面这种做法也是可以的
arr2 = np.arange(16).reshape((4, 4))
print(arr2.flags["F_CONTIGUOUS"])  # False
# 这里面有一个 copy 参数, 表示是否拷贝, 默认为 True, 这样得到的新数组和原数组之前没有任何关系
# 但这里我们显然无需拷贝, 因此指定 copy=False 会更有效率
arr2 = np.array(arr2, order="F", copy=False)
print(arr2.flags["F_CONTIGUOUS"])  # True

# 如果要将一个数组改成 C 连续或者 Fortran 连续, 推荐上面两种做法
# 另外, 我们说一维数组既是 C 连续又是 Fortran 连续
arr3 = np.arange(16)
print(arr3.flags["C_CONTIGUOUS"])  # True
print(arr3.flags["F_CONTIGUOUS"])  # True

但是这并不代表我们在 Cython 中声明 memoryview 对象时,一定非要声明成 C 连续或者 Fortran 连续,当你接收的数组的连续性不确定时,应该采用 full strided 类型,也就是声明的时候不指定 :: 1。因为一旦指定连续,不管是 C 连续、还是 Fortran 连续,那么你的数组必须要满足,否则就会报出我们之前出现的错误:ndarray is not C-contiguous 或者 ndarray is not Fortran contiguous。这个时候你就需要先进行转化,通过手动创建一个 C 连续或者 Fortran 连续的数组(或者说 memoryview)的副本,说白了就是将那些指定步长访问的数组对应的元素拷贝一份,建立一个新的连续数组,但是这会带来额外的开销,甚至会超过连续访问带来的性能收益,举个栗子:

import numpy as np

arr = np.arange(16).reshape((4, 4))

# 默认情况下 copy=True,创建数组时会将元素拷贝一份,此时创建的数组默认是 C 连续
# 但如果是在不拷贝的情况下根据已有的数组创建数组,那么元素不会拷贝,而是共享缓冲区
# 此时创建的数组的连续性和已有的数组保持一致,因此,由于 arr[:: 2] 不是连续的,所以 arr1 也不是
arr1 = np.array(arr[:: 2], copy=False)
arr1[0, 0] = 111
# arr 和 arr1 共享缓冲区,修改 arr1 会改变 arr
print(arr[0, 0])  # 111
# 但 arr1 不是 C 连续,因为它指定了步长,实现了跨步访问,所以它就不再具备连续性
print(arr1.flags["C_CONTIGUOUS"])  # False


# 还是相同的方式,但是我创建的时候强行指定让 arr2 连续
arr2 = np.array(arr[:: 2], copy=False, order="C")
print(arr2.flags["C_CONTIGUOUS"])  # True
arr2[1, 1] = 222
# 但是我们看到 arr 并没有被改变
# 原因就在于将一个不是连续的数组变成连续的数组,会将不是连续的数组中对应的元素拷贝一份
# 以此来构建一个连续的数组,所以此时指定 copy=False 是没有用的
print(arr[1, 1])  # 5
# 所以在不指定连续性的情况下,指定 copy=False 则是和原来的数组共享缓冲区
# 但指定了连续性之后,只能拷贝一份,会有额外的性能开销

对于类型化 memoryview,我们传递 None 也是合法的,因此需要有一步检测,或者使用 not None 子句声明。

混合类型

还记得我们之前提到的混合类型吗?假设我希望某一个参数既可以是接收 list 对象,也可以接收 dict 对象,那么可以这么做。

cdef fused list_dict:
    list
    dict

cpdef func(list_dict var):
    return var

而类型化内存的视图的类型也可以是混合类型,这样可以保证更强的泛化能力和灵活性。但是很明显,所谓的混合类型无非就是创建了多个版本。

cimport cython

cpdef cython.floating generic_summer(cython.floating[:] m):
    cdef cython.floating f, s = 0.0
    for f in m:
        s += f
    return s
import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test


print(
    cython_test.generic_summer(np.array([1, 2, 3], dtype="float64"))
)  # 6.0

print(
    cython_test.generic_summer(np.array([1, 2, 3], dtype="float32"))
)  # 6.0

类型化内存视图对元素类型的要求是很严格的,但是我们通过混合类型的方式可以同时接受 float32 和  float64,也就是 C 中的 float 和double。

但是类型化 memoryview 对混合类型目前支持的其实不是特别友好,个人不太建议使用。

使用类型化 memoryview

一旦声明了类型化 memoryview,就必须给它分配一个支持缓冲区的对象,然后两者共享底层缓冲区。

那么问题来了,类型化 memoryview 都支持哪些操作呢?

首先我们可以像 Numpy 一样,对类型化 memoryview 进行访问和修改。

import numpy as np


cpdef array(int[:, :] numbers):
    print("----------")
    print(np.array(numbers))
    numbers[0, 0] = 3
    print("----------")
    print(np.array(numbers))
    print("----------")
    print(np.array(numbers[1: 3, : 2]))
import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test

arr = np.random.randint(1, 10, (3, 3))
cython_test.array(arr)
"""
----------
[[5 8 1]
 [7 9 9]
 [4 6 2]]
----------
[[3 8 1]
 [7 9 9]
 [4 6 2]]
----------
[[7 9]
 [4 6]]
"""

正如我们之前说的,类型化内存视图可以建立高效的索引,特别是当我们通过 boundscheck、wraparound 关闭检查的时候。

from cython cimport boundscheck, wraparound

cpdef summer(int[:, :] numbers):
    cdef int N, M, i, j
    cdef long s=0
    # 这里类型化 memoryview 的 shape 是一个含有 8 个元素的元组
    # 但我们这里只有两个维度, 所以截取前两位, 至于后面的元素都是 0
    N, M = numbers.shape[: 2]
    with boundscheck(False), wraparound(False):
        for i in range(N):
            for j in range(M):
                s += numbers[i, j]
        return s
import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test

arr = np.random.randint(1, 10, (300, 300))
print(np.sum(arr))  # 450832
print(cython_test.summer(arr))  # 450832

另外类型化 memoryview 和 Numpy 中的 array 一样,也支持 ... 这个语法糖,表示某个维度上、或者整体的全部筛选。

import numpy as np

cdef int[:, :] m = np.zeros((2, 2), dtype="int32")
print("*" * 10)
# 直接打印的话是一个 memoryview 对象, 需要再转成 array 进行打印
print(np.array(m))

# 通过 ... 表示全局修改
# 因此 m[...] 等价于 m[:]
m[...] = 123
print("*" * 10)
print(np.array(m))

# 在某一个维度上使用 ..., 可以实现某个维度上的全局修改
# 等价于 m[0, :] = 456
m[0, ...] = 456
print("*" * 10)
print(np.array(m))
import pyximport
pyximport.install(language_level=3)

import cython_test
"""
**********
[[0 0]
 [0 0]]
**********
[[123 123]
 [123 123]]
**********
[[456 456]
 [123 123]]
"""

因此在用法上,类型化 memoryview 和 Numpy 的 array 是一致的,当然我们也可以指定步长等等。注意:类型化 memoryview 的右边只能跟标量。

但是类型化 memoryview 和 Numpy 的 array 还是有些差别的,类型化 memoryview 不支持通用函数,所以在赋值的时候除了简单的标量赋值之外,无法进行广播操作。

import numpy as np

arr = np.arange(9).reshape((3, 3))
print(arr)
"""
[[0 1 2]
 [3 4 5]
 [6 7 8]]
"""
# 会将所有元素都赋值为 123, 因为是标量赋值, 所以类型化 memoryview 是支持的
arr[:] = 123
print(arr)
"""
[[123 123 123]
 [123 123 123]
 [123 123 123]]
"""

# 这里就涉及到了广播, 因为 (3, 3) 和 (3,) 两个维度明显不一致
# 所以会将 arr 的每一行都替换成 [11 22 33], 但这个是 Numpy 的 array 的功能, memoryview 是不支持的, 因为它在广播的时候右边只能跟标量
arr[:] = np.array([11, 22, 33])
print(arr)
"""
[[11 22 33]
 [11 22 33]
 [11 22 33]]
"""

可能有人觉得这也太不方便了吧,但办法总比困难多,我们可以根据类型化 memoryview 拿到对应的 array,对这个 array 进行操作不就行了。那么这么做的效率会不会低呢?答案是不会的,因为类型化 memoryview 和 array 之间是共享内存的,这么做不会有任何的性能损失。正如 torch 里面的 tensor 一样,它和 Numpy 的 array 之间也是共享内存的,由于 Numpy 的 API 用起来非常方便,也已经习惯了,加上个人也懒得使用 tensor 的一些操作,所以我都会先将 tensor 转成 array,对 array 操作之后再转回 tensor。虽然多了两次转化,但还是那句话,它们是共享内存的,所以完全没问题。

from cython cimport boundscheck, wraparound
import numpy as np

cdef int[:, :] m = np.arange(9).reshape((3, 3))
# 这里一定要指定 copy=False, 否则在拿到缓冲区里面的内容时, 还会拷贝一份新的出来
# 这么做的话, 就会有性能损失了, 因为两者本来就是共享内存的, 获取出来之后直接操作就可以了, 为什么要再创建一个副本呢
# 另外, 如果拷贝了一个副本, 你再操作完毕之后还要赋值回去
np.array(m, copy=False)[:] = [1, 2, 3]
# 以上我们就实现了修改, 这里我们再来打印一下看看, 这里打印就创建一个新的吧, 看看数据有没有修改
print(np.array(m))
import pyximport
pyximport.install(language_level=3)

import cython_test
"""
[[1 2 3]
 [1 2 3]
 [1 2 3]]
"""

因此我们可以把类型化 memoryview 看成是非常灵活 Cython 空间,可以有效地共享、索引定位、以及修改同质数据。它具有很多 Numpy array 的特性,特别是通过索引定位数据。至于那些没有特性,也很容易被两者之间转换的高效性所掩盖。

但是类型化 memoryview 实际上是超越了缓冲区协议,所以它还有额外的特性,那就是它还可以对 C 一级的数组进行 view。

view C 级数组

C 级数组可以是在堆上动态分配的,也可以是在栈上分配的。如果要 view C 一级的数组,那么该数组赋值给 memoryview 即可。如果数组的大小是固定的(或完整的),赋值的右边直接写数组名即可,因为 Cython 有足够的信息来跟踪这个 C 数组的大小。

cdef int a[3][5][7]

cdef int[:, :, :: 1] m = a
m[...] = 123

# 这里我们将 memoryview 声明是 C 连续的, 因为 C 里面的数组当然是 C 连续的
# 然后将其赋值为 123
import numpy as np
# 转成 Numpy 中的数组
arr = np.array(m, copy=False)
print("*" * 10)
print(np.sum(arr), 3 * 5 * 7 * 123)
print("*" * 10)
print(arr[:, 1: 3])
import pyximport
pyximport.install(language_level=3)

import cython_test
"""
**********
12915 12915
**********
[[[123 123 123 123 123 123 123]
  [123 123 123 123 123 123 123]]

 [[123 123 123 123 123 123 123]
  [123 123 123 123 123 123 123]]

 [[123 123 123 123 123 123 123]
  [123 123 123 123 123 123 123]]]
"""

以上 C 数组是在栈区分配的,我们也可以改为堆区分配。

from libc.stdlib cimport malloc, free

cdef int *a = <int *>malloc(3 * 5 * 7 * sizeof(int))
# 但是很明显, 这样做的话, 形状的信息就丢失了, Cython 只知道这个一个 int *
# 所以直接将 a 赋值过去的话, 编译时会出现 "Cannot convert long * to memoryviewslice"
# 因此我们在赋值给 memoryview 的时候, 必须给 Cython 提供更多的信息
# 而 <int[:3, :5, :7]> a 会告诉 Cython 这是一个三维数组, 维度分别是 3 5 7; 当然我们这里变成 7 5 3 也是可以的, 因为形状是由我们决定的
cdef int[:, :, :: 1] m = <int[:3, :5, :7]> a
m[...] = 123

import numpy as np
arr = np.array(m, copy=False)
print("*" * 10)
print(np.sum(arr), 3 * 5 * 7 * 123)
print("*" * 10)
print(arr[:, 1: 3])

# 最后将 a 给释放掉
free(a)

导入这个模块之后,打印的结果和之前是一样的。

在 C 级别,光靠一个头指针是没有办法确定动态分配的 C 数组的长度,而这一点是需要由我们来确定的。因此将一个 C 数组转成类型化 memoryview 时,如果数据不正确,那么可能会导致缓冲区溢出、段错误,或者数据损坏等等。

到此我们就算介绍完了类型化memoryview 的特性,并展示了它如何在支持新缓冲区协议的 Python 对象和 C 级数组中使用它们。如果 Cython 函数中有一个类型化 memoryview 参数,那么可以通过传递一个支持新缓冲区协议的 Python 对象或者 C 数组作为参数进行调用。

另外,当在函数中想返回一个类型化 memoryview 时,Cython 会直接将缓冲区内容(没有拷贝)并转化成 Python 的 memoryview,比如在函数中返回一个由 C 数组构建的类型化 memoryview,假设叫 m。如果这个 C 数组是在堆区通过 malloc 动态分配的,那么返回 m 没有任何问题;但如果它是在栈区分配的,在函数结束后就会被销毁,而我们说 Cython 又不会对缓冲区内容进行拷贝,因此此时就会出现错误。

不过即便是 C 数组在堆区申请,也是存在问题的。那就是当我们不再使用 memoryview 的时候,谁来释放这个在堆上申请的数组呢?如何正确地管理它的内存呢?不过在探讨这个问题之前,我们需要先来说一下另一种 Numpy。

另一种 Numpy

什么叫做另一种 Numpy,因为我们在 Cython 中除了 import numpy 之外,还可以 cimport numpy。

在类型化 memoryview 之前,Cython 可以使用不同的语法来很好地处理 Numpy 数组,这便是 原始缓冲区语法。尽管它已经被类型化 memoryview 所取代,但我们依旧可以正常使用它。

# 这里一定要使用 cimport
cimport numpy as np  # 或者 from numpy cimport ndarray

# 如果是 import numpy as np, 那么不好意思, np.ndarray 是无法作为参数类型和返回值的
# 编译时会报错: 'np' is not a cimported module
# 如果是 from numpy import ndarray, 编译时同样报错: 'ndarray' is not a type identifier
cpdef np.ndarray func(np.ndarray array):
    return array

但是注意:此时我们不能够直接导入,因为它依赖 numpy 的一个头文件,所以我们需要通过编译的方式。

from pathlib import Path
import numpy as np
from distutils.core import Extension, setup
from Cython.Build import cythonize

ext = Extension("cython_test",
                ["cython_test.pyx"],
                # 因为我们 cimport numpy 时会使用 Numpy 提供的一个头文件, 叫做 arrayobject.h
                # 但是很明显, 我们并没有指定它的位置, 因此需要通过 include_dirs 告诉 cython 编译器去哪里找这个头文件
                # 如果没有这个参数, 那么编译时会报错:fatal error: numpy/arrayobject.h: No such file or directory
                include_dirs=[str(Path(np.__file__).parent / "core" / "include")])
                # 当然 numpy 也为我们封装了一个函数, 直接通过 np.get_include() 即可获取改路径
                # 对于我当前的环境来说就是 C:\python38\lib\site-packages\numpy\core\include

setup(
    ext_modules=cythonize(ext, language_level=3),
)

然后导入测试一下:

import numpy as np
import cython_test

print(cython_test.func(np.array([[1, 2], [3, 4]])))
"""
[[1 2]
 [3 4]]
"""

print(cython_test.func(np.array([["xx", None], [(1, 2), {1, 2}]])))
"""
[['xx' None]
 [(1, 2) {1, 2}]]
"""

测试是没有问题的,并且接收的就是 array,返回的也是 array。并且我们看到,这对 array 的类型没有任何限制,但如果我们希望限制的 array 的类型、甚至是维度,这个时候这怎么做呢?

Cython为numpy提供了专门的方法:比如希望参数是一个接收类型为 int64、维度为 2 的数组,就可以使用 ndarray[long, ndim=2] 这种方式,我们演示一下。

cimport numpy as np

# 类型需要统一
# long 对应 np.int64, int 对应 np.int32, short 对应 np.int16, char 对应 int8
# unsigned long 对应 np.uint64, 其它同理
def func1(np.ndarray[long, ndim=2] array):
    print(array)


def func2(np.ndarray[double, ndim=1] array):
    print(array)


def func3(np.ndarray[object, ndim=1] array):
    print(array)
    
# 除了作为函数参数和返回值之外, 还可以用来声明普通的静态变量
# 比如: cdef np.ndarray[double, ndim=2] arr

编译之后并且导入进行测试:

import numpy as np
import cython_test

# 这里我们传递的时候, 参数和维度一定要匹配, 默认是 int64
cython_test.func1(np.array([[1, 2], [3, 4]]))
"""
[[1 2]
 [3 4]]
"""
try:
    cython_test.func1(np.array([1, 2, 3, 4]))
except ValueError as e:
    print(e)  # Buffer has wrong number of dimensions (expected 2, got 1)

try:
    cython_test.func2(np.array([1, 2, 3, 4]))
except ValueError as e:
    print(e)  # Buffer dtype mismatch, expected 'double' but got 'long'

cython_test.func2(np.array([1, 2, 3, 4], dtype="float64"))
"""
[1. 2. 3. 4.]
"""
cython_test.func3(np.array(["a", "b", object], dtype="O"))
"""
['a' 'b' <class 'object'>]
"""

但以上是原始缓冲区语法,现在我们更推荐类型化 memoryview,虽然它比 Numpy 中的 array 少了许多功能,但我们说这两者之间是可以高效转换的。并且如果是通过 np 来调用的话,那么两者是等价的。举个栗子:

import numpy as np

def func(int[:, :: 1] m):
    # m 本身没有 sum 方法, 但是我们可以将它传递给 np.sum
    print(np.sum(m))
    print(np.sum(m, axis=0))
    print(np.sum(m, axis=1))
import pyximport
pyximport.install(language_level=3)

import numpy as np
import cython_test


cython_test.func(np.array([[1, 2], [3, 4]], dtype="int32"))
"""
10
[4 6]
[3 7]
"""

因此这种声明方式是更加推荐的,而且也更加清晰和简洁,以及我们可以通过 pyximport 进行导入了。之前的方式由于依赖一个头文件,你必须要告诉 cython 编译器头文件去哪里找,而直接导入、自动编译的话会找不到头文件。但是现在不需要了,因为我们根本没有使用 cimport numpy。

但是以上这些说实话都不能算是优点,所以肯定还有其它的优点,那么都有哪些呢?

1. 类型化 memoryview 支持的对象种类非常多,只要是支持缓冲区的对象均可,比如:Numpy array、Python memoryview 对象、array.array 对象、支持缓冲区协议任意类型的对象,并且它也适用于 C 数组。所以它比原始缓冲区语法更加通用,原始缓冲区语法只适用于 Numpy array。

2. 类型化 memoryview 有着更多的选择来控制数据到底是 C 连续还是 Fortran 连续,以及直接或者间接访问。并且一些选项可以按照维度逐个控制,而 Numpy 的原始缓冲区语法不提供这种级别的控制。

3. 在任何情况下,类型化 memoryview 都有着超越原始缓冲区语法的性能,这一点才是我们最关注的,至于语法的间接程度其实没有太大的意义,个人这么认为。

包装 C、C++ 数组

回到我们之前的问题,当我们 C 的数组在堆上分配,那么返回之后要如何释放堆区申请的内存呢?我们举个栗子:

// heap_malloc.h
float *make_matrix_c(int nrows, int ncols);


// heap_malloc.c
#include <stdlib.h>

float *make_matrix_c(int nrows, int ncols) {
    float *matrix = (float *) malloc(nrows * ncols * sizeof(float));
    return matrix;
}

以上是返回一个在堆上分配的 C 数组:

import numpy as np

cdef extern from "heap_malloc.h":
    float *make_matrix_c(int nrows, int ncols)


def make_matrix(int nrows, int ncols):
    cdef float[:, :: 1] m = <float[:nrows, :ncols]> make_matrix_c(nrows, ncols)
    # 因为元素都未初始化, 所以里面的值都是不确定的, 虽然这无伤大雅, 但是更优雅的处理方式是将值都初始化为零值
    m[...] = 0.0
    # 如果直接返回 m, 那么 Python 在接收到之后会打印 <MemoryView of 'ndarray' object>
    # 所以这里我们转成 array 之后返回, 当然这里返回 m、然后在 Python 中接收到之后再将 m 转成 array 也是可以的
    # 只不过毕竟涉及到类型化 memoryview, 它是 Cython 的东西, 所以这一步还是在 Cython 中直接完成会比较好, 返回给 Python 的就是一个普通的 array
    return np.array(m, copy=False)

显然这个例子已经无需解释了,然后我们直接编译:

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

ext = Extension("cython_test",
                ["cython_test.pyx", "heap_malloc.c"])
setup(
    ext_modules=cythonize(ext, language_level=3),
)

编译完成之后将文件移到当前目录,导入测试:

import cython_test

arr1 = cython_test.make_matrix(3, 4)
print(arr1.__class__)  # <class 'numpy.ndarray'>
print(arr1)
"""
[[0. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 0. 0. 0.]]
"""

arr2 = cython_test.make_matrix(2, 5)
print(arr2)
"""
[[0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0.]]
"""

整体没有任何问题,很 happy,但是 arr1 和 arr2 对应的 C 数组都是在堆区,那这两个堆区的数组咋办?所以目前这个 make_matrix 是存在致命缺陷的,因为它有内存泄露的风险,而这个风险对于任何一个程序而言都是具有致命性的。

使用 Cython 正确(自动)地管理 C 数组

首先我们是需要对内存负责的,但很明显当和 C 共享数组的时候,合适的解决内存问题就变成了一件很棘手的问题,因为 C 语言没有自动管理内存的特性。通常在这种情况下,最干净利索的做法就是复制数据,来澄清各自对数据的所有权。比如:在返回 array 的时候,我们不加 copy=False,这样就会创建一个新的副本,所以你的是你的,我的是我的,两者之间没有关系。

from libc.stdlib cimport free
import numpy as np

cdef extern from "heap_malloc.h":
    float *make_matrix_c(int nrows, int ncols)


def make_matrix(int nrows, int ncols):
    cdef float[:, :: 1] m = <float[:nrows, :ncols]> make_matrix_c(nrows, ncols)
    m[...] = 0.0
    matrix = np.array(m)
    # 这里直接就可以释放了, 因为此时的 matrix 里面的元素是独立的
    free(m)
    return matrix

但很明显,如果数据量非常大,我们这么做是不是很影响效率呢?所以它虽然是一个解决问题的办法,但不是最好的办法。

应该有一种机制,在 Cython 中返回数组之前应该将这个数组和某个函数关联起来,一旦当数组被回收,那么要触发相应的函数,而在这个函数里面执行释放堆区内存的逻辑。那么我们看看这是如何实现的。

首先 Numpy 中的数组是在底层是一个 PyArrayObject,而 Numpy / C API 在此基础之上定义了一个 base 属性,而这个 base 属性就是为了这个目的而存在的。如果你正在使用 C API 来构建一个数组,并在堆区手动申请内存,那么你应该使用 PyArray_SetBaseObject 函数将 base 属性设置在内存申请在堆区的数组中。但是为了达到同样的目的,我们会使用 Cython 提供的函数,而不是使用 PyArray_SetBaseObject。

下面我们修改一下 pyx 文件,C 的头文件和源文件不需要做改动。

import numpy as np
from libc.stdlib cimport free
cimport numpy as c_np  # 用于访问 Numpy 的 C API

cdef extern from "heap_malloc.h":
    float *make_matrix_c(int nrows, int ncols)


cdef class _finalizer:
    # 定义一个类, 这个 void * 显然是申请在堆区的数组
    cdef void *_data
    def __dealloc__(self):
        # 析构函数, 这里再打印一句话
        print("数组被回收了")
        if self._data != NULL:
            free(self._data)

cdef void *set_base(c_np.ndarray arr, void *c_arr):
    # 函数接收两个参数: ndarray, 堆区申请的数组转成的 void * 指针

    # 实例化一个类
    cdef _finalizer f = _finalizer()
    # 将数组转成 void *, 交给 f._data
    f._data = c_arr
    # 重点来了, 这里便是底层对应的 PyArray_SetBaseObject, 但这里我们可以很轻松的进行调用
    # 因为 Cython 帮我们做好了这一切, 调用这个函数将 ndarray 和 f 关联起来, 而 f 里面存储了指向堆区的指针
    # 当 arr 被销毁时, 会触发 f 内部的 __dealloc__ 方法, 在这个方法将指针指向的内存释放掉
    c_np.set_array_base(arr, f)


def make_matrix(int nrows, int ncols):
    # 获取堆区申请的指针
    cdef float *c_arr = make_matrix_c(nrows, ncols)
    cdef float[:, :: 1] m = <float[:nrows, :ncols]> c_arr
    m[...] = 0.0
    # 这里通过 cdef c_np.ndarray 的方式静态声明
    cdef c_np.ndarray matrix = np.array(m, copy=False)
    # 调用 set_base, 传入 ndarray 和 指向堆区的指针
    set_base(matrix, <void *> c_arr)
    return matrix

然后我们来重新编译一下:

import numpy as np
from distutils.core import Extension, setup
from Cython.Build import cythonize

ext = Extension("cython_test",
                ["cython_test.pyx", "heap_malloc.c"],
                # 在当前路径和 np.get_include() 中寻找头文件
                include_dirs=[".", np.get_include()])
setup(
    ext_modules=cythonize(ext, language_level=3),
)

说实话个人觉得还是比较麻烦的,如果 C 返回的数组不大的话,还是直接拷贝一份吧,这样绝对是最方便的选择。话说只是一个数组而已,如果你不是对性能要求到极致的话,那么直接通过拷贝的方式绝对是最方便、最稳妥的选择。

总结

到目前为止,我们介绍了 Python 中的各种可以转成类型化 memoryview 的对象,但是 Numpy array 绝对是最具普遍性、灵活性以及表现力。而且除了 Python 对象之外,类型化 memoryview 还可以使用 C 级数组,不管是栈上分配,还是堆上分配。

而核心就在于 Cython 的类型化 memoryview,它提供了一个一致的抽象,这个抽象适用于所有支持缓冲区协议的对象,并为我们提供了高效的 C 级访问缓冲区,并且它们是共享内存的。

posted @ 2020-07-12 17:23  古明地盆  阅读(4877)  评论(0编辑  收藏  举报