当元素数量小于3-5百万时，numpy如何提高速度？

Question

当元素数量低于一定数量时，一些 numpy 操作会突然提高性能。

这是一个使用 numpy 求数组最大值的函数。

def np_max(arr):
    return np.max(arr)

为了进行比较，这里还有一个使用 numba 来查找最大值的函数。

@numba.njit
def nb_max(arr):
    a_max = arr[0]
    for a in arr[1:]:
        a_max = max(a_max, a)
    return a_max

这是用 perfplot 测量的运行时间。

代码：

def setup(n):
    rng = np.random.default_rng(0)
    return rng.random(n)


def benchmark_timeit():
    for f in [np_max, nb_max]:
        f(setup(100))

    n_run = 100
    for n in np.logspace(6, 7, num=30, dtype=int).tolist():
        arr = setup(n)
        np_max_time = timeit(lambda: np_max(arr), number=n_run) / n_run
        nb_max_time = timeit(lambda: nb_max(arr), number=n_run) / n_run
        print(f"n={n,}:"
              f" np_max={np_max_time * 1000:.2f}"
              f", nb_max={nb_max_time * 1000:.2f}"
              f", np/nb={np_max_time / nb_max_time:.2f}")


def benchmark_perfplot():
    for f in [np_max, nb_max]:
        f(setup(100))

    data = perfplot.bench(
        n_range=np.logspace(1, 8, num=8 + 7 * 9, dtype=int).tolist(),
        # n_range=np.logspace(6, 7, num=30, dtype=int).tolist(),
        setup=setup,
        kernels=[np_max, nb_max],
        equality_check=np.allclose,
        target_time_per_measurement=1.0,
    )
    data.save("./temp2.png")


if __name__ == "__main__":
    benchmark_perfplot()
    benchmark_timeit()

结果：

正如你所看到的，10^6 和 10^7 之间有一个很大的跳跃。 用timeit测量的结果也证实了这一点。

n=(1000000,): np_max=0.11, nb_max=0.27, np/nb=0.42 <-- np/nb = numpy-runtime / numba-runtime
n=(1082636,): np_max=0.12, nb_max=0.30, np/nb=0.40
n=(1172102,): np_max=0.13, nb_max=0.32, np/nb=0.40
n=(1268961,): np_max=0.15, nb_max=0.46, np/nb=0.31
n=(1373823,): np_max=0.16, nb_max=0.38, np/nb=0.41
n=(1487352,): np_max=0.17, nb_max=0.41, np/nb=0.41
n=(1610262,): np_max=0.18, nb_max=0.45, np/nb=0.41
n=(1743328,): np_max=0.19, nb_max=0.48, np/nb=0.40
n=(1887391,): np_max=0.21, nb_max=0.52, np/nb=0.41
n=(2043359,): np_max=0.23, nb_max=0.56, np/nb=0.41
n=(2212216,): np_max=0.25, nb_max=0.61, np/nb=0.41
n=(2395026,): np_max=0.30, nb_max=0.69, np/nb=0.44
n=(2592943,): np_max=0.31, nb_max=0.72, np/nb=0.42
n=(2807216,): np_max=0.33, nb_max=0.77, np/nb=0.43
n=(3039195,): np_max=0.38, nb_max=0.88, np/nb=0.43 <-- 0.4
n=(3290344,): np_max=0.50, nb_max=0.97, np/nb=0.51      |
n=(3562247,): np_max=0.65, nb_max=1.07, np/nb=0.61      |
n=(3856620,): np_max=0.80, nb_max=1.19, np/nb=0.67      | 2x difference
n=(4175318,): np_max=1.00, nb_max=1.35, np/nb=0.74      |
n=(4520353,): np_max=1.11, nb_max=1.49, np/nb=0.75      |
n=(4893900,): np_max=1.25, nb_max=1.59, np/nb=0.78      |
n=(5298316,): np_max=1.44, nb_max=1.79, np/nb=0.81 <-- 0.8
n=(5736152,): np_max=1.57, nb_max=1.95, np/nb=0.80
n=(6210169,): np_max=1.71, nb_max=2.08, np/nb=0.82
n=(6723357,): np_max=1.85, nb_max=2.27, np/nb=0.81
n=(7278953,): np_max=2.02, nb_max=2.49, np/nb=0.81
n=(7880462,): np_max=2.17, nb_max=2.67, np/nb=0.81
n=(8531678,): np_max=2.44, nb_max=2.91, np/nb=0.84
n=(9236708,): np_max=2.61, nb_max=3.17, np/nb=0.82
n=(10000000,): np_max=2.81, nb_max=3.50, np/nb=0.80

我的问题是：

1. 什么优化带来了如此大的变化？
2. 可以在numba中重现它吗？

请注意，我想要完成的是获得有关此优化的知识，而不是使用并行化等开箱即用的方法来击败 numpy。

---

规格：

• AMD 锐龙 9 5900X
• 内存 64 GB
• Windows 10
• Python 3.10.11
• numpy 1.26.2
• numba 0.58.1

Answer 1

numpy 内部使用 BLAS，它可能使用 SIMD 或某种优化方法来计算大型数组的最大值。最初的高时间可能是调用这些函数的成本。

您可以通过对数组进行切片并对大型数组并行执行最大值来优化 numba 实现。