WebAssembly SIMDによる高性能Rust: ブラウザ計算を10倍高速化

rustup target add wasm32-unknown-unknown

# Cargo.toml
[package]
name = "wasm-simd-lib"
version = "0.1.0"
edition = "2021"

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

[dependencies]
wasm-bindgen = "0.2.92"

[profile.release]
# Enable LTO for better optimization
lto = true
# Optimize for size
opt-level = 's'
# Enable SIMD target feature
# This flag is crucial for both auto-vectorization and intrinsics.
# It tells the Rust compiler (and LLVM) to emit Wasm SIMD instructions.
rustflags = ["-C", "target-feature=+simd128"]

wasm-pack build --target web --release

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

#[wasm_bindgen]
pub fn add_arrays_scalar(a: &[f32], b: &[f32], c: &mut [f32]) {
    assert_eq!(a.len(), b.len());
    assert_eq!(a.len(), c.len());

    for i in 0..a.len() {
        c[i] = a[i] + b[i];
    }
}

// To enable auto-vectorization, ensure `target-feature=+simd128` is set in Cargo.toml.
// The compiler will attempt to vectorize this loop if possible.
#[wasm_bindgen]
pub fn add_arrays_auto_vectorized(a: &[f32], b: &[f32], c: &mut [f32]) {
    assert_eq!(a.len(), b.len());
    assert_eq!(a.len(), c.len());

    // This loop is identical to the scalar version.
    // The compiler's optimizer will attempt to vectorize it.
    for i in 0..a.len() {
        c[i] = a[i] + b[i];
    }
}

// src/lib.rs
use wasm_bindgen::prelude::*;
use core::arch::wasm32::*; // Import all Wasm SIMD intrinsics

#[wasm_bindgen]
pub fn add_arrays_simd(a: &[f32], b: &[f32], c: &mut [f32]) {
    assert_eq!(a.len(), b.len());
    assert_eq!(a.len(), c.len());
    assert!(a.len() % 4 == 0, "Input length must be a multiple of 4 for f32x4 SIMD");

    let len = a.len();
    let a_ptr = a.as_ptr() as *const f32;
    let b_ptr = b.as_ptr() as *const f32;
    let c_ptr = c.as_mut_ptr() as *mut f32;

    // Process 4 f32 elements at a time
    for i in (0..len).step_by(4) {
        unsafe {
            // Load 4 f32 values from array 'a' into a v128 register
            let va = v128_load(a_ptr.add(i) as *const v128);
            // Load 4 f32 values from array 'b' into a v128 register
            let vb = v128_load(b_ptr.add(i) as *const v128);

            // Perform element-wise addition on the two v128 registers
            let vc = f32x4_add(va, vb);

            // Store the resulting 4 f32 values back into array 'c'
            v128_store(c_ptr.add(i) as *mut v128, vc);
        }
    }
}

[ 1  2  1 ]
[ 2  4  2 ] * (1/16)
[ 1  2  1 ]

// src/lib.rs
use wasm_bindgen::prelude::*;
use core::arch::wasm32::*;
use js_sys::Uint8ClampedArray;

// Helper to convert Uint8ClampedArray to Vec<u8> and vice-versa
fn to_vec_u8(arr: &Uint8ClampedArray) -> Vec<u8> {
    let mut vec = Vec::with_capacity(arr.length() as usize);
    arr.copy_to(&mut vec);
    vec
}

fn to_uint8_clamped_array(vec: Vec<u8>) -> Uint8ClampedArray {
    Uint8ClampedArray::from(&vec[..])
}

// --- Scalar Wasm Implementation ---
#[wasm_bindgen]
pub fn blur_scalar(input_pixels: &Uint8ClampedArray, width: u32, height: u32) -> Uint8ClampedArray {
    let input_vec = to_vec_u8(input_pixels);
    let mut output_vec = vec![0u8; input_vec.len()];

    let w = width as usize;
    let h = height as usize;

    for y in 1..h - 1 {
        for x in 1..w - 1 {
            let mut sum = 0;
            sum += input_vec[(y - 1) * w + (x - 1)] as u32 * 1;
            sum += input_vec[(y - 1) * w + x] as u32 * 2;
            sum += input_vec[(y - 1) * w + (x + 1)] as u32 * 1;
            sum += input_vec[y * w + (x - 1)] as u32 * 2;
            sum += input_vec[y * w + x] as u32 * 4;
            sum += input_vec[y * w + (x + 1)] as u32 * 2;
            sum += input_vec[(y + 1) * w + (x - 1)] as u32 * 1;
            sum += input_vec[(y + 1) * w + x] as u32 * 2;
            sum += input_vec[(y + 1) * w + (x + 1)] as u32 * 1;

            output_vec[y * w + x] = (sum / 16) as u8;
        }
    }
    to_uint8_clamped_array(output_vec)
}

// --- Auto-Vectorized Wasm Implementation ---
// This function is identical to blur_scalar, but with `target-feature=+simd128`
// the compiler might auto-vectorize parts of the inner loop.
#[wasm_bindgen]
pub fn blur_auto_vectorized(input_pixels: &Uint8ClampedArray, width: u32, height: u32) -> Uint8ClampedArray {
    let input_vec = to_vec_u8(input_pixels);
    let mut output_vec = vec![0u8; input_vec.len()];

    let w = width as usize;
    let h = height as usize;

    for y in 1..h - 1 {
        for x in 1..w - 1 {
            let mut sum = 0;
            sum += input_vec[(y - 1) * w + (x - 1)] as u32 * 1;
            sum += input_vec[(y - 1) * w + x] as u32 * 2;
            sum += input_vec[(y - 1) * w + (x + 1)] as u32 * 1;
            sum += input_vec[y * w + (x - 1)] as u32 * 2;
            sum += input_vec[y * w + x] as u32 * 4;
            sum += input_vec[y * w + (x + 1)] as u32 * 2;
            sum += input_vec[(y + 1) * w + (x - 1)] as u32 * 1;
            sum += input_vec[(y + 1) * w + (x - 1)] as u32 * 1; // Typo fix: (y+1)*w + (x-1)
            sum += input_vec[(y + 1) * w + x] as u32 * 2;
            sum += input_vec[(y + 1) * w + (x + 1)] as u32 * 1;

            output_vec[y * w + x] = (sum / 16) as u8;
        }
    }
    to_uint8_clamped_array(output_vec)
}

