multithreading - 使用 Numba 对 2 个矩阵求和的最快方法是什么？

Question

我正在尝试找到使用 Numba 对 2 个相同大小的矩阵求和的最快方法。我想出了 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")

结果如下:

function 1 (sum_numpy): t1= 0.1655790419999903

function 2 (sum_numba_simple): t2= 0.3019776669998464

function 3 (sum_numba_vectorized): t3= 0.16486266700030683

function 4 (sum_numba_loop): t4= 0.1862256660001549

如您所见，结果表明在这种情况下使用 Numba 没有任何优势。因此，我的问题是:
是否有任何其他实现可以提高求和速度？

Answer 1

您的代码受 page-faults 约束(有关更多信息，请参阅 here、here 和 there)。发生页面错误是因为数组是新分配的。一种解决方案是对其进行预分配，然后在其中写入，这样就不会导致页面在物理内存中重新映射。 np.add(A, B, out=C) 按照评论中@August 的指示执行此操作。另一种解决方案可能是采用标准分配器，这样就不会以显着的内存使用开销为代价将内存返还给操作系统(例如，AFAIK TC-Malloc 可以做到这一点)。

在大多数平台(尤其是 x86 平台)上还有另一个问题:cache-line write allocations写回缓存在写期间是昂贵的。避免这种情况的典型解决方案是进行非临时存储(如果在目标处理器上可用，在 x86-64 上是这种情况，但在其他处理器上可能不是)。话虽如此，Numpy 和 Numba 都无法做到这一点。对于 Numba，我填写了一个 issue涵盖一个简单的用例。编译器本身(Numpy 的 GCC 和 Numba 的 Clang)往往不会生成这样的指令，因为当数组适合缓存并且编译器在编译时不知道数组的大小时，它们可能会损害性能(它们可能会生成特定的代码他们可以评估计算的数据量，但这并不容易，并且会减慢一些其他代码的速度)。 AFAIK，解决此问题的唯一可能方法是编写 C 代码并使用低级指令或使用编译器指令。在您的案例中，大约 25% 的带宽由于这种影响而丢失，导致速度降低高达 33%。

使用多线程并不总能使内存绑定(bind)代码更快。事实上，它通常几乎无法扩展，因为当 RAM 已经饱和时，使用更多内核并不能加快执行速度。通常需要很少的内核才能使大多数平台上的 RAM 饱和。页面错误可以受益于针对目标系统使用多个内核(Linux 可以很好地并行处理，Windows 通常不能很好地扩展，IDK 用于 MacOS)。

最后，还有一个问题:代码没有矢量化(至少在我的机器上没有，虽然它可以)。解决方案是展平数组 View 并执行一个编译器可以更轻松地矢量化的大循环(基于 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 处理器和 ~42 GiB/s RAM 的结果:

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。使用更多核心并没有帮助(事实上，由于内存 Controller 的争用，它肯定会慢一点)。两者都不是最优的，因为它们不使用可以防止缓存行写入分配的非临时存储指令(导致数据在写回之前从 RAM 中获取)。最佳时间是预期使用此类指令的时间。请注意，由于 RAM 混合读/写，预计只能达到 RAM 带宽的 65-80%。实际上，交错读取和写入会导致低级开销，从而防止 RAM 饱和。有关 RAM 工作原理的更多信息，请考虑阅读 Introduction to High Performance Scientific Computing -- Chapter 1.3和 What Every Programmer Should Know About Memory (可能还有 this )。