使用 NumExpr 提升 NumPy 代码的运行时间：分析答案

【问题标题】：Boosting the runtime of NumPy Code with NumExpr: An analysis使用 NumExpr 提升 NumPy 代码的运行时间：分析
【发布时间】：2019-04-14 12:11:53
【问题描述】：

由于 NumPy 不使用多核，我正在学习使用 NumExpr 加速 NumPy 代码，因为它对多线程有很好的支持。以下是我正在使用的示例：

# input array to work with
x = np.linspace(-1, 1, 1e7)

# a cubic polynomial expr
cubic_poly = 0.25*x**3 + 0.75*x**2 + 1.5*x - 2

%timeit -n 10 cubic_poly = 0.25*x**3 + 0.75*x**2 + 1.5*x - 2
# 657 ms ± 5.04 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

现在，我们可以使用 NumExpr 来做同样的事情：

cubic_poly_str = "0.25*x**3 + 0.75*x**2 + 1.5*x - 2"
# set number of threads to 1 for fair comparison
ne.set_num_threads(1)

%timeit -n 10 ne.evaluate(cubic_poly_str)
# 60.5 ms ± 908 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

从时序中可以看出，NumExpr 的速度快了 10 倍以上，即使我们使用的线程数与 NumPy 使用的线程数相同（即 1）

现在，让我们增加计算并使用所有可用线程并观察：

# use all available threads/cores
ne.set_num_threads(ne.detect_number_of_threads())

%timeit -n 10 ne.evaluate(cubic_poly_str)
# 16.1 ms ± 82.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# sanity check
np.allclose(cubic_poly, ne.evaluate(cubic_poly_str))

不出所料，这比仅使用单线程要快 5 倍。

为什么即使使用相同数量的线程（即 1），NumExpr 也快 10 倍？

【问题讨论】：

使用 numpy 时，计算的一些中间值会存储在内存中（如 1.5*x），对于如此大的 x，这需要时间。另一方面，numexpr 将数据分成小块并执行完整计算，而无需大量存储中间值。

标签： python numpy multidimensional-array polynomial-math numexpr

【解决方案1】：

您认为加速仅/主要来自并行化的假设是错误的。正如@Brenlla 已经指出的那样，numexpr 加速的最大份额通常来自更好地利用缓存。但是还有其他一些原因。

首先，numpy 和 numexpr 不会计算相同的表达式：

numpy 将x**3 和x**2 计算为pow(x,3) 和pow(x,2)。
numexpr 冒昧地将其评估为x**3=x*x*x 和x**2=x*x。

pow 比一两次乘法更复杂，因此要慢得多，比较一下：

ne.set_num_threads(1)
%timeit ne.evaluate("0.25*x**3 + 0.75*x**2 + 1.5*x - 2")
# 60.7 ms ± 1.2 ms, base line on my machine

%timeit 0.25*x**3 + 0.75*x**2 + 1.5*x - 2
# 766 ms ± 4.02 ms
%timeit 0.25*x*x*x + 0.75*x*x + 1.5*x - 2 
# 130 ms ± 692 µs

现在，numexpr 的速度只有两倍。我的猜测是，pow-version 受 CPU 限制，而乘法版本受更多内存限制。

当数据很大时，Numexpr 最亮眼 - 大于 L3 缓存（例如我的机器上的 15Mb），如您的示例中给出的，x 约为 76Mb：

numexp 以块为单位进行评估 - 即所有操作都针对一个块进行评估，并且每个块都适合（至少）L3 缓存，从而最大限度地提高缓存的利用率。只有在完成一个块后，才会评估另一个块。
numpy 对整个数据执行一个又一个操作，因此数据在被重用之前会从缓存中清除。

我们可以使用 valgrind 来查看缓存未命中（参见本文附录中的脚本）：

