如何在 Rust 中使用 SIMD 进行并行处理？

2 July 2024 — developer language

SIMD（单指令多数据）是加速高性能计算中数据密集型操作的强大工具。虽然我们之前探讨了如何使用 Rayon 进行线程级别的并行处理，但 SIMD 可以在单个核心内实现并行处理，同时操作多个数据点。理解并利用 SIMD 对于最大化代码性能至关重要。

截至 2024 年中，Rust 提供了多种 SIMD 开发途径。虽然标准库中的实验性 SIMD 模块（std::simd）仅在 nightly 版本中可用，但稳定版 Rust 提供了几种在生产环境中使用 SIMD 的选择：

Rust 编译器的自动矢量化功能（自动将循环等代码优化为 SIMD 指令）
通过 std::arch 模块使用特定平台的内置函数
Rust 在 std::simd 中的实验性 SIMD 实现

每种方法在性能、可移植性和易用性方面各有优劣，我们将在本文中详细探讨。

我们将重点介绍在稳定版 Rust 中实际可用的 SIMD 技术。本文将涵盖以下内容：

SIMD 操作的基本原理
利用编译器自动矢量化功能
使用 std::arch 中的特定平台内置函数
带有性能考虑的 SIMD 实际应用示例
有效使用 SIMD 的最佳实践和注意事项

本文结束时，希望你能够掌握如何在 Rust 代码中利用 SIMD 提升数据并行操作的性能，使用稳定版特性。你将具备在项目中合理使用 SIMD 的知识，平衡性能提升与代码的可维护性和可移植性。

理解 SIMD

当我们谈论 SIMD（单指令多数据）时，我们是在探讨现代 CPU 架构的一个强大特性。Rust 开发者可以利用这一特性，在从高性能计算到嵌入式系统的各种应用中获得显著的性能提升。

SIMD 范式及其基础

从本质上讲，SIMD 利用 CPU 内部的特殊宽寄存器（register）。这些寄存器并不是普通的寄存器；它们是可以同时存储多个值的大型数据持有者。想象一个 256 位寄存器——它就像一个容器，可以一次性容纳八个 32 位浮点数或十六个 16 位整数。这种能力是 SIMD 强大的基础，也是 Rust 的 SIMD 功能所要充分利用的。

不同架构有各自的 SIMD 指令集：

x86/x64: SSE, AVX
ARM: NEON
RISC-V: 向量扩展

这些专门的 CPU 指令旨在同时操作多个数据点。例如，单个 SIMD 指令可以一次性完成四对数字的加法，而标量方法则需要进行四次单独的加法操作。我们可以通过平台特定的内置函数，或者通过支持 SIMD 的 crates 在 Rust 中访问这些指令。

Rust 开发者的好处和应用

对于处理数据密集型操作的 Rust 项目来说，使用 SIMD 的好处非常明显。在理想情况下，性能提升可以直接与 SIMD 寄存器的宽度成正比。Rust 应用中的实际性能增益通常较为温和，但仍然显著，一般比标量代码快两到四倍。

SIMD 不仅仅是提升速度；它还提高了效率。通过用更少的指令处理更多数据，SIMD 可以减少功耗和内存带宽的使用。这种效率使得 SIMD 在以下 Rust 应用领域特别有吸引力：

高性能计算：科学模拟、机器学习、金融建模
多媒体处理：图像/视频处理、音频分析、3D 图形
系统编程：网络数据包处理、文件系统、数据库引擎
嵌入式系统：实时信号处理、传感器数据融合、控制系统
密码学和安全：加密/解密、哈希、安全通信

在这些领域中，同时处理多个数据点可以显著提升性能和资源利用率。这使得 Rust 开发者能够利用现代硬件功能创建高性能应用程序，从服务器级 CPU 到微控制器。

Rust 实现的考虑事项

虽然 SIMD 提供了显著的好处，但 Rust 开发者需要了解其局限性和注意事项：