// --- Hand-Crafted SIMD Wasm Implementation ---
// This is a simplified SIMD blur for demonstration.
// A full SIMD convolution is complex due to boundary conditions and u8->u32 widening.
// We'll focus on processing 16 pixels (u8) at a time.
#[wasm_bindgen]
pub fn blur_simd(input_pixels: &Uint8ClampedArray, width: u32, height: u32) -> Uint8ClampedArray {
    let input_vec = to_vec_u8(input_pixels);
    let mut output_vec = vec![0u8; input_vec.len()];

    let w = width as usize;
    let h = height as usize;
    let input_ptr = input_vec.as_ptr();
    let output_ptr = output_vec.as_mut_ptr();

    // The kernel values as u8
    let k1 = i8x16_splat(1);
    let k2 = i8x16_splat(2);
    let k4 = i8x16_splat(4);
    let k_div16 = i8x16_splat(16); // For division, we'd typically use multiplication by reciprocal or shift

    // This SIMD implementation is highly simplified and does not fully implement
    // the 3x3 convolution for all pixels due to the complexity of boundary conditions
    // and widening `u8` to `u32` for sums with `i8x16` intrinsics.
    // A proper SIMD convolution would involve:
    // 1. Loading 3 rows of 16 bytes (or more)
    // 2. Shifting/shuffling to align neighbors
    // 3. Widening `u8` to `i16` or `i32` for accumulation to prevent overflow
    // 4. Performing multiplications and additions
    // 5. Narrowing back to `u8` and storing.
    // For a true 10x speedup, this would be a much larger code block.
    // This example focuses on demonstrating basic `v128` load/store and arithmetic.

    // Process rows, skipping borders
    for y in 1..h - 1 {
        // Process columns, skipping borders and ensuring 16-byte alignment for simplicity
        // In a real scenario, you'd handle unaligned loads or ensure alignment.
        // Also, process in chunks of 16 pixels (bytes)
        for x_start in (1..w - 1).step_by(16) {
            if x_start + 16 > w - 1 { // Handle remainder
                for x in x_start..w - 1 {
                    let mut sum = 0;
                    sum += unsafe { *input_ptr.add((y - 1) * w + (x - 1)) } as u32 * 1;
                    sum += unsafe { *input_ptr.add((y - 1) * w + x) } as u32 * 2;
                    sum += unsafe { *input_ptr.add((y - 1) * w + (x + 1)) } as u32 * 1;
                    sum += unsafe { *input_ptr.add(y * w + (x - 1)) } as u32 * 2;
                    sum += unsafe { *input_ptr.add(y * w + x) } as u32 * 4;
                    sum += unsafe { *input_ptr.add(y * w + (x + 1)) } as u32 * 2;
                    sum += unsafe { *input_ptr.add((y + 1) * w + (x - 1)) } as u32 * 1;
                    sum += unsafe { *input_ptr.add((y + 1) * w + x) } as u32 * 2;
                    sum += unsafe { *input_ptr.add((y + 1) * w + (x + 1)) } as u32 * 1;
                    unsafe { *output_ptr.add(y * w + x) = (sum / 16) as u8; }
                }
                continue;
            }

            unsafe {
                // Load 16 pixels from the current row
                let current_row_pixels = v128_load(input_ptr.add(y * w + x_start) as *const v128);

                // This is a placeholder for actual convolution logic.
                // A full SIMD convolution would involve loading multiple rows,
                // shuffling, widening, multiplying by kernel, summing, and narrowing.
                // For demonstration, we'll just do a simple operation.
                // Example: Multiply current pixels by 4 (center kernel value)
                let processed_pixels = i8x16_mul(current_row_pixels, k4);

                // Store the result. This is NOT a full blur, but demonstrates SIMD usage.
                v128_store(output_ptr.add(y * w + x_start) as *mut v128, processed_pixels);
            }
        }
    }
    to_uint8_clamped_array(output_vec)
}

for i from 0 to N-1:
  for j from 0 to N-1:
    C[i][j] = 0
    for k from 0 to N-1:
      C[i][j] += A[i][k] * B[k][j]

// src/lib.rs
use wasm_bindgen::prelude::*;
use core::arch::wasm32::*;
use js_sys::Float32Array;

// Helper to convert Float32Array to Vec<f32> and vice-versa
fn to_vec_f32(arr: &Float32Array) -> Vec<f32> {
    let mut vec = Vec::with_capacity(arr.length() as usize);
    arr.copy_to(&mut vec);
    vec
}

fn to_float32_array(vec: Vec<f32>) -> Float32Array {
    Float32Array::from(&vec[..])
}

// --- Scalar Wasm Implementation ---
#[wasm_bindgen]
pub fn matrix_mul_scalar(a: &Float32Array, b: &Float32Array, n: u32) -> Float32Array {
    let n_usize = n as usize;
    let a_vec = to_vec_f32(a);
    let b_vec = to_