Rust 模拟 SIMD 指令：打造跨平台高性能计算方案

你好！咱们今天来聊聊一个硬核话题：SIMD 指令模拟。别担心，我会尽量用大白话给你讲明白，再配上 Rust 代码示例，保证你能看懂，还能上手实践。

啥是 SIMD？它有啥用？

SIMD，全称 Single Instruction, Multiple Data，单指令多数据。顾名思义，就是一条指令，同时处理多个数据。这就像你有八只手，可以同时干八件事，效率自然蹭蹭往上涨。

举个例子，你要把两个包含 1000 个数字的数组对应位置相加。普通方法，你得循环 1000 次，每次加一个数字。但有了 SIMD，如果你的 CPU 支持 256 位 SIMD 指令，一次就能处理 8 个 32 位整数（256 / 32 = 8）。理论上，速度能提升 8 倍！

所以，SIMD 在图像处理、音视频编解码、科学计算、游戏开发等领域，那可是香饽饽。能用上 SIMD，性能提升效果杠杠的。

为啥要模拟 SIMD？

既然 SIMD 这么好，为啥还要模拟呢？原因有二：

1. 硬件不支持： 不是所有 CPU 都支持 SIMD，或者支持的指令集版本不同。比如，你的 CPU 只支持 SSE2，但你想用 AVX2 的指令，那就得模拟。
2. 跨平台兼容： 你写的程序，可能要跑在不同的平台上，有的支持 AVX2，有的只支持 SSE2，甚至有的啥也不支持。为了保证程序都能跑，就得模拟。

Rust 如何模拟 SIMD？

Rust 提供了强大的底层编程能力，可以让你直接操作内存，编写高效的 SIMD 模拟代码。下面，我们就一步步来实现一个简单的例子：用 SSE2 模拟 AVX2 的 _mm256_add_epi32 指令（256 位整数加法）。

1. 准备工作

首先，你需要安装 Rust 开发环境。这个我就不多说了，网上教程一大堆。然后，创建一个新的 Rust 项目：

cargo new simd_simulation --bin
cd simd_simulation

2. 编写 Rust 代码

打开 src/main.rs 文件，输入以下代码：

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;
 
#[cfg(not(target_arch = "x86_64"))]
use std::arch::x86::*;
 
// 定义一个 256 位整数向量类型
#[derive(Clone, Copy)]
#[repr(C)]
pub struct i32x8 {
    data: [i32; 8],
}
 
// SSE2 实现的 128 位整数加法
fn add_epi32_sse2(a: __m128i, b: __m128i) -> __m128i {
    unsafe { _mm_add_epi32(a, b) }
}
 
// 模拟 AVX2 的 256 位整数加法
fn add_epi32_simulated(a: i32x8, b: i32x8) -> i32x8 {
    // 将 256 位向量拆分成两个 128 位向量
    let a_lo = unsafe { _mm_loadu_si128(a.data.as_ptr() as *const __m128i) };
    let a_hi = unsafe { _mm_loadu_si128(a.data.as_ptr().add(4) as *const __m128i) };
    let b_lo = unsafe { _mm_loadu_si128(b.data.as_ptr() as *const __m128i) };
    let b_hi = unsafe { _mm_loadu_si128(b.data.as_ptr().add(4) as *const __m128i) };
 
    // 分别进行 128 位加法
    let lo = add_epi32_sse2(a_lo, b_lo);
    let hi = add_epi32_sse2(a_hi, b_hi);
 
    // 将结果合并成 256 位向量
    let mut result = i32x8 { data: [0; 8] };
    unsafe {
        _mm_storeu_si128(result.data.as_mut_ptr() as *mut __m128i, lo);
        _mm_storeu_si128(result.data.as_mut_ptr().add(4) as *mut __m128i, hi);
    }
 
    result
}
 
fn main() {
    // 初始化两个 256 位向量
    let a = i32x8 { data: [1, 2, 3, 4, 5, 6, 7, 8] };
    let b = i32x8 { data: [8, 7, 6, 5, 4, 3, 2, 1] };
 
    // 调用模拟函数
    let result = add_epi32_simulated(a, b);
 
    // 打印结果
    println!("{:?}", result.data);
 
    // 检查结果
    let expected = [9, 9, 9, 9, 9, 9, 9, 9];
    assert_eq!(result.data, expected);
 
}