数据对齐（Data Alignment）：SIMD 操作通常需要正确对齐的数据以获得最佳性能。Rust 的类型系统和内存布局控制可以帮助确保正确对齐。
可移植性：不同 CPU 架构支持不同的 SIMD 指令集。Rust 的 cfg 属性和条件编译可以帮助管理这一点，允许在不支持的架构上使用备用实现。
复杂性：SIMD 代码可能更难编写和维护。Rust 的抽象机制，无论是通过标准库的 std::arch 还是第三方 crates，都旨在减少这种复杂性。
适用性：SIMD 对于在大数据集上执行相同操作的算法最为有效。并非所有问题都适合 SIMD 优化。
资源限制：嵌入式系统中的 SIMD 功能可能有限。Rust 的零成本抽象有助于在不增加额外开销的情况下利用 SIMD。
测试和验证：SIMD 优化可能引入微妙的错误。Rust 的强类型系统和测试框架在确保代码正确性方面非常宝贵。

我们将探讨如何有效应对这些考虑事项，利用 Rust 的特性编写高效、可移植和易维护的 SIMD 代码，适用于从高性能服务器到资源受限的嵌入式系统的各种领域。

理解 Rust 中的自动矢量化

虽然显式 SIMD 编程可以对矢量化进行精细控制，但由 LLVM 驱动的 Rust 编译器在某些条件下可以自动矢量化代码。这一功能称为自动矢量化，使开发者可以编写简单的标量代码，编译器则在可能的情况下将其转换为 SIMD 指令。

自动矢量化如何工作

自动矢量化是一种优化技术，编译器会分析循环，并在安全且有益的情况下将其转换为 SIMD 指令。这个过程在编译阶段进行，不需要开发者显式编写 SIMD 代码。

编写有利于自动矢量化的代码

虽然编译器的自动矢量化能力非常强大，但某些编码实践可以增加成功矢量化的可能性：

使用简单、直接的循环，而不是复杂的控制流。
确保数据访问模式是可预测的，最好是连续的。
避免在循环中调用无法内联的函数。
最小化循环迭代之间的依赖性。
启用优化编译（例如，cargo build –release）。

示例：矩阵行累加和

让我们来看一个计算矩阵每行累加和的函数实现：

fn matrix_row_cumsum(matrix: &[&[f64]]) -> Vec<Vec<f64>> {
    matrix.iter().map(|row| {
        let mut cumsum = 0.0;
        row.iter().map(|&x| {
            cumsum += x;
            cumsum
        }).collect()
    }).collect()
}

这个函数使用了迭代器和闭包，通常被认为是典型的 Rust 代码。由于其简单性和没有显式的循环依赖性，这段代码可能是自动矢量化的理想候选。

自动矢量化的现实

为了检查编译器如何处理这段代码，我们可以使用 Rust Playground，这是一个允许我们查看 Rust 代码汇编输出的在线工具。你可以在 Rust Playground 查看这个例子。

在查看汇编输出时，我们可能期望看到矢量化的明显迹象，例如使用 SIMD 指令。然而，自动矢量化的现实往往更加复杂：

矢量指令的存在并不总是意味着有效的矢量化。
编译器可能会以意想不到的方式或对意想不到的代码部分应用矢量化。
自动矢量化的结果在很大程度上依赖于编译器版本、优化级别和目标架构。

对 Rust 开发者的影响

相信编译器：由 LLVM 支持的 Rust 编译器非常强大。它可能会发现开发者无法立即察觉的矢量化机会。
专注于清晰、惯用的代码：与其试图通过“矢量化友好”的代码来智胜编译器，不如专注于编写清晰、惯用的 Rust 代码。编译器通常能够有效地优化写得好的简洁代码。
进行性能基准测试：由于自动矢量化结果可能不可预测，持续使用实际数据集对代码进行基准测试，以衡量实际性能提升。
在必要时使用显式 SIMD：对于需要保证 SIMD 操作的性能关键部分，使用 Rust 的 SIMD 内置函数或库进行显式 SIMD 编程。
了解目标架构：自动矢量化结果可能因 CPU 架构而异。在优化代码时，请考虑你的目标平台。

验证自动矢量化

虽然 Rust Playground 可以通过查看汇编输出提供潜在矢量化的见解，但它并不总是能告诉你全部情况。为了全面了解：

使用性能分析工具分析代码的运行时行为。
使用大规模、现实的数据集对代码进行基准测试。
在不同的架构和不同的编译器版本上进行测试。

记住，有效矢量化的最终衡量标准是在实际应用中的性能提升，而不仅仅是汇编代码中出现的矢量指令。

