我正在尝试整理一个简单的示例来说明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 上运行,活动监视器报告 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] 删除。
我来说两句