rustytorch/benchmarks/wasm/rtx_wasm_bench.js

/**
 * RustyTorch WASM Inference Benchmark
 *
 * This module benchmarks the rtx-wasm-inference compiled WASM module.
 * Usage:
 *   - Node.js: node rtx_wasm_bench.js
 *   - Browser: Import as ES module
 *
 * Build the WASM module first:
 *   wasm-pack build crates/production/rtx-wasm-inference --target web
 */

// For Node.js compatibility
const isNode = typeof window === 'undefined';

/**
 * Benchmark configuration
 */
export const defaultConfig = {
    modelSize: 'small',     // tiny, small, medium
    seqLength: 64,          // Sequence length
    iterations: 100,        // Benchmark iterations
    warmup: 10,             // Warmup iterations
};

/**
 * Benchmark result structure
 */
export class BenchmarkResult {
    constructor(framework, metrics) {
        this.framework = framework;
        this.timestamp = new Date().toISOString();
        this.avgLatencyMs = metrics.avgLatencyMs;
        this.p50LatencyMs = metrics.p50LatencyMs;
        this.p95LatencyMs = metrics.p95LatencyMs;
        this.p99LatencyMs = metrics.p99LatencyMs;
        this.minLatencyMs = metrics.minLatencyMs;
        this.maxLatencyMs = metrics.maxLatencyMs;
        this.throughputTokensPerSec = metrics.throughputTokensPerSec;
        this.memoryMB = metrics.memoryMB;
        this.loadTimeMs = metrics.loadTimeMs;
    }

    toJSON() {
        return {
            framework: this.framework,
            timestamp: this.timestamp,
            metrics: {
                avgLatencyMs: this.avgLatencyMs,
                p50LatencyMs: this.p50LatencyMs,
                p95LatencyMs: this.p95LatencyMs,
                p99LatencyMs: this.p99LatencyMs,
                minLatencyMs: this.minLatencyMs,
                maxLatencyMs: this.maxLatencyMs,
                throughputTokensPerSec: this.throughputTokensPerSec,
                memoryMB: this.memoryMB,
                loadTimeMs: this.loadTimeMs,
            }
        };
    }
}

/**
 * Calculate percentile from sorted array
 */
function percentile(sortedArr, p) {
    const idx = Math.ceil((p / 100) * sortedArr.length) - 1;
    return sortedArr[Math.max(0, Math.min(idx, sortedArr.length - 1))];
}

/**
 * Get memory usage in MB
 */
function getMemoryMB() {
    if (isNode) {
        const usage = process.memoryUsage();
        return usage.heapUsed / 1024 / 1024;
    } else if (performance.memory) {
        return performance.memory.usedJSHeapSize / 1024 / 1024;
    }
    return null;
}

/**
 * High-resolution timer
 */
function now() {
    if (isNode) {
        const [sec, nsec] = process.hrtime();
        return sec * 1000 + nsec / 1e6;
    }
    return performance.now();
}

/**
 * Load RustyTorch WASM module
 */
async function loadRtxWasm() {
    // Path to compiled WASM module
    const wasmPath = isNode
        ? '../../crates/production/rtx-wasm-inference/pkg/rtx_wasm_inference.js'
        : '/pkg/rtx_wasm_inference.js';

    try {
        const rtx = await import(wasmPath);
        await rtx.default(); // Initialize WASM
        return rtx;
    } catch (err) {
        console.warn('RustyTorch WASM not found. Using mock implementation.');
        console.warn('Build with: wasm-pack build crates/production/rtx-wasm-inference --target web');
        return null;
    }
}

/**
 * Create mock inference engine for testing
 */
function createMockEngine(config) {
    const modelSizes = {
        tiny: { hiddenSize: 256, numLayers: 2 },
        small: { hiddenSize: 512, numLayers: 4 },
        medium: { hiddenSize: 1024, numLayers: 8 },
    };

    const spec = modelSizes[config.modelSize];

    return {
        async infer(tokens) {
            // Simulate computation time based on model size
            const baseTime = spec.hiddenSize * spec.numLayers * 0.00001;
            const computeTime = baseTime * tokens.length;
            await new Promise(r => setTimeout(r, computeTime));

            // Return mock logits
            return new Float32Array(tokens.length * 32000).map(() => Math.random() * 2 - 1);
        },
        getMemoryUsage() {
            return spec.hiddenSize * spec.numLayers * 4 / 1024; // MB
        }
    };
}

