显式矩阵乘法比 numpy.matmul 快得多？-解网

问：

在 Python 代码中，我需要在某些时候分别将两个 2x2 矩阵的大列表相乘。在代码中，这两个列表都是形状为（n，2,2）的 numpy 数组。另一个（n，2,2）数组中的预期结果，其中矩阵 1 是第一个列表的矩阵 1 和第二个列表的矩阵 1 相乘的结果，依此类推。

经过一番分析，我发现矩阵乘法是一个性能瓶颈。出于好奇，我尝试“显式”地写矩阵乘法。下面是一个包含测量运行时间的代码示例。

import timeit
import numpy as np

def explicit_2x2_matrices_multiplication(
    mats_a: np.ndarray, mats_b: np.ndarray
) -> np.ndarray:
    matrices_multiplied = np.empty_like(mats_b)
    for i in range(2):
        for j in range(2):
            matrices_multiplied[:, i, j] = (
                mats_a[:, i, 0] * mats_b[:, 0, j] + mats_a[:, i, 1] * mats_b[:, 1, j]
            )

    return matrices_multiplied


matrices_a = np.random.random((1000, 2, 2))
matrices_b = np.random.random((1000, 2, 2))

assert np.allclose( # Checking that the explicit version is correct 
    matrices_a @ matrices_b,
    explicit_2x2_matrices_multiplication(matrices_a, matrices_b),
)

print(  # 1.1814142999992328 seconds
    timeit.timeit(lambda: matrices_a @ matrices_b, number=10000)
)
print(  # 1.1954495010013488 seconds
    timeit.timeit(lambda: np.matmul(matrices_a, matrices_b), number=10000)
)
print(  # 2.2304022700009227 seconds
    timeit.timeit(lambda: np.einsum('lij,ljk->lik', matrices_a, matrices_b), number=10000)
)
print(  # 0.19581600800120214 seconds
    timeit.timeit(
        lambda: explicit_2x2_matrices_multiplication(matrices_a, matrices_b),
        number=10000,
    )
)

如代码中测试的那样，此函数生成的结果与常规矩阵结果相同。然而，速度不一样：在我的机器上，显式表达式的速度要快 10 倍。__matmul__

这对我来说是一个非常令人惊讶的结果。我本来希望 numpy 表达式更快，或者至少与更长的 Python 版本相当，而不是像我们在这里看到的那样慢一个数量级。我很好奇为什么性能差异如此巨大。

我正在运行 numpy 版本 1.25 和 Python 版本 3.10.6。

python numpy 矩阵数学优化

Numpy 的性能针对“任何大小的输入”进行了优化，而您的代码针对 2x2 进行了优化，没有其他任何内容。如果您不需要任意输入，或者您的输入大小有限，则专用（“展开”）代码将始终胜过通用代码，因为没有开销，无论开销有多小。

0赞 TheFamousRat 7/21/2023

@Mike'Pomax'Kamermans：这是一个公平的观点。这是我在 2x2 矩阵的 np.linalg.det 中已经体验过的事情，其中显式版本也比通用版本快 10 倍。我想知道是否还有更多，因为速度快一个数量级似乎很多，尤其是对于较小矩阵的操作。虽然你当然是对的，但更通用的代码总是会有更多的开销。

答：

1赞 hpaulj 7/21/2023 #1

验证第一个“批次”维度的粗略线性度：matmul

您的（1000,2,2）数组：

In [353]: timeit matrices_a@matrices_b
251 µs ± 767 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

以及一半和十分之一尺寸：

In [354]: timeit matrices_a[:500]@matrices_b[:500]
129 µs ± 783 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [355]: timeit matrices_a[:100]@matrices_b[:100]
28.7 µs ± 532 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

你的显式

In [360]: explicit_2x2_matrices_multiplication(matrices_a, matrices_b).shape
Out[360]: (1000, 2, 2)
In [361]: timeit explicit_2x2_matrices_multiplication(matrices_a, matrices_b)
59.9 µs ± 1.91 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

np.einsum不尝试任何重新排序或其他优化：

