numba.prange 表现不佳

Question

我正在尝试组合一个简单的示例来说明使用 numba.prange 对我自己和一些同事的好处，但我无法获得像样的加速。我编写了一个简单的 1D 扩散求解器，它本质上是在一个长数组上循环，组合元素 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秒。我运行使用具有 4 个物理内核的 MacBook Pro，并且 Activity 监视器报告 100% 和大约 700% CPU 使用率有和没有并行化。

我本以为会有接近 4 倍的加速。我做错了什么吗？

Answer 1

糟糕的可扩展性肯定是由于桌面计算机上所有内核共享的 RAM 饱和造成的。实际上，您的代码是 memory-bound，与 CPU（或 GPU）的计算能力相比，现代机器的内存吞吐量非常有限。因此，1 或 2 个内核通常足以使大多数台式机上的 RAM 饱和（计算服务器上需要更多内核）。

在具有 40~43 GiB/s RAM 的 10 核英特尔至强处理器上，代码并行需要 1.32 秒，顺序需要 2.56 秒。这意味着 10 个内核的速度只有 2 倍。也就是说，并行循环每个时间步读取一次完整的 C 数组，并且每个时间步读取+写入一次完整的 C_ 数组（x86 处理器需要默认读取写入的内存，因为写入分配缓存策略）。 C[:] = C_ 做同样的事情。这意味着 (2*3)*(8*10e6)*100/1024**3 = 44.7 GiB 或 RAM 在仅 1.32 秒的并行时间内为 read/written，导致 33.9 GiB/s 内存吞吐量达到 RAM 带宽的 80%（非常适合此 use-case).

要加快此代码的速度，您需要 read/write 更少的数据 from/to RAM 并在缓存中尽可能多地计算数据。要做的第一件事是使用具有两个视图的 double-buffering 方法，以避免非常昂贵的副本。另一个优化是尝试同时并行执行多个时间步长。这在理论上可以使用复杂的 trapezoidal tiling strategy 但在实践中实施起来非常棘手，尤其是在 Numba 中。 High-performance 模板库应该可以为您做到这一点。这种优化不仅应该改进顺序执行，而且还应该改进生成的并行代码的可扩展性。

numba.prange 表现不佳

Poor performance from numba.prange

python

numba