用Cython编写C语言的高效实践

本文详细介绍了如何直接使用Cython编写接近C/C++性能的代码,包括原生C结构体、数组的应用,以及如何通过自定义内存管理避免内存泄漏,同时保持Python的灵活性。

过去两年,几乎所有工作都通过Cython完成。这里并非指先编写Python代码再通过类型声明等方式进行"Cython化",而是直接编写Cython代码。

采用原生C结构体、数组以及偶尔使用的C++向量,配合自行编写的malloc/free薄封装层。这种写法的性能几乎总是与C/C++持平,因为这些代码本质上就是带有语法糖的C/C++——同时保留了随时调用Python的能力。

这与Python等语言的传统承诺形成鲜明对比:理论上可以用Python编写整个应用,仅对热点进行C优化即可获得C的速度和Python的便利。但实践中,数据结构选择会极大影响代码效率与开发体验:数组高效但难用,列表便捷但缓慢。由于Python循环和函数调用开销较大,最终往往需要将大部分代码改用C实现。

某技术社区近期出现关于为Python编写C扩展的讨论帖,作者同时提供了纯Python实现和基于Numpy C API的C实现。为此专门编写了Cython实现进行对比:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import random
from cymem.cymem cimport Pool
from libc.math cimport sqrt
cimport cython

cdef struct Point:
    double x
    double y

cdef class World:
    cdef Pool mem
    cdef int N
    cdef double* m
    cdef Point* r
    cdef Point* v
    cdef Point* F
    cdef readonly double dt

    def __init__(self, N, threads=1, m_min=1, m_max=30.0, 
                 r_max=50.0, v_max=4.0, dt=1e-3):
        self.mem = Pool()
        self.N = N
        self.m = <double*>self.mem.alloc(N, sizeof(double))
        self.r = <Point*>self.mem.alloc(N, sizeof(Point))
        self.v = <Point*>self.mem.alloc(N, sizeof(Point))
        self.F = <Point*>self.mem.alloc(N, sizeof(Point))
        for i in range(N):
            self.m[i] = random.uniform(m_min, m_max)
            self.r[i].x = random.uniform(-r_max, r_max)
            self.r[i].y = random.uniform(-r_max, r_max)
            self.v[i].x = random.uniform(-v_max, v_max)
            self.v[i].y = random.uniform(-v_max, v_max)
            self.F[i].x = 0
            self.F[i].y = 0
        self.dt = dt

@cython.cdivision(True)
def compute_F(World w):
    """计算作用力"""
    cdef int i, j
    cdef double s3, tmp
    cdef Point s, F
    for i in range(w.N):
        w.F[i].x = 0
        w.F[i].y = 0
        for j in range(i+1, w.N):
            s.x = w.r[j].x - w.r[i].x
            s.y = w.r[j].y - w.r[i].y
            s3 = sqrt(s.x * s.x + s.y * s.y)
            s3 *= s3 * s3
            tmp = w.m[i] * w.m[j] / s3
            F.x = tmp * s.x
            F.y = tmp * s.y
            w.F[i].x += F.x
            w.F[i].y += F.y
            w.F[j].x -= F.x
            w.F[j].y -= F.y

@cython.cdivision(True)
def evolve(World w, int steps):
    """模拟演化过程"""
    cdef int _, i
    for _ in range(steps):
        compute_F(w)
        for i in range(w.N):
            w.v[i].x += w.F[i].x * w.dt / w.m[i]
            w.v[i].y += w.F[i].y * w.dt / w.m[i]
            w.r[i].x += w.v[i].x * w.dt
            w.r[i].y += w.v[i].y * w.dt

该Cython实现耗时约30分钟完成,性能与C代码完全相当——因为这些本质上就是带有语法糖的C代码。开发者无需学习复杂的外部C API,直接编写C或C++代码即可(尽管C++支持稍显笨拙)。实测Cython版本和C版本都比使用Numpy数组的纯Python实现快约70倍。

与原生C的区别在于使用了自研的malloc/free封装工具cymem,其自动记录分配地址并在Pool对象垃圾回收时释放对应内存,有效解决了内存泄漏问题。

虽然Cython支持通过类型化内存视图使用Numpy多维数组特性,但对于涉及稀疏数组的应用场景,自定义数据结构仍是更优选择。

comments powered by Disqus
使用 Hugo 构建
主题 StackJimmy 设计