numba.prange 性能不佳

我试图整理一个简单的例子来说明使用的好处numba.prange对于我自己和一些同事来说，但我无法获得像样的加速。我编写了一个简单的一维扩散求解器，它本质上是在一个长数组上循环，组合元素i+1, i, and i-1，并将结果写入elementi第二个数组的。这应该是并行 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 上运行，活动监视器报告无论是否有并行化，CPU 使用率均为 100% 和大约 700%。

我原本预计加速速度会接近 4 倍。难道我做错了什么？

较差的可扩展性肯定是由于台式机上所有内核共享的 RAM 饱和所致。的确，你的代码是受内存限制的与 CPU（或 GPU）的计算能力相比，现代机器的内存吞吐量非常有限。因此，1 或 2 个核心通常足以使大多数台式机上的 RAM 饱和（计算服务器需要更多的核心）。

在具有 40~43 GiB/s RAM 的 10 核 Intel Xeon 处理器上，代码并行耗时 1.32 秒，顺序耗时 2.56 秒。这意味着 10 个核心的速度仅提高 2 倍。话虽这么说，并行循环读取了完整的C每个时间步一次数组，并且还读+写完整的C_每个时间步一次数组（x86 处理器默认需要读取写入的内存，因为写入分配缓存策略). The C[:] = C_做同样的事情。这意味着(2*3)*(8*10e6)*100/1024**3 = 44.7GiB 或 RAM 的并行读/写时间仅为 1.32 秒，从而实现 33.9 GiB/s 的内存吞吐量，达到 RAM 带宽的 80%（非常适合此用例）。

为了加速此代码，您需要从 RAM 读取/向 RAM 写入更少的数据，并在缓存中尽可能多地计算数据。首先要做的是使用双缓冲使用两个视图的方法，以避免非常昂贵的副本。另一个优化是尝试同时并行执行多个时间步。理论上这是可能的，使用复杂的梯形平铺策略 https://crd.lbl.gov/assets/pubs_presos/sirev09-stencil.pdf但在实践中实施起来非常棘手，尤其是在 Numba 中。高性能模板库应该可以为您做到这一点。这样的优化不仅应该提高顺序执行和可扩展性生成的并行代码。