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

How to parallelize this Python for loop when using Numba

我正在使用 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 装饰器告诉 Numba 为 csrMult()函数。

在我的实验中,我的函数 csrMult() 大约 scipy .dot()[= 的两倍 96=] 方法。对于 Numba 来说,这是一个令人印象深刻的结果。

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


问题:

如何在使用 Numba 时并行化外部 for 循环?

Numba 曾经有一个 prange() 函数,这使得并行化变得很简单 for -循环。不幸的是,Numba 不再有 prange() [实际上,这是错误的,请参阅下面的编辑]。 那么现在并行化此 for 循环的正确方法是什么,Numba 的 prange() 功能消失了?

prange() 从 Numba 中删除时,Numba 的开发人员想到了什么替代方案?


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 

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

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


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

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

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


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

感谢您的量化更新,Daniel。
以下几行可能难以下咽,但请相信我,还有更多事情需要考虑。我研究过 / / problems
在尺度中有矩阵~N [TB]; N > 10和它们的稀疏伴奏,所以一些经验可能对你有用更多意见。

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

将一段代码并行化的愿望听起来越来越像现代重新表达的法力。 问题是不是代码,而是这种移动的成本。

经济是头号问题。 Amdahl 定律最初由 Gene Amdahl 制定,没有考虑 [PAR]-processes-setups + [PAR]-processes-finalisations & terminations 的成本,这些成本确实必须在每个真实世界的实现。

The overhead-strict Amdahl's Law depicts the scale of these un-avoidable adverse effects and helps understand a few new aspects that have to be evaluated before one opts to introduce parallelisation(这样做的成本是可以接受的,因为付出比可能获得的收益更多的钱确实非常容易——处理性能下降带来的天真失望更容易故事的一部分)。

如果愿意更好地理解这个主题并预先计算实际,请随时阅读更多关于开销严格的阿姆达尔定律重新制定的帖子"minimum"-subProblem-"size"sum-of-[PAR]-overheads 至少会从现实世界的工具中获得合理的,用于将子问题的并行拆分引入 N_trully_[PAR]_processes(不是任何“只是” -[CONCURRENT],但是 true-[PARALLEL]——它们是不相等的)。


Python 可能会服用一剂类固醇以提高表现:

Python 是一个很棒的原型生态系统,而 numbanumpy和其他编译扩展相比原生的 GIL 步进 python(协同)处理通常提供的性能,对性能的提升有很大帮助。

在这里,你尝试强制numba.jit()安排工作几乎-免费,只是通过它的自动化jit()-time lexical-analyser(你把你的代码放在上面),它应该既“理解”你的全局目标(What to do),也应该提出一些矢量化-技巧(如何最好 assemble一堆CPU指令,以实现此类代码执行的最大效率)。

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

Travis Oliphant 的团队在 numba 工具方面取得了 巨大进步 ,但让我们现实一点,公平地说,不要期望任何形式的自动化巫术在内部实现.jit()-词法分析器 + 代码分析,在尝试转换代码和 assemble 更高效的机器指令流以实现高级任务目标时。

@guvectorize?这里?认真的吗?

由于 [PSPACE] 大小,您可能会立即忘记要求 numba 以某种方式有效地用数据“填充”GPU 引擎,其内存占用量远远落后于 GPU- GDDR 大小(根本不是说太“浅”的 GPU 内核大小,因为这种数学上的“微小”处理只是乘法,可能在 [PAR] 中,但稍后在 [SEQ] 中求和)。

(Re-)-用数据加载 GPU 需要大量时间。如果支付了费用,GPU 内内存延迟对于“微型”GPU 内核经济也不是很友好——您的 GPU-SMX 代码执行将 必须支付 ~ 350-700 [ns]只是为了获取一个数字(很可能不会自动重新对齐以在后续步骤中最好地合并 SM 缓存友好地重用,您可能会注意到,您永远不会,让我重复一遍,永远不会重新 -完全使用单个矩阵单元,因此缓存本身不会在每个矩阵单元 350~700 [ns] 下提供任何东西),而智能纯 numpy-矢量化代码可以处理矩阵-即使在最大的 [PSPACE]-footprints.

上,每个单元格的矢量积也少于 1 [ns]

这是一个比较标准。

( Profiling 会更好地在这里展示铁的事实,但原则是事先众所周知的,没有测试如何将一些 TB 数据移动到 GPU-fabric 上只是为了自己实现这一点。 )


最糟糕的坏消息:

给定矩阵的内存规模A,预期的最坏效果是,矩阵表示存储的稀疏组织将最可能会破坏大部分(如果不是全部的话)可能通过 numba-向量化技巧在密集矩阵表示上实现的性能提升,因为高效内存获取缓存行重用的机会几乎为零,稀疏性也会破坏任何实现矢量化操作的紧凑映射的简单方法,并且这些操作几乎无法轻松转换为高级CPU-硬件矢量处理资源。


可解决问题盘点:

  • 最好预先分配向量 Ax = np.zeros_like( A[:,0] ) 并将其作为另一个参数传递到代码的 numba.jit()- 编译部分,以避免重复支付额外的费用 [PTIME,PSPACE]-创建(再次)新内存分配的成本(如果怀疑向量在外部协调的迭代优化过程中使用,则成本更高)
  • 总是更好地指定(为了缩小通用性,为了最终的代码性能)
    至少 numba.jit( "f8[:]( f4[:], f4[:,:], ... )" ) 调用接口指令
  • 始终查看所有可用的 numba.jit() 选项及其各自的默认值 (可能会更改版本) 针对您的具体情况(禁用 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...(...):
      ...

Numba 已更新,prange() 现在可用! (我正在回答我自己的问题。)

Numba 并行计算能力的改进在 2017 年 12 月 12 日 blog post 中进行了讨论。以下是来自博客的相关片段:

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-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,但未并行化。此外,在我的实验中,并行版本比使用 scipy 的稀疏矩阵向量乘法函数的命令 Ax = A.dot(x) 快大约 5 倍。 Numba 击败了 scipy,我终于有了一个 python 稀疏矩阵向量乘法例程,它 和 MATLAB 一样快.