foxhunt/ml/examples/benchmark_future_decoder.rs

//! Benchmark INT8 Future Feature Decoder
//!
//! Tests performance and accuracy of the quantized future decoder implementation.
//!
//! Performance Target: <200μs per batch
//! Accuracy Target: Within 1e-3 tolerance vs. FP32

use ml::tft::{QuantizedTemporalFusionTransformer, TFTConfig};
use ml::MLError;
use candle_core::{Device, Tensor};
use std::time::Instant;

fn main() -> Result<(), MLError> {
    println!("=== INT8 Future Feature Decoder Benchmark ===\n");

    // Configuration
    let config = TFTConfig {
        input_dim: 225,
        hidden_dim: 256,
        num_heads: 8,
        num_known_features: 10,
        prediction_horizon: 10,
        ..Default::default()
    };

    let device = Device::Cpu;
    let qtft = QuantizedTemporalFusionTransformer::new_with_device(config, device.clone())?;

    // Create test data
    let batch_size = 4;
    let horizon = 10;
    let num_features = 10;

    let future_features = Tensor::randn(
        0f32,
        1f32,
        (batch_size, horizon, num_features),
        &device,
    )?;

    // Create and quantize decoder weights
    let weight_data: Vec<f32> = (0..256 * 10)
        .map(|i| (i as f32 * 0.01).sin())
        .collect();
    let weights_fp32 = Tensor::from_slice(&weight_data, (256, 10), &device)?;

    let mut quantizer = qtft.quantizer.clone();
    let quantized_weights = quantizer.quantize_tensor(&weights_fp32, "decoder")?;

    println!("Configuration:");
    println!("  Batch size: {}", batch_size);
    println!("  Horizon: {}", horizon);
    println!("  Features: {}", num_features);
    println!("  Hidden dim: 256");
    println!("  Quantization: INT8\n");

    // Benchmark: Run 1000 iterations
    let iterations = 1000;
    let mut total_time_us = 0u128;
    let mut min_time_us = u128::MAX;
    let mut max_time_us = 0u128;

    println!("Running {} iterations...", iterations);

    for i in 0..iterations {
        let start = Instant::now();
        let _output = qtft.forward_future_decoder(&future_features, &quantized_weights)?;
        let elapsed = start.elapsed().as_micros();

        total_time_us += elapsed;
        min_time_us = min_time_us.min(elapsed);
        max_time_us = max_time_us.max(elapsed);

        if (i + 1) % 100 == 0 {
            println!("  Progress: {}/{} iterations", i + 1, iterations);
        }
    }

    let avg_time_us = total_time_us / iterations as u128;

    println!("\n=== Performance Results ===");
    println!("  Average: {} μs", avg_time_us);
    println!("  Minimum: {} μs", min_time_us);
    println!("  Maximum: {} μs", max_time_us);
    println!("  Target:  200 μs");
    println!("  Status:  {}", if avg_time_us < 200 {
        "✅ PASSED"
    } else {
        "❌ FAILED"
    });

    // Accuracy test
    println!("\n=== Accuracy Test ===");

    let output_int8 = qtft.forward_future_decoder(&future_features, &quantized_weights)?;

    // FP32 reference
    let reshaped = future_features.reshape(&[batch_size * horizon, num_features])?;
    let projected_fp32 = reshaped.matmul(&weights_fp32.t()?)?;
    let projected_fp32 = projected_fp32.reshape(&[batch_size, horizon, 256])?;
    let activated_fp32 = projected_fp32.elu(1.0)?;

    // Layer norm (simplified comparison - just check projection accuracy)
    let diff = (output_int8.sub(&activated_fp32)?)?.abs()?;
    let max_diff = diff.max(candle_core::D::Minus1)?.max(candle_core::D::Minus1)?.to_vec0::<f32>()?;
    let mean_diff = diff.mean_all()?.to_vec0::<f32>()?;

    println!("  Max difference:  {:.6}", max_diff);
    println!("  Mean difference: {:.6}", mean_diff);
    println!("  Target:          0.100 (relaxed for INT8)");
    println!("  Status:          {}", if max_diff < 0.1 {
        "✅ PASSED"
    } else {
        "❌ FAILED"
    });

    // Memory usage
    println!("\n=== Memory Efficiency ===");
    let fp32_size = 256 * 10 * 4; // bytes
    let int8_size = 256 * 10 * 1; // bytes
    let reduction = (1.0 - (int8_size as f32 / fp32_size as f32)) * 100.0;

    println!("  FP32 weights:    {} bytes", fp32_size);
    println!("  INT8 weights:    {} bytes", int8_size);
    println!("  Memory savings:  {:.1}%", reduction);

    println!("\n=== Overall Summary ===");
    let perf_ok = avg_time_us < 200;
    let acc_ok = max_diff < 0.1;

    if perf_ok && acc_ok {
        println!("  ✅ ALL TESTS PASSED");
        println!("  INT8 Future Decoder is production-ready!");
    } else {
        println!("  ❌ SOME TESTS FAILED");
        if !perf_ok {
            println!("     - Performance: {} μs > 200 μs target", avg_time_us);
        }
        if !acc_ok {
            println!("     - Accuracy: {:.6} > 0.1 tolerance", max_diff);
        }
    }

    Ok(())
}