在 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 进行并行计算。但是，编译器和（在某些平台上）链接器需要与特殊标志一起使用才能激活 OpenMP 指令，这些指令默认情况下不会被激活。

对于 IPython，这可以通过

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

在带有 gcc 的 Linux 上。

在带有 MSVC 的 Windows 上，参数为：

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

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

---

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

即在 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 交互。