应用-理解CPU如何帮助加速Numba和NumPy代码

目录

理解CPU如何帮助加速Numba和NumPy代码

当你需要加速NumPy处理——或者只是减少内存使用——Numba即时编译器是一个很好的工具。它允许你编写Python代码，并在运行时将其编译为机器码，从而获得类似于C、Fortran或Rust等语言的性能提升。

至少，理论上是这样。实际上，你最初的Numba代码可能并不比NumPy等效代码快。

但你可以做得更好，只要你更好地理解CPU的工作原理。这些知识将帮助你更广泛地应用于任何编译语言。

在本文中，我们将：

• 考虑一个简单的图像处理问题。
• 尝试用Numba加速它，但最初失败了。
• 回顾一下现代CPU为何如此快速，以及编译器的局限性。
• 基于我们的新理解，展示如何调整代码使其比原始版本快25倍。

---

问题：去除图像中的噪声

假设你从数字显微镜中获取16位图像，而你只关心图像的明亮部分。暗区没有信息，并且有很多噪声。如果我们想要清理这些区域，一个简单的算法是将所有低于某个阈值的值设置为完全黑色，即0。

以下是使用NumPy的实现：

image[image < 1000] = 0

为了测试运行速度，我生成了一个模拟图像：

import numpy as np
rng = np.random.default_rng(12345)
noise = rng.integers(0, high=1000, size=(4096, 4096), dtype=np.uint16)
signal = rng.integers(0, high=5000, size=(4096, 4096), dtype=np.uint16)
image = noise | signal

然后使用Jupyter和IPython中的%timeit魔法命令进行计时，禁用CPU的Turbo Boost。代码会修改图像，因此我们需要在副本上操作，复制会增加开销，因此我们也测量了复制的时间：

>>> %timeit image2 = image.copy()
5.84 ms ± 189 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
>>> %timeit image2 = image.copy(); image2[image2 < 1000] = 0
54.2 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

使用NumPy运行这个简单算法大约需要48ms（54ms - 6ms）。我们能做得更好吗？

---

尝试#0：将原始算法翻译为Numba

Numba允许我们编写Python代码的子集，并将其编译为本地代码。作为第一次尝试，我们可以实现以下算法：

from numba import njit

@njit
def remove_noise_numba_0(arr, noise_level):
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

remove_noise_numba_0(image.copy(), 1000)

我们确保运行一次函数，以便后续的计时不包括编译代码的一次性时间。

使用%timeit测量，耗时46ms，比NumPy代码快不了多少。

---

尝试#1：选择好的类型

虽然Numba看起来像Python，但最好将其视为C或其他静态类型的编译语言。它会根据传入数据的类型生成特定的代码。

对于数组，我们知道类型：一个16位无符号整数的二维数组。但noise_level的类型是Numba需要猜测的，因为我们传入了一个Python整数，它会将其转换为原始类型，类似于NumPy的dtype。

我们知道它也应该是一个16位无符号整数（噪声级别不能为负数，也不能比最亮的点大！），因此我们可以调整代码以强制执行：

