foxhunt/ml/benches/tft_int8_memory_bench.rs

//! TFT INT8 Memory Profiling Benchmark
//!
//! Comprehensive memory footprint benchmarking comparing FP32 and INT8 TFT models.
//! Uses Criterion for performance testing and custom memory profiling for VRAM tracking.
//!
//! ## Benchmark Scope
//!
//! 1. **FP32 Model Memory Footprint**
//!    - Parameter memory (model weights)
//!    - Activation memory (intermediate tensors)
//!    - Optimizer state memory (Adam: gradients + momentum + variance)
//!    - Total GPU VRAM usage
//!
//! 2. **INT8 Model Memory Footprint**
//!    - Quantized parameter memory (INT8 + scales)
//!    - Activation memory (FP32 dequantized tensors)
//!    - Optimizer state memory (if training enabled)
//!    - Total GPU VRAM usage
//!
//! 3. **INT8 with Weight Caching**
//!    - Cached dequantized weights (trade memory for speed)
//!    - Activation memory
//!    - Total GPU VRAM usage
//!
//! 4. **GPU VRAM Usage (CUDA)**
//!    - Real-time VRAM monitoring via nvidia-smi
//!    - Peak VRAM during inference
//!    - Memory fragmentation analysis
//!
//! ## Performance Targets
//!
//! - **FP32 Baseline**: ~400-500 MB total VRAM
//! - **INT8 Target**: ~100 MB total VRAM (75% reduction)
//! - **INT8 + Cache**: ~150 MB total VRAM (62.5% reduction)
//! - **RTX 3050 Ti Budget**: <256 MB per model (to fit all 4 models in 4GB VRAM)
//!
//! ## Usage
//!
//! ```bash
//! # Run memory benchmarks (requires CUDA)
//! cargo bench --bench tft_int8_memory_bench --features cuda
//!
//! # Generate HTML report
//! cargo bench --bench tft_int8_memory_bench --features cuda -- --save-baseline main
//! ```

#![allow(unused_crate_dependencies)]

use candle_core::{Device, Tensor};
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use ml::benchmark::memory_profiler::MemoryProfiler;
use ml::tft::{QuantizedTemporalFusionTransformer, TFTConfig, TemporalFusionTransformer};
use std::time::Duration;

/// Benchmark configuration constants
const WARMUP_ITERATIONS: usize = 5;
const BATCH_SIZE: usize = 1; // Single inference for memory analysis
const SEQ_LEN: usize = 60;
const HORIZON: usize = 10;

/// Create standard TFT configuration (225 features, Wave C+D)
fn create_tft_config() -> TFTConfig {
    TFTConfig {
        input_dim: 225,
        hidden_dim: 256,
        num_heads: 8,
        num_layers: 3,
        prediction_horizon: HORIZON,
        sequence_length: SEQ_LEN,
        num_quantiles: 3,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 210,
        learning_rate: 0.001,
        batch_size: 32,
        dropout_rate: 0.1,
        l2_regularization: 0.0001,
        use_flash_attention: false,
        mixed_precision: false,
        memory_efficient: true,
        max_inference_latency_us: 3200,
        target_throughput_pps: 10_000,
    }
}

/// Generate synthetic TFT inputs
fn generate_tft_inputs(
    batch_size: usize,
    config: &TFTConfig,
    device: &Device,
) -> Result<(Tensor, Tensor, Tensor), Box<dyn std::error::Error>> {
    let static_features =
        Tensor::randn(0f32, 1f32, (batch_size, config.num_static_features), device)?;

    let historical_features = Tensor::randn(
        0f32,
        1f32,
        (
            batch_size,
            config.sequence_length,
            config.num_unknown_features,
        ),
        device,
    )?;

    let future_features = Tensor::randn(
        0f32,
        1f32,
        (
            batch_size,
            config.prediction_horizon,
            config.num_known_features,
        ),
        device,
    )?;

    Ok((static_features, historical_features, future_features))
}

