multithreading - Cython:使 prange 并行化线程安全

Question

Cython 启动器在这里。我正在尝试通过使用多个线程来加速计算某个成对统计数据(在几个 bin 中)。特别是，我正在使用来自 cython.parallel 的 prange，它在内部使用 openMP。

以下最小示例说明了该问题(通过 Jupyter 笔记本 Cython 魔术进行编译)。

笔记本设置:

%load_ext Cython
import numpy as np

赛通代码:

%%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 上的错误共享，您的性能都会很差。向量。

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

带有手动“2D”索引的平面 numpy 数组效果很好。通过避免将数组的私有(private)部分填充到 64 字节(这是典型的高速缓存行大小)，您可以获得一点额外的性能。这避免了内核之间的错误共享。私有(private)部分在并行区域之外被简单地串行总结。

%%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。剩余的差距可能是由于涡轮模式。我现在无法使用合适的机器进行测试。