CPython、Cython与NumPy数组性能对比

Question

我正在对一种生成素数的程序进行性能测试，这个程序是从http://docs.cython.org/src/tutorial/numpy.html上来的。下面是我在kmax=1000时的性能测试结果。

用纯Python在CPython上运行的结果是：0.15秒

用纯Python在Cython上运行的结果是：0.07秒

def primes(kmax):
    p = []
    k = 0
    n = 2
    while k < kmax:
        i = 0
        while i < k and n % p[i] != 0:
            i = i + 1
        if i == k:
            p.append(n)
            k = k + 1
        n = n + 1
    return p

用纯Python加Numpy在CPython上运行的结果是：1.25秒

import numpy

def primes(kmax):
    p = numpy.empty(kmax, dtype=int)
    k = 0
    n = 2
    while k < kmax:
        i = 0
        while i < k and n % p[i] != 0:
            i = i + 1
        if i == k:
            p[k] = n
            k = k + 1
        n = n + 1
    return p

用Cython实现，使用int*的结果是：0.003秒

from libc.stdlib cimport malloc, free

def primes(int kmax):
    cdef int n, k, i
    cdef int *p = <int *>malloc(kmax * sizeof(int))
    result = []
    k = 0
    n = 2
    while k < kmax:
        i = 0
        while i < k and n % p[i] != 0:
            i = i + 1
        if i == k:
            p[k] = n
            k = k + 1
            result.append(n)
        n = n + 1
    free(p)
    return result

上面的结果表现得很好，但看起来很糟糕，因为它保存了两份数据...所以我尝试重新实现了一下：

Cython加Numpy的结果是：1.01秒

import numpy as np
cimport numpy as np
cimport cython

DTYPE = np.int
ctypedef np.int_t DTYPE_t

@cython.boundscheck(False)
def primes(DTYPE_t kmax):
    cdef DTYPE_t n, k, i
    cdef np.ndarray p = np.empty(kmax, dtype=DTYPE)
    k = 0
    n = 2
    while k < kmax:
        i = 0
        while i < k and n % p[i] != 0:
            i = i + 1
        if i == k:
            p[k] = n
            k = k + 1
        n = n + 1
    return p

我有几个问题：

1. 为什么在CPython上，Numpy数组的速度比Python列表慢得多？
2. 我在Cython加Numpy的实现中做错了什么？Cython显然没有把Numpy数组当作int[]来处理。

3. 我该如何将Numpy数组转换为int*？下面的代码不起作用。

   cdef numpy.nparray a = numpy.zeros(100, dtype=int)
   cdef int * p = <int *>a.data
   

Answer 1

为什么在CPython上，numpy数组的速度比Python列表慢得多呢？

因为你没有完全指定类型。使用

cdef np.ndarray[dtype=np.int, ndim=1] p = np.empty(kmax, dtype=DTYPE)

我怎么把numpy数组转换成int*？

要用 np.intc 作为数据类型，而不是 np.int（后者是C语言中的 long）。这样做是对的。

cdef np.ndarray[dtype=int, ndim=1] p = np.empty(kmax, dtype=np.intc)

（不过其实，使用内存视图会更好，它们更简洁，而且Cython的开发者希望最终能摆脱NumPy数组的语法。）

Answer 2

cdef DTYPE_t [:] p_view = p

用这个代替计算中的 p，让我的运行时间从580毫秒降到了2.8毫秒。这差不多和用 *int 实现的运行时间一样快。而这大概就是你能期待的最大效果了。

DTYPE = np.int
ctypedef np.int_t DTYPE_t

@cython.boundscheck(False)
def primes(DTYPE_t kmax):
    cdef DTYPE_t n, k, i
    cdef np.ndarray p = np.empty(kmax, dtype=DTYPE)
    cdef DTYPE_t [:] p_view = p
    k = 0
    n = 2
    while k < kmax:
        i = 0
        while i < k and n % p_view[i] != 0:
            i = i + 1
        if i == k:
            p_view[k] = n
            k = k + 1
        n = n + 1
    return p

Answer 3

我找到的最佳语法如下：

import numpy
cimport numpy
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def primes(int kmax):
    cdef int n, k, i
    cdef numpy.ndarray[int] p = numpy.empty(kmax, dtype=numpy.int32)
    k = 0
    n = 2
    while k < kmax:
        i = 0
        while i < k and n % p[i] != 0:
            i = i + 1
        if i == k:
            p[k] = n
            k = k + 1
        n = n + 1
    return p

注意我在这里用了numpy.int32而不是int。在cdef的左边用的都是C语言的数据类型（所以int相当于int32，float相当于float32），而在cdef的右边（或者说在cdef外面）用的都是Python的数据类型（int相当于int64，float相当于float64）。