3. 代码解读

• i32x8 结构体：自定义的 256 位整数向量类型，方便操作。
• add_epi32_sse2 函数：使用 SSE2 指令实现的 128 位整数加法。
• add_epi32_simulated 函数：模拟 AVX2 的 256 位整数加法。它将 256 位向量拆分成两个 128 位向量，分别调用 add_epi32_sse2 函数进行计算，最后将结果合并。
• main 函数：测试代码，初始化两个向量，调用模拟函数，打印并验证结果。
• unsafe：由于直接操作了底层指令和内存，所以需要使用 unsafe 块。

4. 编译和运行

在命令行中运行以下命令：

cargo run

如果一切正常，你应该能看到输出：

[9, 9, 9, 9, 9, 9, 9, 9]

这表明我们的模拟函数 কাজ করছে！

如何在 JavaScript 中调用？

要把 Rust 写的 SIMD 模拟函数给 JavaScript 用，我们需要用到 WebAssembly。

1. 安装 wasm-pack

cargo install wasm-pack

2. 修改 Cargo.toml

在 Cargo.toml 文件中添加以下内容：

[lib]
crate-type = ["cdylib", "rlib"]
 
[dependencies]
wasm-bindgen = "0.2"

3. 修改 Rust 代码

修改 src/lib.rs 文件：

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;
 
#[cfg(not(target_arch = "x86_64"))]
use std::arch::x86::*;
 
use wasm_bindgen::prelude::*;
 
// 定义一个 256 位整数向量类型
#[derive(Clone, Copy)]
#[repr(C)]
pub struct i32x8 {
    data: [i32; 8],
}
 
// SSE2 实现的 128 位整数加法
fn add_epi32_sse2(a: __m128i, b: __m128i) -> __m128i {
    unsafe { _mm_add_epi32(a, b) }
}
 
// 模拟 AVX2 的 256 位整数加法
#[wasm_bindgen]
pub fn add_epi32_simulated(a: &[i32], b: &[i32]) -> Vec<i32> {
    // 将输入切片转换为 i32x8 向量
     let a_vec = i32x8 { data: [a[0],a[1],a[2],a[3],a[4],a[5],a[6],a[7]] };
    let b_vec = i32x8 { data: [b[0],b[1],b[2],b[3],b[4],b[5],b[6],b[7]] };
    // 将 256 位向量拆分成两个 128 位向量
    let a_lo = unsafe { _mm_loadu_si128(a_vec.data.as_ptr() as *const __m128i) };
    let a_hi = unsafe { _mm_loadu_si128(a_vec.data.as_ptr().add(4) as *const __m128i) };
    let b_lo = unsafe { _mm_loadu_si128(b_vec.data.as_ptr() as *const __m128i) };
    let b_hi = unsafe { _mm_loadu_si128(b_vec.data.as_ptr().add(4) as *const __m128i) };
 
    // 分别进行 128 位加法
    let lo = add_epi32_sse2(a_lo, b_lo);
    let hi = add_epi32_sse2(a_hi, b_hi);
 
    // 将结果合并成 256 位向量
    let mut result = i32x8 { data: [0; 8] };
    unsafe {
        _mm_storeu_si128(result.data.as_mut_ptr() as *mut __m128i, lo);
        _mm_storeu_si128(result.data.as_mut_ptr().add(4) as *mut __m128i, hi);
    }
 
    result.data.to_vec()
}

4. 编译 WebAssembly

wasm-pack build --target web

5. 在 JavaScript 中使用

import init, { add_epi32_simulated } from './pkg/simd_simulation.js';
 
async function run() {
  await init();
 
  const a = [1, 2, 3, 4, 5, 6, 7, 8];
  const b = [8, 7, 6, 5, 4, 3, 2, 1];
 
  const result = add_epi32_simulated(a, b);
 
  console.log(result);
}
 
run();

总结

今天，咱们一起学习了 SIMD 的基本概念，以及如何用 Rust 模拟 SIMD 指令。我们还通过 WebAssembly，实现了 Rust 与 JavaScript 的互操作。这只是 SIMD 模拟的冰山一角，还有很多更复杂的指令和优化技巧等你探索。希望这篇文章能给你带来启发，让你在高性能计算的道路上更进一步！