rustytorch/crates/production/rtx-wasm-inference/examples/browser_inference.js

/**
 * RustyTorch WASM Inference Example
 *
 * This example demonstrates using the WASM tensor operations
 * in a browser environment with automatic WebGPU/CPU fallback.
 *
 * Usage:
 *   1. Build the WASM package: wasm-pack build --target web
 *   2. Serve this directory: python -m http.server 8080
 *   3. Open http://localhost:8080/examples/browser_inference.html
 */

import init, {
    WasmTensor,
    WasmBackend,
    getBackendInfo,
    isSIMDAvailable,
    getSIMDWidth,
} from '../pkg/rtx_wasm_inference.js';

// Initialize WASM module
async function initialize() {
    await init();
    console.log('WASM module initialized');
    console.log('Backend info:', getBackendInfo());
    console.log('SIMD available:', isSIMDAvailable());
    console.log('SIMD width:', getSIMDWidth());
}

// Basic tensor operations
function basicOperations() {
    console.log('\n=== Basic Tensor Operations ===');

    // Create tensors
    const zeros = WasmTensor.zeros([2, 3]);
    console.log('Zeros shape:', zeros.shape());
    console.log('Zeros data:', zeros.toArray());

    const ones = WasmTensor.ones([2, 3]);
    console.log('Ones shape:', ones.shape());

    // From array
    const data = new Float32Array([1, 2, 3, 4, 5, 6]);
    const custom = WasmTensor.fromArray(data, [2, 3]);
    console.log('Custom tensor:', custom.toArray());

    // Reshape
    const reshaped = custom.reshape([3, 2]);
    console.log('Reshaped to [3, 2]:', reshaped.shape());

    // Element access
    console.log('Element at index 0:', custom.getFlat(0));
    console.log('Element at index 5:', custom.getFlat(5));
}

// Arithmetic operations
function arithmeticOperations() {
    console.log('\n=== Arithmetic Operations ===');

    const a = WasmTensor.fromArray(new Float32Array([1, 2, 3, 4]), [2, 2]);
    const b = WasmTensor.fromArray(new Float32Array([5, 6, 7, 8]), [2, 2]);

    // Addition
    const sum = a.add(b);
    console.log('a + b:', sum.toArray());  // [6, 8, 10, 12]

    // Subtraction
    const diff = a.sub(b);
    console.log('a - b:', diff.toArray());  // [-4, -4, -4, -4]

    // Multiplication
    const prod = a.mul(b);
    console.log('a * b:', prod.toArray());  // [5, 12, 21, 32]

    // Scalar operations
    const scaled = a.scale(2.0);
    console.log('a * 2:', scaled.toArray());  // [2, 4, 6, 8]

    const shifted = a.addScalar(10.0);
    console.log('a + 10:', shifted.toArray());  // [11, 12, 13, 14]
}

// Activation functions
function activationFunctions() {
    console.log('\n=== Activation Functions ===');

    const x = WasmTensor.fromArray(new Float32Array([-2, -1, 0, 1, 2]), [5]);
    console.log('Input:', x.toArray());

    // ReLU
    const relu = x.relu();
    console.log('ReLU:', relu.toArray());  // [0, 0, 0, 1, 2]

    // GELU
    const gelu = x.gelu();
    console.log('GELU:', gelu.toArray());

    // Sigmoid
    const sigmoid = x.sigmoid();
    console.log('Sigmoid:', sigmoid.toArray());

    // Tanh
    const tanh = x.tanh();
    console.log('Tanh:', tanh.toArray());

    // SiLU (Swish)
    const silu = x.silu();
    console.log('SiLU:', silu.toArray());
}

// Reduction operations
function reductionOperations() {
    console.log('\n=== Reduction Operations ===');

    const x = WasmTensor.fromArray(new Float32Array([1, 2, 3, 4, 5]), [5]);
    console.log('Input:', x.toArray());

    console.log('Sum:', x.sum());    // 15
    console.log('Mean:', x.mean());  // 3
    console.log('Max:', x.max());    // 5
    console.log('Min:', x.min());    // 1
}

// Matrix operations
function matrixOperations() {
    console.log('\n=== Matrix Operations ===');

    // Matrix multiplication: [2,3] @ [3,2] = [2,2]
    const a = WasmTensor.fromArray(new Float32Array([1, 2, 3, 4, 5, 6]), [2, 3]);
    const b = WasmTensor.fromArray(new Float32Array([1, 2, 3, 4, 5, 6]), [3, 2]);

    console.log('A shape:', a.shape());
    console.log('B shape:', b.shape());

    const c = a.matmul(b);
    console.log('A @ B shape:', c.shape());  // [2, 2]
    console.log('A @ B:', c.toArray());      // [22, 28, 49, 64]

    // Dot product
    const v1 = WasmTensor.fromArray(new Float32Array([1, 2, 3]), [3]);
    const v2 = WasmTensor.fromArray(new Float32Array([4, 5, 6]), [3]);
    const dot = v1.dot(v2);
    console.log('Dot product:', dot);  // 1*4 + 2*5 + 3*6 = 32
}

