518 lines
14 KiB
Rust
518 lines
14 KiB
Rust
//! WASM-compatible Tensor Operations
|
|
//!
|
|
//! CPU-only tensor operations optimized for WebAssembly execution.
|
|
//! Supports optional SIMD acceleration when available.
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
use wasm_bindgen::prelude::*;
|
|
|
|
/// Data type for tensor elements
|
|
#[wasm_bindgen]
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum WasmDType {
|
|
/// 32-bit floating point
|
|
Float32,
|
|
/// 16-bit floating point (stored as f32)
|
|
Float16,
|
|
/// 8-bit integer (quantized)
|
|
Int8,
|
|
/// 4-bit integer (quantized, packed)
|
|
Int4,
|
|
}
|
|
|
|
impl WasmDType {
|
|
/// Get bytes per element
|
|
pub fn bytes_per_element(&self) -> usize {
|
|
match self {
|
|
WasmDType::Float32 => 4,
|
|
WasmDType::Float16 => 2,
|
|
WasmDType::Int8 => 1,
|
|
WasmDType::Int4 => 1, // Packed 2 per byte
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A tensor for WASM computation
|
|
#[wasm_bindgen]
|
|
#[derive(Debug, Clone)]
|
|
pub struct WasmTensor {
|
|
/// Data storage
|
|
data: Vec<f32>,
|
|
/// Shape
|
|
shape: Vec<usize>,
|
|
/// Data type
|
|
dtype: WasmDType,
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl WasmTensor {
|
|
/// Create a new tensor filled with zeros
|
|
#[wasm_bindgen(constructor)]
|
|
pub fn zeros(shape: Vec<usize>) -> Self {
|
|
let size: usize = shape.iter().product();
|
|
Self {
|
|
data: vec![0.0; size],
|
|
shape,
|
|
dtype: WasmDType::Float32,
|
|
}
|
|
}
|
|
|
|
/// Create a tensor from a flat array
|
|
#[wasm_bindgen]
|
|
pub fn from_array(data: Vec<f32>, shape: Vec<usize>) -> Result<WasmTensor, JsError> {
|
|
let expected_size: usize = shape.iter().product();
|
|
if data.len() != expected_size {
|
|
return Err(JsError::new(&format!(
|
|
"Data size {} doesn't match shape {:?} (expected {})",
|
|
data.len(),
|
|
shape,
|
|
expected_size
|
|
)));
|
|
}
|
|
Ok(Self {
|
|
data,
|
|
shape,
|
|
dtype: WasmDType::Float32,
|
|
})
|
|
}
|
|
|
|
/// Create a tensor filled with a value
|
|
#[wasm_bindgen]
|
|
pub fn full(shape: Vec<usize>, value: f32) -> Self {
|
|
let size: usize = shape.iter().product();
|
|
Self {
|
|
data: vec![value; size],
|
|
shape,
|
|
dtype: WasmDType::Float32,
|
|
}
|
|
}
|
|
|
|
/// Create a tensor with random values
|
|
#[wasm_bindgen]
|
|
pub fn randn(shape: Vec<usize>) -> Self {
|
|
let size: usize = shape.iter().product();
|
|
let mut data = Vec::with_capacity(size);
|
|
|
|
// Box-Muller transform for normal distribution
|
|
for _ in 0..size {
|
|
let u1: f32 = rand::random();
|
|
let u2: f32 = rand::random();
|
|
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
|
|
data.push(z);
|
|
}
|
|
|
|
Self {
|
|
data,
|
|
shape,
|
|
dtype: WasmDType::Float32,
|
|
}
|
|
}
|
|
|
|
/// Get tensor shape
|
|
#[wasm_bindgen(getter)]
|
|
pub fn shape(&self) -> Vec<usize> {
|
|
self.shape.clone()
|
|
}
|
|
|
|
/// Get number of dimensions
|
|
#[wasm_bindgen(getter)]
|
|
pub fn ndim(&self) -> usize {
|
|
self.shape.len()
|
|
}
|
|
|
|
/// Get total number of elements
|
|
#[wasm_bindgen(getter)]
|
|
pub fn numel(&self) -> usize {
|
|
self.data.len()
|
|
}
|
|
|
|
/// Get data type
|
|
#[wasm_bindgen(getter)]
|
|
pub fn dtype(&self) -> WasmDType {
|
|
self.dtype
|
|
}
|
|
|
|
/// Get data as array
|
|
#[wasm_bindgen]
|
|
pub fn to_array(&self) -> Vec<f32> {
|
|
self.data.clone()
|
|
}
|
|
|
|
/// Get element at index
|
|
#[wasm_bindgen]
|
|
pub fn get(&self, index: usize) -> Option<f32> {
|
|
self.data.get(index).copied()
|
|
}
|
|
|
|
/// Set element at index
|
|
#[wasm_bindgen]
|
|
pub fn set(&mut self, index: usize, value: f32) -> Result<(), JsError> {
|
|
if index >= self.data.len() {
|
|
return Err(JsError::new("Index out of bounds"));
|
|
}
|
|
self.data[index] = value;
|
|
Ok(())
|
|
}
|
|
|
|
/// Reshape tensor
|
|
#[wasm_bindgen]
|
|
pub fn reshape(&self, new_shape: Vec<usize>) -> Result<WasmTensor, JsError> {
|
|
let new_size: usize = new_shape.iter().product();
|
|
if new_size != self.data.len() {
|
|
return Err(JsError::new(&format!(
|
|
"Cannot reshape tensor of size {} to shape {:?}",
|
|
self.data.len(),
|
|
new_shape
|
|
)));
|
|
}
|
|
Ok(Self {
|
|
data: self.data.clone(),
|
|
shape: new_shape,
|
|
dtype: self.dtype,
|
|
})
|
|
}
|
|
|
|
/// Element-wise addition
|
|
#[wasm_bindgen]
|
|
pub fn add(&self, other: &WasmTensor) -> Result<WasmTensor, JsError> {
|
|
if self.shape != other.shape {
|
|
return Err(JsError::new("Shape mismatch for addition"));
|
|
}
|
|
let data: Vec<f32> = self
|
|
.data
|
|
.iter()
|
|
.zip(other.data.iter())
|
|
.map(|(a, b)| a + b)
|
|
.collect();
|
|
Ok(Self {
|
|
data,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
})
|
|
}
|
|
|
|
/// Element-wise multiplication
|
|
#[wasm_bindgen]
|
|
pub fn mul(&self, other: &WasmTensor) -> Result<WasmTensor, JsError> {
|
|
if self.shape != other.shape {
|
|
return Err(JsError::new("Shape mismatch for multiplication"));
|
|
}
|
|
let data: Vec<f32> = self
|
|
.data
|
|
.iter()
|
|
.zip(other.data.iter())
|
|
.map(|(a, b)| a * b)
|
|
.collect();
|
|
Ok(Self {
|
|
data,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
})
|
|
}
|
|
|
|
/// Scalar multiplication
|
|
#[wasm_bindgen]
|
|
pub fn scale(&self, scalar: f32) -> WasmTensor {
|
|
let data: Vec<f32> = self.data.iter().map(|x| x * scalar).collect();
|
|
Self {
|
|
data,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
}
|
|
}
|
|
|
|
/// Matrix multiplication (2D tensors)
|
|
#[wasm_bindgen]
|
|
pub fn matmul(&self, other: &WasmTensor) -> Result<WasmTensor, JsError> {
|
|
if self.shape.len() != 2 || other.shape.len() != 2 {
|
|
return Err(JsError::new("matmul requires 2D tensors"));
|
|
}
|
|
if self.shape[1] != other.shape[0] {
|
|
return Err(JsError::new(&format!(
|
|
"Shape mismatch: {:?} @ {:?}",
|
|
self.shape, other.shape
|
|
)));
|
|
}
|
|
|
|
let m = self.shape[0];
|
|
let k = self.shape[1];
|
|
let n = other.shape[1];
|
|
|
|
let mut result = vec![0.0; m * n];
|
|
|
|
// Basic GEMM (can be optimized with SIMD)
|
|
for i in 0..m {
|
|
for j in 0..n {
|
|
let mut sum = 0.0;
|
|
for l in 0..k {
|
|
sum += self.data[i * k + l] * other.data[l * n + j];
|
|
}
|
|
result[i * n + j] = sum;
|
|
}
|
|
}
|
|
|
|
Ok(Self {
|
|
data: result,
|
|
shape: vec![m, n],
|
|
dtype: self.dtype,
|
|
})
|
|
}
|
|
|
|
/// Softmax along last dimension
|
|
#[wasm_bindgen]
|
|
pub fn softmax(&self) -> WasmTensor {
|
|
let last_dim = *self.shape.last().unwrap_or(&1);
|
|
let batch_size = self.data.len() / last_dim;
|
|
|
|
let mut result = vec![0.0; self.data.len()];
|
|
|
|
for b in 0..batch_size {
|
|
let start = b * last_dim;
|
|
let end = start + last_dim;
|
|
let slice = &self.data[start..end];
|
|
|
|
// Find max for numerical stability
|
|
let max_val = slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
|
|
|
// Compute exp and sum
|
|
let exp_vals: Vec<f32> = slice.iter().map(|x| (x - max_val).exp()).collect();
|
|
let sum: f32 = exp_vals.iter().sum();
|
|
|
|
// Normalize
|
|
for (i, val) in exp_vals.iter().enumerate() {
|
|
result[start + i] = val / sum;
|
|
}
|
|
}
|
|
|
|
Self {
|
|
data: result,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
}
|
|
}
|
|
|
|
/// ReLU activation
|
|
#[wasm_bindgen]
|
|
pub fn relu(&self) -> WasmTensor {
|
|
let data: Vec<f32> = self.data.iter().map(|x| x.max(0.0)).collect();
|
|
Self {
|
|
data,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
}
|
|
}
|
|
|
|
/// GELU activation
|
|
#[wasm_bindgen]
|
|
pub fn gelu(&self) -> WasmTensor {
|
|
let data: Vec<f32> = self
|
|
.data
|
|
.iter()
|
|
.map(|x| {
|
|
0.5 * x
|
|
* (1.0
|
|
+ ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
|
|
})
|
|
.collect();
|
|
Self {
|
|
data,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
}
|
|
}
|
|
|
|
/// SiLU (Swish) activation
|
|
#[wasm_bindgen]
|
|
pub fn silu(&self) -> WasmTensor {
|
|
let data: Vec<f32> = self.data.iter().map(|x| x / (1.0 + (-x).exp())).collect();
|
|
Self {
|
|
data,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
}
|
|
}
|
|
|
|
/// Layer normalization
|
|
#[wasm_bindgen]
|
|
pub fn layer_norm(&self, eps: f32) -> WasmTensor {
|
|
let last_dim = *self.shape.last().unwrap_or(&1);
|
|
let batch_size = self.data.len() / last_dim;
|
|
|
|
let mut result = vec![0.0; self.data.len()];
|
|
|
|
for b in 0..batch_size {
|
|
let start = b * last_dim;
|
|
let end = start + last_dim;
|
|
let slice = &self.data[start..end];
|
|
|
|
// Compute mean
|
|
let mean: f32 = slice.iter().sum::<f32>() / last_dim as f32;
|
|
|
|
// Compute variance
|
|
let var: f32 = slice.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / last_dim as f32;
|
|
|
|
// Normalize
|
|
let std = (var + eps).sqrt();
|
|
for (i, val) in slice.iter().enumerate() {
|
|
result[start + i] = (val - mean) / std;
|
|
}
|
|
}
|
|
|
|
Self {
|
|
data: result,
|
|
shape: self.shape.clone(),
|
|
dtype: self.dtype,
|
|
}
|
|
}
|
|
|
|
/// Sum all elements
|
|
#[wasm_bindgen]
|
|
pub fn sum(&self) -> f32 {
|
|
self.data.iter().sum()
|
|
}
|
|
|
|
/// Mean of all elements
|
|
#[wasm_bindgen]
|
|
pub fn mean(&self) -> f32 {
|
|
self.sum() / self.data.len() as f32
|
|
}
|
|
|
|
/// Max of all elements
|
|
#[wasm_bindgen]
|
|
pub fn max(&self) -> f32 {
|
|
self.data.iter().copied().fold(f32::NEG_INFINITY, f32::max)
|
|
}
|
|
|
|
/// Argmax (index of maximum element)
|
|
#[wasm_bindgen]
|
|
pub fn argmax(&self) -> usize {
|
|
self.data
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
|
|
.map_or(0, |(i, _)| i)
|
|
}
|
|
|
|
/// Memory usage in bytes
|
|
#[wasm_bindgen]
|
|
pub fn memory_usage(&self) -> usize {
|
|
self.data.len() * std::mem::size_of::<f32>()
|
|
}
|
|
}
|
|
|
|
/// KV-Cache for transformer inference
|
|
#[wasm_bindgen]
|
|
#[derive(Debug, Clone)]
|
|
pub struct WasmKvCache {
|
|
/// Key cache: [num_layers, num_heads, max_seq_len, head_dim]
|
|
keys: Vec<f32>,
|
|
/// Value cache: [num_layers, num_heads, max_seq_len, head_dim]
|
|
values: Vec<f32>,
|
|
/// Configuration
|
|
max_seq_len: usize,
|
|
num_layers: usize,
|
|
num_heads: usize,
|
|
head_dim: usize,
|
|
/// Current position
|
|
position: usize,
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl WasmKvCache {
|
|
/// Create a new KV cache
|
|
#[wasm_bindgen(constructor)]
|
|
pub fn new(max_seq_len: usize, num_layers: usize, num_heads: usize, head_dim: usize) -> Self {
|
|
let cache_size = num_layers * num_heads * max_seq_len * head_dim;
|
|
Self {
|
|
keys: vec![0.0; cache_size],
|
|
values: vec![0.0; cache_size],
|
|
max_seq_len,
|
|
num_layers,
|
|
num_heads,
|
|
head_dim,
|
|
position: 0,
|
|
}
|
|
}
|
|
|
|
/// Get current position
|
|
#[wasm_bindgen(getter)]
|
|
pub fn position(&self) -> usize {
|
|
self.position
|
|
}
|
|
|
|
/// Advance position by one
|
|
pub fn step(&mut self) {
|
|
self.position = (self.position + 1).min(self.max_seq_len - 1);
|
|
}
|
|
|
|
/// Clear the cache
|
|
#[wasm_bindgen]
|
|
pub fn clear(&mut self) {
|
|
self.keys.fill(0.0);
|
|
self.values.fill(0.0);
|
|
self.position = 0;
|
|
}
|
|
|
|
/// Memory usage in bytes
|
|
#[wasm_bindgen]
|
|
pub fn memory_usage(&self) -> usize {
|
|
(self.keys.len() + self.values.len()) * std::mem::size_of::<f32>()
|
|
}
|
|
|
|
/// Get cache utilization (0.0 - 1.0)
|
|
#[wasm_bindgen]
|
|
pub fn utilization(&self) -> f32 {
|
|
self.position as f32 / self.max_seq_len as f32
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_tensor_zeros() {
|
|
let t = WasmTensor::zeros(vec![2, 3]);
|
|
assert_eq!(t.shape(), vec![2, 3]);
|
|
assert_eq!(t.numel(), 6);
|
|
assert_eq!(t.sum(), 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_from_array() {
|
|
let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
|
|
let t = WasmTensor::from_array(data, vec![2, 3]).unwrap();
|
|
assert_eq!(t.get(0), Some(1.0));
|
|
assert_eq!(t.get(5), Some(6.0));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_matmul() {
|
|
// 2x3 @ 3x2 = 2x2
|
|
let a = WasmTensor::from_array(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
|
|
let b = WasmTensor::from_array(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]).unwrap();
|
|
let c = a.matmul(&b).unwrap();
|
|
|
|
assert_eq!(c.shape(), vec![2, 2]);
|
|
// [1,2,3] @ [1,2] = 1+6+15 = 22
|
|
assert!((c.get(0).unwrap() - 22.0).abs() < 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_softmax() {
|
|
let t = WasmTensor::from_array(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
|
|
let s = t.softmax();
|
|
assert!((s.sum() - 1.0).abs() < 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_kv_cache() {
|
|
let mut cache = WasmKvCache::new(1024, 32, 32, 128);
|
|
assert_eq!(cache.position(), 0);
|
|
cache.step();
|
|
assert_eq!(cache.position(), 1);
|
|
cache.clear();
|
|
assert_eq!(cache.position(), 0);
|
|
}
|
|
}
|