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

Posted 2023-02-16

【中文标题】使用 Numba 时如何并行化此 Python for 循环

【英文标题】：How to parallelize this Python for loop when using Numba

【发布时间】：2018-04-06 01:14:37

【问题描述】：

我正在使用 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 数组。这是调用上述函数的一段sn-p代码：

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() 方法的两倍。这对 Numba 来说是一个非常令人印象深刻的结果。

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

---

问题：

使用 Numba 时如何并行化外部 for-loop？

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

当prange() 从 Numba 中移除时，Numba 的开发人员想到了什么替代方案？

---

编辑 1： 我更新到 Numba 的最新版本，即 0.35，prange() 又回来了！它不包含在我一直使用的版本 .33 中。 这是个好消息，但不幸的是，当我尝试使用 prange() 并行化我的 for 循环时收到一条错误消息。这是 Numba 文档中的并行 for 循环 example（请参阅第 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 

当我使用上面给出的代码 sn-p 调用此函数时，我收到以下错误：

AttributeError：在 nopython 处失败（转换为 parfors）'SetItem' 对象没有属性“get_targets”

---

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

正确的方法是什么（使用prange 或替代方法）并行化此 Python for-loop？

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

---

编辑 2： 这是我的 C++ 版本的csrMult()。在 C++ 版本中并行化 for() 循环使我的测试中的代码快了大约 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;
        
    

【问题讨论】：

您是否尝试过 parallel=True 装饰器的 parallel=True 关键字参数？我的意思是用@jit(parallel=True)注释它？

@fxx 我刚刚尝试用@jit(parallel=True) 替换@jit，当我运行我的测试代码sn-p 时，我收到以下错误消息： KeyError: " 不支持选项：'parallel'"

是的，这是一项实验性功能（取决于您的 numba 版本可能尚不可用）。好的，删除该选项后，我接下来尝试将实现移植到@vectorize 或@guvectorize（以生成ufunc）。也许您甚至必须为此将内部循环展开到另一个函数中。

@littleO 让我们在问题表述中更加量化。 A 矩阵（rows, cols, dtype）+ a（sparse/dense）占用率有多大和有多稀疏？注意：试图将 MATLAB 代码执行与 Py3/Numba 生态系统工具进行比较可能会产生很大的误导。

@user3666197 我用一些重要的新信息更新了这个问题。 A 有 56,469 行和 39,279 列以及 142,258,302 个非零条目（因此大约 6.4% 的条目是非零的）。 type(A[0,0]) 的输出是 numpy.float32。我在 C++ 中编写了一个非常相似的 csrMult 函数，其中并行化 for 循环很简单（因为 C++ 本身支持 openMP），并且我的函数快了大约 6 或 7 倍。我希望在使用 Numba 时通过并行化 Python 中的 for 循环来实现类似的加速。

【参考方案1】：

Numba 已更新，prange() 现在可以使用了！ _{（我正在回答我自己的问题。）}

2017 年 12 月 12 日的 blog post 讨论了 Numba 并行计算能力的改进。以下是博客中的相关 sn-p：

很久以前（超过 20 个版本！），Numba 曾经支持 编写并行 for 循环的习惯用法称为 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，但未并行化。此外，在我的实验中，并行版本比使用 scipy 的稀疏矩阵向量乘法函数的命令 Ax = A.dot(x) 快大约 5 倍。 Numba 已经碾压了 scipy，我终于有了一个 与 MATLAB 一样快的 Python 稀疏矩阵向量乘法例程。

【讨论】：

一条很酷的新闻。如果这在 Intel、AMD、ARM 等架构上通用，那么代码重新设计确实是一个绝妙的举措。如果诀窍只是使用新的可能性，来自基于硬件的扩展寄存器和矢量化操作指令，而不存在于其他处理器架构上，那么 ARM 和可能还有 AMD 端口将无法享受您所享受的性能观察.无论如何，请享受可用于进一步扩展您有价值的研究的新功能。

感谢您向我指出这一点！我已将链接转发给 Numba 团队以鼓励他们。

@MichaelGrant 如果您不介意，我有一个问题要问您。你知道 Numba 在使用prange() 并行化for-loop 时是否提供了一种指定“块大小”的方法？

仔细想想，A * x 在 MATLAB 中会比A' * x 慢是有道理的。使用 CSC 存储，A' * x，更容易并行化，因为每一行都有自己的线程。

@GeoffreyNegiar 我很犹豫是否接受我自己的答案并撤销对不同答案的接受，但你是对的。我只是把这个作为公认的答案。

【参考方案2】：

_{感谢您的 quant 更新，Daniel。以下几行可能难以理解，但请相信我，还有更多事情需要考虑。我曾在 hpc / parallel-processing / parallelism-amdahl problems在音阶中有矩阵 ~ N [TB]; N > 10 及其稀疏的伴奏工作过，所以一些经验可能有用供您进一步的意见。}

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

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