/// Benchmark 1: FP32 model memory footprint
fn bench_fp32_memory_footprint(c: &mut Criterion) {
    let mut group = c.benchmark_group("tft_fp32_memory_footprint");
    group.sample_size(10); // Small sample for memory-focused benchmarks
    group.measurement_time(Duration::from_secs(15));

    let config = create_tft_config();
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    // Skip benchmark if CUDA not available
    if matches!(device, Device::Cpu) {
        println!("⚠️  CUDA not available, skipping FP32 memory benchmark");
        return;
    }

    group.bench_function("model_creation", |b| {
        b.iter(|| {
            let mut profiler = MemoryProfiler::new(0);

            // Baseline snapshot
            let baseline = profiler.take_snapshot().expect("Baseline snapshot failed");

            // Create FP32 model
            let _model = black_box(
                TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
                    .expect("Model creation failed"),
            );

            // Wait for VRAM allocation to stabilize
            std::thread::sleep(Duration::from_millis(100));

            // Measure memory after model creation
            let after_create = profiler
                .take_snapshot()
                .expect("Post-creation snapshot failed");

            let vram_used_mb = after_create.vram_used_mb - baseline.vram_used_mb;

            black_box(vram_used_mb)
        });
    });

    group.bench_function("inference_memory", |b| {
        let mut profiler = MemoryProfiler::new(0);
        let baseline = profiler.take_snapshot().expect("Baseline snapshot failed");

        let mut model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
            .expect("Model creation failed");

        let (static_features, historical_features, future_features) =
            generate_tft_inputs(BATCH_SIZE, &config, &device).expect("Input generation failed");

        // Warmup
        for _ in 0..WARMUP_ITERATIONS {
            let _ = model.forward(&static_features, &historical_features, &future_features);
        }

        b.iter(|| {
            // Run inference
            let _ = black_box(
                model
                    .forward(&static_features, &historical_features, &future_features)
                    .expect("Forward pass failed"),
            );

            // Measure peak memory
            let snapshot = profiler.take_snapshot().expect("Snapshot failed");
            let vram_used_mb = snapshot.vram_used_mb - baseline.vram_used_mb;

            black_box(vram_used_mb)
        });
    });

    // Print summary statistics
    let mut profiler = MemoryProfiler::new(0);
    let baseline = profiler.take_snapshot().expect("Baseline failed");

    let _model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
        .expect("Model creation failed");

    let after_create = profiler.take_snapshot().expect("Post-create failed");
    let param_memory_mb = after_create.vram_used_mb - baseline.vram_used_mb;

    println!("\n=== FP32 Memory Footprint ===");
    println!("Parameter Memory: {:.0} MB", param_memory_mb);
    println!(
        "Estimated Optimizer Memory: {:.0} MB (2x params for Adam)",
        param_memory_mb * 2.0
    );
    println!(
        "Total Budget (params + optimizer): {:.0} MB",
        param_memory_mb * 3.0
    );

    group.finish();
}

