WebAssembly+Rust：前端高性能计算（矩阵乘法）

2401_83640238

242人浏览 · 2025-11-01 20:42:41

2401_83640238 · 2025-11-01 20:42:41 发布

WebAssembly + Rust：前端高性能矩阵乘法实现

矩阵乘法是计算密集型操作，传统JavaScript实现效率较低。通过WebAssembly（Wasm）和Rust的组合，可在浏览器中实现接近原生性能的矩阵运算。以下是完整实现方案：

数学原理

设矩阵$A$为$m \times n$，矩阵$B$为$n \times p$，则乘积$C = AB$的元素定义为： $$c_{ij} = \sum_{k=1}^{n} a_{ik} \times b_{kj}$$ 其中$1 \leq i \leq m$，$1 \leq j \leq p$

实现步骤

1. Rust端实现（高性能计算核心）

// Cargo.toml
[package]
name = "wasm-matrix"
version = "0.1.0"
edition = "2021"

[lib]
crate-type = ["cdylib"]

[dependencies]
wasm-bindgen = "0.2"

// src/lib.rs
use wasm_bindgen::prelude::*;

#[wasm_bindgen]
pub fn matrix_multiply(a: &[f32], b: &[f32], m: usize, n: usize, p: usize) -> Vec<f32> {
    let mut result = vec![0.0; m * p];
    
    // 三层循环实现矩阵乘法
    for i in 0..m {
        for k in 0..n {
            let a_temp = a[i * n + k];
            for j in 0..p {
                result[i * p + j] += a_temp * b[k * p + j];
            }
        }
    }
    result
}

2. 编译为WebAssembly

# 安装编译工具
cargo install wasm-pack

# 编译为Wasm
wasm-pack build --target web

3. 前端调用（JavaScript）

<!DOCTYPE html>
<script type="module">
  import init, { matrix_multiply } from './pkg/wasm_matrix.js';

  async function run() {
    await init();
    
    // 创建测试矩阵 (2x3 和 3x2)
    const a = new Float32Array([1, 2, 3, 4, 5, 6]); // 2x3
    const b = new Float32Array([7, 8, 9, 10, 11, 12]); // 3x2
    
    // 执行Wasm矩阵乘法
    const result = matrix_multiply(a, b, 2, 3, 2);
    
    // 输出结果 (2x2)
    console.log("Result:", result); 
    // 预期: [58, 64, 139, 154]
  }

  run();
</script>

性能优化技巧

内存布局优化

// 使用列优先存储提升缓存命中率
for j in 0..p {
    for k in 0..n {
        let b_temp = b[k * p + j];
        for i in 0..m {
            result[i * p + j] += a[i * n + k] * b_temp;
        }
    }
}

SIMD并行计算（需浏览器支持）

#[cfg(target_feature = "simd128")]
unsafe fn simd_multiply(a: &[f32], b: &[f32], m: usize, n: usize, p: usize) -> Vec<f32> {
    use std::arch::wasm32::*;
    let mut result = vec![0.0; m * p];
    
    for i in 0..m {
        for j in (0..p).step_by(4) {
            let mut sum = f32x4_splat(0.0);
            for k in 0..n {
                let a_val = f32x4_splat(a[i * n + k]);
                let b_vec = v128_load(b.as_ptr().add(k * p + j));
                sum = f32x4_add(sum, f32x4_mul(a_val, b_vec));
            }
            v128_store(result.as_mut_ptr().add(i * p + j), sum);
        }
    }
    result
}