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

Posted

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了在cython中并行化for循环:超越prange相关的知识,希望对你有一定的参考价值。

我正在努力使用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]

这段代码适用于一些简单的并行化(valuesdist非常大):需要将第一个for循环分割为单独的进程,每个进程都在自己的countvarg数组版本上运行。完成后,必须通过在第二个for循环(更短)之前将countvarg的不同版本相加来将所有内容组合在一起。

也就是说,这是两天我试图理解如何在cython中有效地实现这一点,我开始怀疑使用当前版本的语言是不可能的。请注意,仅使用来自prangecython.parallel作为第一个循环并不能提供正确的结果,因为(我假设)同时访问来自不同线程的nicountvarg

cython并行支持真的如此有限吗?我得到了这么好的加速单线程,我只希望我能继续......

答案

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

  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完成的工作相当大,这很容易也不会太糟糕
  2. 制作countvarg 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中的想法做类似的事情。
  3. (注意 - 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*s,但它应该在使用标准C数字类型时生成正确的C代码。(检查代码是否确定)

以上是关于在cython中并行化for循环:超越prange的主要内容,如果未能解决你的问题,请参考以下文章

如何在异步函数中并行化 for 循环并跟踪 for 循环执行状态?

在 Python 中通过线程/核心/节点并行化 for 循环

在 OpenMP 中并行化嵌套循环并使用更多线程执行内部循环

使用 OpenMP 在 C、C++ 中并行化嵌套 for 循环的几种方法之间的区别

如何在 C++ 中并行化一个 for 循环,只创建一次线程池

使用 CUDA 在 python 中展开一个可并行化的 for 循环