/**
 * Run RustyTorch WASM benchmark
 */
export async function runRtxBenchmark(config = defaultConfig) {
    console.log('=== RustyTorch WASM Benchmark ===');
    console.log(`Model: ${config.modelSize}, Seq: ${config.seqLength}, Iterations: ${config.iterations}`);

    // Load WASM module
    const loadStart = now();
    const rtx = await loadRtxWasm();
    const loadTime = now() - loadStart;

    let engine;
    if (rtx) {
        // Use real RustyTorch WASM
        const modelConfig = rtx.ModelConfig[config.modelSize] || rtx.ModelConfig.tiny();
        engine = new rtx.WasmInferenceEngine(modelConfig);
    } else {
        // Use mock for testing
        engine = createMockEngine(config);
    }

    // Create input tokens
    const inputTokens = new Uint32Array(config.seqLength);
    for (let i = 0; i < config.seqLength; i++) {
        inputTokens[i] = Math.floor(Math.random() * 32000);
    }

    // Warmup
    console.log('Warming up...');
    for (let i = 0; i < config.warmup; i++) {
        await engine.infer(inputTokens);
    }

    // Benchmark
    console.log('Benchmarking...');
    const latencies = [];
    const memoryBefore = getMemoryMB();

    for (let i = 0; i < config.iterations; i++) {
        const start = now();
        await engine.infer(inputTokens);
        const elapsed = now() - start;
        latencies.push(elapsed);

        if ((i + 1) % 20 === 0) {
            process.stdout?.write(`  Progress: ${i + 1}/${config.iterations}\r`);
        }
    }

    const memoryAfter = getMemoryMB();

    // Calculate statistics
    const sortedLatencies = [...latencies].sort((a, b) => a - b);
    const avgLatency = latencies.reduce((a, b) => a + b, 0) / latencies.length;
    const throughput = (config.seqLength * 1000) / avgLatency;

    const result = new BenchmarkResult('RustyTorch WASM', {
        avgLatencyMs: avgLatency,
        p50LatencyMs: percentile(sortedLatencies, 50),
        p95LatencyMs: percentile(sortedLatencies, 95),
        p99LatencyMs: percentile(sortedLatencies, 99),
        minLatencyMs: sortedLatencies[0],
        maxLatencyMs: sortedLatencies[sortedLatencies.length - 1],
        throughputTokensPerSec: throughput,
        memoryMB: memoryAfter - memoryBefore || engine.getMemoryUsage?.() || null,
        loadTimeMs: loadTime,
    });

    console.log('\n');
    console.log('Results:');
    console.log(`  Avg Latency: ${result.avgLatencyMs.toFixed(3)} ms`);
    console.log(`  P50 Latency: ${result.p50LatencyMs.toFixed(3)} ms`);
    console.log(`  P99 Latency: ${result.p99LatencyMs.toFixed(3)} ms`);
    console.log(`  Throughput:  ${result.throughputTokensPerSec.toFixed(0)} tokens/sec`);
    console.log(`  Load Time:   ${result.loadTimeMs.toFixed(0)} ms`);

    return result;
}

/**
 * Run all benchmarks
 */
export async function runAllBenchmarks(config = defaultConfig) {
    const results = [];

    // RustyTorch WASM
    results.push(await runRtxBenchmark(config));

    return results;
}

// CLI entry point
if (isNode && import.meta.url === `file://${process.argv[1]}`) {
    const config = {
        ...defaultConfig,
        iterations: parseInt(process.argv[2]) || 100,
    };

    runRtxBenchmark(config)
        .then(result => {
            console.log('\nJSON Output:');
            console.log(JSON.stringify(result.toJSON(), null, 2));
        })
        .catch(err => {
            console.error('Benchmark failed:', err);
            process.exit(1);
        });
}