当计算成本太高时,你会怎么做?

最近我有了一个绝妙的商业点子:Mandelbrot 即服务!公司不用自己计算分形,我将为他们在云端实时计算,他们无需做任何工作。通过使用云计算,我将能够扩展到无疑会为我这项新颖服务付费的大量客户。

我有两个目标:

1. 加快结果速度: 我能越快返回分形图,客户就会越满意。

2. 降低成本: 如果我能为云提供商支付更少的计算费用,我的利润就会上升!

不幸的是,由于我只销售刚刚计算出来的、CPU 还温热的 Mandelbrot 图,我不能依赖缓存。

在这种情况下,你会怎么做?

一个显而易见的方法是并行化:使用线程或多进程。这确实会加快结果速度,所以绝对值得去做,但它 不会降低我的成本^[1] 。如果我们使用 10 个核心而不是 1 个核心,服务将返回结果的速度快 10 倍,但我们将不得不支付大约 10 倍的费用,因为我们将使用 10 倍大小的实例。

然而,如果我们能想办法在单个核心上加速计算,这将有助于实现我们的两个目标。我们既能获得更快的结果(与任何多核处理相结合),又能降低计算成本。

在本文中,我们将:

• 快速回顾一个用 Rust 编写的标准 Mandelbrot 实现。
• 讨论为什么在单个 CPU 核心上优化 Mandelbrot 算法可能很棘手。
• 演示如何通过使用掩码 SIMD 操作来实现这一点。
• 使用 Rust 的 Rayon 库轻松添加多核并行性。
• 盈利!

Rust 中的 Mandelbrot

Mandelbrot 计算用 Rust 编写有两个原因:

1. 它是 编写 Python 扩展的优秀语言^[2] 。
2. 我找到了一个预先编写的版本,其中包括标量(普通)实现和 SIMD 实现。

本页中的代码示例和图像版权归 Rust 开发者所有,并根据 MIT 或 Apache 2.0 许可证授权,由你选择。我已将原始代码移植到较新的 Rust SIMD API,并为了可读性进行了简化。我将在文章末尾链接到完整项目。

我们想得到这样漂亮的图像:

涉及的计算相当简单。对于每个像素,我们有一个起始复数,我们反复变换它,直到它从一个常数阈值(默认为 4.0)"发散"。变换的次数被转换为相应像素的颜色:

fn get_count(mut current: Complex<f64>, max_iter: usize) -> usize {
    let mut z = Complex { re: 0.0, im: 0.0 };
    for i in 0..max_iter {
        if z.norm_sqr() > 4.0 {
            return i;
        }
        z = z * z + current;
    }
    max_iter
}

其他未显示的代码为每个像素创建一个复数,对其调用 get_count(),并将结果转换为该像素的颜色。

优化这段代码的一些问题

正如我在即将出版的 优化计算代码的书^[3] 中所介绍的,我们可以尝试不同的方法来优化代码。这里我们不会采用一些方法:

• 消除不必要的工作: 在这种情况下会很困难 - 计算中没有太多浪费的工作。真正的 Mandelbrot 服务可以使用缓存,这样就不必重新计算常见的输出。
• 放宽结果精度: 我们可以稍微调整计算,使其不那么精确,以换取更快的结果,但目前我们想要最高质量的 Mandelbrot。

我们剩下的是尝试利用现代 CPU 功能,在单个核心内实现并行性。重要的是,这不涉及线程或进程,而是让单个线程运行得更快,因为 CPU 某种程度上会同时做多件事。

CPU 透明应用的一种机制是指令级并行性。不幸的是,算法的工作方式使这种并行性变得困难。在每次迭代中,current 计算依赖于先前计算的 current 值,所以我们最多只能在单个循环迭代内获得并行性,而不能跨迭代。

我们能用 SIMD 优化 Mandelbrot 吗?

我们将用来优化 Mandelbrot 的方法是 SIMD:单指令多数据。这些是专门的 CPU 指令,允许你例如用一条指令一次添加 4 个值,而不是通常的一次一个值。

例如,如果我们一次操作 4 个值,我们可以将两个包含 4 个 int64 的向量打包到两个特殊的 SIMD "寄存器"(用于计算的变量的特殊 CPU 内存槽)中,并用一条特殊指令将它们全部相加:

let a = i64x4::from_array([1, 2, 3, 4]);
let b = i64x4::from_array([10, 20, 30, 40]);
let sum = a + b;
assert_eq!(sum, i64x4::from_array([11, 22, 33, 44]));

现在,如果你看我们的 Mandelbrot 算法,仍然不清楚批量操作如何有帮助。我们可以一次处理 4 或 8 个像素...但对于不同的起始值,for 循环中的迭代次数会不同。所以看起来我们不能只是对所有像素不断进行相同的操作。

这就是 SIMD "掩码"发挥作用的地方:通过使用掩码,可以对不同的索引(在 SIMD 术语中称为"通道")进行不同的操作。因此,我们可以创建一个掩码,指示哪些通道我们已经处理完毕(因为它们已经发散并达到阈值),哪些通道我们需要继续处理。这意味着我们可以对不同的值执行略有不同的操作,同时仍然每个 CPU 指令处理 8 个值。