>>> valgrind --tool=cachegrind python np_version.py
...
...
==5676== D   refs:      1,144,572,370  (754,717,376 rd   + 389,854,994 wr)
==5676== D1  misses:      220,844,716  (181,436,970 rd   +  39,407,746 wr)
==5676== LLd misses:      217,056,340  (178,062,890 rd   +  38,993,450 wr)
==5676== D1  miss rate:          19.3% (       24.0%     +        10.1%  )
==5676== LLd miss rate:          19.0% (       23.6%     +        10.0%  )
....

对我们来说有趣的部分是 LLd-misses（即 L3 未命中，有关输出解释的信息，请参阅 here） - 大约 25% 的读取访问是未命中的。

对 numexpr 的相同分析显示：

>>> valgrind --tool=cachegrind python ne_version.py 
...
==5145== D   refs:      2,612,495,487  (1,737,673,018 rd   + 874,822,469 wr)
==5145== D1  misses:      110,971,378  (   86,949,951 rd   +  24,021,427 wr)
==5145== LLd misses:       29,574,847  (   15,579,163 rd   +  13,995,684 wr)
==5145== D1  miss rate:           4.2% (          5.0%     +         2.7%  )
==5145== LLd miss rate:           1.1% (          0.9%     +         1.6%  )
...

只有 5% 的读取失败！

然而，numpy 也有一些优点：在底层 numpy 使用 mkl-routines（至少在我的机器上），而 numexpr 没有。因此 numpy 最终使用打包的 SSE 操作 (movups+mulpd+addpd)，而 numexpr 最终使用标量版本 (movsd+mulsd)。

这解释了 numpy 版本的 25% 未命中率：一次读取是 128 位 (movups)，这意味着在 4 次读取后处理缓存行（64 字节）并产生未命中。可以在配置文件中看到（例如 Linux 上的perf）：

 32,93 │       movups 0x10(%r15,%rcx,8),%xmm4                                                                               
  1,33 │       movups 0x20(%r15,%rcx,8),%xmm5                                                                               
  1,71 │       movups 0x30(%r15,%rcx,8),%xmm6                                                                               
  0,76 │       movups 0x40(%r15,%rcx,8),%xmm7                                                                               
 24,68 │       movups 0x50(%r15,%rcx,8),%xmm8                                                                               
  1,21 │       movups 0x60(%r15,%rcx,8),%xmm9                                                                               
  2,54 │       movups 0x70(%r15,%rcx,8),%xmm10

每四个movups 需要更多时间，因为它等待内存访问。

Numpy 适合较小的数组大小，适合 L1 缓存（但最大的份额是开销而不是计算本身，这在 numpy 中更快 - 但这并没有起到很大的作用）：

x = np.linspace(-1, 1, 10**3)
%timeit ne.evaluate("0.25*x*x*x + 0.75*x*x + 1.5*x - 2")
# 20.1 µs ± 306 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit 0.25*x*x*x + 0.75*x*x + 1.5*x - 2
# 13.1 µs ± 125 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

附带说明：将函数评估为((0.25*x + 0.75)*x + 1.5)*x - 2 会更快。

都是因为 CPU 使用率较低：

# small x - CPU bound
x = np.linspace(-1, 1, 10**3)
%timeit ((0.25*x + 0.75)*x + 1.5)*x - 2
#  9.02 µs ± 204 ns

更少的内存访问：

# large x - memory bound
x = np.linspace(-1, 1, 10**7)
%timeit ((0.25*x + 0.75)*x + 1.5)*x - 2
#  73.8 ms ± 3.71 ms

列表：

A np_version.py:

import numpy as np

x = np.linspace(-1, 1, 10**7)
for _ in range(10):
    cubic_poly = 0.25*x*x*x + 0.75*x*x + 1.5*x - 2

B ne_version.py:

import numpy as np
import numexpr as ne

x = np.linspace(-1, 1, 10**7)
ne.set_num_threads(1)
for _ in range(10):
    ne.evaluate("0.25*x**3 + 0.75*x**2 + 1.5*x - 2")

【讨论】：