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

2022-01-10 00:00:00 python sparse-matrix anaconda parallel-processing numba

问题描述

我正在使用 Python 的 Anaconda 发行版以及 Numba，并且我编写了以下 Python 函数，该函数乘以稀疏矩阵 A(存储在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是一个很大的scipy稀疏矩阵，

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

和 x 是一个 numpy 数组.这是调用上述函数的代码片段:

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 对 csrMult() 进行即时编译 功能.

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

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

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 执行矩阵向量乘法的速度仍然比 csrMult() 快 6 倍.我相信这是因为 MATLAB 在执行稀疏矩阵向量乘法时使用了多线程.

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 时如何并行化外部 for 循环?

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

Numba 曾经有一个 prange() 函数，这使得并行化变得简单，令人尴尬的并行 for-循环.不幸的是，Numba 不再具有 prange() [实际上，这是错误的，请参阅下面的编辑].那么现在并行化这个 for 循环的正确方法是什么，Numba 的 prange() 函数不见了?

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?

当 prange() 从 Numba 中移除时，Numba 的开发人员想到了哪些替代方案?

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

编辑 1:
我更新到 Numba 的最新版本，即 .35，prange() 又回来了！它不包含在我一直使用的版本 .33 中.
这是个好消息，但不幸的是，当我尝试使用 prange() 并行化我的 for 循环时收到一条错误消息.这是 Numba 文档中的一个并行 for 循环 示例(请参阅第 1.9.2 节显式并行循环")，下面是我的新代码:

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

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

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

AttributeError:在 nopython 处失败(转换为 parfors)'SetItem'对象没有属性get_targets"

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

<小时>

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

正确的方法是什么(使用 prange 或替代方法)并行化这个 Python for-loop?

Given
the above attempt to use prange crashes, my question stands:

What is the correct way ( using prange or an alternative method ) to parallelize this Python for-loop?

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

As noted below, it was trivial to parallelize a similar for loop in C++ and obtain an 8x speedup, having been run on 20-omp-threads. There must be a way to do it using Numba, since the for loop is embarrassingly parallel (and since sparse matrix-vector multiplication is a fundamental operation in scientific computing).

编辑 2:
这是我的 csrMult() 的 C++ 版本.在 C++ 版本中并行化 for() 循环使我的测试中的代码快了大约 8 倍.这向我表明，在使用 Numba 时，Python 版本应该可以实现类似的加速.

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

解决方案

Numba 已经更新，prange() 现在可以使用了！ (我在回答我自己的问题.)

Numba has been updated and prange() works now! (I'm answering my own question.)

本博文，日期为 2017 年 12 月 12 日.以下是博客的相关片段:

The improvements to Numba's parallel computing capabilities are discussed in this blog post, dated December 12, 2017. Here is a relevant snippet from the blog:

很久以前(超过 20 个版本！)，Numba 曾经支持编写名为 prange() 的并行循环的习惯用法.大一之后在 2014 年重构代码库，这个特性不得不被移除，但它一直是最常被请求的 Numba 功能之一从那之后.英特尔开发人员并行化阵列后表达，他们意识到带回 prange 将是公平的容易

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-循环:

Using Numba version 0.36.1, I can parallelize my embarrassingly parallel for-loop using the following simple code:

@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，我终于有了一个 与 MATLAB 一样快的 Python 稀疏矩阵向量乘法例程.

In my experiments, parallelizing the for-loop made the function execute about eight times faster than the version I posted at the beginning of my question, which was already using Numba, but which was not parallelized. Moreover, in my experiments the parallelized version is about 5x faster than the command Ax = A.dot(x) which uses scipy's sparse matrix-vector multiplication function. Numba has crushed scipy and I finally have a python sparse matrix-vector multiplication routine that is as fast as MATLAB.

相关文章