In [362]: print(np.einsum_path('ijk,ikl->ijl',matrices_a, matrices_b, optimize='greedy
     ...: ')[1])
  Complete contraction:  ijk,ikl->ijl
         Naive scaling:  4
     Optimized scaling:  4
      Naive FLOP count:  1.600e+04
  Optimized FLOP count:  1.600e+04
   Theoretical speedup:  1.000
  Largest intermediate:  4.000e+03 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   4                ikl,ijk->ijl                                 ijl->ijl

3赞 Jérôme Richard 7/21/2023 #2

TL;DR：问题中提供的所有方法都非常低效。事实上，Numpy显然是次优的，仅使用Numpy就无法有效地计算这一点。尽管如此，还是有比问题中提供的解决方案更快的解决方案。

解释和更快的实施

Numpy 代码利用强大的泛型迭代器将给定的计算（如矩阵乘法）应用于多个数组切片。这样的迭代器对于实现广播和生成相对简单的实现很有用。问题是，当迭代次数很大而数组很小时，它们也相当昂贵。这正是您的用例中发生的情况。虽然可以通过优化 Numpy 代码来减轻 Numpy 迭代器的开销，但在这个特定用例中，没有办法将开销减少到可以忽略不计的时间。事实上，每个矩阵只有 12 个浮点运算要执行。相对较新的 x86-64 处理器可以使用标量 FMA 单元在不到 10 纳秒的时间内计算每个矩阵。事实上，人们可以使用 SIMD 指令，因此只需几纳秒即可计算出每个矩阵。einsum

首先，我们可以通过对沿第一个轴的向量进行矩阵乘法操作来消除内部 Numpy 迭代器的开销。这正是这样做的！explicit_2x2_matrices_multiplication

虽然应该快得多，但它仍然是次优的：它执行非连续的内存读取，创建几个无用的临时数组，并且每个 Numpy 调用都会引入很小的开销。更快的解决方案是编写 Numba/Cython 代码，以便底层编译器可以生成针对 2x2 矩阵优化的非常好的指令序列。在这种情况下，好的编译器甚至可以生成 SIMD 指令，这对 Numpy 来说是不可能的。下面是生成的代码：explicit_2x2_matrices_multiplication

import numba as nb
@nb.njit('(float64[:,:,::1], float64[:,:,::1])')
def compute_fastest(matrices_a, matrices_b):
    assert matrices_a.shape == matrices_b.shape
    assert matrices_a.shape[1] == 2 and matrices_a.shape[2] == 2

    n = matrices_a.shape[0]
    result_matrices = np.empty((n, 2, 2))
    for k in range(n):
        for i in range(2):
            for j in range(2):
                result_matrices[k,i,j] = matrices_a[k,i,0] * matrices_b[k,0,j] + matrices_a[k,i,1] * matrices_b[k,1,j]

    return result_matrices

性能结果

以下是我的机器在使用 i5-9600KF CPU 处理 1000x2x2 矩阵时的性能结果：

Naive einsum:                           214   µs
matrices_a @ matrices_b:                102   µs
explicit_2x2_matrices_multiplication:    24   µs
compute_fastest:                          2.7 µs   <-----

讨论

Numba 实现达到 4.5 GFlops。每个矩阵的计算时间仅为 12 个 CPU 周期（2.7 ns）！我的机器在实践中能够达到 300 GFlops（理论上为 432 GFlops），但只有 50 GFlops 使用 1 个内核和 12.5 GFlops 使用标量代码（理论上为 18 GFlops）。操作的粒度太小，多线程无法使用（生成线程的开销至少为几微秒）。此外，SIMD 代码很难使 de FMA 单位饱和，因为输入数据布局需要 SIMD 随机排序，因此 50 GFlops 实际上是一个乐观的上限。因此，我们可以肯定地说，Numba 实现是一个非常有效的实现。尽管如此，由于 SIMD 指令，可以编写更快的代码（我预计在实践中可以提高大约 x2 的速度）。话虽如此，使用 SIMD 内部函数编写原生代码以帮助编译器生成快速的 SIMD 代码确实不是一件容易的事（更不用说最终代码会很丑陋且难以维护）。因此，SIMD 实现可能不值得付出努力。

上一个：Scipy Rotvec 到 Extrinsic Euler

下一个：浮点数学坏了吗？

显式矩阵乘法比 numpy.matmul 快得多？

Explicit matrix multiplication much faster than numpy.matmul?

评论

解释和更快的实施

性能结果

讨论