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

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.

这篇关于使用 Numba 时如何并行化此 Python for 循环的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持跟版网！