Rust 中 Mandelbrot 的 SIMD 实现

Rust 中的可移植 SIMD 仍然不稳定,所以你必须做更多工作才能使用它。首先,我们需要使用 nightly(不稳定)编译器;我在仓库中添加了一个 rust-toolchain.toml 文件:

[toolchain]
channel = "nightly"

接下来,在 src/lib.rs 中,我们可以通过在文件顶部添加以下内容来启用实验性的可移植 SIMD 功能:

#![feature(portable_simd)]

现在我们可以编写与上面看到的相同算法的 SIMD 实现:

use std::simd::*;

const LANES: usize = 8;
type F64x8 = Simd<f64, LANES>;

fn get_count(mut current: Complex<F64x8>, max_iter: usize) -> Simd<u32, LANES> {
    let mut z = Complex {
        re: F64x8::splat(0.0),
        im: F64x8::splat(0.0),
    };
    let mut count = Simd::splat(0u32);
    let max_iter = u32::try_from(max_iter).unwrap();
    let mut mask = Mask::splat(true);

    for i in 0..max_iter {
        mask &= z.norm_sqr().simd_lt(F64x8::splat(4.0));
        if !mask.any() {
            break;
        }
        z = z * z + current;
        count = mask.select(Simd::splat(i + 1), count);
    }
    count
}

同样,这里没有显示其他代码,它创建了 Complex 实例,每个实例包含 8 个值,然后将 get_count() 的结果转换为相应像素的颜色值。

接下来,我们想编译这段代码。这是一个可移植的 SIMD 库,所以它可以使用不同的 SIMD 指令。这很重要,因为不同的 CPU 型号有不同的 SIMD 指令。例如,AVX-512 SIMD 指令在一些 x86-64 CPU 上可用,但不是所有(我的电脑就缺少它们!)。

所以当我们编译代码时,第一步是告诉编译器针对当前 CPU。这意味着二进制文件不会是可移植的;在我的机器上编译的东西可能无法在 10 年前的计算机上运行。但它确实使 Rust 能够使用本地可用的最佳 SIMD 指令。

[build]
rustflags = ["-C", "target-cpu=native"]

这种可移植性缺失是一个问题,但也是可以解决的。我希望写一篇后续文章,解释如何制作可移植的二进制文件,同时仍然可以利用任何可用的 CPU 功能。

比较性能:标量 vs. SIMD

所以我们已经重写了我们的算法以使用 SIMD - 它真的更快吗?为了找出答案,我生成了 3200×3200 的图像,并测量了两个版本的运行时间:

SIMD 版本确实快得多。

使用 Rayon 利用多核并行性

到目前为止,我们一直在单核上运行。但我们确实希望尽快向用户返回结果,所以我们想利用多个核心。当然,还有一个问题是,一旦使用多个核心,我们看到的速度提升是否会持续。

由于我们使用的是 Rust, Rayon^[4]  crate 使并行化变得容易。我们将其添加为依赖项:

[dependencies]
rayon = "1.7"

现在,我们只需确保 API 可用:

use rayon::prelude::*;

我们找到驱动逻辑的核心迭代器,例如 SIMD 实现中的这个循环:

for chunk in buffer.chunks_mut(LANES) {
    // ...
}

我们将 chunks_mut() 改为 par_chunks_mut(),以并行迭代,自动使用所有 CPU 核心:

for chunk in buffer.par_chunks_mut(LANES) {
    // ...
}

就这样!我们现在有了多核并行性,使用 Rayon 管理的线程池。

以下是我在 i7-12700K CPU 上获得的性能(禁用超线程和 turboboost):

在多核情况下,我们从 SIMD 获得的加速并不像单核那么高,但它仍然给我们带来了显著的性能提升。

想看完整的、可运行的项目吗? 它在 GitHub 上^[5] 。

大局观:你可以在单核上更快

我们最初的实现是用一种快速的编译语言编写的 - 但通过切换到 SIMD,我们设法让它更快,同时仍然在单个 CPU 核心上的单个线程中运行。使用 SIMD 将我们的成本削减了 78% 到 62% 之间,即使我们切换到多核并行,它也能为我们带来好处。

我们用相当通用的 SIMD 代码做到了这一点,从标量到 SIMD 代码的转换相当简单。很可能有人可以做得更好。

使用 SIMD 当然只是一个例子。在其他情况下,其他优化技术会给你带来相同的好处:更快的结果和更低的成本。多核并行非常有用,但它不应该是你性能工具包中唯一的技巧。

参考链接

1. 不会降低我的成本: https://pythonspeed.com/articles/do-you-need-cluster-or-multiprocessing/
2. 编写 Python 扩展的优秀语言: https://pythonspeed.com/articles/intro-rust-python-extensions/
3. 优化计算代码的书: https://pythonspeed.com/products/lowlevelcode/
4. Rayon: https://docs.rs/rayon/
5. 它在 GitHub 上: https://github.com/pythonspeed/mandelbrot-simd