Question

这里是Cython的首发。我试图通过使用多个线程来加速某个成对统计（在几个箱中）的计算。特别是，我使用cython.parallel中的prange，它在内部使用openMP。

以下最小的例子说明了这个问题（通过Jupyter笔记本Cython魔术编译）。

笔记本设置：

%load_ext Cython
import numpy as np

Cython代码：

%%cython --compile-args=-fopenmp --link-args=-fopenmp -a

from cython cimport boundscheck
import numpy as np
from cython.parallel cimport prange, parallel

@boundscheck(False)
def my_parallel_statistic(double[:] X, double[:,::1] bins, int num_threads):

    cdef: 
        int N = X.shape[0]
        int nbins = bins.shape[0]
        double Xij,Yij
        double[:] Z = np.zeros(nbins,dtype=np.float64)
        int i,j,b

    with nogil, parallel(num_threads=num_threads):
        for i in prange(N,schedule='static',chunksize=1):
            for j in range(i):
                #some pairwise quantities
                Xij = X[i]-X[j]
                Yij = 0.5*(X[i]+X[j])
                #check if in bin
                for b in range(nbins):
                    if (Xij < bins[b,0]) or (Xij > bins[b,1]):
                        continue
                    Z[b] += Xij*Yij

    return np.asarray(Z)

模拟数据和箱子

X = np.random.rand(10000)
bin_edges = np.linspace(0.,1,11)
bins = np.array([bin_edges[:-1],bin_edges[1:]]).T
bins = bins.copy(order='C')

通过时间

%timeit my_parallel_statistic(X,bins,1)
%timeit my_parallel_statistic(X,bins,4)

产量

1 loop, best of 3: 728 ms per loop
1 loop, best of 3: 330 ms per loop

这不是一个完美的缩放，但这不是问题的主要观点。（但是，除了添加通常的装饰器或微调prange参数之外，如果你有建议，请告诉我。）

但是，这个计算显然不是线程安全的：

Z1 = my_parallel_statistic(X,bins,1)
Z4 = my_parallel_statistic(X,bins,4)
np.allclose(Z1,Z4)

显示两个结果之间的显着差异（在此示例中高达20％）。

我强烈怀疑问题是多线程可以做到

Z[b] += Xij*Yij

同时。但我不知道的是如何在不牺牲加速的情况下解决这个问题。

在我的实际使用案例中，Xij和Yij的计算更加昂贵，因此我想每对只执行一次。此外，为所有对预先计算和存储Xij和Yij然后简单地循环通过bin也不是一个好的选择，因为N可以变得非常大，而且我不能在内存中存储100,000 x 100,000个numpy数组（这是实际上是在Cython中重写它的主要动机！）。

系统信息（在评论中添加以下建议）：

CPU(s): 8
Model name: Intel(R) Core(TM) i7-4790K CPU @ 4.00GHz
OS: Red Hat Linux v6.8
Memory: 16 GB

Answer 1

是的，Z[b] += Xij*Yij确实是竞争条件。

有两种方法可以制作atomic或critical。除了Cython的实现问题之外，由于共享Z向量上的错误共享，您在任何情况下都会有糟糕的性能。

所以更好的选择是为每个线程保留一个私有数组。再次有几个（非）选项。可以使用私人malloc指针，但我想坚持使用np。内存片不能指定为私有变量。二维(num_threads, nbins)数组有效，但由于某种原因生成非常复杂的低效数组索引代码。这有效，但速度较慢，无法扩展。

带有手动＆＃34; 2D＆＃34;平坦的numpy阵列索引效果很好。通过避免将数组的私有部分填充为64字节（这是典型的高速缓存行大小），您可以获得一些额外的性能。这避免了核心之间的错误共享。私有部分简单地在并行区域之外连续求和。

%%cython --compile-args=-fopenmp --link-args=-fopenmp -a
from cython cimport boundscheck
import numpy as np
from cython.parallel cimport prange, parallel
cimport openmp

@boundscheck(False)
def my_parallel_statistic(double[:] X, double[:,::1] bins, int num_threads):

    cdef: 
        int N = X.shape[0]
        int nbins = bins.shape[0]
        double Xij,Yij
        # pad local data to 64 byte avoid false sharing of cache-lines
        int nbins_padded = (((nbins - 1) // 8) + 1) * 8
        double[:] Z_local = np.zeros(nbins_padded * num_threads,dtype=np.float64)
        double[:] Z = np.zeros(nbins)
        int i,j,b, bb, tid

    with nogil, parallel(num_threads=num_threads):
        tid = openmp.omp_get_thread_num()
        for i in prange(N,schedule='static',chunksize=1):
            for j in range(i):
                #some pairwise quantities
                Xij = X[i]-X[j]
                Yij = 0.5*(X[i]+X[j])
                #check if in bin
                for b in range(nbins):
                    if (Xij < bins[b,0]) or (Xij > bins[b,1]):
                        continue
                    Z_local[tid * nbins_padded + b] += Xij*Yij
    for tid in range(num_threads):
        for bb in range(nbins):
            Z[bb] += Z_local[tid * nbins_padded + bb]


    return np.asarray(Z)

这在我的4核计算机上运行良好，720 ms / 191 ms，加速3.6。剩余的差距可能是由于涡轮模式。我现在无法使用合适的机器进行测试。

Answer 2

你是对的，Z的访问是在竞争条件下。

最好定义Z的num_threads个副本，作为cdef double[:] Z = np.zeros((num_threads, nbins), dtype=np.float64)，并在prange循环后沿轴0执行求和。

return np.sum(Z, axis=0)

Cython代码在并行区域中可以有with gil语句，但只记录错误处理。您可以查看一般的C代码，看看是否会触发原子OpenMP操作，但我对此表示怀疑。

Cython：使prange并行化线程安全

2 个答案: