問(wèn)題描述

我正在使用 Python 的 Anaconda 發(fā)行版以及 Numba，并且我編寫(xiě)了以下 Python 函數(shù)，該函數(shù)乘以稀疏矩陣 A(存儲(chǔ)在CSR 格式)由密集向量 x:

I'm using the Anaconda distribution of Python, together with Numba, and I've written the following Python function that multiplies a sparse matrix A (stored in a CSR format) by a dense vector 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是一個(gè)很大的scipy稀疏矩陣，

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

和 x 是一個(gè) numpy 數(shù)組.這是調(diào)用上述函數(shù)的代碼片段:

and x is a numpy array. Here is a snippet of code that calls the above function:

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

注意 @jit 裝飾器，它告訴 Numba 對(duì) csrMult() 進(jìn)行即時(shí)編譯 功能.

Notice the @jit-decorator that tells Numba to do a just-in-time compilation for the csrMult() function.

在我的實(shí)驗(yàn)中，我的函數(shù) csrMult() 大約是 scipy .dot() 方法.這對(duì) Numba 來(lái)說(shuō)是一個(gè)非常令人印象深刻的結(jié)果.

In my experiments, my function csrMult() is about twice as fast as the scipy .dot() method. That is a pretty impressive result for Numba.

但是，MATLAB 執(zhí)行矩陣向量乘法的速度仍然比 csrMult() 快 6 倍.我相信這是因?yàn)?MATLAB 在執(zhí)行稀疏矩陣向量乘法時(shí)使用了多線程.

However, MATLAB still performs this matrix-vector multiplication about 6 times faster than csrMult(). I believe that is because MATLAB uses multithreading when performing sparse matrix-vector multiplication.

使用 Numba 時(shí)如何并行化外部 for 循環(huán)?

How can I parallelize the outer for-loop when using Numba?

Numba 曾經(jīng)有一個(gè) prange() 函數(shù)，這使得并行化變得簡(jiǎn)單，令人尷尬的并行 for-循環(huán).不幸的是，Numba 不再具有 prange() [實(shí)際上，這是錯(cuò)誤的，請(qǐng)參閱下面的編輯].那么現(xiàn)在并行化這個(gè) for 循環(huán)的正確方法是什么，Numba 的 prange() 函數(shù)不見(jiàn)了?

Numba used to have a prange() function, that made it simple to parallelize embarassingly parallel for-loops. Unfortunately, Numba no longer has prange() [actually, that is false, see the edit below]. So what is the correct way to parallelize this for-loop now, that Numba's prange() function is gone?

當(dāng) prange() 從 Numba 中移除時(shí)，Numba 的開(kāi)發(fā)人員想到了哪些替代方案?

When prange() was removed from Numba, what alternative did the developers of Numba have in mind?

編輯 1:
我更新到 Numba 的最新版本，即 .35，prange() 又回來(lái)了！它不包含在我一直使用的版本 .33 中.
這是個(gè)好消息，但不幸的是，當(dāng)我嘗試使用 prange() 并行化我的 for 循環(huán)時(shí)收到一條錯(cuò)誤消息.這是 Numba 文檔中的一個(gè)并行 for 循環(huán) 示例(請(qǐng)參閱第 1.9.2 節(jié)顯式并行循環(huán)")，下面是我的新代碼:

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 

當(dāng)我使用上面給出的代碼片段調(diào)用此函數(shù)時(shí)，我收到以下錯(cuò)誤:

When I call this function, using the code snippet given above, I receive the following error:

AttributeError:在 nopython 處失敗(轉(zhuǎn)換為 parfors)'SetItem'對(duì)象沒(méi)有屬性get_targets"

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

<小時(shí)>

鑒于
上述使用 prange 的嘗試崩潰，我的問(wèn)題是:

正確的方法是什么(使用 prange 或替代方法)并行化這個(gè) Python for-loop?

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;
        }
    }
}

@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