python - 加速特定矩阵和向量大小的 einsum

Question

我有 2 个数组，其中一个的大小为:

A = np.random.uniform(size=(48, 1000000, 2))

另一个是

B = np.random.uniform(size=(48))

我想做以下总结:

np.einsum("i, ijk -> jk", B, A)

尽可能快。

求和需要进行数千万次，所以速度很重要。在每次迭代中，B 都会发生变化，而 A 会保持不变。我尝试了 2 个选项(见下文)，但它们都具有可比性。有什么办法可以加快速度吗？

代码:

import numpy as np

def computation_einsum(x,y):
    z = np.einsum("i, ijk -> jk", y, x)
    return z

def computation_dot(x,y):
    z = y @ x
    return z

A = np.random.uniform(size=(48, 1000000, 2))
B = np.random.uniform(size=(48))
C = A.transpose(1,0,2)

时间:

%timeit -n 10 computation_einsum(A, B)
100 ms ± 238 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit -n 10 computation_dot(C, B)
107 ms ± 2.11 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

我知道对于较大的矩阵，有些选项会有所不同，但我有这些特定的形状和大小。

Answer 1

默认情况下，np.einsum 不会优化已完成的计算，因为执行优化的成本有时可能比计算本身更大(尤其是对于小型数组)。您可以使用以下方法解决此问题:

np.einsum("i, ijk -> jk", B, A, optimize='optimal')

有了这个，np.einsum 在内部使用一个 BLAS 实现(像默认情况下在大多数系统上的 OpenBLAS，包括在我的机器上)这是更有效的:它应该并行运行并使用更快的 SIMD 指令(尤其是在使用 AVX 而不是 SSE 的 x86 上)。

此操作主要是内存限制，上述解决方案成功达到了约 30 GiB/s 的内存吞吐量，而峰值约为 ~40/s GiB。这非常好(对于通用 Numpy 代码)但次优。

请注意@MichaelSzczesny 的解决方案给出了类似的性能结果(使用 64 位 float )。

---

击败 OpenBLAS

虽然 OpenBLAS 在这里不是最优的，但编写更快的代码远非易事，因为这样的库积极优化以达到高性能(AFAIK，他们利用特定于体系结构的内核和甚至一些手写的汇编代码)。话虽这么说，它们是为一般情况编写的，并假设计算主要受计算限制，但这里不是这种情况。

击败 OpenBLAS 的解决方案是编写一个并行 Numba 代码来解决您的特定情况，其中:B 具有较小的固定大小，A 的最后一个维度 始终为 2，A 相对较大。 Numba 是一个 Python 模块，可以使用即时编译器将特定的 Python 代码编译为 native 二进制文件。

主要思想是沿着第二个维度虚拟地将 A 分成 block ，然后沿着第一个维度(不是连续存储的)执行缓存友好的加权和在内存中)。该代码使用断言和编译时常量来帮助编译器主动展开和向量化循环。输出矩阵在函数的参数中传递，因此无需支付昂贵的成本page-faults .这是最终代码:

import numpy as np
import numba as nb

block_size = 1024
B_size = 48

@nb.njit(fastmath=True, inline='never')
def compute_block(A, B, out, iStart, iEnd, complete):
    n = A.shape[1]

    # Check everything is Ok and help the compiler 
    # to optimize the loops more aggressively.
    assert A.shape[0] == B_size
    assert A.shape[2] == 2
    assert out.shape[1] == 2
    assert B.size == B_size
    assert out.shape[0] == n

    for i in range(iStart, iEnd):
        out[i, 0] = out[i, 1] = 0.0
    for j in range(B_size):
        factor = B[j]

        # Help the compiler to perform (better) loop unrolling
        if complete:
            for i in range(iStart, iStart+block_size):
                out[i, 0] += A[j, i, 0] * factor
                out[i, 1] += A[j, i, 1] * factor
        else:
            for i in range(iStart, iEnd):
                out[i, 0] += A[j, i, 0] * factor
                out[i, 1] += A[j, i, 1] * factor

@nb.njit('void(float64[:,:,::1], float64[::1], float64[:,::1])', parallel=True)
def compute(A, B, out):
    n = A.shape[1]

    # Parallel computation of many blocks
    for bi in nb.prange(n // block_size):
        iStart = bi * block_size
        iEnd = iStart + block_size
        compute_block(A, B, out, iStart, iEnd, True)

    # Remaining part
    bi = n // block_size
    iStart = bi * block_size
    if iStart < n:
        compute_block(A, B, out, iStart, n, False)

# Example of use:
A = np.random.uniform(size=(48, 1000000, 2))
B = np.random.uniform(size=(48))

C = np.empty((A.shape[1], 2), dtype=np.float64)
compute(A, B, C)

---

基准

以下是我的 i5-9600KF 处理器和约 40 GiB DRAM 的结果:

np.einsum("i, ijk -> jk", B, A):                        54.4 ms
np.einsum("i, ijk -> jk", B, A, optimize='optimal'):    25.2 ms
compute(A, B, C)                                        19.5 ms
Theoretical optimal time (lower bound):                 18.6 ms

最终实现成功达到了 38.2 GiB/s 的内存吞吐量，这非常接近我机器上的最佳值。

请注意，正如@MichaelSzczesny 所指出的，使用 32 位 float 会使时序快大约 2 倍，因为 32 位 float 占用的内存空间少两倍，而且操作受内存限制。不过，32 位 float 不太精确。