/// Benchmark 2: INT8 model memory footprint
fn bench_int8_memory_footprint(c: &mut Criterion) {
    let mut group = c.benchmark_group("tft_int8_memory_footprint");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(15));

    let config = create_tft_config();
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    if matches!(device, Device::Cpu) {
        println!("⚠️  CUDA not available, skipping INT8 memory benchmark");
        return;
    }

    group.bench_function("model_creation", |b| {
        b.iter(|| {
            let mut profiler = MemoryProfiler::new(0);
            let baseline = profiler.take_snapshot().expect("Baseline snapshot failed");

            // Create FP32 source model (required for quantization)
            let fp32_model =
                TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
                    .expect("FP32 model creation failed");

            // Quantize to INT8
            let _int8_model = black_box(
                QuantizedTemporalFusionTransformer::new_from_fp32(&fp32_model)
                    .expect("Quantization failed"),
            );

            std::thread::sleep(Duration::from_millis(100));

            let after_create = profiler.take_snapshot().expect("Post-create failed");
            let vram_used_mb = after_create.vram_used_mb - baseline.vram_used_mb;

            black_box(vram_used_mb)
        });
    });

    group.bench_function("inference_memory", |b| {
        let mut profiler = MemoryProfiler::new(0);
        let baseline = profiler.take_snapshot().expect("Baseline failed");

        let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
            .expect("FP32 model creation failed");

        let int8_model = QuantizedTemporalFusionTransformer::new_from_fp32(&fp32_model)
            .expect("Quantization failed");

        let (static_features, historical_features, future_features) =
            generate_tft_inputs(BATCH_SIZE, &config, &device).expect("Input generation failed");

        // Warmup
        for _ in 0..WARMUP_ITERATIONS {
            let _ = int8_model.forward(&static_features, &historical_features, &future_features);
        }

        b.iter(|| {
            let _ = black_box(
                int8_model
                    .forward(&static_features, &historical_features, &future_features)
                    .expect("Forward pass failed"),
            );

            let snapshot = profiler.take_snapshot().expect("Snapshot failed");
            let vram_used_mb = snapshot.vram_used_mb - baseline.vram_used_mb;

            black_box(vram_used_mb)
        });
    });

    // Print summary
    let mut profiler = MemoryProfiler::new(0);
    let baseline = profiler.take_snapshot().expect("Baseline failed");

    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
        .expect("FP32 model creation failed");

    let _int8_model = QuantizedTemporalFusionTransformer::new_from_fp32(&fp32_model)
        .expect("Quantization failed");

    let after_create = profiler.take_snapshot().expect("Post-create failed");
    let param_memory_mb = after_create.vram_used_mb - baseline.vram_used_mb;

    println!("\n=== INT8 Memory Footprint ===");
    println!("Parameter Memory: {:.0} MB", param_memory_mb);
    println!(
        "Estimated Optimizer Memory: {:.0} MB",
        param_memory_mb * 2.0
    );
    println!("Total Budget: {:.0} MB", param_memory_mb * 3.0);

    group.finish();
}

/// Benchmark 3: INT8 with weight caching (trade memory for speed)
fn bench_int8_with_caching_memory(c: &mut Criterion) {
    let mut group = c.benchmark_group("tft_int8_cached_memory");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(15));

    let config = create_tft_config();
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    if matches!(device, Device::Cpu) {
        println!("⚠️  CUDA not available, skipping INT8 caching benchmark");
        return;
    }

    // Note: This benchmark simulates cached dequantized weights by pre-loading them
    group.bench_function("cached_inference_memory", |b| {
        let mut profiler = MemoryProfiler::new(0);
        let baseline = profiler.take_snapshot().expect("Baseline failed");

        let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
            .expect("FP32 model creation failed");

        let int8_model = QuantizedTemporalFusionTransformer::new_from_fp32(&fp32_model)
            .expect("Quantization failed");

        let (static_features, historical_features, future_features) =
            generate_tft_inputs(BATCH_SIZE, &config, &device).expect("Input generation failed");

        // Warmup (pre-cache weights via inference)
        for _ in 0..WARMUP_ITERATIONS {
            let _ = int8_model.forward(&static_features, &historical_features, &future_features);
        }

        // Measure memory with cached weights
        let after_warmup = profiler.take_snapshot().expect("Post-warmup failed");
        let cached_memory_mb = after_warmup.vram_used_mb - baseline.vram_used_mb;

        b.iter(|| {
            let _ = black_box(
                int8_model
                    .forward(&static_features, &historical_features, &future_features)
                    .expect("Forward pass failed"),
            );

            black_box(cached_memory_mb)
        });
    });

    println!("\n=== INT8 with Weight Caching ===");
    println!("Note: Cached weights stored in FP32 for faster inference");
    println!("Trade-off: +25% memory for -50% latency (estimated)");

    group.finish();
}

