Python 中不那么随意的性能优化

2023年8月1日

我的上一篇文章（老实说，主要是为了测试这个网站的主题而创建的）提供了几个计算 $N$ 项平方倒数和的代码片段。我用我最喜欢的四种语言——Python、C、Rust 和 Haskell——编写了代码，但当我运行 Python 代码时，它的速度慢得令人尴尬。与 Rust 顺序执行所需的约 950 毫秒相比，Python 花了 70 秒！因此，在这篇文章中，我们将尝试让 Python 获得一些更合理的数字。

原始方案

def basel(N: int) -> float:
    return sum(x**(-2) for x in range(1,N))

这是完成此任务最 Pythonic 的方式。一行易于阅读的代码，使用了内置的生成器。让我们来计时，使用 $N = 1 0^{8}$ 而不是上一篇文章中的 $1 0^{9}$ （因为我不想等那么久）：

import time

# 计时输入的函数
def time_function(func):
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        print(f"函数 '{func.__name__}' 执行耗时 {execution_time:.6f} 秒。")
        return result
    return wrapper

>>> f = time_function(basel)
>>> f(100000000) # 10^8
函数 'basel' 执行耗时 6.641589 秒。
1.644934057834575

让我们尝试用不那么 Pythonic 的方式重写它，使用 for 循环：

def basel_less_pythonic(N: int) -> float:
    s = 0.0
    for x in range(1, N):
        s += x**(-2)
    return s

>>> f(100000000) # 10^8
函数 'basel_less_pythonic' 执行耗时 5.908466 秒。
1.644934057834575

有意思。不那么地道的写法实际上提升了性能。为什么？让我们使用 Python 反汇编模块 dis 来探究一下。

这是原始版本。我添加了一些有用的注释：

# 加载变量
00 LOAD_GLOBAL              0 (sum)
02 LOAD_CONST               1 (<code object <genexpr> at 0x104f42e40>)
04 LOAD_CONST               2 ('basel.<locals>.<genexpr>')

# 创建函数
06 MAKE_FUNCTION            0

# 创建 `range` 对象
08 LOAD_GLOBAL              1 (range)
10 LOAD_CONST               3 (1)
12 LOAD_FAST                0 (N)
14 CALL_FUNCTION            2

# 转换为迭代器
16 GET_ITER

# 调用生成器
18 CALL_FUNCTION            1

# 调用 sum 函数
20 CALL_FUNCTION            1

# 返回
22 RETURN_VALUE

f <code object <genexpr> at 0x104f42e40>:
# 本质上是一个带 yield 的 `for` 循环
00 GEN_START                0
02 LOAD_FAST                0 (.0)
04 FOR_ITER                 7 (to 20)
06 STORE_FAST               1 (x)
08 LOAD_FAST                1 (x)
10 LOAD_CONST               0 (-2)
12 BINARY_POWER
14 YIELD_VALUE
16 POP_TOP
18 JUMP_ABSOLUTE            2 (to 4)
20 LOAD_CONST               1 (None)
22 RETURN_VALUE

这是 for 循环版本：

# 加载变量
00 LOAD_CONST               1 (0.0)
02 STORE_FAST               1 (s)

# 创建 `range` 对象
04 LOAD_GLOBAL              0 (range)
06 LOAD_CONST               2 (1)
08 LOAD_FAST                0 (N)
10 CALL_FUNCTION            2
12 GET_ITER

# 声明 `for` 循环
14 FOR_ITER                 8 (to 32)
16 STORE_FAST               2 (x)

# 幂运算、加法、存储
18 LOAD_FAST                1 (s)
20 LOAD_FAST                2 (x)
22 LOAD_CONST               3 (-2)
24 BINARY_POWER
26 INPLACE_ADD
28 STORE_FAST               1 (s)

# 跳回 `for` 循环顶部
30 JUMP_ABSOLUTE            7 (to 14)

# 返回 `s`
32 LOAD_FAST                1 (s)
34 RETURN_VALUE

在这个 特定情况 下，使用生成器减慢了程序速度。正如我们所见，生成器循环的代码与第二个版本完全相同。第一个版本只是额外处理了生成器对象，这需要（缓慢的）CALL_FUNCTION 指令。请注意，在实践中，生成器由于其惰性特性通常更快。只是在这种情况下，额外的开销并不值得。

尽管如此，两个版本之间的性能差异（约 1 秒）与 Python 和 C/Rust 之间的整体差异相比是微不足道的。即使使用更快的版本，Rust 代码也比 Python 快约 65 倍。

