使用Cython中prange函数实现for循环的并行

上一篇文章我们探讨了 GIL 的原理以及如何释放 GIL 实现并行做法是将函数声明为 nogil然后使用 with nogil 上下文管理器即可。在使用上非常简单但如果我们想让循环也能够并行执行那么该方式就不太方便了为此 Cython 提供了一个 prange 函数专门用于循环的并行执行。这个 prange 的特殊功能是 Cython 独一无二的并且 prange 只能与 for 循环搭配使用不能独立存在。Cython 使用 OpenMP API 实现 prange用于多平台共享内存的处理。但 OpenMP 需要 C 或者 C 编译器支持并且编译时需要指定特定的编译参数来启动。例如当我们使用 gcc 时必须在编译和链接二进制文件的时候指定一个 -fopenmp以确保启用 OpenMP。许多编译器均支持 OpenMP 包括免费的和商业的。但 Clang/LLVM 则是一个最显著的例外它只在一个单独的分支中得到了初步的支持而为它完全实现的 OpenMP 还在开发当中。而使用 prange需要从 cython.parallel 中进行导入。但是在这之前我们先来看一个例子123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657importnumpy as npfromcython cimport boundscheck, wraparoundcdef inline double norm2(doublecomplexz) nogil:接收一个复数 z, 计算它的模的平方由于 norm2 要被下面的 escape 函数多次调用这里通过 inline 声明成内联函数:param z::return:returnz.real*z.realz.imag*z.imagcdefintescape(doublecomplexz,doublecomplexc,double z_max,intn_max) nogil:这个函数具体做什么, 不是我们的重点我们不需要关心cdef:inti0double z_max2z_max*z_maxwhilenorm2(z) z_max2andi n_max:zz*zci1returniboundscheck(False)wraparound(False)defcalc_julia(intresolution,doublecomplexc,double bound1.5,double z_max4.0,intn_max1000):我们将要在 Python 中调用的函数cdef:double step2.0*bound/resolutioninti, jdoublecomplexzdouble real, imagint[:, ::1] countscountsnp.zeros((resolution1, resolution1), dtypeint32)foriinrange(resolution1):real-boundi*stepforjinrange(resolution1):imag-boundj*stepzrealimag*1jcounts[i, j]escape(z, c, z_max, n_max)returnnp.array(counts, copyFalse)我们手动编译一下然后调用 calc_julia 函数这个函数做什么不需要关心我们只需要将注意力放在那两层 for 循环准确的说是外层循环上即可这里我们采用手动编译的形式。123456importcython_testimportnumpy as npimportmatplotlib.pyplot as pltarrcython_test.calc_julia(1000,0.3220.05j)plt.imshow(np.log(arr))plt.show()那么 calc_julia 这个函数耗时多少呢我们来测试一下使用 prange对于上面的代码来说外层循环里面的逻辑是彼此独立的即当前循环不依赖上一层循环的结果因此这非常适合并行执行。所以 prange 便闪亮登场了我们只需要做简单的修改即可1234567891011121314151617181920212223242526272829303132333435363738394041424344importnumpy as npfromcython cimport boundscheck, wraparoundfromcython.parallel cimport prangecdef inline double norm2(doublecomplexz) nogil:returnz.real*z.realz.imag*z.imagcdefintescape(doublecomplexz,doublecomplexc,double z_max,intn_max) nogil:cdef:inti0double z_max2z_max*z_maxwhilenorm2(z) z_max2andi n_max:zz*zci1returniboundscheck(False)wraparound(False)defcalc_julia(intresolution,doublecomplexc,double bound1.5,double z_max4.0,intn_max1000):cdef:double step2.0*bound/resolutioninti, jdoublecomplexzdouble real, imagint[:, ::1] countscountsnp.zeros((resolution1, resolution1), dtypeint32)# 只需要将外层的 range 换成 prangeforiinprange(resolution1, nogilTrue):real-boundi*stepforjinrange(resolution1):imag-boundj*stepzrealimag*1jcounts[i, j]escape(z, c, z_max, n_max)returnnp.array(counts, copyFalse)我们只需要将外层循环的 range 换成 prange 即可里面指定 nogilTrue便可实现并行的效果至于这个函数的其它参数以及用法后面会说。而且一旦使用了 prange那么在编译的时候必须启用 OpenMP下面看一下编译脚本。123456789fromdistutils.coreimportsetup, ExtensionfromCython.Buildimportcythonizeext[Extension(cython_test,sources[cython_test.pyx],extra_compile_args[-fopenmp],extra_link_args[-fopenmp])]setup(ext_modulescythonize(ext, language_level3))编译测试一下我们看到效率大概是提升了两倍因为我 Windows 上使用的不是 gcc所以这里是在 CentOS 上演示的。而我的 CentOS 服务器只有两个核因此效率提升大概两倍左右。所以只是做了一些非常简单的修改便可带来如此巨大的性能提升简直妙啊。prange 是要搭配 for 循环来使用的如果 for 循环内部的逻辑彼此独立即第二层循环不依赖第一层循环的某些结果那么不妨使用 prange 吧。注意还没完我们还能做得更好下面就来看看 prange 里面的其它的参数这样我们能更好利用 prange 的并行特性。prange 的其它参数prange 函数的原型如下123456789# 第一个参数 self 我们不需要管# prange 实际上是类 CythonDotParallel 的成员函数# 因为 Cython 内部执行了下面这行逻辑# sys.modules[cython.parallel] CythonDotParallel()# 所以它将一个实例对象变成了一个模块defprange(self, start0, stopNone, step1,nogilFalse, scheduleNone,chunksizeNone, num_threadsNone):我们先来看前三个参数start、stop、step。prange(3): 相当于 start0、stop3prange(1, 3): 相当于 start1、stop3prange(1, 3, 2): 相当于 start1、stop3、step2类似于 range同样不包含结尾 stop。然后是第四个参数 nogil它默认是 False但事实上我们必须将其设置为 True否则会报出编译错误。然后剩下的三个参数如果我们不指定的话那么 Cython 编译器采取的策略是将整个循环分成多个大小相同的连续块然后给每一个可用线程一个块。然而这个策略实际上并不是最好的因为每一层循环用的时间不一定一样如果一个线程很快就完成了那么不就造成资源上的浪费了吗我们修改一下将 schedule 指定为 staticchunksize 指定为 112foriinprange(resolution1, nogilTrue,schedulestatic, chunksize1):其它地方不变只是加两个参数然后重新测试一下。我们看到效率上是差不多的原因是我的机器只有两个核如果核数再多一些的话那么速度就会明显地提升。下面来解释一下剩余的三个参数的含义首先是 schedule它有以下几个选项static整个循环在编译时会以一种固定的方式分配给多个线程如果 chunksize 没有指定那么会分成 num_threads 个连续块一个线程一个块。如果指定了 chunksize那么每一块会以轮询调度算法Round Robin交给线程进行处理适用于任务均匀分布的情况。dynamic线程在运行时动态地向调度器申请下一个块chunksize 默认为 1当任务负载不均时动态调度是最佳的选择。guided块是动态分布的就像 dynamic 一样但这与 dynamic 还不同chunksize 的比例不是固定的而是和 剩余迭代次数 / 线程数 成比例关系。runtime不常用。控制 schedule 和 chunksize 可以方便地探索不同的并行执行策略、以及工作负载分配通常指定 schedule 为 static加上设置一个合适的 chunksize 是最好的选择。而 dynamic 和 guided 适用于动态变化的执行上下文但会导致运行时开销。当然还有最后一个参数 num_threads很明显不需要解释就是使用的线程数量。如果不指定那么 prange 会使用尽可能多的线程。所以我们只是做了一点修改便可以带来巨大的性能提升这种性能提升与 Cython 在纯 Python 上带来的性能提升成倍增关系。