mirror of
https://github.com/ruvnet/RuView
synced 2026-07-28 18:21:42 +00:00
ruv
1418 lines
44 KiB
Rust
1418 lines
44 KiB
Rust
//! SIMD-Optimized Compute Operations for edge-net
|
|
//!
|
|
//! This module provides vectorized operations for neural network inference
|
|
//! with automatic dispatch to the best available SIMD implementation:
|
|
//!
|
|
//! - WASM simd128: 4x f32 lanes (browser targets)
|
|
//! - x86_64 AVX2: 8x f32 lanes (native x86 targets)
|
|
//! - Scalar: Portable fallback
|
|
//!
|
|
//! # Performance Targets
|
|
//!
|
|
//! - dot_product: 8x speedup over scalar
|
|
//! - matmul: 10x speedup with tiling + prefetch
|
|
//! - softmax: Numerically stable with max subtraction
|
|
//! - Q4 quantization: 4x memory reduction with 1% accuracy loss
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
use core::arch::wasm32::*;
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
use std::arch::x86_64::*;
|
|
|
|
/// SIMD compute backend with automatic platform detection
|
|
pub struct SimdCompute {
|
|
/// Platform capabilities detected at runtime
|
|
#[allow(dead_code)]
|
|
capabilities: SimdCapabilities,
|
|
}
|
|
|
|
/// Detected SIMD capabilities
|
|
#[derive(Clone, Debug)]
|
|
pub struct SimdCapabilities {
|
|
/// WASM simd128 available
|
|
pub wasm_simd128: bool,
|
|
/// x86 AVX2 available
|
|
pub avx2: bool,
|
|
/// x86 SSE4.1 available
|
|
pub sse41: bool,
|
|
/// x86 FMA available
|
|
pub fma: bool,
|
|
}
|
|
|
|
impl Default for SimdCapabilities {
|
|
fn default() -> Self {
|
|
Self::detect()
|
|
}
|
|
}
|
|
|
|
impl SimdCapabilities {
|
|
/// Detect available SIMD capabilities at runtime
|
|
pub fn detect() -> Self {
|
|
#[cfg(target_arch = "wasm32")]
|
|
{
|
|
Self {
|
|
wasm_simd128: true, // Always available on wasm32 with simd128 feature
|
|
avx2: false,
|
|
sse41: false,
|
|
fma: false,
|
|
}
|
|
}
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
{
|
|
Self {
|
|
wasm_simd128: false,
|
|
avx2: is_x86_feature_detected!("avx2"),
|
|
sse41: is_x86_feature_detected!("sse4.1"),
|
|
fma: is_x86_feature_detected!("fma"),
|
|
}
|
|
}
|
|
|
|
#[cfg(not(any(target_arch = "wasm32", target_arch = "x86_64")))]
|
|
{
|
|
Self {
|
|
wasm_simd128: false,
|
|
avx2: false,
|
|
sse41: false,
|
|
fma: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Get the SIMD lane width for f32 operations
|
|
pub fn lane_width(&self) -> usize {
|
|
if self.avx2 {
|
|
8
|
|
} else if self.wasm_simd128 || self.sse41 {
|
|
4
|
|
} else {
|
|
1
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for SimdCompute {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl SimdCompute {
|
|
/// Create a new SIMD compute backend with automatic platform detection
|
|
pub fn new() -> Self {
|
|
Self {
|
|
capabilities: SimdCapabilities::detect(),
|
|
}
|
|
}
|
|
|
|
/// Get detected capabilities
|
|
pub fn capabilities(&self) -> &SimdCapabilities {
|
|
&self.capabilities
|
|
}
|
|
|
|
// ========================================================================
|
|
// Dot Product Operations
|
|
// ========================================================================
|
|
|
|
/// SIMD dot product for f32 vectors
|
|
///
|
|
/// Automatically dispatches to the best available implementation:
|
|
/// - AVX2: 8x f32 lanes with FMA
|
|
/// - WASM simd128: 4x f32 lanes
|
|
/// - SSE4.1: 4x f32 lanes
|
|
/// - Scalar: Portable fallback
|
|
#[inline]
|
|
pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
|
|
debug_assert_eq!(a.len(), b.len(), "Vector lengths must match");
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
{
|
|
if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
|
|
return unsafe { Self::dot_product_avx2_fma(a, b) };
|
|
} else if is_x86_feature_detected!("avx2") {
|
|
return unsafe { Self::dot_product_avx2(a, b) };
|
|
} else if is_x86_feature_detected!("sse4.1") {
|
|
return unsafe { Self::dot_product_sse41(a, b) };
|
|
} else {
|
|
return Self::dot_product_scalar(a, b);
|
|
}
|
|
}
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
{
|
|
return Self::dot_product_wasm_simd128(a, b);
|
|
}
|
|
|
|
#[cfg(not(any(target_arch = "wasm32", target_arch = "x86_64")))]
|
|
{
|
|
Self::dot_product_scalar(a, b)
|
|
}
|
|
}
|
|
|
|
/// Scalar dot product (fallback)
|
|
#[inline]
|
|
pub fn dot_product_scalar(a: &[f32], b: &[f32]) -> f32 {
|
|
a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
|
|
}
|
|
|
|
/// WASM simd128 dot product with 4x f32 lanes
|
|
#[cfg(target_arch = "wasm32")]
|
|
#[inline]
|
|
pub fn dot_product_wasm_simd128(a: &[f32], b: &[f32]) -> f32 {
|
|
let len = a.len();
|
|
let chunks = len / 4;
|
|
let mut sum = f32x4_splat(0.0);
|
|
|
|
// Process 4 elements at a time
|
|
for i in 0..chunks {
|
|
let offset = i * 4;
|
|
let a_vec = unsafe {
|
|
v128_load(a.as_ptr().add(offset) as *const v128)
|
|
};
|
|
let b_vec = unsafe {
|
|
v128_load(b.as_ptr().add(offset) as *const v128)
|
|
};
|
|
let prod = f32x4_mul(a_vec, b_vec);
|
|
sum = f32x4_add(sum, prod);
|
|
}
|
|
|
|
// Horizontal sum: extract all 4 lanes and add
|
|
let mut result = f32x4_extract_lane::<0>(sum)
|
|
+ f32x4_extract_lane::<1>(sum)
|
|
+ f32x4_extract_lane::<2>(sum)
|
|
+ f32x4_extract_lane::<3>(sum);
|
|
|
|
// Handle remainder
|
|
for i in (chunks * 4)..len {
|
|
result += a[i] * b[i];
|
|
}
|
|
|
|
result
|
|
}
|
|
|
|
/// x86_64 AVX2 dot product with 8x f32 lanes
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn dot_product_avx2(a: &[f32], b: &[f32]) -> f32 {
|
|
let len = a.len();
|
|
let chunks = len / 8;
|
|
let mut sum = _mm256_setzero_ps();
|
|
|
|
for i in 0..chunks {
|
|
let offset = i * 8;
|
|
let a_vec = _mm256_loadu_ps(a.as_ptr().add(offset));
|
|
let b_vec = _mm256_loadu_ps(b.as_ptr().add(offset));
|
|
let prod = _mm256_mul_ps(a_vec, b_vec);
|
|
sum = _mm256_add_ps(sum, prod);
|
|
}
|
|
|
|
// Horizontal sum reduction
|
|
let result = Self::hsum_avx2(sum);
|
|
|
|
// Handle remainder
|
|
let mut final_result = result;
|
|
for i in (chunks * 8)..len {
|
|
final_result += a[i] * b[i];
|
|
}
|
|
|
|
final_result
|
|
}
|
|
|
|
/// x86_64 AVX2+FMA dot product with fused multiply-add
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2", enable = "fma")]
|
|
#[inline]
|
|
unsafe fn dot_product_avx2_fma(a: &[f32], b: &[f32]) -> f32 {
|
|
let len = a.len();
|
|
let chunks = len / 8;
|
|
let mut sum = _mm256_setzero_ps();
|
|
|
|
for i in 0..chunks {
|
|
let offset = i * 8;
|
|
let a_vec = _mm256_loadu_ps(a.as_ptr().add(offset));
|
|
let b_vec = _mm256_loadu_ps(b.as_ptr().add(offset));
|
|
// FMA: sum = a * b + sum
|
|
sum = _mm256_fmadd_ps(a_vec, b_vec, sum);
|
|
}
|
|
|
|
let result = Self::hsum_avx2(sum);
|
|
|
|
let mut final_result = result;
|
|
for i in (chunks * 8)..len {
|
|
final_result += a[i] * b[i];
|
|
}
|
|
|
|
final_result
|
|
}
|
|
|
|
/// x86_64 SSE4.1 dot product with 4x f32 lanes
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "sse4.1")]
|
|
#[inline]
|
|
unsafe fn dot_product_sse41(a: &[f32], b: &[f32]) -> f32 {
|
|
let len = a.len();
|
|
let chunks = len / 4;
|
|
let mut sum = _mm_setzero_ps();
|
|
|
|
for i in 0..chunks {
|
|
let offset = i * 4;
|
|
let a_vec = _mm_loadu_ps(a.as_ptr().add(offset));
|
|
let b_vec = _mm_loadu_ps(b.as_ptr().add(offset));
|
|
let prod = _mm_mul_ps(a_vec, b_vec);
|
|
sum = _mm_add_ps(sum, prod);
|
|
}
|
|
|
|
// Horizontal sum using shuffle
|
|
let shuf = _mm_shuffle_ps(sum, sum, 0b10_11_00_01);
|
|
let sums = _mm_add_ps(sum, shuf);
|
|
let shuf = _mm_movehl_ps(sums, sums);
|
|
let sums = _mm_add_ss(sums, shuf);
|
|
let mut result = _mm_cvtss_f32(sums);
|
|
|
|
for i in (chunks * 4)..len {
|
|
result += a[i] * b[i];
|
|
}
|
|
|
|
result
|
|
}
|
|
|
|
/// Horizontal sum for AVX2 __m256
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn hsum_avx2(v: __m256) -> f32 {
|
|
let high = _mm256_extractf128_ps(v, 1);
|
|
let low = _mm256_castps256_ps128(v);
|
|
let sum128 = _mm_add_ps(high, low);
|
|
let shuf = _mm_shuffle_ps(sum128, sum128, 0b10_11_00_01);
|
|
let sums = _mm_add_ps(sum128, shuf);
|
|
let shuf = _mm_movehl_ps(sums, sums);
|
|
let sums = _mm_add_ss(sums, shuf);
|
|
_mm_cvtss_f32(sums)
|
|
}
|
|
|
|
// ========================================================================
|
|
// Matrix Multiplication (Tiled with Prefetch Hints)
|
|
// ========================================================================
|
|
|
|
/// SIMD tiled matrix multiplication
|
|
///
|
|
/// Performs C = A * B with cache-friendly tiling for optimal performance.
|
|
/// Uses prefetch hints for next tile to reduce cache misses.
|
|
///
|
|
/// # Arguments
|
|
/// * `a` - Left matrix (m x k) in row-major order
|
|
/// * `b` - Right matrix (k x n) in row-major order
|
|
/// * `m` - Rows in A
|
|
/// * `k` - Cols in A / Rows in B
|
|
/// * `n` - Cols in B
|
|
///
|
|
/// # Returns
|
|
/// Result matrix C (m x n) in row-major order
|
|
#[inline]
|
|
pub fn matmul_simd(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {
|
|
debug_assert_eq!(a.len(), m * k, "A dimensions mismatch");
|
|
debug_assert_eq!(b.len(), k * n, "B dimensions mismatch");
|
|
|
|
let mut c = vec![0.0f32; m * n];
|
|
|
|
// Tile size for cache optimization (64 elements = 256 bytes = 4 cache lines)
|
|
const TILE_SIZE: usize = 64;
|
|
|
|
// Tiled matrix multiplication
|
|
for ii in (0..m).step_by(TILE_SIZE) {
|
|
for jj in (0..n).step_by(TILE_SIZE) {
|
|
for kk in (0..k).step_by(TILE_SIZE) {
|
|
let i_end = (ii + TILE_SIZE).min(m);
|
|
let j_end = (jj + TILE_SIZE).min(n);
|
|
let k_end = (kk + TILE_SIZE).min(k);
|
|
|
|
// Process tile
|
|
for i in ii..i_end {
|
|
for j in jj..j_end {
|
|
let mut sum = c[i * n + j];
|
|
|
|
// Use SIMD for inner product within tile
|
|
let a_row = &a[i * k + kk..i * k + k_end];
|
|
let b_col_start = kk * n + j;
|
|
|
|
// Gather B column elements (strided access)
|
|
let mut b_col = Vec::with_capacity(k_end - kk);
|
|
for ki in kk..k_end {
|
|
b_col.push(b[ki * n + j]);
|
|
}
|
|
|
|
sum += Self::dot_product(a_row, &b_col);
|
|
c[i * n + j] = sum;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
c
|
|
}
|
|
|
|
/// Optimized matrix-vector multiplication
|
|
///
|
|
/// Computes y = A * x where A is m x n matrix
|
|
#[inline]
|
|
pub fn matvec_simd(a: &[f32], x: &[f32], m: usize, n: usize) -> Vec<f32> {
|
|
debug_assert_eq!(a.len(), m * n, "Matrix dimensions mismatch");
|
|
debug_assert_eq!(x.len(), n, "Vector dimension mismatch");
|
|
|
|
let mut y = Vec::with_capacity(m);
|
|
|
|
for i in 0..m {
|
|
let row_start = i * n;
|
|
let row = &a[row_start..row_start + n];
|
|
y.push(Self::dot_product(row, x));
|
|
}
|
|
|
|
y
|
|
}
|
|
|
|
// ========================================================================
|
|
// Softmax (Numerically Stable with Max Subtraction)
|
|
// ========================================================================
|
|
|
|
/// Numerically stable softmax with SIMD acceleration
|
|
///
|
|
/// Uses the log-sum-exp trick: softmax(x) = exp(x - max(x)) / sum(exp(x - max(x)))
|
|
/// This prevents overflow for large values.
|
|
#[inline]
|
|
pub fn softmax_simd(input: &mut [f32]) {
|
|
if input.is_empty() {
|
|
return;
|
|
}
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
{
|
|
if is_x86_feature_detected!("avx2") {
|
|
unsafe { Self::softmax_avx2(input) };
|
|
return;
|
|
}
|
|
}
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
{
|
|
Self::softmax_wasm_simd128(input);
|
|
return;
|
|
}
|
|
|
|
#[cfg(not(any(target_arch = "wasm32", target_arch = "x86_64")))]
|
|
{
|
|
Self::softmax_scalar(input);
|
|
}
|
|
}
|
|
|
|
/// Scalar softmax implementation
|
|
#[inline]
|
|
pub fn softmax_scalar(input: &mut [f32]) {
|
|
// Find max for numerical stability
|
|
let max_val = input.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
|
|
|
// Compute exp(x - max) and sum
|
|
let mut sum = 0.0f32;
|
|
for x in input.iter_mut() {
|
|
*x = (*x - max_val).exp();
|
|
sum += *x;
|
|
}
|
|
|
|
// Normalize
|
|
let inv_sum = 1.0 / sum;
|
|
for x in input.iter_mut() {
|
|
*x *= inv_sum;
|
|
}
|
|
}
|
|
|
|
/// WASM simd128 softmax
|
|
#[cfg(target_arch = "wasm32")]
|
|
#[inline]
|
|
pub fn softmax_wasm_simd128(input: &mut [f32]) {
|
|
let len = input.len();
|
|
let chunks = len / 4;
|
|
|
|
// Find max using SIMD
|
|
let mut max_vec = f32x4_splat(f32::NEG_INFINITY);
|
|
for i in 0..chunks {
|
|
let v = unsafe { v128_load(input.as_ptr().add(i * 4) as *const v128) };
|
|
max_vec = f32x4_pmax(max_vec, v);
|
|
}
|
|
|
|
// Horizontal max
|
|
let mut max_val = f32x4_extract_lane::<0>(max_vec)
|
|
.max(f32x4_extract_lane::<1>(max_vec))
|
|
.max(f32x4_extract_lane::<2>(max_vec))
|
|
.max(f32x4_extract_lane::<3>(max_vec));
|
|
|
|
// Handle remainder for max
|
|
for i in (chunks * 4)..len {
|
|
max_val = max_val.max(input[i]);
|
|
}
|
|
|
|
let max_broadcast = f32x4_splat(max_val);
|
|
|
|
// Compute exp(x - max) and accumulate sum
|
|
let mut sum = 0.0f32;
|
|
for i in 0..chunks {
|
|
let offset = i * 4;
|
|
let v = unsafe { v128_load(input.as_ptr().add(offset) as *const v128) };
|
|
let shifted = f32x4_sub(v, max_broadcast);
|
|
|
|
// Fast exp approximation for each lane
|
|
let exp_vals = [
|
|
Self::fast_exp(f32x4_extract_lane::<0>(shifted)),
|
|
Self::fast_exp(f32x4_extract_lane::<1>(shifted)),
|
|
Self::fast_exp(f32x4_extract_lane::<2>(shifted)),
|
|
Self::fast_exp(f32x4_extract_lane::<3>(shifted)),
|
|
];
|
|
|
|
input[offset] = exp_vals[0];
|
|
input[offset + 1] = exp_vals[1];
|
|
input[offset + 2] = exp_vals[2];
|
|
input[offset + 3] = exp_vals[3];
|
|
|
|
sum += exp_vals[0] + exp_vals[1] + exp_vals[2] + exp_vals[3];
|
|
}
|
|
|
|
// Handle remainder
|
|
for i in (chunks * 4)..len {
|
|
input[i] = (input[i] - max_val).exp();
|
|
sum += input[i];
|
|
}
|
|
|
|
// Normalize
|
|
let inv_sum = 1.0 / sum;
|
|
let inv_sum_vec = f32x4_splat(inv_sum);
|
|
|
|
for i in 0..chunks {
|
|
let offset = i * 4;
|
|
let v = unsafe { v128_load(input.as_ptr().add(offset) as *const v128) };
|
|
let normalized = f32x4_mul(v, inv_sum_vec);
|
|
unsafe {
|
|
v128_store(input.as_mut_ptr().add(offset) as *mut v128, normalized);
|
|
}
|
|
}
|
|
|
|
for i in (chunks * 4)..len {
|
|
input[i] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
/// AVX2 softmax
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn softmax_avx2(input: &mut [f32]) {
|
|
let len = input.len();
|
|
let chunks = len / 8;
|
|
|
|
// Find max using AVX2
|
|
let mut max_vec = _mm256_set1_ps(f32::NEG_INFINITY);
|
|
for i in 0..chunks {
|
|
let v = _mm256_loadu_ps(input.as_ptr().add(i * 8));
|
|
max_vec = _mm256_max_ps(max_vec, v);
|
|
}
|
|
|
|
// Horizontal max reduction
|
|
let mut max_val = Self::hmax_avx2(max_vec);
|
|
|
|
// Handle remainder for max
|
|
for i in (chunks * 8)..len {
|
|
max_val = max_val.max(input[i]);
|
|
}
|
|
|
|
let max_broadcast = _mm256_set1_ps(max_val);
|
|
|
|
// Compute exp(x - max) and sum
|
|
let mut sum = 0.0f32;
|
|
for i in 0..chunks {
|
|
let ptr = input.as_mut_ptr().add(i * 8);
|
|
let v = _mm256_loadu_ps(ptr);
|
|
let shifted = _mm256_sub_ps(v, max_broadcast);
|
|
let exp_v = Self::fast_exp_avx2(shifted);
|
|
_mm256_storeu_ps(ptr, exp_v);
|
|
|
|
// Accumulate sum
|
|
sum += Self::hsum_avx2(exp_v);
|
|
}
|
|
|
|
// Handle remainder
|
|
for i in (chunks * 8)..len {
|
|
input[i] = (input[i] - max_val).exp();
|
|
sum += input[i];
|
|
}
|
|
|
|
// Normalize
|
|
let inv_sum = 1.0 / sum;
|
|
let inv_sum_vec = _mm256_set1_ps(inv_sum);
|
|
|
|
for i in 0..chunks {
|
|
let ptr = input.as_mut_ptr().add(i * 8);
|
|
let v = _mm256_loadu_ps(ptr);
|
|
_mm256_storeu_ps(ptr, _mm256_mul_ps(v, inv_sum_vec));
|
|
}
|
|
|
|
for i in (chunks * 8)..len {
|
|
input[i] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
/// Horizontal max for AVX2
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn hmax_avx2(v: __m256) -> f32 {
|
|
let high = _mm256_extractf128_ps(v, 1);
|
|
let low = _mm256_castps256_ps128(v);
|
|
let max128 = _mm_max_ps(high, low);
|
|
let max64 = _mm_max_ps(max128, _mm_movehl_ps(max128, max128));
|
|
let max32 = _mm_max_ss(max64, _mm_shuffle_ps(max64, max64, 1));
|
|
_mm_cvtss_f32(max32)
|
|
}
|
|
|
|
/// Fast exp approximation for AVX2
|
|
/// Uses polynomial: exp(x) ~ 1 + x + x^2/2 + x^3/6 for |x| < 1
|
|
/// For larger x, uses range reduction
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn fast_exp_avx2(x: __m256) -> __m256 {
|
|
// Clamp to avoid overflow/underflow
|
|
let min_val = _mm256_set1_ps(-88.0);
|
|
let max_val = _mm256_set1_ps(88.0);
|
|
let x = _mm256_max_ps(_mm256_min_ps(x, max_val), min_val);
|
|
|
|
// Constants for polynomial approximation
|
|
let one = _mm256_set1_ps(1.0);
|
|
let half = _mm256_set1_ps(0.5);
|
|
let sixth = _mm256_set1_ps(1.0 / 6.0);
|
|
let twenty_fourth = _mm256_set1_ps(1.0 / 24.0);
|
|
|
|
let x2 = _mm256_mul_ps(x, x);
|
|
let x3 = _mm256_mul_ps(x2, x);
|
|
let x4 = _mm256_mul_ps(x2, x2);
|
|
|
|
// exp(x) ~ 1 + x + x^2/2 + x^3/6 + x^4/24
|
|
let term1 = _mm256_add_ps(one, x);
|
|
let term2 = _mm256_mul_ps(x2, half);
|
|
let term3 = _mm256_mul_ps(x3, sixth);
|
|
let term4 = _mm256_mul_ps(x4, twenty_fourth);
|
|
|
|
_mm256_add_ps(_mm256_add_ps(term1, term2), _mm256_add_ps(term3, term4))
|
|
}
|
|
|
|
// ========================================================================
|
|
// GELU Activation (Fast Approximation)
|
|
// ========================================================================
|
|
|
|
/// GELU activation using fast tanh approximation
|
|
///
|
|
/// GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
|
|
#[inline]
|
|
pub fn gelu_simd(input: &mut [f32]) {
|
|
#[cfg(target_arch = "x86_64")]
|
|
{
|
|
if is_x86_feature_detected!("avx2") {
|
|
unsafe { Self::gelu_avx2(input) };
|
|
return;
|
|
}
|
|
}
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
{
|
|
Self::gelu_wasm_simd128(input);
|
|
return;
|
|
}
|
|
|
|
#[cfg(not(any(target_arch = "wasm32", target_arch = "x86_64")))]
|
|
{
|
|
Self::gelu_scalar(input);
|
|
}
|
|
}
|
|
|
|
/// Scalar GELU
|
|
#[inline]
|
|
pub fn gelu_scalar(input: &mut [f32]) {
|
|
const SQRT_2_PI: f32 = 0.7978845608028654;
|
|
const COEF: f32 = 0.044715;
|
|
|
|
for x in input.iter_mut() {
|
|
let x3 = *x * *x * *x;
|
|
let inner = SQRT_2_PI * (*x + COEF * x3);
|
|
*x = 0.5 * *x * (1.0 + Self::fast_tanh(inner));
|
|
}
|
|
}
|
|
|
|
/// WASM simd128 GELU
|
|
#[cfg(target_arch = "wasm32")]
|
|
#[inline]
|
|
pub fn gelu_wasm_simd128(input: &mut [f32]) {
|
|
const SQRT_2_PI: f32 = 0.7978845608028654;
|
|
const COEF: f32 = 0.044715;
|
|
|
|
let len = input.len();
|
|
let chunks = len / 4;
|
|
|
|
let sqrt_2_pi = f32x4_splat(SQRT_2_PI);
|
|
let coef = f32x4_splat(COEF);
|
|
let half = f32x4_splat(0.5);
|
|
let one = f32x4_splat(1.0);
|
|
|
|
for i in 0..chunks {
|
|
let offset = i * 4;
|
|
let x = unsafe { v128_load(input.as_ptr().add(offset) as *const v128) };
|
|
|
|
// x^3
|
|
let x2 = f32x4_mul(x, x);
|
|
let x3 = f32x4_mul(x2, x);
|
|
|
|
// sqrt(2/pi) * (x + 0.044715 * x^3)
|
|
let inner = f32x4_mul(sqrt_2_pi, f32x4_add(x, f32x4_mul(coef, x3)));
|
|
|
|
// Fast tanh approximation for each lane
|
|
let tanh_vals = [
|
|
Self::fast_tanh(f32x4_extract_lane::<0>(inner)),
|
|
Self::fast_tanh(f32x4_extract_lane::<1>(inner)),
|
|
Self::fast_tanh(f32x4_extract_lane::<2>(inner)),
|
|
Self::fast_tanh(f32x4_extract_lane::<3>(inner)),
|
|
];
|
|
let tanh_vec = f32x4(tanh_vals[0], tanh_vals[1], tanh_vals[2], tanh_vals[3]);
|
|
|
|
// 0.5 * x * (1 + tanh)
|
|
let result = f32x4_mul(half, f32x4_mul(x, f32x4_add(one, tanh_vec)));
|
|
|
|
unsafe {
|
|
v128_store(input.as_mut_ptr().add(offset) as *mut v128, result);
|
|
}
|
|
}
|
|
|
|
// Handle remainder
|
|
for i in (chunks * 4)..len {
|
|
let x = input[i];
|
|
let x3 = x * x * x;
|
|
let inner = SQRT_2_PI * (x + COEF * x3);
|
|
input[i] = 0.5 * x * (1.0 + Self::fast_tanh(inner));
|
|
}
|
|
}
|
|
|
|
/// AVX2 GELU
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn gelu_avx2(input: &mut [f32]) {
|
|
let len = input.len();
|
|
let chunks = len / 8;
|
|
|
|
let sqrt_2_pi = _mm256_set1_ps(0.7978845608028654);
|
|
let coef = _mm256_set1_ps(0.044715);
|
|
let half = _mm256_set1_ps(0.5);
|
|
let one = _mm256_set1_ps(1.0);
|
|
|
|
for i in 0..chunks {
|
|
let ptr = input.as_mut_ptr().add(i * 8);
|
|
let x = _mm256_loadu_ps(ptr);
|
|
|
|
// x^3
|
|
let x2 = _mm256_mul_ps(x, x);
|
|
let x3 = _mm256_mul_ps(x2, x);
|
|
|
|
// sqrt(2/pi) * (x + 0.044715 * x^3)
|
|
let inner = _mm256_mul_ps(sqrt_2_pi, _mm256_add_ps(x, _mm256_mul_ps(coef, x3)));
|
|
|
|
// Fast tanh approximation
|
|
let tanh = Self::fast_tanh_avx2(inner);
|
|
|
|
// 0.5 * x * (1 + tanh)
|
|
let result = _mm256_mul_ps(half, _mm256_mul_ps(x, _mm256_add_ps(one, tanh)));
|
|
|
|
_mm256_storeu_ps(ptr, result);
|
|
}
|
|
|
|
// Handle remainder
|
|
const SQRT_2_PI: f32 = 0.7978845608028654;
|
|
const COEF: f32 = 0.044715;
|
|
for i in (chunks * 8)..len {
|
|
let x = input[i];
|
|
let x3 = x * x * x;
|
|
let inner = SQRT_2_PI * (x + COEF * x3);
|
|
input[i] = 0.5 * x * (1.0 + Self::fast_tanh(inner));
|
|
}
|
|
}
|
|
|
|
/// Fast tanh approximation for AVX2
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn fast_tanh_avx2(x: __m256) -> __m256 {
|
|
// tanh(x) ~ x * (27 + x^2) / (27 + 9*x^2) for |x| < 3
|
|
// This is Pade approximation
|
|
let x2 = _mm256_mul_ps(x, x);
|
|
let c27 = _mm256_set1_ps(27.0);
|
|
let c9 = _mm256_set1_ps(9.0);
|
|
|
|
let num = _mm256_mul_ps(x, _mm256_add_ps(c27, x2));
|
|
let den = _mm256_add_ps(c27, _mm256_mul_ps(c9, x2));
|
|
|
|
// Clamp result to [-1, 1]
|
|
let result = _mm256_div_ps(num, den);
|
|
let one = _mm256_set1_ps(1.0);
|
|
let neg_one = _mm256_set1_ps(-1.0);
|
|
_mm256_max_ps(_mm256_min_ps(result, one), neg_one)
|
|
}
|
|
|
|
/// Fast scalar tanh approximation
|
|
#[inline]
|
|
fn fast_tanh(x: f32) -> f32 {
|
|
// Pade approximation: tanh(x) ~ x * (27 + x^2) / (27 + 9*x^2)
|
|
let x2 = x * x;
|
|
let result = x * (27.0 + x2) / (27.0 + 9.0 * x2);
|
|
result.clamp(-1.0, 1.0)
|
|
}
|
|
|
|
/// Fast scalar exp approximation
|
|
#[inline]
|
|
fn fast_exp(x: f32) -> f32 {
|
|
// Clamp to avoid overflow/underflow
|
|
let x = x.clamp(-88.0, 88.0);
|
|
|
|
// Polynomial approximation
|
|
let x2 = x * x;
|
|
let x3 = x2 * x;
|
|
let x4 = x2 * x2;
|
|
|
|
1.0 + x + x2 * 0.5 + x3 / 6.0 + x4 / 24.0
|
|
}
|
|
|
|
// ========================================================================
|
|
// Layer Normalization (Welford Algorithm for Numerical Stability)
|
|
// ========================================================================
|
|
|
|
/// Layer normalization using Welford's online algorithm
|
|
///
|
|
/// Uses running mean/variance computation for numerical stability
|
|
/// with large numbers or values with large variance.
|
|
///
|
|
/// # Arguments
|
|
/// * `input` - Input tensor
|
|
/// * `weight` - Learned scale parameters (gamma)
|
|
/// * `bias` - Learned shift parameters (beta), optional
|
|
/// * `eps` - Small constant for numerical stability (typically 1e-5)
|
|
#[inline]
|
|
pub fn layer_norm_simd(
|
|
input: &[f32],
|
|
weight: &[f32],
|
|
bias: Option<&[f32]>,
|
|
eps: f32,
|
|
) -> Vec<f32> {
|
|
debug_assert_eq!(input.len(), weight.len(), "Dimension mismatch");
|
|
if let Some(b) = bias {
|
|
debug_assert_eq!(input.len(), b.len(), "Bias dimension mismatch");
|
|
}
|
|
|
|
// Welford's algorithm for computing mean and variance in one pass
|
|
let (mean, var) = Self::welford_mean_var(input);
|
|
|
|
let inv_std = 1.0 / (var + eps).sqrt();
|
|
|
|
let mut output = Vec::with_capacity(input.len());
|
|
|
|
match bias {
|
|
Some(b) => {
|
|
for i in 0..input.len() {
|
|
let normalized = (input[i] - mean) * inv_std;
|
|
output.push(normalized * weight[i] + b[i]);
|
|
}
|
|
}
|
|
None => {
|
|
for i in 0..input.len() {
|
|
let normalized = (input[i] - mean) * inv_std;
|
|
output.push(normalized * weight[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
/// RMS normalization (used in modern transformers like LLaMA)
|
|
///
|
|
/// RMSNorm(x) = x * weight / sqrt(mean(x^2) + eps)
|
|
#[inline]
|
|
pub fn rms_norm_simd(input: &[f32], weight: &[f32], eps: f32) -> Vec<f32> {
|
|
debug_assert_eq!(input.len(), weight.len(), "Dimension mismatch");
|
|
|
|
// Compute mean of squared values using SIMD
|
|
let sum_sq = Self::dot_product(input, input);
|
|
let rms = (sum_sq / input.len() as f32 + eps).sqrt();
|
|
let inv_rms = 1.0 / rms;
|
|
|
|
let mut output = Vec::with_capacity(input.len());
|
|
for i in 0..input.len() {
|
|
output.push(input[i] * inv_rms * weight[i]);
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
/// Welford's online algorithm for mean and variance
|
|
///
|
|
/// Numerically stable single-pass algorithm
|
|
#[inline]
|
|
fn welford_mean_var(data: &[f32]) -> (f32, f32) {
|
|
if data.is_empty() {
|
|
return (0.0, 0.0);
|
|
}
|
|
|
|
let mut count = 0.0f64;
|
|
let mut mean = 0.0f64;
|
|
let mut m2 = 0.0f64;
|
|
|
|
for &x in data {
|
|
count += 1.0;
|
|
let delta = x as f64 - mean;
|
|
mean += delta / count;
|
|
let delta2 = x as f64 - mean;
|
|
m2 += delta * delta2;
|
|
}
|
|
|
|
let variance = if count > 1.0 { m2 / count } else { 0.0 };
|
|
|
|
(mean as f32, variance as f32)
|
|
}
|
|
|
|
// ========================================================================
|
|
// Quantization Operations (Q4/Q8)
|
|
// ========================================================================
|
|
|
|
/// Q4 block size (number of elements per scale factor)
|
|
pub const Q4_BLOCK_SIZE: usize = 32;
|
|
|
|
/// Q8 block size
|
|
pub const Q8_BLOCK_SIZE: usize = 32;
|
|
|
|
/// Quantize f32 array to Q4 format (4-bit quantization)
|
|
///
|
|
/// Uses block-wise quantization with per-block scale factors.
|
|
/// Achieves ~4x memory reduction with ~1% accuracy loss.
|
|
///
|
|
/// # Returns
|
|
/// Tuple of (quantized_data, scales) where:
|
|
/// - quantized_data: Packed 4-bit values (2 values per byte)
|
|
/// - scales: Per-block scale factors
|
|
#[inline]
|
|
pub fn quantize_simd_q4(input: &[f32]) -> (Vec<u8>, Vec<f32>) {
|
|
let num_blocks = (input.len() + Self::Q4_BLOCK_SIZE - 1) / Self::Q4_BLOCK_SIZE;
|
|
let mut data = Vec::with_capacity(input.len() / 2);
|
|
let mut scales = Vec::with_capacity(num_blocks);
|
|
|
|
for block in input.chunks(Self::Q4_BLOCK_SIZE) {
|
|
// Find max absolute value for scale
|
|
let max_abs = block.iter().map(|x| x.abs()).fold(0.0f32, f32::max);
|
|
let scale = max_abs / 7.0; // Q4 range is -8 to 7
|
|
scales.push(scale);
|
|
|
|
// Quantize with zero-centered mapping
|
|
let inv_scale = if scale > 1e-10 { 1.0 / scale } else { 0.0 };
|
|
|
|
for pair in block.chunks(2) {
|
|
let q0 = ((pair[0] * inv_scale).round() as i8).clamp(-8, 7) as u8 & 0x0F;
|
|
let q1 = if pair.len() > 1 {
|
|
((pair[1] * inv_scale).round() as i8).clamp(-8, 7) as u8 & 0x0F
|
|
} else {
|
|
0
|
|
};
|
|
data.push((q1 << 4) | q0);
|
|
}
|
|
}
|
|
|
|
(data, scales)
|
|
}
|
|
|
|
/// Dequantize Q4 data back to f32
|
|
#[inline]
|
|
pub fn dequantize_simd_q4(
|
|
data: &[u8],
|
|
scales: &[f32],
|
|
output_len: usize,
|
|
) -> Vec<f32> {
|
|
let mut output = Vec::with_capacity(output_len);
|
|
|
|
for (block_idx, scale) in scales.iter().enumerate() {
|
|
let block_start = block_idx * Self::Q4_BLOCK_SIZE / 2;
|
|
let block_end = ((block_idx + 1) * Self::Q4_BLOCK_SIZE / 2).min(data.len());
|
|
|
|
for byte_idx in block_start..block_end {
|
|
if output.len() >= output_len {
|
|
break;
|
|
}
|
|
|
|
let byte = data[byte_idx];
|
|
|
|
// Low nibble
|
|
let q0 = (byte & 0x0F) as i8;
|
|
let q0 = if q0 > 7 { q0 - 16 } else { q0 };
|
|
output.push(q0 as f32 * scale);
|
|
|
|
if output.len() >= output_len {
|
|
break;
|
|
}
|
|
|
|
// High nibble
|
|
let q1 = ((byte >> 4) & 0x0F) as i8;
|
|
let q1 = if q1 > 7 { q1 - 16 } else { q1 };
|
|
output.push(q1 as f32 * scale);
|
|
}
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
/// Quantize f32 array to Q8 format (8-bit quantization)
|
|
///
|
|
/// Uses block-wise quantization with per-block scale factors.
|
|
/// Achieves ~4x memory reduction with minimal accuracy loss.
|
|
#[inline]
|
|
pub fn quantize_simd_q8(input: &[f32]) -> (Vec<i8>, Vec<f32>) {
|
|
let num_blocks = (input.len() + Self::Q8_BLOCK_SIZE - 1) / Self::Q8_BLOCK_SIZE;
|
|
let mut data = Vec::with_capacity(input.len());
|
|
let mut scales = Vec::with_capacity(num_blocks);
|
|
|
|
for block in input.chunks(Self::Q8_BLOCK_SIZE) {
|
|
// Find max absolute value for scale
|
|
let max_abs = block.iter().map(|x| x.abs()).fold(0.0f32, f32::max);
|
|
let scale = max_abs / 127.0; // Q8 range is -128 to 127
|
|
scales.push(scale);
|
|
|
|
// Quantize
|
|
let inv_scale = if scale > 1e-10 { 1.0 / scale } else { 0.0 };
|
|
for &x in block {
|
|
let q = (x * inv_scale).round() as i8;
|
|
data.push(q);
|
|
}
|
|
}
|
|
|
|
(data, scales)
|
|
}
|
|
|
|
/// Dequantize Q8 data back to f32
|
|
#[inline]
|
|
pub fn dequantize_simd_q8(data: &[i8], scales: &[f32], output_len: usize) -> Vec<f32> {
|
|
let mut output = Vec::with_capacity(output_len);
|
|
|
|
for (block_idx, scale) in scales.iter().enumerate() {
|
|
let block_start = block_idx * Self::Q8_BLOCK_SIZE;
|
|
let block_end = ((block_idx + 1) * Self::Q8_BLOCK_SIZE).min(data.len());
|
|
|
|
for idx in block_start..block_end {
|
|
if output.len() >= output_len {
|
|
break;
|
|
}
|
|
output.push(data[idx] as f32 * scale);
|
|
}
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
/// Quantized matrix-vector multiplication (Q4 * f32 -> f32)
|
|
///
|
|
/// Efficient implementation that dequantizes on-the-fly without
|
|
/// allocating full dequantized matrix.
|
|
#[inline]
|
|
pub fn matvec_q4(
|
|
data: &[u8],
|
|
scales: &[f32],
|
|
x: &[f32],
|
|
m: usize,
|
|
n: usize,
|
|
) -> Vec<f32> {
|
|
let mut y = vec![0.0f32; m];
|
|
let total_elements = m * n;
|
|
let num_blocks = (total_elements + Self::Q4_BLOCK_SIZE - 1) / Self::Q4_BLOCK_SIZE;
|
|
|
|
for row in 0..m {
|
|
let mut sum = 0.0f32;
|
|
let row_offset = row * n;
|
|
|
|
for col in 0..n {
|
|
let idx = row_offset + col;
|
|
// Find which block this element belongs to
|
|
let block_idx = idx / Self::Q4_BLOCK_SIZE;
|
|
let scale = if block_idx < scales.len() {
|
|
scales[block_idx]
|
|
} else {
|
|
// Fallback for last partial block
|
|
scales.last().copied().unwrap_or(1.0)
|
|
};
|
|
|
|
let byte = data[idx / 2];
|
|
let q = if idx % 2 == 0 {
|
|
(byte & 0x0F) as i8
|
|
} else {
|
|
((byte >> 4) & 0x0F) as i8
|
|
};
|
|
let q = if q > 7 { q - 16 } else { q };
|
|
sum += q as f32 * scale * x[col];
|
|
}
|
|
|
|
y[row] = sum;
|
|
}
|
|
|
|
y
|
|
}
|
|
|
|
// ========================================================================
|
|
// Additional Activation Functions
|
|
// ========================================================================
|
|
|
|
/// SiLU (Swish) activation: x * sigmoid(x)
|
|
#[inline]
|
|
pub fn silu_simd(input: &mut [f32]) {
|
|
for x in input.iter_mut() {
|
|
*x = *x / (1.0 + (-*x).exp());
|
|
}
|
|
}
|
|
|
|
/// ReLU activation: max(0, x)
|
|
#[inline]
|
|
pub fn relu_simd(input: &mut [f32]) {
|
|
#[cfg(target_arch = "x86_64")]
|
|
{
|
|
if is_x86_feature_detected!("avx2") {
|
|
unsafe { Self::relu_avx2(input) };
|
|
return;
|
|
}
|
|
}
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
{
|
|
Self::relu_wasm_simd128(input);
|
|
return;
|
|
}
|
|
|
|
for x in input.iter_mut() {
|
|
*x = x.max(0.0);
|
|
}
|
|
}
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
#[inline]
|
|
fn relu_wasm_simd128(input: &mut [f32]) {
|
|
let len = input.len();
|
|
let chunks = len / 4;
|
|
let zero = f32x4_splat(0.0);
|
|
|
|
for i in 0..chunks {
|
|
let offset = i * 4;
|
|
let v = unsafe { v128_load(input.as_ptr().add(offset) as *const v128) };
|
|
let result = f32x4_pmax(v, zero);
|
|
unsafe {
|
|
v128_store(input.as_mut_ptr().add(offset) as *mut v128, result);
|
|
}
|
|
}
|
|
|
|
for i in (chunks * 4)..len {
|
|
input[i] = input[i].max(0.0);
|
|
}
|
|
}
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
#[target_feature(enable = "avx2")]
|
|
#[inline]
|
|
unsafe fn relu_avx2(input: &mut [f32]) {
|
|
let len = input.len();
|
|
let chunks = len / 8;
|
|
let zero = _mm256_setzero_ps();
|
|
|
|
for i in 0..chunks {
|
|
let ptr = input.as_mut_ptr().add(i * 8);
|
|
let v = _mm256_loadu_ps(ptr);
|
|
let result = _mm256_max_ps(v, zero);
|
|
_mm256_storeu_ps(ptr, result);
|
|
}
|
|
|
|
for i in (chunks * 8)..len {
|
|
input[i] = input[i].max(0.0);
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Quantized Weight Storage
|
|
// ============================================================================
|
|
|
|
/// Q4 quantized weight matrix for memory-efficient inference
|
|
#[derive(Clone)]
|
|
pub struct Q4Weights {
|
|
/// Packed 4-bit quantized data
|
|
data: Vec<u8>,
|
|
/// Per-block scale factors
|
|
scales: Vec<f32>,
|
|
/// Matrix dimensions
|
|
rows: usize,
|
|
cols: usize,
|
|
}
|
|
|
|
impl Q4Weights {
|
|
/// Create Q4 weights from f32 matrix (row-major)
|
|
pub fn from_f32(weights: &[f32], rows: usize, cols: usize) -> Self {
|
|
debug_assert_eq!(weights.len(), rows * cols);
|
|
|
|
let (data, scales) = SimdCompute::quantize_simd_q4(weights);
|
|
|
|
Self {
|
|
data,
|
|
scales,
|
|
rows,
|
|
cols,
|
|
}
|
|
}
|
|
|
|
/// Matrix-vector multiplication with on-the-fly dequantization
|
|
pub fn matvec(&self, x: &[f32]) -> Vec<f32> {
|
|
debug_assert_eq!(x.len(), self.cols);
|
|
SimdCompute::matvec_q4(&self.data, &self.scales, x, self.rows, self.cols)
|
|
}
|
|
|
|
/// Get matrix dimensions
|
|
pub fn dims(&self) -> (usize, usize) {
|
|
(self.rows, self.cols)
|
|
}
|
|
|
|
/// Memory usage in bytes
|
|
pub fn memory_bytes(&self) -> usize {
|
|
self.data.len() + self.scales.len() * 4
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_dot_product_scalar() {
|
|
let a = vec![1.0, 2.0, 3.0, 4.0];
|
|
let b = vec![1.0, 1.0, 1.0, 1.0];
|
|
let result = SimdCompute::dot_product_scalar(&a, &b);
|
|
assert!((result - 10.0).abs() < 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dot_product_simd() {
|
|
let a: Vec<f32> = (0..256).map(|i| i as f32 * 0.1).collect();
|
|
let b: Vec<f32> = (0..256).map(|i| (255 - i) as f32 * 0.1).collect();
|
|
|
|
let scalar_result = SimdCompute::dot_product_scalar(&a, &b);
|
|
let simd_result = SimdCompute::dot_product(&a, &b);
|
|
|
|
assert!(
|
|
(scalar_result - simd_result).abs() < 0.1,
|
|
"Scalar: {}, SIMD: {}",
|
|
scalar_result,
|
|
simd_result
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_softmax_scalar() {
|
|
let mut values = vec![1.0, 2.0, 3.0];
|
|
SimdCompute::softmax_scalar(&mut values);
|
|
|
|
let sum: f32 = values.iter().sum();
|
|
assert!((sum - 1.0).abs() < 1e-5);
|
|
assert!(values[2] > values[1]);
|
|
assert!(values[1] > values[0]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_softmax_numerical_stability() {
|
|
// Test with large values that would overflow without max subtraction
|
|
let mut values = vec![1000.0, 1001.0, 1002.0];
|
|
SimdCompute::softmax_simd(&mut values);
|
|
|
|
let sum: f32 = values.iter().sum();
|
|
assert!((sum - 1.0).abs() < 1e-5);
|
|
assert!(values.iter().all(|&x| x.is_finite()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_gelu() {
|
|
let mut values = vec![-2.0, -1.0, 0.0, 1.0, 2.0];
|
|
SimdCompute::gelu_scalar(&mut values);
|
|
|
|
// GELU(0) = 0
|
|
assert!(values[2].abs() < 1e-5);
|
|
// GELU(-2) is very small negative, GELU(-1) is also small negative
|
|
// For large negative inputs, GELU approaches 0 from below
|
|
// GELU(-2) ~ -0.045, GELU(-1) ~ -0.158
|
|
// So GELU(-2) > GELU(-1) (less negative)
|
|
// For x > 0, GELU is monotonically increasing and positive
|
|
assert!(values[1] < values[2]); // GELU(-1) < GELU(0)
|
|
assert!(values[2] < values[3]); // GELU(0) < GELU(1)
|
|
assert!(values[3] < values[4]); // GELU(1) < GELU(2)
|
|
// GELU(-2) > GELU(-1) because GELU(-2) is closer to 0
|
|
assert!(values[0] > values[1]); // GELU(-2) > GELU(-1)
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_norm() {
|
|
let input = vec![1.0, 2.0, 3.0, 4.0];
|
|
let weight = vec![1.0, 1.0, 1.0, 1.0];
|
|
let bias = vec![0.0, 0.0, 0.0, 0.0];
|
|
|
|
let output = SimdCompute::layer_norm_simd(&input, &weight, Some(&bias), 1e-5);
|
|
|
|
// Mean of output should be ~0
|
|
let mean: f32 = output.iter().sum::<f32>() / output.len() as f32;
|
|
assert!(mean.abs() < 1e-5);
|
|
|
|
// Variance should be ~1
|
|
let var: f32 = output.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / output.len() as f32;
|
|
assert!((var - 1.0).abs() < 0.1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_rms_norm() {
|
|
let input = vec![1.0, 2.0, 3.0, 4.0];
|
|
let weight = vec![1.0, 1.0, 1.0, 1.0];
|
|
|
|
let output = SimdCompute::rms_norm_simd(&input, &weight, 1e-5);
|
|
|
|
assert_eq!(output.len(), input.len());
|
|
// RMS normalized values should be smaller for larger inputs
|
|
assert!(output[0].abs() < input[0].abs());
|
|
}
|
|
|
|
#[test]
|
|
fn test_q4_quantization() {
|
|
let input: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.1).collect();
|
|
|
|
let (data, scales) = SimdCompute::quantize_simd_q4(&input);
|
|
let output = SimdCompute::dequantize_simd_q4(&data, &scales, input.len());
|
|
|
|
assert_eq!(output.len(), input.len());
|
|
|
|
// Check that dequantized values are close to original
|
|
let max_error: f32 = input
|
|
.iter()
|
|
.zip(output.iter())
|
|
.map(|(a, b)| (a - b).abs())
|
|
.fold(0.0, f32::max);
|
|
|
|
// Q4 should have reasonable accuracy (within 10% of range)
|
|
let range = 6.4; // -3.2 to 3.2
|
|
assert!(max_error < range * 0.15, "Max error: {}", max_error);
|
|
}
|
|
|
|
#[test]
|
|
fn test_q8_quantization() {
|
|
let input: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.1).collect();
|
|
|
|
let (data, scales) = SimdCompute::quantize_simd_q8(&input);
|
|
let output = SimdCompute::dequantize_simd_q8(&data, &scales, input.len());
|
|
|
|
assert_eq!(output.len(), input.len());
|
|
|
|
// Q8 should be more accurate than Q4
|
|
let max_error: f32 = input
|
|
.iter()
|
|
.zip(output.iter())
|
|
.map(|(a, b)| (a - b).abs())
|
|
.fold(0.0, f32::max);
|
|
|
|
let range = 6.4;
|
|
assert!(max_error < range * 0.02, "Max error: {}", max_error);
|
|
}
|
|
|
|
#[test]
|
|
fn test_q4_weights() {
|
|
let weights: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.01).collect();
|
|
let q4 = Q4Weights::from_f32(&weights, 8, 8);
|
|
|
|
assert_eq!(q4.dims(), (8, 8));
|
|
|
|
// Test matvec
|
|
let x = vec![1.0; 8];
|
|
let y = q4.matvec(&x);
|
|
assert_eq!(y.len(), 8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_matvec() {
|
|
// 2x3 matrix times 3-vector
|
|
let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
|
|
let x = vec![1.0, 1.0, 1.0];
|
|
|
|
let y = SimdCompute::matvec_simd(&a, &x, 2, 3);
|
|
|
|
assert_eq!(y.len(), 2);
|
|
assert!((y[0] - 6.0).abs() < 1e-5); // 1+2+3
|
|
assert!((y[1] - 15.0).abs() < 1e-5); // 4+5+6
|
|
}
|
|
|
|
#[test]
|
|
fn test_matmul() {
|
|
// 2x2 * 2x2
|
|
let a = vec![1.0, 2.0, 3.0, 4.0];
|
|
let b = vec![5.0, 6.0, 7.0, 8.0];
|
|
|
|
let c = SimdCompute::matmul_simd(&a, &b, 2, 2, 2);
|
|
|
|
assert_eq!(c.len(), 4);
|
|
// [[1,2],[3,4]] * [[5,6],[7,8]] = [[19,22],[43,50]]
|
|
assert!((c[0] - 19.0).abs() < 1e-4, "c[0]={}", c[0]);
|
|
assert!((c[1] - 22.0).abs() < 1e-4, "c[1]={}", c[1]);
|
|
assert!((c[2] - 43.0).abs() < 1e-4, "c[2]={}", c[2]);
|
|
assert!((c[3] - 50.0).abs() < 1e-4, "c[3]={}", c[3]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_relu() {
|
|
let mut values = vec![-2.0, -1.0, 0.0, 1.0, 2.0];
|
|
SimdCompute::relu_simd(&mut values);
|
|
|
|
assert_eq!(values, vec![0.0, 0.0, 0.0, 1.0, 2.0]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_silu() {
|
|
let mut values = vec![0.0, 1.0, -1.0];
|
|
SimdCompute::silu_simd(&mut values);
|
|
|
|
// SiLU(0) = 0
|
|
assert!(values[0].abs() < 1e-5);
|
|
// SiLU(1) ~ 0.731
|
|
assert!((values[1] - 0.731).abs() < 0.01);
|
|
// SiLU(-1) ~ -0.269
|
|
assert!((values[2] + 0.269).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_welford() {
|
|
let data = vec![2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];
|
|
let (mean, var) = SimdCompute::welford_mean_var(&data);
|
|
|
|
assert!((mean - 5.0).abs() < 1e-5);
|
|
assert!((var - 4.0).abs() < 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_capabilities_detection() {
|
|
let caps = SimdCapabilities::detect();
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
assert!(caps.wasm_simd128);
|
|
|
|
// lane_width should be at least 1
|
|
assert!(caps.lane_width() >= 1);
|
|
}
|
|
}
|