numba.prange表现差

Question

我正在尝试编写一个简单的示例来说明使用numba.prange的好处，以便我和同事们更好地理解。但是，我无法得到有效的加速。我编写了一个简单的一维扩散求解器，它基本上循环遍历一个长数组，将元素i+1、i和i-1组合起来，然后将结果写入第二个数组的元素i中。这应该是一个非常适合并行for循环的完美用例，类似于Fortran或C中的OpenMP。

以下是我的完整示例：

import numpy as np
from numba import jit, prange

@jit(nopython=True, parallel=True)
def diffusion(Nt):
    alpha = 0.49
    x = np.linspace(0, 1, 10000000)
    # Initial condition
    C = 1/(0.25*np.sqrt(2*np.pi)) * np.exp(-0.5*((x-0.5)/0.25)**2)
    # Temporary work array
    C_ = np.zeros_like(C)
    # Loop over time (normal for-loop)
    for j in range(Nt):
        # Loop over array elements (space, parallel for-loop)
        for i in prange(1, len(C)-1):
            C_[i] = C[i] + alpha*(C[i+1] - 2*C[i] + C[i-1])
        C[:] = C_
    return C

# Run once to just-in-time compile
C = diffusion(1)

# Check timing
%timeit C = diffusion(100)

当使用parallel=False运行时，需要约2秒钟，而使用parallel=True则需要约1.5秒钟。我在MacBook Pro上运行，该计算机有4个物理核心，Activity Monitor报告在并行化和非并行化的情况下分别占用100％和约700％的CPU使用率。

我本来希望能够获得接近4倍的速度提升。我做错了什么吗？

- Tor

---

运行您的代码，使用parallel=False，我得到了1.52秒的时间；使用parallel=True，我得到了524毫秒的时间；使用parallel=True并将其他范围更改为prange，我得到了326毫秒的时间。 - Nin17

有意思，谢谢。也许我应该重启一下我的电脑或者做些其他的尝试。一会回来。（顺便提一下，从扩散求解器的角度看，我不确定把第一个范围改为 prange 是否是个好主意，因为每次迭代都依赖于完整的上一次迭代） - Tor

重启后，我的Mac笔记本电脑仍然得到相同的结果，但是当我尝试在Linux桌面上运行时，我看到了更合理的加速。奇怪。但无论如何，还是谢谢，我想这意味着我的代码没有问题。 - Tor

@Nin17 并行化Nt循环只会导致由于竞争条件而产生错误的结果。因此，在实践中这不是一个好主意。 - Jérôme Richard

@JérômeRichard 我知道这更多是为了展示使用prange比range可以提高性能。 - Nin17

最简单的方法是避免使用临时数组。https://pastebin.com/Ps4bJqt0 尽管是单线程，但已经比您的并行版本更快。通过一些努力，这也可以并行化，但如果没有进一步优化缓存使用，由于内存带宽限制，它不会很好地扩展。 - max9111

Answer 1

这段话的意思是：性能扩展性不佳的原因在于台式机上所有内核共享的RAM已经饱和。事实上，您的代码是内存限制型的，与CPU（或GPU）的计算能力相比，现代机器的内存吞吐量相当有限。因此，在大多数台式机上，1或2个内核通常足以使RAM饱和（计算服务器需要更多的内核）。

在一台配备10个Intel Xeon处理器、40~43 GiB/s RAM的计算机上，该代码并行运行需要1.32秒，串行运行需要2.56秒。这意味着使用10个核心只能提高2倍速度。话虽如此，并行循环每个时间步骤读取一次完整的C数组，同时每个时间步骤也会读取+写入完整的C_数组（由于写分配缓存策略，x86处理器默认需要读取写入的内存）。而C[:] = C_执行了同样的操作。这意味着在仅1.32秒的并行运行中，读/写44.7 GiB或RAM，达到33.9 GiB/s的内存吞吐量，接近80%的RAM带宽（对于此用例非常好）。

为了加速这段代码，你需要减少从/向 RAM 读写的数据量，并尽可能多地在缓存中计算数据。首先要做的是使用双缓冲方法，利用两个视图避免非常昂贵的复制。另一个优化是尝试同时进行多个时间步骤的并行计算。理论上可以使用复杂的梯形瓦片策略实现，但在 Numba 中实现起来非常棘手。高性能的阶梯库应该可以为您完成这项工作。这样的优化不仅应该提高顺序执行的效率，还应该提高结果并行代码的可伸缩性。

- Jérôme Richard

---

为了实现双缓冲方法，最简单的解决方案是在每个时间步骤结束时执行 C_, C = C, C_（它交换了数组的引用，因此可以在常数时间内非常快速地完成）。在函数结束时，根据 Nt 是奇数还是偶数，您可以返回 C 或 C_。这个技巧可以使您的代码运行速度提高两倍。 - Jérôme Richard

要小心，C = C_ 绝对不会做你想的那样：它不会复制这个数组。在Python中，（几乎）一切都是引用。C = C_ 会导致 C 引用 C_ 所引用的对象（即数组）。因此，修改C将修改C_，因为两者都引用同一个数组。C[:] = C_ 会导致复制，因为C[:]是视图，而视图上的赋值操作有所不同：它会复制数组内容。使用C = C_会导致源和目标相同，从而导致结果错误。经验之谈是始终检查结果 ;)。 - Jérôme Richard

啊，是的，当然。现在想起来了，在那种情况下，我曾经遇到过很多次奇怪的结果让我感到困惑。 - Tor

@JérômeRichard 我有什么遗漏吗？用C_, C = C, C_替换C[:] = C_会得到不同的结果，而且两者都使用会使它变慢吗？ - Nin17

哈！我忘记了你的范围不是从0开始。因此，第一个和最后一个项目保持不变，可以来自“坏”数组。你需要将两个边框项目从C复制到C_中。这是一个很好的例子，说明为什么上述经验法则如此重要:D! - Jérôme Richard