使用 std::arch 的特定平台内置函数

虽然自动矢量化提供了一种免动手的 SIMD 方法，但 Rust 还通过特定平台的内置函数（intrinsics）提供了更直接的控制。标准库中的 std::arch 模块提供了对特定 CPU 架构 SIMD 指令的低级访问。这种方法提供了最大的性能，但需要仔细处理跨平台兼容性。

理解 std::arch

std::arch 模块包含不同架构的子模块。对于基于 ARM 的系统（如 M1/M2 MacBook），我们特别关注 std::arch::aarch64 模块，它提供对 ARM NEON SIMD 指令的访问。

我们将探索一个使用 ARM NEON 内置函数实现音频回声效果的例子。这个例子展示了如何利用特定平台的 SIMD 指令进行音频处理任务。请记住，大多数数据中心硬件使用 x86。我出于兴趣探索 ARM NEON，因为我使用的是搭载 ARM 处理器的 MacBook。

use std::arch::aarch64::*;

// Echo parameters
const DELAY_SAMPLES: usize = 11025; // 0.25 second at 44.1kHz
const ECHO_ATTENUATION: f32 = 0.6;

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn process_samples_neon(
    input: &[f32],
    output: &mut [f32],
    delay_line: &mut [f32],
    delay_index: &mut usize,
) {
    let attenuation = vdupq_n_f32(ECHO_ATTENUATION);

    for (i, &sample) in input.iter().enumerate() {
        let current = vdupq_n_f32(sample);
        let delayed = vdupq_n_f32(delay_line[*delay_index]);

        let echo = vmulq_f32(delayed, attenuation);
        let result = vaddq_f32(current, echo);

        let mut result_array = [0.0f32; 4];
        vst1q_f32(result_array.as_mut_ptr(), result);

        output[i] = result_array[0];
        delay_line[*delay_index] = sample;
        *delay_index = (*delay_index + 1) % DELAY_SAMPLES;
    }
}

让我们分解这个 NEON 优化函数的关键元素：

配置和安全性

#[cfg(target_arch = "aarch64")] 确保此函数仅在 ARM64 架构上编译。
#[target_feature(enable = "neon")] 表示该函数需要 NEON 指令支持。
由于我们使用了底层 SIMD 指令（low-level SIMD instructions），因此必须使用 unsafe 关键字。

NEON 内置函数

vdupq_n_f32: 创建一个所有通道都设置为相同值的向量。
vmulq_f32: 执行两个向量的元素逐个相乘。
vaddq_f32: 执行两个向量的元素逐个相加。
vst1q_f32: 将一个向量存储到内存中。

处理循环

每个样本单独处理，对当前和延迟样本的向量化版本应用 NEON 操作。

为了安全地使用这个 NEON 优化函数，我们提供了一个运行时检查 NEON 支持的包装器：

fn process_samples(
    input: &[f32],
    output: &mut [f32],
    delay_line: &mut [f32],
    delay_index: &mut usize,
) {
    #[cfg(target_arch = "aarch64")]
    {
        if std::arch::is_aarch64_feature_detected!("neon") {
            unsafe {
                return process_samples_neon(input, output, delay_line, delay_index);
            }
        }
    }
    process_samples_fallback(input, output, delay_line, delay_index);
}

这个包装器使用 is_aarch64_feature_detected!("neon") 在运行时检测 NEON 支持，如果 NEON 不可用，则回退到标量实现。

何时使用特定平台内置函数

在以下情况下考虑使用 std::arch 内置函数：

需要在特定架构上保证 SIMD 性能。
自动矢量化未能提供所需的性能。
你正在处理性能关键代码，并且复杂性的权衡是值得的。
针对特定平台，可以充分利用其 SIMD 功能。

在我们的音频处理示例中，使用 NEON 内置函数可以在回声效果计算中带来潜在的性能提升。然而，重要的是要将这种实现与标量版本进行基准测试，以确保增加的复杂性能够带来实际的性能提升。

以下是完整的示例：

use std::arch::aarch64::*;
use std::error::Error;

