Raf*_*rad 2 performance time multithreading numpy numba
我试图找到使用 Numba 对两个相同大小的矩阵求和的最快方法。我想出了 3 种不同的方法，但没有一种能打败 Numpy。这是我的代码：
import numpy as np
from numba import njit,vectorize, prange,float64
import timeit
import time
# function 1:
def sum_numpy(A,B):
return A+B
# function 2:
sum_numba_simple= njit(cache=True,fastmath=True) (sum_numpy)
# function 3:
@vectorize([float64(float64, float64)])
def sum_numba_vectorized(A,B):
return A+B
# function 4:
@njit('(float64[:,:],float64[:,:])', cache=True, fastmath=True, parallel=True)
def sum_numba_loop(A,B):
n=A.shape[0]
m=A.shape[1]
C = np.empty((n, m), A.dtype)
for i in prange(n):
for j in prange(m):
C[i,j]=A[i,j]+B[i,j]
return C
#Test the functions with 2 matrices of size 1,000,000x3:
N=1000000
np.random.seed(123)
A=np.random.uniform(low=-10, high=10, size=(N,3))
B=np.random.uniform(low=-5, high=5, size=(N,3))
t1=min(timeit.repeat(stmt='sum_numpy(A,B)',timer=time.perf_counter,repeat=3, number=100,globals=globals()))
t2=min(timeit.repeat(stmt='sum_numba_simple(A,B)',timer=time.perf_counter,repeat=3, number=100,globals=globals()))
t3=min(timeit.repeat(stmt='sum_numba_vectorized(A,B)',timer=time.perf_counter,repeat=3, number=100,globals=globals()))
t4=min(timeit.repeat(stmt='sum_numba_loop(A,B)',timer=time.perf_counter,repeat=3, number=100,globals=globals()))
print("function 1 (sum_numpy): t1= ",t1,"\n")
print("function 2 (sum_numba_simple): t2= ",t2,"\n")
print("function 3 (sum_numba_vectorized): t3= ",t3,"\n")
print("function 4 (sum_numba_loop): t4= ",t4,"\n")
结果如下：
函数 1 (sum_numpy): t1= 0.1655790419999903
函数 2 (sum_numba_simple): t2= 0.3019776669998464
函数 3 (sum_numba_vectorized): t3= 0.16486266700030683
函数 4 (sum_numba_loop): t4= 0.1862256660001549
正如您所看到的，结果表明在这种情况下使用 Numba 没有任何优势。因此，我的问题是：
是否有其他实现可以提高求和速度？
您的代码受到页面错误的约束（有关此问题的更多信息，请参阅此处、此处和那里）。发生页错误是因为数组是新分配的。解决方案是预先分配它，然后在其中写入，这样就不会导致页面在物理内存中重新映射。np.add(A, B, out=C)按照评论中@August 的指示执行此操作。另一种解决方案可能是调整标准分配器，这样就不会以显着的内存使用开销为代价将内存返还给操作系统（例如，AFAIK TC-Malloc 可以做到这一点）。
大多数平台（尤其是 x86 平台）上还存在另一个问题：写回缓存的缓存行写入分配在写入过程中非常昂贵。避免这种情况的典型解决方案是进行非临时存储（如果在目标处理器上可用，则在 x86-64 处理器上是这种情况，但在其他处理器上可能不是）。话虽如此，Numpy 和 Numba 都还无法做到这一点。对于 Numba，我填写了一个涵盖简单用例的问题。编译器本身（Numpy 的 GCC 和 Numba 的 Clang）往往不会生成此类指令，因为当数组适合缓存并且编译器在编译时不知道数组的大小时，它们可能会损害性能（它们可以在以下情况下生成特定代码）：他们可以评估计算的数据量，但这并不容易，并且可能会减慢其他一些代码的速度）。AFAIK，解决此问题的唯一可能方法是编写 C 代码并使用低级指令或使用编译器指令。在您的情况下，由于这种影响，大约有 25% 的带宽丢失，导致速度下降高达 33%。
使用多个线程并不总是能让内存绑定的代码更快。事实上，它通常几乎无法扩展，因为当 RAM 已经饱和时，使用更多核心并不能加快执行速度。在大多数平台上，通常需要很少的内核来使 RAM 饱和。页面错误可以受益于在目标系统上使用多个内核（Linux 在并行方面做得很好，Windows 通常不能很好地扩展，IDK 对于 MacOS）。
最后，还有另一个问题：代码没有矢量化（至少在我的机器上没有矢量化）。一种解决方案是展平数组视图并执行一个大循环，编译器可以更轻松地进行矢量化（基于 j 的循环太小，SIMD 指令无法有效）。还应指定输入数组的连续性，以便编译器生成快速的 SIMD 代码。下面是生成的 Numba 代码：
@njit('(float64[:,::1], float64[:,::1], float64[:,::1])', cache=True, fastmath=True, parallel=True)
def sum_numba_fast_loop(A, B, C):
n, m = A.shape
assert C.shape == A.shape
A_flat = A.reshape(n*m)
B_flat = B.reshape(n*m)
C_flat = C.reshape(n*m)
for i in prange(n*m):
C_flat[i]=A_flat[i]+B_flat[i]
return C
以下是我的 6 核 i5-9600KF 处理器（RAM 约为 42 GiB/s）上的结果：
sum_numpy: 0.642 s 13.9 GiB/s
sum_numba_simple: 0.851 s 10.5 GiB/s
sum_numba_vectorized: 0.639 s 14.0 GiB/s
sum_numba_loop serial: 0.759 s 11.8 GiB/s
sum_numba_loop parallel: 0.472 s 18.9 GiB/s
Numpy "np.add(A, B, out=C)": 0.281 s 31.8 GiB/s <----
Numba fast: 0.288 s 31.0 GiB/s <----
Optimal time: 0.209 s 32.0 GiB/s
Numba 代码和 Numpy 代码使我的 RAM 饱和。使用更多的核心并没有帮助（事实上，由于内存控制器的争用，它肯定会慢一些）。两者都不是最优的，因为它们不使用可阻止缓存行写入分配的非临时存储指令（导致数据在写回之前从 RAM 中获取）。最佳时间是使用此类指令的预期时间。请注意，由于 RAM 混合读/写，预计只能达到 RAM 带宽的 65-80%。事实上，交错读取和写入会导致低级开销，从而防止 RAM 饱和。有关 RAM 工作原理的更多信息，请考虑阅读《高性能科学计算简介 - 第 1.3 章》和《每个程序员应该了解的内存知识》（可能还有此）。