为什么？主要是因为 Python 是一种弱类型的解释型语言。这意味着 Python 解释器（CPython）必须即时将 Python 代码翻译成计算机易于执行的东西。它通过快速将 Python 代码编译成一种通用的汇编形式（称为 Python 字节码）来实现这一点，我们刚刚查看的就是这个。

然后由解释器执行这些字节码。这个编译步骤是 Python 遭受的第一个性能损失。由于像 Rust 这样的语言只需要编译一次程序，这个时间不计入程序的 运行时间。但最大的性能损失来自于字节码质量不高、与架构无关，而不是原生优化的。这仅仅是解释型语言本质的一部分，因为它们无法花费大量时间进行高质量的编译。

那么，如何才能写出快速的 Python 代码呢？嗯，你不能。但是，你可以通过调用像 Numpy 这样高度优化的库，进入快速的 C 代码。这些库包含预编译的、向量化的 C 函数，让你可以完全绕过 Python 解释器。

让我们使用 NumPy

import numpy as np

def basel_np(N: int) -> float:
    # [1, 1, ..., 1]
    ones = np.ones(N - 1)
    # [1, 2, ..., N]
    r = np.arange(1, N)
    # [1, 1/2, ..., 1/N]
    div = ones / r
    # [1, 1/4, ..., 1/N^2]
    inv_squares = np.square(div)
    # ~ pi^2/6
    return float(np.sum(inv_squares))

>>> f(100000000) # 10^8
函数 'basel_np' 在 0.460317 秒内执行完毕。
1.6449340568482196

哇，这运行速度提升了约 13 倍！查看这种情况下的字节码不会告诉我们太多信息，因为它只包含几个对 NumPy 的 CALL_FUNCTION 指令，实际工作由 NumPy 完成。让我们看看哪一行代码耗时最多：

def basel_np(N: int) -> tuple[float, list[float]]:
    times = []

    # 计时单个步骤
    start = time.perf_counter()
    ones = np.ones(N - 1)
    end = time.perf_counter()
    step_time = end - start
    times.append(step_time)

    # 其余计时代码已省略
    r = np.arange(1, N)
    div = ones / r
    square = np.square(div)
    ret = np.sum(square)

    return ret, times

操作步骤	单步耗时 (毫秒)	累计耗时 (毫秒)
创建全1数组	97	97
创建范围数组	79	176
对范围数组求平方	98	274
执行除法 (全1数组/平方数组)	112	387
最终求和	58	444

让我们思考一下这些数字。创建全1数组/范围数组 的步骤并不涉及繁重的计算工作。然而，它们花费的时间几乎与“更艰巨”的步骤（如求平方和除法）相同。令人惊讶的是，对最终数组求和是最快的步骤！这表明这里的瓶颈不是 CPU，而是内存访问。让我们看看性能如何随 $N$ 从 $1 0^{7}$ 到 $1 0^{9}$ 变化而扩展。

在此图中，最顶部线的边缘代表总耗时，相邻线之间的区域代表每个步骤的耗时。

我们可以看到：

总耗时非线性增长，尽管算法是 $O (n)$ 的。
divide 步骤耗时最多，在 $\approx 3 \times 1 0^{8}$ 和 $\approx 7 \times 1 0^{8}$ 处有急剧增加。

这很合理，因为与其他操作不同，np.divide 需要同时访问两个大型数组，并将结果写入一个新数组。这意味着大量的主内存访问，并且可能从 SSD 读取数据，而 SSD 的读取速度极慢。

广播机制

实际上，NumPy 针对这类问题内置了一种称为广播机制的优化，它能在向量-向量运算时，将较小尺寸的向量虚拟地“投射”到较大尺寸。

def basel_np_broadcast(N) -> float:
    ones = 1
    r = np.arange(1, N)
    div = ones / r
    square = np.square(div)
    # 效果如同 [1, 1, ..., 1] / [1, 4, ..., N^2]
    return float(np.sum(square))

这让我们节省了大量宝贵的缓存空间。让我们运行与之前相同的性能指标测试。取 $N = 1 0^{8}$ ：

操作	单步耗时 (ms)	累计耗时 (ms)
创建全一标量	0.00	0
创建范围数组	68.56	70.48
对范围数组求平方	105.14	180.74
执行标量/平方数组除法	133.08	271.30
最终求和	71.08	310.41