// Echo parameters
const DELAY_SAMPLES: usize = 11025; // 0.25 second at 44.1kHz
const ECHO_ATTENUATION: f32 = 0.6;

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn process_samples_neon(
    input: &[f32],
    output: &mut [f32],
    delay_line: &mut [f32],
    delay_index: &mut usize,
) {
    let attenuation = vdupq_n_f32(ECHO_ATTENUATION);

    for (i, &sample) in input.iter().enumerate() {
        let current = vdupq_n_f32(sample);
        let delayed = vdupq_n_f32(delay_line[*delay_index]);

        let echo = vmulq_f32(delayed, attenuation);
        let result = vaddq_f32(current, echo);

        let mut result_array = [0.0f32; 4];
        vst1q_f32(result_array.as_mut_ptr(), result);

        output[i] = result_array[0];
        delay_line[*delay_index] = sample;
        *delay_index = (*delay_index + 1) % DELAY_SAMPLES;
    }
}

fn process_samples_fallback(
    input: &[f32],
    output: &mut [f32],
    delay_line: &mut [f32],
    delay_index: &mut usize,
) {
    for (i, &sample) in input.iter().enumerate() {
        let echo = delay_line[*delay_index] * ECHO_ATTENUATION;
        output[i] = sample + echo;
        delay_line[*delay_index] = sample;
        *delay_index = (*delay_index + 1) % DELAY_SAMPLES;
    }
}

fn process_samples(
    input: &[f32],
    output: &mut [f32],
    delay_line: &mut [f32],
    delay_index: &mut usize,
) {
    #[cfg(target_arch = "aarch64")]
    {
        if std::arch::is_aarch64_feature_detected!("neon") {
            unsafe {
                println!("Using NEON instructions");
                return process_samples_neon(input, output, delay_line, delay_index);
            }
        }
    }
    println!("Using fallback instructions");
    process_samples_fallback(input, output, delay_line, delay_index);
}

fn main() -> Result<(), Box<dyn Error>> {
    // Open the input WAV file
    let mut reader = hound::WavReader::open("input.wav")?;
    let spec = reader.spec();

    // Read samples and convert to f32
    let samples: Vec<f32> = match spec.sample_format {
        hound::SampleFormat::Float => reader.samples::<f32>().map(|s| s.unwrap()).collect(),
        hound::SampleFormat::Int => reader
            .samples::<i16>()
            .map(|s| s.unwrap() as f32 / i16::MAX as f32)
            .collect(),
    };

    // Prepare output buffer and delay line
    let mut output = vec![0.0f32; samples.len()];
    let mut delay_line = vec![0.0f32; DELAY_SAMPLES];
    let mut delay_index = 0;

    // Process samples
    process_samples(&samples, &mut output, &mut delay_line, &mut delay_index);

    // Prepare the output WAV file
    let mut writer = hound::WavWriter::create("output.wav", spec)?;

    // Write processed samples
    match spec.sample_format {
        hound::SampleFormat::Float => {
            for &sample in &output {
                writer.write_sample(sample)?;
            }
        }
        hound::SampleFormat::Int => {
            for &sample in &output {
                writer.write_sample((sample.clamp(-1.0, 1.0) * i16::MAX as f32) as i16)?;
            }
        }
    }

    writer.finalize()?;

    println!("Echo effect applied and saved to output.wav!");
    Ok(())
}

我对这两种实现进行了基准测试，发现 SIMD 优化代码和标量版本之间的性能差异微乎其微。在我的 M1 MacBook 上，这两种实现都能快速处理音频。虽然这个例子并不实用——在这种情况下，我建议让编译器自行优化——但希望它能展示 SIMD 提升性能关键代码速度的潜力，这在更密集的任务中可能更有用。

探索 std::simd：Rust 中可移植 SIMD 的未来

在使用特定平台内置函数之后，我对 Rust 中更可移植的 SIMD 解决方案产生了兴趣。这引导我探索 std::simd，这是 Rust 标准库中的一个实验模块，旨在提供 SIMD 操作的可移植抽象。

截至 Rust 1.79（我当前的版本），std::simd 仍然是一个不稳定的特性，这意味着它仅在 nightly 频道可用，需要明确选择加入。尽管它处于实验阶段，std::simd 代表了 Rust SIMD 编程的一个令人兴奋的方向，承诺将 SIMD 的性能优势与 Rust 对可移植性和安全性的承诺结合起来。

要使用 std::simd，我需要切换到 nightly 频道：

rustup default nightly

