python - 使用 Numba 时如何并行化此 Python for 循环

Question

我正在使用 Python 的 Anaconda 发行版和 Numba，我编写了以下 Python 函数来乘以稀疏矩阵 A(存储在CSR 格式)通过密集向量 x:

@jit
def csrMult( x, Adata, Aindices, Aindptr, Ashape ):

    numRowsA = Ashape[0]
    Ax       = numpy.zeros( numRowsA )

    for i in range( numRowsA ):
        Ax_i = 0.0
        for dataIdx in range( Aindptr[i], Aindptr[i+1] ):

            j     = Aindices[dataIdx]
            Ax_i +=    Adata[dataIdx] * x[j]

        Ax[i] = Ax_i

    return Ax 

这里A是一个很大的scipy稀疏矩阵，

>>> A.shape
( 56469, 39279 )
#                  having ~ 142,258,302 nonzero entries (so about 6.4% )
>>> type( A[0,0] )
dtype( 'float32' )

和 x 是一个 numpy 数组。下面是调用上述函数的代码片段:

x       = numpy.random.randn( A.shape[1] )
Ax      = A.dot( x )   
AxCheck = csrMult( x, A.data, A.indices, A.indptr, A.shape )

注意 @jit 装饰器，它告诉 Numba 为 csrMult() 进行即时编译 功能。

在我的实验中，我的函数 csrMult() 大约是 两倍 scipy .dot( ) 方法。对于 Numba 来说，这是一个令人印象深刻的结果。

但是，MATLAB 执行此矩阵向量乘法的速度仍然比 csrMult() 快 6 倍。我认为这是因为 MATLAB 在执行稀疏矩阵向量乘法时使用了多线程。

---

问题:

如何在使用 Numba 时并行化外部 for 循环？

Numba 曾经有一个 prange() 函数，这使得并行化 embarrassingly 并行化变得简单 for -循环。不幸的是，Numba 不再有 prange() [实际上，这是错误的，请参阅下面的编辑]。 那么现在并行化这个 for 循环的正确方法是什么，Numba 的 prange() 函数已经消失了？

当从 Numba 中删除 prange() 时，Numba 的开发人员想到了什么替代方案？

---

Edit 1:
I updated to the latest version of Numba, which is .35, and prange() is back! It was not included in version .33, the version I had been using.
That is good news, but unfortunately I am getting an error message when I attempt to parallelize my for loop using prange(). Here is a parallel for loop example from the Numba documentation (see section 1.9.2 "Explicit Parallel Loops"), and below is my new code:

from numba import njit, prange
@njit( parallel=True )
def csrMult_numba( x, Adata, Aindices, Aindptr, Ashape):

    numRowsA = Ashape[0]    
    Ax       = np.zeros( numRowsA )

    for i in prange( numRowsA ):
        Ax_i = 0.0        
        for dataIdx in range( Aindptr[i],Aindptr[i+1] ):

            j     = Aindices[dataIdx]
            Ax_i +=    Adata[dataIdx] * x[j]

        Ax[i] = Ax_i            

    return Ax 

当我使用上面给出的代码片段调用此函数时，我收到以下错误:

AttributeError: Failed at nopython (convert to parfors) 'SetItem' object has no attribute 'get_targets'

---

鉴于
上述尝试使用prange 崩溃，我的问题是:

什么是正确的方法(使用 prange 或替代方法)并行化此 Python for 循环？

如下所述，在 20-omp-threads 上运行后，在 C++ 中并行化类似的 for 循环并获得 8x 加速是微不足道的。必须有一种方法可以使用 Numba 来完成它，因为 for 循环是令人尴尬的并行(并且因为稀疏矩阵向量乘法是科学计算中的基本操作)。

---

Edit 2:
Here is my C++ version of csrMult(). Parallelizing the for() loop in the C++ version makes the code about 8x faster in my tests. This suggests to me that a similar speedup should be possible for the Python version when using Numba.

void csrMult(VectorXd& Ax, VectorXd& x, vector<double>& Adata, vector<int>& Aindices, vector<int>& Aindptr)
{
    // This code assumes that the size of Ax is numRowsA.
    #pragma omp parallel num_threads(20)
    {       
        #pragma omp for schedule(dynamic,590) 
        for (int i = 0; i < Ax.size(); i++)
        {
            double Ax_i = 0.0;
            for (int dataIdx = Aindptr[i]; dataIdx < Aindptr[i + 1]; dataIdx++)
            {
                Ax_i += Adata[dataIdx] * x[Aindices[dataIdx]];
            }

            Ax[i] = Ax_i;
        }
    }
}

Answer 1

Numba 已更新，prange() 现在可以使用了! _{(我正在回答我自己的问题。)}

本 blog post 中讨论了对 Numba 并行计算能力的改进。 ，日期为 2017 年 12 月 12 日。以下是博客中的相关片段:

Long ago (more than 20 releases!), Numba used to have support for an idiom to write parallel for loops called prange(). After a major refactoring of the code base in 2014, this feature had to be removed, but it has been one of the most frequently requested Numba features since that time. After the Intel developers parallelized array expressions, they realized that bringing back prange would be fairly easy

使用 Numba 版本 0.36.1，我可以使用以下简单代码并行化我的令人尴尬的并行 for 循环:

@numba.jit(nopython=True, parallel=True)
def csrMult_parallel(x,Adata,Aindices,Aindptr,Ashape): 

    numRowsA = Ashape[0]    
    Ax = np.zeros(numRowsA)

    for i in numba.prange(numRowsA):
        Ax_i = 0.0        
        for dataIdx in range(Aindptr[i],Aindptr[i+1]):

            j = Aindices[dataIdx]
            Ax_i += Adata[dataIdx]*x[j]

        Ax[i] = Ax_i            

    return Ax

在我的实验中，并行化 for 循环使函数执行速度比我在问题开头发布的版本快八倍，该版本已经在使用 Numba，但未并行化.此外，在我的实验中，并行版本比使用 scipy 的稀疏矩阵向量乘法函数的命令 Ax = A.dot(x) 快大约 5 倍。 Numba 击败了 scipy，我终于有了一个 python 稀疏矩阵向量乘法例程，它和 MATLAB 一样快。