numpy.einsum('ij,ji', a, b) 性能问题

Question

有人能解释一下，为什么 numpy.einsum('ij,ji', A, B) 比 numpy.einsum('ij,ij', A, B) 慢得多，因为它是如下所示？

In [1]: import numpy as np
In [2]: a = np.random.rand(1000,1000)                                                                                                                                
In [3]: b = np.random.rand(1000,1000)                                                                                                                                
In [4]: timeit np.einsum('ij,ij', a, b)                                                                                                                               
532 µs ± 5.36 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [5]: timeit np.einsum('ij,ji', a, b)                                                                                                                               
1.28 ms ± 20.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

此致， 马立克

Answer 1

第一嫌疑人：https://en.wikipedia.org/wiki/Locality_of_reference#Spatial_and_temporal_locality_usage

在第一种情况下，a 和 b 的内存都是以连续的顺序模式访问的。

在第二种情况下，b 的内存在内部循环中以 1000 个元素（8000 字节）的步幅访问。

大多数现代 x86 处理器都有 32KiB L1 缓存和 64 字节缓存行，这意味着总共有 512 条缓存行。所以整个 L1 缓存在外循环迭代之间被驱逐了~两次。

请注意，如果您 运行

timeit np.einsum('ij,ji', a, b.T)

你应该得到与第一个例子大致相同的时间。

同样

timeit np.einsum('ij,ij', a, b.T)

应该给出与第二个示例相同的时间。