multithreading - cython : beyond prange中的for循环并行处理

Question

我正在努力使用cython正确并行化一个函数。基本上，问题是要对某些数据进行装箱。实际的代码有点长，但最终它会执行以下操作:

def bin_var(double[:] dist,
            double[:] values,
            double[:] bin_def,
            double[:] varg, long[:] count):

    dbin = (bin_def[1] - bin_def[0]) / bin_def[2]

    for n1 in range(values.size):
            if (dist[n1] < bin_def[0]) or (dist[n1] >= bin_def[1]):
                continue
            else:
                ni = int((dist - bin_def[0]) / dbin)
                count[ni] += 1
                varg[ni] += calc_something(values[ni])

    # compute the mean
    for n1 in range(int(bin_def[2])):
        varg[ni] /= count[ni]

此代码使自己可以进行一些简单的并行化( values和 dist很大):一个人需要将第一个 for循环拆分为单独的进程，每个进程都在自己的 count和 varg数组版本上工作。完成此操作后，您必须通过在第二个for循环(要短得多)之前将不同版本的 count和 varg相加来将所有内容组合在一起。

就是说，两天来我试图了解如何在 cython中有效地实现这一点，并且我开始怀疑当前版本的语言是不可能的。请注意，仅将 prange中的 cython.parallel用于第一个循环并不能提供正确的结果，因为(我假设)是从不同线程同时访问 ni， count和 varg。
cython并行支持真的那么有限吗？我获得了如此出色的单线程加速，我只是希望我可以继续...

Answer 1

我可以在这里想到三个选择:

使用GIL来确保+=是单线程完成的:

varg_ni = calc_something(values[ni]) # keep this out 
               # of the single threaded block...
with gil:
    count[ni] += 1
    varg[ni] += varg_ni

只要在calc_something中完成的工作相当大，那么这很容易并且不会太糟

每个线程写入不同的列，以生成count和varg 2D数组。然后沿第二维求和:

# rough, untested outline....

# might need to go in a `with parallel()` block
num_threads = openmp.omp_get_num_threads()

cdef double[:,:] count_tmp = np.zeros((count.shape[0],num_threads))
cdef double[:,:] varg_tmp = np.zeros((varg.shape[0],num_threads))

# then in the loop:
count_tmp[ni,cython.parallel.threadid()] += 1
varg_tmp[ni,cython.parallel.threadid()] += calc_something(values[ni])

# after the loop:
count[:] = np.sum(count_tmp,axis=1)
varg[:] = np.sum(varg_tmp,axis=1)

您还可以使用 local_buf example in the documentation中的想法执行类似的操作。

(注意-GCC当前为此提供了一个“内部编译器错误”-我认为它应该可以工作，但是目前看来似乎不起作用，因此，尝试方法3后果自负。 )使用openmp atomic directive原子地添加。这需要一些工作来规避Cython，但应该不会太困难。使用add_inplace宏创建一个简短的C头文件:

#define add_inplace(x,y) _Pragma("omp atomic") x+=y

_Pragma是C99的一项功能，应允许您在预处理器语句中放入编译指示。然后，向Cython告知该头文件(就好像它是一个函数一样):

cdef extern from "header.h":
    void add_inplace(...) nogil # just use varargs to make Cython think it accepts anything

然后在循环中执行:

add_inplace(count[ni], 1)
add_inplace(varg[ni], calc_something(values[ni]))

因为这使用了宏技巧，所以它可能有点脆弱(即，绝对不能与PyObject*一起使用，但在使用标准C数值类型时，它应该生成正确的C代码。(请检查代码以确保正确)