从现在起，我将把广播版本称为“NumPy 解决方案”。

尽管已有显著改进，我们仍能看到那些延迟尖峰。让我们尝试解决这个问题。

关于内存的思考

为了改善函数的内存使用，我们可以采用分块处理的方式，每次对数据范围的一小段执行相同的操作，最后将所有结果累加。这样我们只需在内存中保留一个数据块，有望在利用 NumPy 向量化代码优势的同时，获得更好的缓存利用率。

让我们修改函数以处理 $N$ 的某个区间：

# [N1, N2]，包含两端
def basel_np_range(N1: int, N2: int) -> float:
    # 计时代码已省略
    ones = 1
    r = np.arange(N1, N2 + 1)
    div = ones / r
    squares = np.square(div)
    return float(np.sum(squares))

并编写一个函数来汇总所有数据块。

def basel_chunks(N: int, chunk_size: int) -> float:
    # 计时代码已省略
    s = 0.0
    num_chunks = N // chunk_size
    for i in range(num_chunks):
        s += basel_np_range(i * chunk_size + 1, (i + 1) * chunk_size)
    return s

当 $N = 1 0^{8}$ 时：

函数 'basel_chunks' 执行耗时 0.108557 秒。
1.6449340568482258

很好！它比 NumPy 方案快了约 3 倍。

此外，由于 IEEE 754 浮点数的特性，结果实际上会更精确一些。让我们像之前一样观察性能如何随输入规模变化，同时保持块大小不变。

首先，请注意 y 轴。我们现在的时间尺度是秒，而不是分钟。我们还看到运行时间线性增长，这与算法 $O (n)$ 的时间复杂度一致。我们可以推测，在 chunk_size 为 20000 时，所有数据都能放入缓存，因此不会出现因缓存未命中导致的运行时间尖峰。

让我们尝试在 $[5 \times 1 0^{2}, 1 0^{6}]$ 范围内改变 chunk_size，同时保持 $N$ 固定在 $1 0^{9}$ 。

从图中可以看出，性能提升一直持续到 chunk_size 约为 51000，之后出现的延迟尖峰很可能由缓存未命中引起。

粗略估算

每个 float64 占用 64 位，即 8 字节。我们操作的是 3 个长度为 $N$ 的 float64 数组，因此单次调用的内存使用量为 $51000 \times 8 \times 3 = 1224000 \approx 1.2$ MB。

我的 M1 Pro 芯片规格如下：

内存类型	大小
L1	128KB (数据缓存，每核心)
L2	24MB (由 6 个性能核心共享)
L3	24MB
主内存	16GB

这表明该函数可用的最大 $L 1 + L 2$ 缓存空间约为 $1.2$ MB，如果超过这个值，我们就必须等待 $L 3$ 缓存。

多进程处理

现在，尝试让计算机将更多数组数据放入缓存。一种可能性是尝试在所有核心上运行该函数，这样我们就可以为函数使用更多的 $L 2$ 缓存，让我们试试看。

首先，修改 basel_chunks 函数，使其接受一个 $N$ 的范围作为输入。

# (N1, N2]
def basel_chunks_range(N1: int, N2: int, chunk_size: int):
    # 计时代码已省略
    s = 0.0
    num_chunks = (N2 - N1) // chunk_size
    for i in range(num_chunks):
        s += basel_np_range(N1 + i * chunk_size + 1, N1 + (i + 1) * chunk_size)
    return s

现在进行实际的多进程处理：

from multiprocessing import Pool

def basel_multicore(N: int, chunk_size: int):
    # 我的机器有10个核心
    NUM_CORES = 10
    N_PER_CORE = N // NUM_CORES
    Ns = [
        (i * N_PER_CORE, (i + 1) * N_PER_CORE, chunk_size)
        for i in range(NUM_CORES)
    ]
    # 并行处理10个批次
    with Pool(NUM_CORES) as p:
        result = p.starmap(basel_chunks_range, Ns)
    return sum(result)

在底层，Python 的 multiprocessing 模块会 pickle 函数对象，并生成 9 个新的 python3.10 实例，所有这些实例都使用 $N = 1 0^{9}$ 和 num_chunks = 50000 来执行函数。

让我们比较一下性能。

函数 'basel_multicore' 执行耗时 1.156558 秒。
1.6449340658482263
函数 'basel_chunks' 执行耗时 1.027350 秒。
1.6449340658482263

