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

Question

在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,
    )
)

根据代码测试，这个函数产生的结果与常规矩阵的__matmul__结果相同。然而不同的是速度：在我的机器上，显式表达式要快10倍。

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

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

- TheFamousRat

---

我的初步猜测是你提供的是“最坏情况”，其中Numpy广播开销是主要成本。我建议你也包括一个einsum比较。 - Reinderien

我的初步猜测是，你提出了“最坏情况”，其中Numpy广播开销是主要的成本。我建议你也包括一个einsum的比较。 - Reinderien

numpy的性能针对“任意大小的输入”进行了优化，而你的代码只针对2x2进行了优化，其他情况下并不适用。如果你不需要任意输入，或者你只有有限数量的输入大小，专门（“展开”）编写的代码将始终胜过通用代码，因为没有任何开销，无论这个开销多么小。 - Mike 'Pomax' Kamermans

numpy的性能针对“任意大小的输入”进行了优化，而你的代码只针对2x2进行了优化，其他情况并不适用。如果你不需要任意输入，或者只有有限数量的输入大小，专门（展开）的代码将始终胜过通用代码，因为没有任何额外开销，无论这个开销多么小。 - Mike 'Pomax' Kamermans

numpy的性能针对“任意大小的输入”进行了优化，而你的代码只针对2x2进行了优化，其他情况则不适用。如果你不需要任意输入，或者只有有限数量的输入大小，专门（展开）的代码将始终胜过通用代码，因为没有任何开销，无论这个开销有多小。 - undefined

Answer 1

TL;DR：问题中提供的所有方法都非常低效。事实上，NumPy明显不是最优的，仅使用NumPy无法高效计算此问题。然而，有比问题中提供的解决方案更快的解决方法。

---

解释和更快的实现

Numpy代码利用强大的通用迭代器来对多个数组切片应用给定的计算（如矩阵乘法）。这些迭代器对于实现广播以及生成相对简单的einsum实现非常有用。然而，当迭代次数巨大而数组很小的时候，它们也是非常昂贵的。这正是您的使用情况。虽然可以通过优化Numpy代码来减少迭代器的开销，但在这种特定的使用情况下，无法将开销降低到可忽略的时间。事实上，每个矩阵只需要执行12次浮点运算。一个相对较新的x86-64处理器可以使用标量FMA单元在不到10纳秒的时间内计算每个矩阵。实际上，可以使用SIMD指令，在几个纳秒内计算每个矩阵。

首先，我们可以通过自己进行矩阵乘法，沿着第一个轴操作向量，从而基本消除内部Numpy迭代器的开销。这正是explicit_2x2_matrices_multiplication所做的！

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

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纳秒）计算！我的机器在实际中能够达到最高300 GFlops（理论上432 GFlops），但只有使用1个核心时才能达到50 GFlops，使用标量代码时只有12.5 GFlops（理论上18 GFlops）。操作的粒度对于多线程来说太小，无法发挥作用（生成线程的开销至少几微秒）。此外，SIMD代码很难充分利用FMA单元，因为输入数据布局需要进行SIMD洗牌，所以50 GFlops实际上是一个乐观的上限。因此，我们可以放心地说Numba的实现非常高效。尽管如此，通过SIMD指令编写更快的代码是可能的（我预计实际上可以加速约两倍）。话虽如此，使用SIMD内嵌函数编写本地代码或者帮助编译器生成快速的SIMD代码并不容易（更不用说最终的代码将会丑陋且难以维护）。因此，实施SIMD可能并不值得努力。

- Jérôme Richard

Answer 2

验证第一个“批次”维度上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

- hpaulj