在 Cython 中使用 prange 的并行性失败

Question

我尝试使用 prange 来加速代码，但时间成本与初始版本几乎相同。初始版本如下：

%%cython -a --cplus
import cython

cdef extern from "<complex.h>" namespace "std" nogil:
    double complex exp(double complex)

@cython.boundscheck(False)
@cython.wraparound(False)    
def cphaseshift(double px,double py,double[:,:] kX,double[:,:] kY,double complex[:,:] f):
    cdef double complex I = 1j
    cdef double pi = 3.141592653589793
    cdef int i
    cdef int j
    for i in range(kX.shape[0]):
        for j in range(kY.shape[1]):
            f[i,j] = f[i,j]*exp(-2*pi*I*(px*kX[i,j]+py*kY[i,j]))

初始版本速度：

px = 1
py = 1
kx = np.linspace(1,256,256)
kX,kY = np.meshgrid(kx,kx)
f0 = np.ones_like(kX,dtype='complex128')

%timeit cphaseshift(px,py,kX,kY,f0)

2.07 ms ± 23.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

平行版本：

%%cython -a --cplus
import cython
from cython.parallel import prange

cdef extern from "<complex.h>" namespace "std" nogil:
    double complex exp(double complex)

@cython.boundscheck(False)
@cython.wraparound(False)    
def cphaseshift_para(double px,double py,double[:,:] kX,double[:,:] kY,double complex[:,:] f):
    cdef double complex I = 1j
    cdef double pi = 3.141592653589793
    cdef int i
    cdef int j
    for i in prange(kX.shape[0],nogil=True,num_threads=6):
        for j in range(kY.shape[1]):
            f[i,j] = f[i,j]*exp(-2*pi*I*(px*kX[i,j]+py*kY[i,j]))

这个版本的速度：

px = 1
py = 1
kx = np.linspace(1,256,256)
kX,kY = np.meshgrid(kx,kx)
f0 = np.ones_like(kX,dtype='complex128')

%timeit cphaseshift_para(px,py,kX,kY,f0)

2.12 ms ± 28.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

我注意到 for i in prange(kX.shape[0],nogil=True,num_threads=6): 提示 python 互动。如何正确使用并行来加速代码？谢谢！

Answer 1

prange 使用 OpenMP 进行并行计算。但是，编译器和（在某些平台上）linker 需要与特殊标志一起使用才能激活 OpenMP 指令，默认情况下不会激活这些指令。

对于IPython，这可以通过

完成

%%cython -c=-fopenmp --link-args=-fopenmp
....

在 Linux 上使用 gcc。

在 Windows 使用 MSVC 时，参数将是：

%%cython -c=/openmp 
....

（在这种情况下不需要额外的 link 参数）。

如果使用安装文件，则 Extension 需要 -fopenmp/copenmp 通过 extra_compile_args 和 extra_link_args 传递。

即在 Linux+gcc:

...
ext_modules = [
    Extension(
        "myextension",
        ["myextension.pyx"],
        extra_compile_args=['-fopenmp'],
        extra_link_args=['-fopenmp'],
    )
]
...

在 Windows+MSVC 上只需要 extra_compile_args：

...
ext_modules = [
    Extension(
        "myextension",
        ["myextension.pyx"],
        extra_compile_args=['/openmp'],
    )
]
...

在我的机器上启用 openMP 进行编译导致加速因子约为 3：

%timeit cphaseshift(px,py,kX,kY,f0)
# 2.87 ms ± 115 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit cphaseshift_para(px,py,kX,kY,f0)
# 912 µs ± 36.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

您无法避免 Python-

中的互动

for i in prange(kX.shape[0],nogil=True,num_threads=6):

行，因为需要释放GIL然后再获取-这两个动作都需要与Python-interpreter交互。

在 Cython 中使用 prange 的并行性失败

Parallelism using prange in Cython fails

python

cython