使用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 -decorator告诉Numba对 csrMult() 函数进行即时编译。

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

然而，MATLAB仍然执行有关此矩阵-向量乘法 6倍的速度
比csrMult()。我相信这是因为MATLAB在执行稀疏矩阵向量乘法时会使用多线程。

题：

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

Numba曾经有一个 prange() 函数，可以轻松地使尴尬的并行 for
循环并行化。不幸的是，Numba不再具有prange()[ 实际上，这是错误的，请参见下面的编辑 ]。
那么，for现在使Numbaprange()函数消失的并行此循环的正确方法是什么？

当prange()从Numba被删除，有什么其他没Numba的开发人员有什么想法？

编辑1：
我更新到Numba的最新版本，即0.35，prange()现在又回来了！它没有包含在我一直使用的.33版本中。

这是个好消息，但是不幸的是，当我尝试使用并行化我的for循环时，我收到一条错误消息prange()。这是Numba文档中的并行for循环示例（请参见第1.9.2节“显式并行循环”），以下是我的新代码：

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：在nopython失败（转换为parfors）’SetItem’对象没有属性’get_targets’

鉴于

以上使用prange崩溃的尝试，我的问题是：

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

如下所述，在 20个* -omp-threads上运行时，并行化C ++中类似的for循环并获得 8倍的
加速是微不足道的。必须使用Numba进行处理，因为for循环令人尴尬地是并行的（并且稀疏矩阵向量乘法是科学计算中的基本运算）。 *

编辑2：
这是我的C 版本 csrMult()。 for()在我的测试中，并行执行C
版本中的循环可使代码快8倍。这向我暗示，使用Numba时，Python版本应该可以实现类似的加速。

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() 现在可以工作！ （我正在回答自己的问题。）

日期为2017年12月12日的此博客文章讨论了Numba并行计算功能的改进。这是该博客的相关摘要：

很久以前（超过20个发行版！），Numba曾经支持惯用语来编写称为的并行for循环prange()。在2014年对代码库进行了重大重构之后，必须删除此功能，但自那以来，它一直是Numba最常请求的功能之一。英特尔开发人员将数组表达式并行化之后，他们意识到带回prange是很容易的

使用Numba版本0.36.1，我可以for使用以下简单代码并行化尴尬的parallel -loop：

@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-loop使该函数的执行速度比我在问题开头发布的版本快了八倍，该版本已经在使用Numba了，但尚未并行化。此外，在我的实验中，并行版本比Ax = A.dot(x)使用scipy的稀疏矩阵矢量乘法函数的命令快大约5倍。 Numba压碎了scipy
，我终于有了一个python稀疏的矩阵向量乘法例程，它的 速度与MATLAB一样快 。