cython并行性能-计算滚动求和 rolling function

cython通过编译为C程序提高性能有很多例子，通过OpenMP并行的性能没那么多。

今天尝试了一下似乎gcc对parallelism reduction优化的很厉害，加上OpenMP并行可以提高20倍性能（相对于pandas rolling），这不是简单的2 core带来的性能提高。

滚动求和 rolling sum的例子

最简单的实现pandas.rolling，通过操作numpy array，速度也还算能接受。

# test_para.py
import numpy as np
#import pyximport; pyximport.install(reload_support=True, setup_args={"include_dirs":np.get_include()})
import timeit
import pandas as pd
import para.cpara as cpara

X = -1 + 2*np.random.rand(100000)
ss = pd.Series(X)
ss.rolling(100).apply(np.sum,raw=True)

print('==============')
print('multi thread')
start_time = timeit.default_timer()
sum_cython=pd.Series(cpara.window_sum(X, 100))
print(timeit.default_timer() - start_time)

print('single thread')
start_time = timeit.default_timer()
sum_pandas=ss.rolling(100).apply(np.sum,raw=True)
print(timeit.default_timer() - start_time)
print(np.max(np.abs(sum_cython - sum_pandas)))

cython源文件

# cpara.pyx
cimport cython
import numpy as np
from cython.parallel import prange,parallel
cimport numpy as cnp
from libc.stdlib cimport malloc

@cython.boundscheck(False)
def window_sum(cnp.ndarray[double, ndim=1] arr, int window):
    cdef h = np.zeros_like(arr)
    cdef int imax = arr.shape[0]
    cdef double *buffer = <double *>malloc(imax * sizeof(double))
    cdef double result = 0.0
    cdef int i, j
    with nogil, parallel():
        for i in prange(imax, schedule='dynamic'):
            buffer[i] = 0.0
            if i >= window-1:
                for j in range(window):
                    buffer[i] += arr[i-j]

    for i in range(imax):
        if i < window -1:
            h[i] = np.nan
        else:
            h[i] = buffer[i]

    return h

setup.py中要加入openmp的编译链接参数

EXT = Extension("*",
                ["para/*.pyx"],
                define_macros=[('CYTHON_TRACE', CYTHON_DEBUG),
                               ('CYTHON_TRACE_NOGIL', CYTHON_DEBUG),
                               ('CYTHON_BINDING', CYTHON_DEBUG),
                               ("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION"),
                               ('CYTHON_FAST_PYCCALL', '1')],
                extra_compile_args = ["-fopenmp" ],
                extra_link_args=['-fopenmp'],
                include_dirs=[".", np.get_include()])

性能比较

%timeit pd.Series(cpara.window_sum(X, 100))
23.4 ms ± 325 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit ss.rolling(100).apply(np.sum,raw=True)
536 ms ± 3.96 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

536/23.4=22.9

测试环境：i3-7100U 2core 2T CPU, ubuntu 18.04 LTS