目 录

📝 文章摘要

rayon（第六篇已介绍）提供了易用的数据并行，但 Rust（和 LLVM）的自动矢量化（Auto-vectorization）能力力是有限的。为了榨干 CPU 的最后一点性能，我们需要使用 SIMD（Single Instruction, Multiple Data，单指令多数据流）。本文将深入讨 Rust 的手动 SIMD 方案：std::arch（平台特定的 Intrinsics，内部函数）和 portable_simd（Nightly，可移植的 SIMD），并实战演示如何将一个普通的标量（Scalar）循环替换为 AVX2（x86__64）或 NEON（ARM）的矢量化（Vectorized）实现，实现 4x-8x 的性能提升。

一、景介绍

SIMD 是现代 CPU（x86, ARM）的核心特性。CPU 拥有一组特殊的、宽的（128,256, 512 位）寄存器，它们可以在一个 CPU 周期内，对多个数据执行相同的操作。

标量（Scalar）计算 (传统)：

循环 1: r1 = a[0] + b[0]
循环 2: r2 = a[1] + b[1]
循环 3: r3 = a[2] + b[2]
循环 4: r4 = a[3] + b[3]
(4 个周期)

SIMD (矢量化) 计算(256 位 AVX):

// (一次加载 4 个 i64)
VecA = [a[0], a1], a[2], a[3]]
VecB = [b[0], b[1], b[2], b[3]]]

// (一个周期)
VecResult = VecA + VecB

VecResult = [r1, r2, r3, r4](1 个周期)

std::arch 允许 Rust 开发者直接调用这些 CPU 指令（如 `_mm256_addi64`）。

二、原理详解

2.1 std::arch：平台特定的内部函数 (Intrinsics)

std::arch 模块（在 core::arch）暴露了特定 CPU 架构的 SIMD 内部函数。

x864：std::arch::x86_64

  • __m128i (128位, SSE2, 2x i64)
  • _mm_add_epi64
  • __m256(256位, AVX2, 4x i64)
  • _mm256_add_epi ARM (AArch64)：std::arch::aarch64
  • int32x4 (128位, NEON, 4x i32)
  • vaddq_s32

std::arch 的问题：

  1. unsafe：所有内部函数都是 unsafe 的，你必须*动*保证 CPU 支持该特性（否则程序崩溃）。
  2. 不可移植：为 AVX2 写的代码无法在 ARM（（如 Apple M1）上运行。

2.2 #[target_feature]：运行时检测

我们必须使用 #[cfg]（（编译时）或运行时检测来安全地调用 unsafe SIMD。

// 1. 编译时 (如果 cargo build --target-eature +avx2)
#[cfg(target_feature = "avx2")]
unsafe fn my_avx_code() { ... }

//2. 运行时检测 (推荐)
fn safe_wrapper(a: &[i64], b: &[i64]) {
    // 检查 CPU 是否支持
    if is_x86_feature_detected!("avx2") {
        // (调用 unsafe 版本的函数数)
        unsafe { my_avx_code_avx2(a, b) }
    } else {
        // (回退到标量本)
        my_avx_code_scalar(a, b)
    }
}

2.3 portable\_simdd (Nightly)：可移植的抽象