糟糕，它反而更慢了。这可能是因为多进程处理的初始工作开销很大，因此只有当函数本身需要相对较长时间来完成时，它才会带来好处。

将 $N$ 增加到 $1 0^{10}$ ：

函数 'basel_multicore' 执行耗时 2.314828 秒。
1.6449340667482235
函数 'basel_chunks' 执行耗时 10.221904 秒。
1.6449340667482235

这就对了！速度提高了约 $\approx 5$ 倍。试试 $1 0^{11}$

函数 'basel_multicore' 执行耗时 13.844876 秒。
multicore 1.6449340668379773
函数 'basel_chunks' 执行耗时 102.480372 秒。
chunks 1.6449340668379773

速度提高了约 $\approx 8$ 倍！随着 $N$ 的增加，单核与10核的运行时间比应趋近于 10:1。

x轴是对数刻度，这就是为什么线性增长运行时间看起来像指数增长。

我们可以看到，使用多进程处理有约 $\approx 1$ 秒的固定开销，因此只有当分块运行时间高于此开销时，才会有性能优势。这大约发生在 $N = 1.3 \times 1 0^{9}$ 时。然而，在此之后，性能差异就非常显著了。

通过实验（此处未展示），我发现对于多进程代码，chunk_size 在 50000 左右也是最优的，这告诉我们操作系统对单个核心的 L2 缓存使用施加了一些限制。

结果摘要

Python 函数	$N = 1 0^{8}$ (秒)	$N = 1 0^{9}$ (秒)	$N = 1 0^{10}$ (秒)
`basel_multicore`	1.04	1.17	2.33
`basel_chunks`	0.09	0.91	9.18
`basel_np_broadcast`	0.27	9.68	OOM
`basel_less_pythonic`	5.99	59.25	592 (估)
`basel` (惯用写法)	7.64	68.28	682 (估)
`basel_np`	0.618	44.27	OOM

与上一篇文章中的其他语言相比如何？

语言	$N = 1 0^{9}$ (毫秒)
Rust (rc, –release, 多核 w/ rayon)	112.65
Python3.10 (`basel_chunks`)	913
Rust (rc, –release)	937.94
C (clang, -O3)	995.38
Python3.10 (`basel_multicore`)	1170
Python3.10 (`basel_np_broadcast`)	9680
Haskell (ghc, -O3)	13454
Python3.10 (`basel_np`)	44270
Python3.10 (`basel_less_pythonic`)	59250
Python3.10 (`basel`)	68280

结果极好！分块的 Numpy 代码是最快的顺序执行解决方案！请记住，Python 在进行计算时，背后运行着整个解释器。

比较 Rust 与 Python 的多核性能：

语言	$N = 1 0^{8}$ (毫秒)	$N = 1 0^{9}$ (毫秒)	$N = 1 0^{10}$ (毫秒)	$N = 1 0^{11}$ (毫秒)
Rust	12.3	111.2	1083	10970
Python	1040	1173	2330	12629

显然，对于较小的 $N$ ，Rust 的并行化方法（通过操作系统线程）远比 Python 基于进程的并行化高效。但随着 $N$ 增大，两者之间的差距会缩小。

总结

如果还不明显的话，整篇文章其实只是一次了解Python底层工作原理的练习。请千万不要试图通过过度优化Python代码来给你的经理留下深刻印象。几乎总有更好的解决方案，比如：

from math import pi

def basel_smart() -> float:
    return pi*pi / 6

看！比其他所有函数都快几个数量级，而且无限简洁。

在考虑性能时，你需要格外小心。不要假设性能有问题，除非你经过测量并确认它确实是个问题。如果是这样，在匆忙输入import multiprocessing并用蛮力解决一个原本简单的问题之前，务必先看看是否有更聪明的解决方法。

另外，如果你的应用真的对性能如此敏感，你或许本就不该用Python来编写它。Python生来就是一种擅长快速原型设计的工具，而非执行速度。但正如你所见，要获得出色的性能也并非不可能。

无论如何，希望你喜欢我在这个网站上的第一篇正式文章。如果你找到了针对这个“问题”的更快的Python“解决方案”，欢迎发邮件告诉我！

更新日志

日期	描述	致谢
2023年8月2日	新增广播部分，Python 成为最快解决方案	u/mcmcmcmcmcmcmcmcmc_ & u/Allanon001
2023年8月2日	将 `time.time()` 改为 `time.perf_counter()`	u/james_pic

✦ 本文的构思、研究、撰写和编辑均未使用大语言模型。