// Normalization operations
function normalizationOperations() {
    console.log('\n=== Normalization Operations ===');

    // Softmax
    const logits = WasmTensor.fromArray(new Float32Array([1, 2, 3, 4]), [4]);
    const probs = logits.softmax();
    console.log('Logits:', logits.toArray());
    console.log('Softmax:', probs.toArray());
    console.log('Sum (should be 1):', probs.sum());

    // Layer normalization
    const x = WasmTensor.fromArray(new Float32Array([1, 2, 3, 4, 5, 6, 7, 8]), [2, 4]);
    const normalized = x.layerNorm(1e-5);
    console.log('Input shape:', x.shape());
    console.log('Layer norm:', normalized.toArray());
}

// Backend-aware operations
function backendOperations() {
    console.log('\n=== Backend Operations ===');

    // Create backend (auto-selects WebGPU or CPU)
    const backend = new WasmBackend();
    console.log('Backend name:', backend.name());
    console.log('Is WebGPU:', backend.isWebGpu());
    console.log('Is fallback:', backend.isFallback());
    console.log('SIMD available:', backend.simdAvailable());
    console.log('SIMD width:', backend.simdWidth());

    // Use backend for operations
    const a = WasmTensor.ones([3, 3]);
    const b = WasmTensor.full([3, 3], 2.0);

    const sum = backend.tensorAdd(a, b);
    console.log('Backend add:', sum.toArray());

    const prod = backend.tensorMul(a, b);
    console.log('Backend mul:', prod.toArray());

    // CPU-only backend
    const cpuBackend = WasmBackend.cpuOnly();
    console.log('CPU backend name:', cpuBackend.name());
}

// Neural network simulation
function neuralNetworkDemo() {
    console.log('\n=== Neural Network Demo ===');

    // Simple 2-layer MLP forward pass
    // Input: [batch=2, features=4]
    // Layer 1: Linear(4, 8) + ReLU
    // Layer 2: Linear(8, 3) + Softmax

    const batchSize = 2;
    const inputFeatures = 4;
    const hiddenSize = 8;
    const outputSize = 3;

    // Create random-ish input and weights
    const input = WasmTensor.fromArray(
        new Float32Array([
            0.1, 0.2, 0.3, 0.4,
            0.5, 0.6, 0.7, 0.8
        ]),
        [batchSize, inputFeatures]
    );

    // Weight matrices (normally these would be learned)
    const w1Data = new Float32Array(inputFeatures * hiddenSize);
    for (let i = 0; i < w1Data.length; i++) {
        w1Data[i] = (i % 10) * 0.1 - 0.5;
    }
    const w1 = WasmTensor.fromArray(w1Data, [inputFeatures, hiddenSize]);

    const w2Data = new Float32Array(hiddenSize * outputSize);
    for (let i = 0; i < w2Data.length; i++) {
        w2Data[i] = (i % 10) * 0.1 - 0.5;
    }
    const w2 = WasmTensor.fromArray(w2Data, [hiddenSize, outputSize]);

    console.log('Input shape:', input.shape());
    console.log('W1 shape:', w1.shape());
    console.log('W2 shape:', w2.shape());

    // Forward pass
    const hidden = input.matmul(w1).relu();
    console.log('Hidden shape:', hidden.shape());

    const output = hidden.matmul(w2).softmax();
    console.log('Output shape:', output.shape());
    console.log('Output (probabilities):', output.toArray());

    // Check probabilities sum to 1 for each sample
    const outputData = output.toArray();
    for (let b = 0; b < batchSize; b++) {
        let sum = 0;
        for (let i = 0; i < outputSize; i++) {
            sum += outputData[b * outputSize + i];
        }
        console.log(`Sample ${b} probability sum:`, sum.toFixed(6));
    }
}

// Performance benchmark
function performanceBenchmark() {
    console.log('\n=== Performance Benchmark ===');

    const sizes = [100, 500, 1000];
    const iterations = 10;

    for (const size of sizes) {
        const a = WasmTensor.ones([size, size]);
        const b = WasmTensor.full([size, size], 2.0);

        // Warm up
        a.add(b);
        a.matmul(b);

        // Benchmark add
        const addStart = performance.now();
        for (let i = 0; i < iterations; i++) {
            a.add(b);
        }
        const addTime = (performance.now() - addStart) / iterations;
        console.log(`Add [${size}x${size}]: ${addTime.toFixed(2)}ms`);

        // Benchmark matmul
        const matmulStart = performance.now();
        for (let i = 0; i < iterations; i++) {
            a.matmul(b);
        }
        const matmulTime = (performance.now() - matmulStart) / iterations;
        console.log(`Matmul [${size}x${size}]: ${matmulTime.toFixed(2)}ms`);

        // Calculate GFLOPS for matmul
        const flops = 2 * size * size * size;  // 2*n^3 for matmul
        const gflops = flops / (matmulTime * 1e6);
        console.log(`Matmul GFLOPS: ${gflops.toFixed(2)}`);
    }
}

// Main entry point
async function main() {
    try {
        await initialize();

        basicOperations();
        arithmeticOperations();
        activationFunctions();
        reductionOperations();
        matrixOperations();
        normalizationOperations();
        backendOperations();
        neuralNetworkDemo();
        performanceBenchmark();

        console.log('\n=== All examples completed successfully! ===');
    } catch (error) {
        console.error('Error:', error);
    }
}

// Run if in browser
if (typeof window !== 'undefined') {
    main();
}

// Export for module usage
export { main, initialize };