cython prange 没有单线程那么快

cython prange not so fast than single thread

我刚刚写了一个简单的程序来测试 cythonprange 是如何执行的,这里是代码:

from cython.parallel import prange
import numpy as np

def func(int r, int c):
  cdef:
    double[:,:] a = np.arange(r*c, dtype=np.double).reshape(r,c)
    double total = 0
    int i, j

  for i in prange(r, nogil=True, schedule='static', chunksize=1):
    for j in range(c):
      total += a[i,j]

  return total

在Mac Book pro上,OMP_NUM_THREADS=3,上面的代码(r,c) = (10000, 100000)用了将近18秒,单线程用了大约21秒。

为什么性能提升这么小?我使用这个 prange 正确吗?

你计算过分配 a 需要多长时间吗?一个 10000 x 100000 的 float64 数组占用 8GB 内存。

a = np.ones((10000, 100000), np.double)

在我的 16GB RAM 笔记本电脑上需要六秒多的时间。如果您没有 8GB 可用空间,那么您将进行交换,这将花费 很多 更长的时间。由于 func 几乎所有的时间都花在了分配 a 上,因此并行化你的外部 for 循环只能让你在总运行时间上获得一小部分改进。

为了证明这一点,我修改了您的函数以接受 a 作为输入。在 tmp.pyx:

#cython: boundscheck=False, wraparound=False, initializedcheck=False

from cython.parallel cimport prange

def serial(double[:, :] a):
    cdef:
        double total = 0
        int i, j
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            total += a[i, j]
    return total

def parallel(double[:, :] a):
    cdef:
        double total = 0
        int i, j
    for i in prange(a.shape[0], nogil=True, schedule='static', chunksize=1):
        for j in range(a.shape[1]):
            total += a[i, j]
    return total

例如:

In [1]: import tmp

In [2]: r, c = 10000, 100000

In [3]: a = np.random.randn(r, c)   # this takes ~6.75 sec

In [4]: %timeit tmp.serial(a)
1 loops, best of 3: 1.25 s per loop

In [5]: %timeit tmp.parallel(a)
1 loops, best of 3: 450 ms per loop

在我的 4 核笔记本电脑上,并行化该函数的时间约为 2.8 倍 speed-up*,但这只是分配 a.

所用时间的一小部分

这里的教训是,在深入优化之前,始终分析您的代码以了解它花费最多时间的地方。


* 您可以通过将更大的 a 块传递给每个工作进程来做得更好,例如通过增加 chunksize 或使用 schedule='guided'