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

Question

我正在使用Anaconda分发的Python和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-loop 并行化的正确方法prange()是什么呢,Numba的功能已经消失了？

当prange()从Numba移除时,Numba的开发人员有什么选择？

---

编辑1:
我更新到Numba的最新版本,即.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 -loop 的正确方法(使用prange或替代方法)是for什么？

如下所述,在C++中并行化类似的for循环并获得8x加速,在20 -omp-threads 上运行是微不足道的.必须有一种方法可以使用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过去常常支持一个习惯用来为被称为循环的并行写入prange().在2014年对代码库进行了重大重构之后,必须删除此功能,但这是自那时以来最常请求的Numba功能之一.在英特尔开发人员对数组表达式进行并行化之后,他们意识到回归prange将非常容易

使用Numba版本0.36.1,我可以for使用以下简单代码并行化我令人尴尬的并行-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一样快.

Answer 2

_{感谢您的量化更新，丹尼尔。
以下几行可能难以理解，但请相信我，还有更多需要考虑的事情。我已经研究了矩阵规模较小的HPC /并行计算问题，N [TB]; N > 10以及它们的稀疏伴随，因此一些经验可能对您的进一步观点很有用。}

警告：不要指望免费提供任何晚餐

希望并行处理一段代码听起来像是越来越现代的重新表达的法术力。问题不在于代码，而在于此举的成本。

经济是头号问题。由基因阿姆达尔（Gene Amdahl）最初制定的阿姆达尔定律（Amdahl's Law）并没有考虑[PAR]-processes-setups + [PAR]-processes-finalisations和终止的非常昂贵的成本，而这些成本实际上必须在每个实际实现中都支付。

该开销严格阿姆达尔定律描述的这些未避免不良反应的规模，并有助于了解那些被评估人之前选择只推出的并行计算的几个新的方面（在这样的一个可接受的成本，因为它是非常，确实很轻松支付不止一个可以从中获得的收益不菲-在这种情况下，较简单的做法是降低处理性能而天真的失望）。

如果愿意更好地理解此主题并预先计算实际的“ 最小 “ -subProblem-” 大小 ”（开销总和[PAR]将为之），请随时阅读有关开销限制的阿姆达尔定律重新制定的更多帖子。从用于将subProblem的并行拆分引入到现实世界的工具 中至少可以证明是合理的N_trully_[PAR]_processes（不是任何“ just”- [CONCURRENT]，但是true-- [PARALLEL]这些方式并不相等）。

Python可能会增加剂量的类固醇以提高性能：

Python是一种伟大的原型的生态系统，而numba，numpy和其他已编译的扩展有很大的帮助，以提高性能的方式远高于本地人，GIL阶梯蟒蛇（CO - ） -处理通常提供。

在这里，你尝试强制numba.jit()安排工作几乎仅仅通过它的自动-用于自由，jit()-时间词法分析器（你要把你的代码），它应该既“懂”你的全球目标（什么做的），和还提出了一些矢量化技巧（如何最好地组装CPU指令堆以最大程度地提高这种代码执行的效率）。

这听起来很容易，但事实并非如此。

Travis Oliphant的团队在工具方面取得了巨大的进步numba，但是现实而公平的是.jit()，当试图转换代码并组装更高效的流程时，不要期望在-lexer +代码分析中实现任何形式的自动wizardry机器指令以实现高级任务的目标。

@guvectorize？这里？认真吗

由于[PSPACE]大小问题，您可能会立即忘记以numba某种方式有效地用数据“填充” GPU引擎，其内存占用量远远落后于GPU-GDDR大小调整（根本谈不上“太”浅的GPU-此类数学上的“微小”处理的内核大小只是相乘，可能乘以，[PAR]而后累加[SEQ]。

（重新）向GPU加载数据需要花费大量时间。如果付出了这些，则In-GPU内存延迟对于“微型” GPU内核经济也不是很友好-您的GPU-SMX代码执行将不得不付出〜350-700 [ns]只是获取一个数字（很可能不是自动获取）重新调整为最佳组合的SM-cache-友好的下一步使用，您可能会注意到，您永远不会，让我重复一遍，永远不要重复使用单个矩阵单元，因此缓存本身不会在350~700 [ns]每个矩阵单元下提供任何东西），而智能的纯numpy矢量化代码可以处理矩阵向量乘积，1 [ns]即使在最大[PSPACE]占用空间上也比每个单元更少。

这是一个可以比较的准绳。

（剖析可以更好地展示这些事实，但是该原理是众所周知的，无需测试如何将一些TB数据移至GPU架构上就可以自己实现。）

最坏的消息：

给定矩阵的内存规模A，更糟糕的效果是，矩阵表示存储的稀疏组织很可能会破坏（即使不是全部）通过numba密集矩阵表示中的矢量化技巧可获得的大部分（如果不是全部）可能的性能提升。，因为高效内存获取缓存行重用的可能性几乎为零，稀疏性也将破坏实现矢量化操作的紧凑映射的任何简单方法，并且这些将很难保持容易地转换为高级CPU硬件的状态。向量处理资源。

可解决问题的清单：

• 总是最好预先分配向量Ax = np.zeros_like( A[:,0] )，并将其作为另一个参数传递到numba.jit()代码的-编译部分，以避免重复[PTIME,PSPACE]花费额外的成本来创建（再次）新的内存分配（如果怀疑向量是向量，则更多）在外部精心设计的迭代优化过程中使用）
• 总是最好
  至少指定numba.jit( "f8[:]( f4[:], f4[:,:], ... )" )-calling接口指令（以缩小通用性，以提高代码性能）
• 始终查看所有numba.jit()可用的-options及其各自的默认值_{（可能会因版本而异）（}针对特定情况）（禁用GIL并使用numba+硬件功能更好地调整目标将始终有助于代码中数字密集的部分）

@jit(   signature = [    numba.float32( numba.float32, numba.int32 ),                                   #          # [_v41] @decorator with a list of calling-signatures
                         numba.float64( numba.float64, numba.int64 )                                    #
                         ],    #__________________ a list of signatures for prepared alternative code-paths, to avoid a deferred lazy-compilation if undefined
        nopython = False,      #__________________ forces the function to be compiled in nopython mode. If not possible, compilation will raise an error.
        nogil    = False,      #__________________ tries to release the global interpreter lock inside the compiled function. The GIL will only be released if Numba can compile the function in nopython mode, otherwise a compilation warning will be printed.
        cache    = False,      #__________________ enables a file-based cache to shorten compilation times when the function was already compiled in a previous invocation. The cache is maintained in the __pycache__ subdirectory of the directory containing the source file.
        forceobj = False,      #__________________ forces the function to be compiled in object mode. Since object mode is slower than nopython mode, this is mostly useful for testing purposes.
        locals   = {}          #__________________ a mapping of local variable names to Numba Types.
        ) #____________________# [_v41] ZERO <____ TEST *ALL* CALLED sub-func()-s to @.jit() too >>>>>>>>>>>>>>>>>>>>> [DONE]
 def r...(...):
      ...