Cython内联函数与numpy数组作为参数
考虑一下这样的代码：
import numpy as np
cimport numpy as np
cdef inline inc(np.ndarray[np.int32_t] arr, int i):
arr[i]+= 1
def test1(np.ndarray[np.int32_t] arr):
cdef int i
for i in xrange(len(arr)):
inc(arr, i)
def test2(np.ndarray[np.int32_t] arr):
cdef int i
for i in xrange(len(arr)):
arr[i] += 1
我用ipython来测量test1和test2的速度：
In [7]: timeit ttt.test1(arr)
100 loops, best of 3: 6.13 ms per loop
In [8]: timeit ttt.test2(arr)
100000 loops, best of 3: 9.79 us per loop
有没有办法优化test1？为什么cython没有像说的那样内联这个函数呢？
更新：其实我需要的是像这样的多维代码：
# cython: infer_types=True
# cython: boundscheck=False
# cython: wraparound=False
import numpy as np
cimport numpy as np
cdef inline inc(np.ndarray[np.int32_t, ndim=2] arr, int i, int j):
arr[i, j] += 1
def test1(np.ndarray[np.int32_t, ndim=2] arr):
cdef int i,j
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
inc(arr, i, j)
def test2(np.ndarray[np.int32_t, ndim=2] arr):
cdef int i,j
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
arr[i,j] += 1
对它的计时：
In [7]: timeit ttt.test1(arr)
1 loops, best of 3: 647 ms per loop
In [8]: timeit ttt.test2(arr)
100 loops, best of 3: 2.07 ms per loop
显式内联让速度提升了300倍。而且我的真实函数相当大，所以内联会让代码的可维护性变得更糟
更新2：
# cython: infer_types=True
# cython: boundscheck=False
# cython: wraparound=False
import numpy as np
cimport numpy as np
cdef inline inc(np.ndarray[np.float32_t, ndim=2] arr, int i, int j):
arr[i, j]+= 1
def test1(np.ndarray[np.float32_t, ndim=2] arr):
cdef int i,j
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
inc(arr, i, j)
def test2(np.ndarray[np.float32_t, ndim=2] arr):
cdef int i,j
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
arr[i,j] += 1
cdef class FastPassingFloat2DArray(object):
cdef float* data
cdef int stride0, stride1
def __init__(self, np.ndarray[np.float32_t, ndim=2] arr):
self.data = <float*>arr.data
self.stride0 = arr.strides[0]/arr.dtype.itemsize
self.stride1 = arr.strides[1]/arr.dtype.itemsize
def __getitem__(self, tuple tp):
cdef int i, j
cdef float *pr, r
i, j = tp
pr = (self.data + self.stride0*i + self.stride1*j)
r = pr[0]
return r
def __setitem__(self, tuple tp, float value):
cdef int i, j
cdef float *pr, r
i, j = tp
pr = (self.data + self.stride0*i + self.stride1*j)
pr[0] = value
cdef inline inc2(FastPassingFloat2DArray arr, int i, int j):
arr[i, j]+= 1
def test3(np.ndarray[np.float32_t, ndim=2] arr):
cdef int i,j
cdef FastPassingFloat2DArray tmparr = FastPassingFloat2DArray(arr)
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
inc2(tmparr, i,j)
计时结果：
In [4]: timeit ttt.test1(arr)
1 loops, best of 3: 623 ms per loop
In [5]: timeit ttt.test2(arr)
100 loops, best of 3: 2.29 ms per loop
In [6]: timeit ttt.test3(arr)
1 loops, best of 3: 201 ms per loop
3 个回答
7
你把数组作为一种Python对象（类型是numpy.ndarray）传给了inc()。这样传递Python对象会比较耗费资源，因为涉及到引用计数的问题，而且这似乎还会影响到内联优化。如果你用C语言的方式传递数组，也就是作为一个指针传递，那么在我的机器上，test1()的运行速度会比test2()更快：
cimport numpy as np
cdef inline inc(int* arr, int i):
arr[i] += 1
def test1(np.ndarray[np.int32_t] arr):
cdef int i
for i in xrange(len(arr)):
inc(<int*>arr.data, i)
21
这个问题已经发布超过3年了，在这段时间里有了很大的进展。在这段代码（问题的更新2）中：
# cython: infer_types=True
# cython: boundscheck=False
# cython: wraparound=False
import numpy as np
cimport numpy as np
cdef inline inc(np.ndarray[np.int32_t, ndim=2] arr, int i, int j):
arr[i, j]+= 1
def test1(np.ndarray[np.int32_t, ndim=2] arr):
cdef int i,j
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
inc(arr, i, j)
def test2(np.ndarray[np.int32_t, ndim=2] arr):
cdef int i,j
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
arr[i,j] += 1
我得到了以下的时间记录：
arr = np.zeros((1000,1000), dtype=np.int32)
%timeit test1(arr)
%timeit test2(arr)
1 loops, best of 3: 354 ms per loop
1000 loops, best of 3: 1.02 ms per loop
所以这个问题在3年后依然存在。Cython现在有了类型化内存视图，我记得这是在Cython 0.16版本中引入的，所以在提问时是没有的。用这个：
# cython: infer_types=True
# cython: boundscheck=False
# cython: wraparound=False
import numpy as np
cimport numpy as np
cdef inline inc(int[:, ::1] tmv, int i, int j):
tmv[i, j]+= 1
def test3(np.ndarray[np.int32_t, ndim=2] arr):
cdef int i,j
cdef int[:, ::1] tmv = arr
for i in xrange(tmv.shape[0]):
for j in xrange(tmv.shape[1]):
inc(tmv, i, j)
def test4(np.ndarray[np.int32_t, ndim=2] arr):
cdef int i,j
cdef int[:, ::1] tmv = arr
for i in xrange(tmv.shape[0]):
for j in xrange(tmv.shape[1]):
tmv[i,j] += 1
我得到了：
arr = np.zeros((1000,1000), dtype=np.int32)
%timeit test3(arr)
%timeit test4(arr)
1000 loops, best of 3: 977 µs per loop
1000 loops, best of 3: 838 µs per loop
我们几乎达到了目标，速度已经比以前的方式快了！现在，inc()函数可以声明为nogil，所以我们来声明一下！但是，哎呀：
Error compiling Cython file:
[...]
cdef inline inc(int[:, ::1] tmv, int i, int j) nogil:
^
[...]
Function with Python return type cannot be declared nogil
啊，我完全忘了返回类型void没有写！再来一次，这次加上void：
cdef inline void inc(int[:, ::1] tmv, int i, int j) nogil:
tmv[i, j]+= 1
最后我得到了：
%timeit test3(arr)
%timeit test4(arr)
1000 loops, best of 3: 843 µs per loop
1000 loops, best of 3: 853 µs per loop
速度和手动内联一样快！
现在，出于好玩，我在这段代码上试了Numba：
import numpy as np
from numba import autojit, jit
@autojit
def inc(arr, i, j):
arr[i, j] += 1
@autojit
def test5(arr):
for i in xrange(arr.shape[0]):
for j in xrange(arr.shape[1]):
inc(arr, i, j)
我得到了：
arr = np.zeros((1000,1000), dtype=np.int32)
%timeit test5(arr)
100 loops, best of 3: 4.03 ms per loop
虽然它比Cython慢了4.7倍，可能是因为JIT编译器没有成功内联inc()，但我觉得这太棒了！我只需要加上@autojit，就不用在代码里搞那些复杂的类型声明；几乎没有成本就实现了88倍的速度提升！
我还尝试了其他一些Numba的功能，比如
@jit('void(i4[:],i4,i4)')
def inc(arr, i, j):
arr[i, j] += 1
或者nopython=True，但没有进一步改善。
改进内联功能在Numba开发者的计划中，我们只需要提交更多请求，让它优先级更高。；)
7
问题在于，给一个numpy数组赋值（或者把它作为函数参数传入）并不是简单的赋值，而是一个“缓冲区提取”的过程。这会填充一个结构体，并将需要的步幅和指针信息提取到本地变量中，以便快速索引。如果你在循环中处理的元素数量适中，这种O(1)的开销在循环中是可以忽略的，但对于小函数来说情况就不一样了。
很多人都希望能改善这个问题，但这并不是一个简单的改动。具体可以参考这个讨论：http://groups.google.com/group/cython-users/browse_thread/thread/8fc8686315d7f3fe