在 cython 中并行化 for 循环：超越 prange

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 中有效地实现它，我开始怀疑当前版本的语言不可能做到这一点。请注意，仅在第一个循环中使用 cython.parallel 中的 prange 并不能提供正确的结果，因为（我假设）同时访问 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 二维数组，每个线程写入不同的列。之后沿第二个维度求和：

# 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 目前为此给我一个 "internal compiler error" - 我觉得它应该工作，但目前它没有'似乎没有用，所以尝试选项 3 需要您自担风险...）使用 openmp atomic directive 以原子方式进行加法。这需要一些工作来规避 Cython，但应该不会太困难。使用 add_inplace 宏创建一个简短的 C 头文件：
```
#define add_inplace(x,y) _Pragma("omp atomic") x+=y
```
_Pragma 是一个 C99 特性，应该允许您将 pragma 放在预处理器语句中。然后告诉 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*s，但它应该在使用标准 C 数字类型时生成正确的 C 代码。（检查代码可以肯定）

在 cython 中并行化 for 循环：超越 prange

Parallelising a for loop in cython: beyond prange

multithreading

multiprocessing

cython

python-multiprocessing