std::simd 的核心理念是提供跨不同架构工作的 SIMD 向量类型和操作。与其编写特定架构的内置函数，你可以编写更通用的 SIMD 代码，由编译器针对目标架构进行优化。

让我们重新审视我们的音频回声效果示例，这次使用 std::simd：

#![feature(portable_simd)]
use std::simd::*;

use std::error::Error;

const DELAY_SAMPLES: usize = 11025; // 0.25 second at 44.1kHz
const ECHO_ATTENUATION: f32 = 0.6;

fn process_samples(
    input: &[f32],
    output: &mut [f32],
    delay_line: &mut [f32],
    delay_index: &mut usize,
) {
    let attenuation = f32x4::splat(ECHO_ATTENUATION);

    for (i, &sample) in input.iter().enumerate() {
        let current = f32x4::splat(sample);
        let delayed = f32x4::splat(delay_line[*delay_index]);

        let echo = delayed * attenuation;
        let result = current + echo;

        output[i] = result[0];
        delay_line[*delay_index] = sample;
        *delay_index = (*delay_index + 1) % DELAY_SAMPLES;
    }
}

fn main() -> Result<(), Box<dyn Error>> {
    // Open the input WAV file
    let mut reader = hound::WavReader::open("input.wav")?;
    let spec = reader.spec();

    // Read samples and convert to f32
    let samples: Vec<f32> = match spec.sample_format {
        hound::SampleFormat::Float => reader.samples::<f32>().map(|s| s.unwrap()).collect(),
        hound::SampleFormat::Int => reader
            .samples::<i16>()
            .map(|s| s.unwrap() as f32 / i16::MAX as f32)
            .collect(),
    };

    // Prepare output buffer, delay line and previous sample
    let mut output = vec![0.0f32; samples.len()];
    let mut delay_line = vec![0.0f32; DELAY_SAMPLES];
    let mut delay_index = 0;

    process_samples(&samples, &mut output, &mut delay_line, &mut delay_index);

    // Prepare the output WAV file
    let mut writer = hound::WavWriter::create("output.wav", spec)?;

    // Write processed samples
    match spec.sample_format {
        hound::SampleFormat::Float => {
            for &sample in &output {
                writer.write_sample(sample)?;
            }
        }
        hound::SampleFormat::Int => {
            for &sample in &output {
                writer.write_sample((sample.clamp(-1.0, 1.0) * i16::MAX as f32) as i16)?;
            }
        }
    }

    writer.finalize()?;

    println!("Echo effect applied and saved to output.wav!");
    Ok(())
}

在这个实现中，我使用了 f32x4，它表示一个包含四个 32 位浮点数的向量。与内置函数版本相比，这些操作更像标准的 Rust 代码，我觉得这种方式更直观，也更容易阅读。

然而，需要注意的是，由于 std::simd 是一个不稳定特性，这段代码不能在稳定版 Rust 上编译。这只是对未来 Rust 中 SIMD 编程可能样子的一个预览，而不是我们今天可以在生产环境中使用的东西。

结论：Rust 中 SIMD 的现状

在总结我对 Rust 中 SIMD 编程的探索时，我有一些混合的印象和见解。Rust 编译器在没有显式 SIMD 指令的情况下优化代码的能力让我印象深刻。在许多情况下，特别是对于简单操作，编译器的优化效果与手写的 SIMD 代码相当。这让我更确信编写清晰、惯用的 Rust 代码，并相信编译器的优化能力的重要性。

特定平台内置函数则稍微复杂一些。虽然内置函数提供了对 SIMD 操作的精细控制，但我发现性能提升并不总是如预期的那样显著，特别是对于简单计算。这次经历强调了基准测试的重要性，以及在潜在性能提升与增加的复杂性之间仔细权衡的必要性。在大多数情况下，编写 SIMD 代码并不是必要的。

尽管 std::simd 仍然处于实验阶段，这次探索让我一窥 Rust 中 SIMD 的光明未来。编写可移植的 SIMD 代码并在不同架构上优化的前景令人兴奋，尽管这种方法目前还不能用于生产。

尽管 Rust 中的 SIMD 编程提供了强大的性能优化工具，但它并不是万能的。最有效的方法通常是结合相信编译器的自动矢量化，并在特定平台内置函数提供明确优势的地方有选择地使用它们。