经济是头号问题。最初由 Gene Amdahl 制定的阿姆达尔定律没有考虑到 [PAR]-processes-setups + [PAR]-processes-finalisations 和终止的成本，这些成本确实必须在每个现实世界中支付实施。

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（这样做的成本是可以接受的，因为它非常，确实非常容易支付远远超过一个人可能从中获得的收益 - 处理性能下降导致的天真失望是故事中更容易的部分)。

如果愿意更好地理解这个主题并预先计算实际的“最小值

，请随意阅读更多关于开销严格的阿姆达尔定律重新制定的帖子em>"-subProblem-"size"，

sum-of-[PAR]-overheads 至少可以从现实世界的工具中得到合理的

用于将 subProblem 的并行拆分引入

N_trully_[PAR]_processes

（不是任何“just”-[CONCURRENT]，而是 true-

[PARALLEL]

- 这些都不是相等）。

---

Python 可能会获得一定剂量的类固醇以提高性能：

Python 是一个很棒的原型生态系统，而 numba、numpy 和其他编译扩展对提高性能的帮助远大于原生, GIL-stepped python (co-)-processing 通常会提供。

在这里，您尝试强制 numba.jit() 安排工作几乎免费，只需通过其自动化的jit()-time 词法分析器（即你把你的代码扔在上面），这应该既“理解”你的全球目标（做什么要做什么），也应该提出一些矢量化技巧（How best组装一堆CPU 指令，以实现此类代码执行的最大效率）。

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

Travis Oliphant 的团队在 numba 工具上取得了巨大的进步，但让我们现实和公平地不要期望任何形式的自动化向导会在 .jit()-lexer + 代码中实现-分析，当尝试转换代码并组装更有效的机器指令流以实现高级任务的目标时。

@guvectorize?这里？认真的吗？

由于[PSPACE] 的大小调整，您可能会立即忘记要求numba 以某种方式有效地向 GPU 引擎“填充”数据，其内存占用量远远落后于 GPU-GDDR 大小调整（不在所有关于这种数学上的“浅”GPU内核大小-“微小”处理只是相乘，可能在[PAR]中，但后来在[SEQ]中求和）。

（重新）-用数据加载 GPU 需要大量时间。如果已经支付了，GPU 内存延迟对于“微小”GPU 内核经济也不是很友好——你的 GPU-SMX 代码执行将必须支付 ~350-700 [ns] 只是为了获取一个number （很可能不会自动重新对齐，以便在接下来的步骤中实现最佳合并 SM 缓存友好的重用，您可能会注意到，您永远不会，让我重复一遍，永远不要重新使用单个矩阵单元完全如此，因此缓存本身不会在每个矩阵单元的 350~700 [ns] 下提供任何东西），而智能纯 numpy-vectorised 代码可以在少于 1 [ns] 每个单元上处理矩阵向量乘积即使是最大的[PSPACE]-footprints。

这是一个比较的标准。

（Profiling 最好在此处显示事实，但原理是事先众所周知的，无需测试如何将一些 TB 的数据移动到 GPU-fabric 上只是为了自己实现这一点。）

---

最糟糕的坏消息：

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

---

可解决问题清单：

总是更好地预先分配向量Ax = np.zeros_like( A[:,0] )并将它作为另一个参数传递到代码的numba.jit()-编译部分，以避免重复支付额外的[PTIME,PSPACE]-创建（再次）新内存的成本-allocations（如果向量被怀疑在外部协调的迭代优化过程中使用，则更多） 总是更好地指定（以缩小通用性，以提高代码性能）至少 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...(...):
      ...

【讨论】：

我不认为指定签名是一个好的建议，它会阻止基于数据连续性的优化（有时会导致性能明显下降）。另外我不确定你为什么在这里提到 GPU。问题中没有提到 GPU。

但我喜欢关于并行处理成本的部分，尤其是经常被忽略的部分，即“支付远远超过可能获得的收益非常非常容易”！

Ad GPU) 实际上在上面的 cmets 中提到了尝试 numba @guvectorize 工具，所以我添加了一些关于隐藏的极端成本的评论（也确实非常经常误用 ) GPU-latency-masking-SMX 玩具来解决这类问题。 GPU 可以帮助“数学上”的大型 GPU 内核在非常紧凑和小型数据区域上运行 + 具有最小、最好没有 SIMT 同步，但不能用于其他任何事情。不惜一切代价进行并行化现在如此频繁。 “Ó Tempóra, ó Mórés ...” :o)

感谢您的详细回答。要记住的一件事是，我在 C++ 中编写了一个非常相似的 csrMult 函数，其中并行化 for 循环很简单（因为 C++ 本身支持 openMP），并且通过并行化 for 循环，我观察到了 6 倍或 7 倍的加速，使用同一个矩阵。我希望这里有类似的加速。无论如何，我认为至少应该可以使用prange() 并行化我的 for 循环而不会导致代码崩溃。在C++中，我只需要在for循环上面写#pragma omp parallel for就可以让循环并行执行。

如果我没看错的话，似乎有一个错误的假设，即 guvectorize 装饰器意味着 GPU 计算，但这是不正确的。事实上，我一直在 CPU 目标上使用这样的结构。