/// Benchmark 4: GPU VRAM usage comparison (CUDA-specific)
fn bench_gpu_vram_usage(c: &mut Criterion) {
    let mut group = c.benchmark_group("tft_gpu_vram_comparison");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(15));

    let config = create_tft_config();
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    if matches!(device, Device::Cpu) {
        println!("⚠️  CUDA not available, skipping VRAM comparison");
        return;
    }

    // FP32 VRAM measurement
    group.bench_function("fp32_vram", |b| {
        b.iter(|| {
            let mut profiler = MemoryProfiler::new(0);
            let baseline = profiler.take_snapshot().expect("Baseline failed");

            let mut model =
                TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
                    .expect("Model creation failed");

            let (static_features, historical_features, future_features) =
                generate_tft_inputs(BATCH_SIZE, &config, &device).expect("Input generation failed");

            // Run multiple inferences to measure peak VRAM
            let mut peak_vram = 0.0f64;
            for _ in 0..10 {
                let _ = model.forward(&static_features, &historical_features, &future_features);

                let snapshot = profiler.take_snapshot().expect("Snapshot failed");
                let vram_mb = snapshot.vram_used_mb - baseline.vram_used_mb;
                peak_vram = peak_vram.max(vram_mb);
            }

            black_box(peak_vram)
        });
    });

    // INT8 VRAM measurement
    group.bench_function("int8_vram", |b| {
        b.iter(|| {
            let mut profiler = MemoryProfiler::new(0);
            let baseline = profiler.take_snapshot().expect("Baseline failed");

            let fp32_model =
                TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
                    .expect("FP32 model creation failed");

            let int8_model = QuantizedTemporalFusionTransformer::new_from_fp32(&fp32_model)
                .expect("Quantization failed");

            let (static_features, historical_features, future_features) =
                generate_tft_inputs(BATCH_SIZE, &config, &device).expect("Input generation failed");

            let mut peak_vram = 0.0f64;
            for _ in 0..10 {
                let _ =
                    int8_model.forward(&static_features, &historical_features, &future_features);

                let snapshot = profiler.take_snapshot().expect("Snapshot failed");
                let vram_mb = snapshot.vram_used_mb - baseline.vram_used_mb;
                peak_vram = peak_vram.max(vram_mb);
            }

            black_box(peak_vram)
        });
    });

    // Print comparison summary
    println!("\n=== GPU VRAM Usage Summary ===");

    // FP32 measurement
    let mut profiler_fp32 = MemoryProfiler::new(0);
    let baseline_fp32 = profiler_fp32.take_snapshot().expect("Baseline failed");

    let mut model_fp32 = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
        .expect("Model creation failed");

    let (static_features, historical_features, future_features) =
        generate_tft_inputs(BATCH_SIZE, &config, &device).expect("Input generation failed");

    let mut fp32_peak = 0.0f64;
    for _ in 0..10 {
        let _ = model_fp32.forward(&static_features, &historical_features, &future_features);
        let snap = profiler_fp32.take_snapshot().expect("Snapshot failed");
        fp32_peak = fp32_peak.max(snap.vram_used_mb - baseline_fp32.vram_used_mb);
    }

    // INT8 measurement
    let mut profiler_int8 = MemoryProfiler::new(0);
    let baseline_int8 = profiler_int8.take_snapshot().expect("Baseline failed");

    let fp32_src = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
        .expect("FP32 model creation failed");

    let model_int8 =
        QuantizedTemporalFusionTransformer::new_from_fp32(&fp32_src).expect("Quantization failed");

    let mut int8_peak = 0.0f64;
    for _ in 0..10 {
        let _ = model_int8.forward(&static_features, &historical_features, &future_features);
        let snap = profiler_int8.take_snapshot().expect("Snapshot failed");
        int8_peak = int8_peak.max(snap.vram_used_mb - baseline_int8.vram_used_mb);
    }

    let reduction_mb = fp32_peak - int8_peak;
    let reduction_pct = (reduction_mb / fp32_peak) * 100.0;

    println!("FP32 Peak VRAM: {:.0} MB", fp32_peak);
    println!("INT8 Peak VRAM: {:.0} MB", int8_peak);
    println!(
        "Memory Reduction: {:.0} MB ({:.1}%)",
        reduction_mb, reduction_pct
    );
    println!(
        "75% Target: {}",
        if reduction_pct >= 75.0 {
            "✅ ACHIEVED"
        } else {
            "❌ NOT MET"
        }
    );

    // RTX 3050 Ti budget validation
    let rtx3050ti_vram_mb = 4096.0;
    let num_models = 4; // DQN, PPO, MAMBA-2, TFT
    let budget_per_model = rtx3050ti_vram_mb / num_models as f64;

    println!("\n=== RTX 3050 Ti Budget Validation ===");
    println!("Total VRAM: {:.0} MB", rtx3050ti_vram_mb);
    println!("Budget per model (4 models): {:.0} MB", budget_per_model);
    println!("FP32 Usage: {:.0} MB", fp32_peak);
    println!("INT8 Usage: {:.0} MB", int8_peak);
    println!(
        "FP32 fits budget: {}",
        if fp32_peak <= budget_per_model {
            "✅ YES"
        } else {
            "❌ NO"
        }
    );
    println!(
        "INT8 fits budget: {}",
        if int8_peak <= budget_per_model {
            "✅ YES"
        } else {
            "❌ NO"
        }
    );

    group.finish();
}

