WASM SIMD 高性能计算实战

(module
  (func $vadd (param $a v128) (param $b v128) (result v128)
    local.get $a
    local.get $b
    f32x4.add)             ;; 4 个 f32 同时相加
  (export "vadd" (func $vadd)))

// 运行时检测 SIMD 支持
import { simd } from "wasm-feature-detect";
if (await simd()) {
  console.log("当前引擎支持 WASM SIMD");
}

[dependencies]
# 使用 nightly 的 portable-simd
# RUSTFLAGS="-C target-feature=+simd128"

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

#[target_feature(enable = "simd128")]
pub fn sum_f32x4(a: &[f32], b: &[f32], out: &mut [f32]) {
    assert!(a.len() % 4 == 0);
    let chunks = a.len() / 4;
    for i in 0..chunks {
        let va = Simd::<f32, 4>::from_slice(&a[i * 4..]);
        let vb = Simd::<f32, 4>::from_slice(&b[i * 4..]);
        let vc = va + vb;
        vc.copy_to_slice(&mut out[i * 4..]);
    }
    // 开启 +simd128 时，编译为 f32x4.add
}

#[cfg(target_feature = "simd128")]
fn fast_path(x: &[f32]) -> f32 {
    // 向量版本
}
#[cfg(not(target_feature = "simd128"))]
fn fast_path(x: &[f32]) -> f32 {
    // 标量版本
}

RUSTFLAGS="-C target-feature=+simd128" cargo build --target wasm32-unknown-unknown --release
# 或用于 WASI
RUSTFLAGS="-C target-feature=+simd128" cargo build --target wasm32-wasip1 --release

#include <wasm_simd128.h>
#include <stdint.h>

// 4 个 f32 点积加速
v128_t dot4(const float *a, const float *b) {
    v128_t va = wasm_v128_load(a);
    v128_t vb = wasm_v128_load(b);
    return wasm_f32x4_mul(va, vb);
}

// 16 个 u8 的逐字节 clamp（图像处理核心）
v128_t clamp_u8(v128_t v) {
    return wasm_u8x16_max(v, wasm_i8x16_splat(0));
}

// shuffle：任意通道重排
v128_t swap_pairs(v128_t v) {
    return wasm_i8x16_shuffle(v, v, 1, 0, 3, 2, 5, 4, 7, 6,
                                    9, 8, 11, 10, 13, 12, 15, 14);
}

clang --target=wasm32-wasi \
  -msimd128 \
  -O3 \
  -o image.wasm image.c
# -msimd128 启用 SIMD 指令生成

// 使用整数定点逼近，避免浮点：Y ≈ (77*R + 150*G + 29*B) >> 8
#include <wasm_simd128.h>

void rgba_to_gray_u8(const uint8_t *src, uint8_t *dst, size_t n_pixels) {
    // 每 v128 = 4 像素（RGBA×4）
    const v128_t w_r = wasm_i16x8_splat(77);
    const v128_t w_g = wasm_i16x8_splat(150);
    const v128_t w_b = wasm_i16x8_splat(29);

    for (size_t i = 0; i < n_pixels; i += 4) {
        v128_t px = wasm_v128_load(src + i * 4);

        // 拆出 R/G/B 通道（每通道 16 个 u8，shuffle 提取）
        v128_t r = wasm_i8x16_shuffle(px, px, 0,4,8,12, 0,0,0,0, 0,0,0,0, 0,0,0,0);
        v128_t g = wasm_i8x16_shuffle(px, px, 1,5,9,13, 0,0,0,0, 0,0,0,0, 0,0,0,0);
        v128_t b = wasm_i8x16_shuffle(px, px, 2,6,10,14,0,0,0,0, 0,0,0,0, 0,0,0,0);

        // 扩到 u16 计算，避免溢出
        v128_t r16 = wasm_i16x8_extend_low_u8x16(r);
        v128_t g16 = wasm_i16x8_extend_low_u8x16(g);
        v128_t b16 = wasm_i16x8_extend_low_u8x16(b);

        v128_t y = wasm_i16x8_add(wasm_i16x8_add(wasm_i16x8_mul(r16, w_r),
                                                 wasm_i16x8_mul(g16, w_g)),
                                  wasm_i16x8_mul(b16, w_b));
        y = wasm_u16x8_shr(y, 8);   // >> 8

        // 写回（shr 后每通道 16 位，pack 成 8 位）
        v128_t packed = wasm_u8x16_narrow_i16x8(y, y);
        // 存前 4 字节
        uint32_t val = wasm_i32x4_extract_lane(packed, 0);
        __builtin_memcpy(dst + i, &val, 4);
    }
}

use std::simd::prelude::*;

#[target_feature(enable = "simd128")]
pub fn mix(a: &[f32], b: &[f32], out: &mut [f32], gain_b: f32) {
    let g = Simd::<f32, 4>::splat(gain_b);
    let mut i = 0;
    while i + 4 <= a.len() {
        let va = Simd::<f32, 4>::from_slice(&a[i..]);
        let vb = Simd::<f32, 4>::from_slice(&b[i..]);
        let vout = va + vb * g;
        vout.copy_to_slice(&mut out[i..]);
        i += 4;
    }
    // 尾数（非 4 倍数）走标量
    for j in i..a.len() {
        out[j] = a[j] + b[j] * gain_b;
    }
}

// BLAKE3 的压缩函数核心——4 路并行
#[target_feature(enable = "simd128")]
pub unsafe fn compress4(
    state: &[u32; 16],
    block: &[u8; 64],
    counter: u64,
    block_len: u32,
    flags: u32,
    out: &mut [u32; 16],
) {
    // 把 4 个不同块的 state/block 按列装入 4 个 i32x4 向量
    // 一次 i32x4.add 同时推进 4 个独立哈希
    // ...（G 函数 → 对角化 → 置换）
}

use std::simd::prelude::*;

#[target_feature(enable = "simd128")]
pub fn matmul4(a: &[[f32; 4]; 4], b: &[[f32; 4]; 4]) -> [[f32; 4]; 4] {
    let b_cols: [Simd<f32, 4>; 4] = (0..4).map(|j| {
        Simd::<f32, 4>::from_slice(&[b[0][j], b[1][j], b[2][j], b[3][j]])
    }).collect::<Vec<_>>().try_into().unwrap();

    let mut out = [[0.0f32; 4]; 4];
    for i in 0..4 {
        let row = Simd::<f32, 4>::from_slice(&a[i]);
        for j in 0..4 {
            // 一行 × 一列 = 一个 f32x4 点积
            let product = row * b_cols[j];
            out[i][j] = product.reduce_sum();
        }
    }
    out
}