使用 cython 并行化 python 循环 numpy.searchsorted

Question

我使用 cython 编写了一个包含以下循环的函数。对数组 A1 的每一行进行二进制搜索以查找数组 A2 中的所有值。所以每次循环迭代 returns 一个索引值的二维数组。数组 A1 和 A2 作为函数参数输入，类型正确。

数组 C 按照 cython 的要求预先分配在最高缩进级别。

我对这个问题做了一些简化。

...
cdef np.ndarray[DTYPEint_t, ndim=3] C = np.zeros([N,M,M], dtype=DTYPEint)

for j in range(0,N):
    C[j,:,:]  = np.searchsorted(A1[j,:], A2, side='left' )

到目前为止一切正常，编译正常，运行符合预期。但是，为了获得更快的速度，我想并行化 j 循环。第一次尝试只是写

for j in prange(0,N, nogil=True):
    C[j,:,:]  = np.searchsorted(A1[j,:], A2, side='left' )

我尝试了很多编码变体，例如将东西放在一个单独的 nogil_function 中，将结果分配给中间数组并编写嵌套循环以避免分配给 C 的切片部分。

错误的形式通常是"Accessing Python attribute not allowed without gil"

我无法让它工作。对我如何做到这一点有什么建议吗？

编辑：

这是我的setup.py

try:
    from setuptools import setup
    from setuptools import Extension
except ImportError:
    from distutils.core import setup
    from distutils.extension import Extension


from Cython.Build import cythonize

import numpy

extensions = [Extension("matchOnDistanceVectors",
                    sources=["matchOnDistanceVectors.pyx"],
                    extra_compile_args=["/openmp", "/O2"],
                    extra_link_args=[]
                   )]


setup(
ext_modules = cythonize(extensions),
include_dirs=[numpy.get_include()]
)

我正在 windows 7 使用 msvc 进行编译。我确实指定了 /openmp 标志，我的数组大小为 200*200。所以一切似乎都井井有条...

Answer 1

我相信 searchsorted 发布了 GIL 本身（参见 https://github.com/numpy/numpy/blob/e2805398f9a63b825f4a2aab22e9f169ff65aae9/numpy/core/src/multiarray/item_selection.c，第 1664 行“NPY_BEGIN_THREADS_DEF”）。

因此，你可以做到

for j in prange(0,N, nogil=True):
    with gil:
      C[j,:,:]  = np.searchsorted(A1[j,:], A2, side='left' )

临时要求 GIL 对 Python 对象做必要的工作（希望很快），然后它应该在 searchsorted 允许在运行中很大程度上并行。

为了更新，我对此进行了快速测试（A1.shape==(105,100)、A2.shape==(302,302)，数字是任意选择的）。对于 10 次重复，串行版本花费了 4.5 秒，并行版本花费了 1.4 秒（在 4 核 CPU 上测试运行）。你没有获得 4 倍的全速提升，但你已经接近了。

这被编译为 described in the documentation。我怀疑如果您没有看到加速，那么它可能是以下任何一种：1）您的数组足够小，以至于 function-call/numpy 检查类型和大小的开销占主导地位； 2) 您没有在启用 OpenMP 的情况下编译它；或 3) 您的编译器不支持 OpenMP。

Answer 2

您遇到了一些问题 22。您需要 GIL 来调用 numpy.searchsorted 但 GIL 会阻止任何类型的并行处理。最好的办法是编写自己的 nogil 版本的 searchsorted:

cdef mySearchSorted(double[:] array, double target) nogil:
    # binary search implementation

for j in prange(0,N, nogil=True):
    for k in range(A2.shape[0]):
        for L in range(A2.shape[1]):
            C[j, k, L]  = mySearchSorted(A1[j, :], A2[k, L])

numpy.searchsorted 也有大量开销，因此如果 N 很大，使用您自己的 searchsorted 只是为了减少开销是有意义的。

使用 cython 并行化 python 循环 numpy.searchsorted

Parallelize python loop numpy.searchsorted using cython

python

arrays

parallel-processing

numpy

cython