/// Memory reduction validation (75% target)
fn bench_memory_reduction_validation(c: &mut Criterion) {
    let mut group = c.benchmark_group("tft_memory_reduction_validation");
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(10));

    let config = create_tft_config();
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    if matches!(device, Device::Cpu) {
        println!("⚠️  CUDA not available, skipping validation benchmark");
        return;
    }

    group.bench_function("validate_75_percent_reduction", |b| {
        b.iter(|| {
            // Measure FP32
            let mut profiler_fp32 = MemoryProfiler::new(0);
            let baseline_fp32 = profiler_fp32.take_snapshot().expect("FP32 baseline failed");

            let _model_fp32 =
                TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
                    .expect("FP32 model creation failed");

            std::thread::sleep(Duration::from_millis(100));

            let after_fp32 = profiler_fp32.take_snapshot().expect("FP32 snapshot failed");
            let fp32_mb = after_fp32.vram_used_mb - baseline_fp32.vram_used_mb;

            // Measure INT8
            let mut profiler_int8 = MemoryProfiler::new(0);
            let baseline_int8 = profiler_int8.take_snapshot().expect("INT8 baseline failed");

            let fp32_src =
                TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
                    .expect("FP32 source creation failed");

            let _model_int8 = QuantizedTemporalFusionTransformer::new_from_fp32(&fp32_src)
                .expect("Quantization failed");

            std::thread::sleep(Duration::from_millis(100));

            let after_int8 = profiler_int8.take_snapshot().expect("INT8 snapshot failed");
            let int8_mb = after_int8.vram_used_mb - baseline_int8.vram_used_mb;

            // Calculate reduction
            let reduction_pct = ((fp32_mb - int8_mb) / fp32_mb) * 100.0;

            black_box((fp32_mb, int8_mb, reduction_pct))
        });
    });

    println!("\n=== Memory Reduction Validation ===");
    println!("Target: 75% reduction (FP32 → INT8)");
    println!("Expected: FP32 ~400 MB → INT8 ~100 MB");

    group.finish();
}

criterion_group!(
    benches,
    bench_fp32_memory_footprint,
    bench_int8_memory_footprint,
    bench_int8_with_caching_memory,
    bench_gpu_vram_usage,
    bench_memory_reduction_validation
);
criterion_main!(benches);