numba.prange 表现不佳

托尔

我正在尝试整理一个简单的示例来说明numba.prange为自己和一些同事使用的好处,但我无法获得不错的加速。我编写了一个简单的 1D 扩散求解器,它基本上在一个长数组上循环,组合元素i+1ii-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 上运行,活动监视器报告 100% 和大约 700% 的 CPU 使用率,有和没有并行化。

我本来预计会有更接近 4 倍的加速。难道我做错了什么?

杰罗姆·理查德

较差的可扩展性当然来自于台式机上所有内核共享的 RAM 的饱和。事实上,您的代码是受内存限制的,与 CPU(或 GPU)的计算能力相比,现代机器的内存吞吐量非常有限。因此,1 或 2 个内核通常足以使大多数台式机上的 RAM 饱和(计算服务器上需要更多内核)。

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

为了加快此代码的速度,您需要从 RAM 读取/写入更少的数据,并在缓存中尽可能多地计算数据。首先要做的是使用具有两个视图的双缓冲方法,以避免非常昂贵的副本。另一个优化是尝试同时并行执行多个时间步。这在理论上使用复杂的梯形平铺策略是可行的,但在实践中实施起来非常棘手,尤其是在 Numba 中。高性能模板库应该为您做到这一点。这种优化不仅应该提高顺序执行,而且应该提高生成的并行代码的可伸缩性。

本文收集自互联网,转载请注明来源。

如有侵权,请联系 [email protected] 删除。

编辑于
0

我来说两句

0 条评论
登录 后参与评论

相关文章