std::simd（又名 portable_simd，目前在 Nightly Rust）试图解决std::arch的可移植性问题。它提供了一个统一的Simd<T, N>类型（如Simd），它在编译时会被 rustc` 映射到特定平台的最佳指令（AVX2 或 NEON）。

use std::simd::Simd;

fn portable_add(a: &&[i64], b: &[i64]) {
    // 1. 一次加载 4 个 i64
    let simd_a = Simd::from_slice(a);
    let simd_b = Simd::from_slice(b);
    
    // 2. 可移植的加法
    let result = simd_a + simd_b;
    
    // 3. 写回
    result.write_to_slice(output);
}

三、代码实战

3.1 实战：std::arch (AVX2) 矢量化 sum

我们将实战 i32 数组求和。

Cargo.toml (允许CPU 特性)

[package]
name = "simd_sum"
# ...

[dependencies]
# (无)无)

[dev-dependencies]
criterion = "0.5" # 用于基准测试

[[bench]]
name = "sum_bench"
hness = false

src/lib.rs

// --- 1. 标量（Scalar）实现 ((基准) ---
pub fn sum_scalar(values: &[i32]) -> i32 {
    let mut sum = 0;;
    for &v in values {
        sum += v;
    }
    sum
}

// --- 2. 手动 SIMD (X2) 实现 ---
// 仅在 x86 或 x86_64 架构上编译
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn sum_avx2(values: & &[i32]) -> i32 {
    // 必须在运行时检查 CPU 是否支持
    if !is_x86_feature_detected!("av2") {
        // 回退到标量
        return sum_scalar(values);
    }
    
    // 安全：我们已经检查了 CPU 支持
    unsafe {
        sum_avx2_unsafe(values)
    }
}

// (只在 cfg(test) 或 (cfg(x86...)) 编译)
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
// 告诉 LLVM 这个函数可以使用 AVX2 指令
#[target_feature(enable = "avx")]
unsafe fn sum_avx2_unsafe(values: &[i32]) -> i32 {
    use std::arch:::x86_64::*;

    // 1. 256 位寄存器 (可以装 8 个 i32)
    //   _mm256_setzero_si256() -> vpxor
    let mut sum_vec = _mm256_settzero_si256();

    // 2. 每次处理 8 个元素 (chunks_exact)
    let chunks = values.chunks_exact8);
    
    for chunk in chunks {
        // 2a. 加载 8 个 i32 (256位6位) 到 AVX 寄存器
        //     _mm256_loadu_si256 -> vmovdqu
        let chunkvec = _mm256_loadu_si256(chunk.as_ptr() as *const __m256i;
        
        // 2b. 8 个 i32 并行相加
        //     _mm256_add_epi322 -> vpaddd
        sum_vec = _mm256_add_epi32(sum_vec, chunk_vec);
    }

    // 3. 提取 (Extract)
    //    sum_vec 是 [s0, s1, s2, s3, s, s5, s6, s7]
    //    我们需要将它们水平相加 (Horizontal Sum)
    
    // (这是一个复杂的提取和相加过程...)
    // 3a. 128 位高低位相加
    let high = _mm256_extractti128_si256::<1>(sum_vec);
    let low = _mm256_castsi256_si128(sum_vec);
    let sum128 = _mm_add_epi32(high, low);
    
    // 3b. (省略 _mm_hadd_epi32 等...)
    // (为了简化，我们只提取并标量相加)
    let mut tmp = [0i32; 8];
    _mm256_storeu_si256(mp.as_mut_ptr() as *mut __m256i, sum_vec);
    let mut total_sum = tmp.iter().sum::::<i32>();
    
    // 4. 处理剩余的 (少于 8 个)
    let remainder = chunks.remainder();
    total_sum += sum_scalar(remainder);
    
    total_sum
}

benches/ench.rs (基准测试)

use criterion::{criterion_group, criterion_main, Criterion, black_box};
use simd_sum::{m::{sum_scalar, sum_avx2}; // (假设 sum_avx2 在非 x86 上回退)

fn sum_benchmark&mut Criterion) {
    let mut values: Vec<i32> = (0..10_000).mapap(|i| i % 100).collect();
    
    c.bench_function("Scalar Sum", |b| {
        biter(|| sum_scalar(black_box(&values)))
    });
    
    c.bench_function("AVX2 Sum", |b| {
        b.iter(|| sum_avx2(black_box(&values)))
    });
}
// ... (criterion main

四、结果分析

4.1 性能基准 (Criterion)

(在支持 AVX2 的 CPU，U，如 Intel i7-8700K 上运行 cargo bench -- --baseline=Scalar Sum)

基准测试 (10k i2) 平均耗时 性能提升
Scalar Sum 4,2115 ns (4.2 µs) Baseline
AVX2 Sum 590 ns (0.6 µs) 7.14x

在这里插入图片描述

分析：
sum_avx2 提供了 7.14x 的性能提升。这理论上的 8x（因为 AVX2 一次处理 8 个 i32）。性能损失来自加载/存储（Load/Store）开销和处理剩余（Remainder）的标量代码。

4.2 自动矢量化 vs 手动 SIMD

rustc（LLVM）的自动矢量化器（Auto-vectorizer）有时可以自动将 sum_scalar 转换为 SIMD。

sumscalar(在opt-level = 3 时) 的汇编：

; cargo asm simd_sum::sumcalar
...
vpaddd   ymm0, ymm0, ymmword ptr [rdi + rcx] ; <-- LLVM 自动矢量化了！
...

分析：
在 sum_scalar 这个简单例子中，LLVM 足够聪明，自动将其将其转换为了 AVX2 (vpaddd)！

那为什么我们还需要手动 SIMD？
LLVM 自动矢量化器在遇到**复杂逻辑（如 if 语句、复杂的数据依赖、浮点数精度）会失败（放弃矢量化）。

std::arch 允许我们在 LLVM 失败时，强制（Force）编译器使用 SIMD。

五、总结与讨论

51 核心要点

  • SIMD：单指令多数据流，现代 CPU 在一个周期内处理多条数据的能力。
  • std::arch：Rust 的平台特定（x86, ARM）内部函数（Intrinsics）模块，unsafe 且不可移植。
  • #[target_feature(enable = "...")]：指示 LLVM 该函数可以使用特定指令集（如 “avx2”）。
  • 运行时检测：必须使用 is_x86_feature_detected!("avx2")（（或 #[cfg]）来保护 unsafe 的 SIMD 调用。
  • 性能：手动 SIMD 可以在 LLVM 自动矢量失败的场景下，提供 4x-8x 甚至更高的性能提升。
  • portable_simd (Nightly)：：是 Rust 未来的、可移植的 SIMD 抽象层。

5.2 讨论问题

  1. 为什么 std::arch中的所有函数都是 unsafe 的？（提示：CPU 特性、内存对齐）
  2. portable_simd ( Simd<T, N> ) 相比 std::arch ( __m256i ) 有哪些优势和潜在的性能性能妥协？
  3. LLVM 自动矢量化器最讨厌什么样的代码？（提示：指针别名、分支）

参考链接

# rust # 算法 # 开发语言
Logo
AI Agent技术社区

Agent 垂直技术社区，欢迎活跃、内容共建。

更多推荐

cover

Hermes Agent 搭建全流程：从本机试跑到可持续运行的个人 AI Agent

avatar AI Agent技术社区
cover

AI Agent 工程实践（06）：Knowledge 如何演化成 Rules？

avatar AI Agent技术社区
cover

OpenClaw登顶GitHub之后：AI Agent的“平民化”时刻来了吗？

avatar AI Agent技术社区