@njit
def remove_noise_numba_1(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

remove_noise_numba_1(image.copy(), 1000)

虽然这个更改的目的不是为了加速——目标是正确性——但这个更改恰好将速度提升到了33ms。我们能做得更好吗？

为了回答这个问题，我们需要考虑CPU的工作原理。

---

现代CPU的更好心智模型

首先，我们在这里讨论的所有内容都涉及单个CPU核心。我们完全忽略了并行化到多个核心的可能性，只假设在单个核心上运行的单个线程。

CPU的第一个心智模型是它们读取机器指令并逐个执行。

虽然这个模型易于理解，并且在许多情况下足够用，但如果你想编写更快的代码，你需要一个更准确的心智模型。单指令模型意味着执行更多指令会使代码变慢，基于添加的指令数量。但这并不总是正确的。不同的指令可能对性能产生非常不同的影响。

现代CPU可以并行运行指令

事实上，你在台式机、笔记本电脑和服务器上遇到的现代CPU可以同时运行多个指令，只要这些指令不相互干扰。考虑以下代码；假设它是用C或其他编译语言编写的，因此表达式可以直接映射到CPU指令：

a = a + 1;
b = b + 1;

这两行代码彼此无关。如果翻译为机器码，CPU很可能能够同时运行这两行代码。再次强调，这与多核CPU或多线程无关：这都是在单个核心上的单个线程。

现在考虑以下代码会发生什么：

a = a + 1;
if (a > 100) {
    b = b + 1;
}

在这个例子中，递增b依赖于a的值。这意味着第三行代码在第一行代码完成并且比较完成后才能执行。不再有并行性，我们的代码突然变得慢了很多！

现代CPU会进行推测执行和分支预测

为了恢复一些丢失的并行性，现代CPU会基于假设大多数代码实际上是可预测的，执行后续代码——并预测条件中的哪个分支将被执行。考虑以下代码：

int a = 0;
int b = 0;
while (true) {
    a = a + 1
    if (a > 100) {
        b = b + 1;
    }
}

在前一百次迭代中，b永远不会递增。之后，b将始终递增。除了过渡期外，CPU很可能能够预测接下来会发生什么，并推测性地选择要执行的代码。

如果CPU猜错了会发生什么？它会撤销工作，这可能会很昂贵，这意味着失败的分支预测可能会对性能产生重大影响。

现代CPU有特殊的CPU指令用于并行工作：SIMD

除了透明地并行化指令外，现代CPU还有SIMD指令：单指令多数据。其思想是，你可以使用单个CPU指令对数组中的多个条目执行相同的操作。

例如，这里我们手动将100添加到四个数组条目中，使用循环：

int array[4] = {1, 2, 3, 4};
for (int i=0; i < 4; i++) {
    array[i] += 100;
}

相反，使用SIMD，我们可以使用特殊的CPU指令一次性将100添加到所有四个数组项中。即使指令稍微慢一些，由于我们将四个指令（可能还有部分循环）替换为单个指令，总体而言，我们可以预期代码运行得更快。

支持这些特殊指令是因为我们对所有四个值执行相同的操作，在这种情况下是添加一个常数。将两个数组相加也将是对所有项执行相同的操作。关键在于名称：单指令，多数据。

你可以显式使用SIMD指令，但Numba不支持，而且这通常会使你的代码依赖于特定的CPU。ARM CPU（如较新的Mac上使用的）与x86-64 CPU有不同的SIMD指令，甚至x86-64 CPU的SIMD指令也因型号而异。

或者，你可以希望你的编译器足够智能，在相关的地方自动使用SIMD。在Numba的情况下，因为它为你的计算机即时编译，所以它可以使用你的CPU支持的任何SIMD指令。

帮助编译器：线性代码是快速代码

无论是自动CPU指令并行化还是基于SIMD的并行化，if语句和其他条件都是一个问题。

• 对于指令并行化，条件代码意味着CPU需要猜测接下来会发生什么，它可能会猜错。
• 对于SIMD，这些指令通常希望对小数组中的所有项执行相同的操作。条件使得执行相同操作变得更加困难。

如果我们想让代码更快，我们希望尽可能使其线性化，没有if语句或其他分支条件。 这将使编译器更有可能识别出使用SIMD的机会，并使CPU更容易并行运行生成的指令。

---

使用perf检查我们的代码

我们已经了解到，条件代码使编译器更难编写快速代码，也使CPU更难快速运行生成的代码。这是我们当前的Numba代码：

@njit
def remove_noise_numba_1(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

这里肯定有一个条件！代码有时会写入数组，有时不会，这取决于比较结果。

我们可以使用Linux上的perf工具来测量这种影响，使用以下装饰器连接到当前进程并获取CPU在附加期间的数据：

from subprocess import Popen
from contextlib import contextmanager
from os import getpid
from time import sleep
from signal import SIGINT

@contextmanager
def perf_stat():
    p = Popen(["perf", "stat", "-p", str(getpid())])
    sleep(0.5)
    yield
    p.send_signal(SIGINT)

我们可以在代码上运行这个：

image2 = image.copy()
with perf_stat():
    remove_noise_numba_1(image2, 1000)

结果是一大堆信息，但对我们来说，我们只看两个数字：错误推测的百分比和分支预测错误的百分比。

...
52.8% Bad Speculation
52.8% Branch Mispredict
...

这是一个很大的预测错误！一旦我们想到我们的算法在做什么，这就不足为奇了：将一堆噪声数字与一个值进行比较，看看它们是高于还是低于阈值。对于杂乱的图像来说，这很难预测。

---

尝试#2：去除条件

我们想要去除条件。一种方法是提出一个等效的计算，该计算始终运行相同的代码，没有分支或条件。这里的一个关键词是“无分支”代码，通过维基百科的搜索，我找到了这个有用的页面，介绍了无分支饱和算术。也就是说，我们想要减去一个值，如果它低于零，则坚持为零，而不是环绕，这是无符号整数的默认行为。

经过几分钟的调试，我想出了以下代码。它给出了相同的结果，没有任何分支。不幸的是，它更难理解：它依赖于一些与CPU内存中整数表示方式相关的技巧。

@njit
def remove_noise_numba_2(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            res = arr[i, j] - noise_level
            mask_to_zero_if_wrapped = -(res <= arr[i, j])
            res = (res & mask_to_zero_if_wrapped) + (
                   noise_level & mask_to_zero_if_wrapped)
            arr[i, j] = res

remove_noise_numba_2(image.copy(), 1000)

如果CPU真的逐个执行指令，这可能会比我们之前的尝试慢：例如，我们对所有条目进行写入。但现代CPU并非如此工作，因此这个版本要快得多。

事实上，我们降到了4ms。

如果我们用perf测量它，我们可以看到我们不再受到分支预测错误的影响：

...
2.4% Bad Speculation
2.4% Branch Mispredict
...

---

尝试#3：帮助编译器帮助我们

我们通过去除条件使代码更快——但代价是晦涩、丑陋且难以维护的代码。我们能做得更好吗？

在数学上，我们的两个版本是相同的。我们可以假设编译器会尝试在明显会减慢CPU速度时避免条件。编译器当然非常擅长将低效的数学表达式转换为高效的表达式。

那么为什么编译器在这种情况下不能这样做呢？让我们再次看看低效的版本：

@njit
def remove_noise_numba_1(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            if arr[i, j] < noise_level:
                arr[i, j] = 0

问题在于我们只进行条件写入：对于高于噪声级别的数字，我们不执行任何操作。我对编译器作者的心智模型是，虽然他们乐于删除不必要的内存读取，但他们会保守地拒绝添加你没有要求的额外内存写入。

那么，如果我们在两种情况下都写入内存，但其他情况下只保留条件呢？

@njit
def remove_noise_numba_3(arr, noise_level):
    noise_level = arr.dtype.type(noise_level)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            arr[i, j] = (
                arr[i, j] if arr[i, j] >= noise_level
                else 0
            )

remove_noise_numba_3(image.copy(), 1000)

是的，仍然有一个条件，但它是在一个相当简单的计算中，不会影响内存写入。我们仍然在条件的两种情况下写入相同的内存地址。聪明的编译器会注意到有一个等效的数学表达式不需要条件，并切换到使用它，如果它认为这会加速代码。

当我们运行这个版本时，运行时间为2ms，这是迄今为止最快的版本！无论编译器找到了什么表达式，它都比我们上面的手动位操作更高效。

---

我们学到的内容及进一步讨论

以下是我们考虑的所有版本的性能：

版本	运行时间（越低越好）
NumPy	48 ms
Numba #1，条件写入	33 ms
Numba #2，丑陋的位操作	4 ms
Numba #3，无条件写入	2 ms

我们成功获得了比原始版本快25倍的代码，但仍然非常易读！我们通过以下方式实现了这一点：

1. 切换到Numba，以便更好地控制执行。
2. 确保所有类型都明确指定并相同（在这种情况下是np.uint16）。
3. 去除条件写入；这使得编译器可以切换到没有分支的线性实现。

这没有涉及线程或多进程：这都是在单个CPU核心上的单线程代码。

结果可能因编译器版本和CPU而异

一方面，所有编译器都尝试为目标CPU编写高效的代码，因此类似的方法可能适用于其他语言，而不仅仅是Numba。

也就是说，不同的编译器版本和实现行为不同，改变了它们进行优化的方式。因此，这可能会产生影响。因此，在一个版本中自动优化的代码可能在另一个版本中退化；如果速度至关重要，拥有基准测试以在部署之前捕获退化是很重要的。

对于C、Rust、Cython等使用的提前编译，编译器通常针对可用功能的子集，以便代码可以在不同的CPU模型之间移植。你可以选择针对当前CPU或特定功能集进行优化，以在你知道只会在特定硬件上运行的情况下最大化性能。Numba是即时编译的，针对你的计算机，因此它始终针对当前计算机CPU的特定功能。这意味着不同的计算机可能会给出不同的结果。

例如，在特定设置中，Numba能够在初始版本remove_noise_numba_0()上获得巨大的加速，与NumPy相比，无需进一步更改。这种加速似乎是Numba使用AVX-512 SIMD指令的结果，而不是在我的计算机上生成的分支机器代码。

哪些Numba变体最终使用了SIMD？

对于浮点SIMD，我的计算机的Intel CPU可以报告运行了多少条指令（例如通过perf stat访问），但在这个例子中我们使用的是整数SIMD，该统计信息不可用。然而，你可以通过运行以下命令获取Numba即时编译的SIMD使用诊断：

import llvmlite.binding as llvm
llvm.set_option('', '--debug-only=loop-vectorize')

据我所知，所有Numba尝试都使用了某种程度的SIMD。但SIMD只是加速代码的一种方式，还有其他瓶颈，版本1（至少在我的计算机上）涉及大量分支。

---