foxhunt/ml/examples/test_memory_optimization.rs

//! Standalone memory optimization test for 4GB GPU
//!
//! Tests quantization and mixed precision features to verify
//! 4GB VRAM compatibility.

use candle_core::{Device, DType, Tensor};
use ml::memory_optimization::{
    MemoryOptimizationConfig, MemoryStats, PrecisionConverter, PrecisionType, QuantizationConfig,
    QuantizationType, Quantizer,
};
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Memory Optimization Test for 4GB GPU ===\n");

    let device = Device::cuda_if_available(0)?;
    println!("Device: {:?}\n", device);

    // Test 1: INT8 Quantization
    test_int8_quantization(&device)?;

    // Test 2: INT4 Quantization
    test_int4_quantization(&device)?;

    // Test 3: FP16 Precision
    test_fp16_precision(&device)?;

    // Test 4: BF16 Precision
    test_bf16_precision(&device)?;

    // Test 5: Full Pipeline
    test_full_optimization_pipeline(&device)?;

    // Test 6: 4GB Compatibility
    test_4gb_compatibility(&device)?;

    println!("\n=== ALL MEMORY OPTIMIZATION TESTS PASSED ===");
    Ok(())
}

fn test_int8_quantization(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 1: INT8 Quantization");
    println!("----------------------------");

    let start = Instant::now();

    // Create test tensor (256x256 = 262,144 elements)
    let tensor = Tensor::randn(0.0f32, 1.0f32, (256, 256), device)?;
    let original_size = tensor.dims().iter().product::<usize>() * 4; // 4 bytes per f32

    println!("Original tensor: {:?}, size: {} bytes ({:.2} MB)",
             tensor.dims(), original_size, original_size as f64 / 1_048_576.0);

    // Configure INT8 quantization
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };

    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize tensor
    let quantized = quantizer.quantize_tensor(&tensor, "test_layer")?;

    println!("Quantization type: {:?}", quantized.quant_type);
    println!("Scale: {}, Zero point: {}", quantized.scale, quantized.zero_point);

    // Check memory savings
    let quantized_size = quantized.memory_bytes();
    let savings_percent = (1.0 - (quantized_size as f64 / original_size as f64)) * 100.0;

    println!("Quantized size: {} bytes ({:.2} MB)",
             quantized_size, quantized_size as f64 / 1_048_576.0);
    println!("Memory savings: {:.1}%", savings_percent);

    // Dequantize and verify
    let _dequantized = quantizer.dequantize_tensor(&quantized)?;

    let elapsed = start.elapsed();
    println!("✓ INT8 quantization test passed ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_int4_quantization(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 2: INT4 Quantization");
    println!("----------------------------");

    let start = Instant::now();

    let tensor = Tensor::randn(0.0f32, 1.0f32, (512, 512), device)?;
    let original_size = tensor.dims().iter().product::<usize>() * 4;

    println!("Original tensor: {:?}, size: {:.2} MB",
             tensor.dims(), original_size as f64 / 1_048_576.0);

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int4,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };

    let mut quantizer = Quantizer::new(config, device.clone());
    let quantized = quantizer.quantize_tensor(&tensor, "int4_layer")?;

    let quantized_size = quantized.memory_bytes();
    let savings_percent = (1.0 - (quantized_size as f64 / original_size as f64)) * 100.0;

    println!("Quantized size: {:.2} MB", quantized_size as f64 / 1_048_576.0);
    println!("Memory savings: {:.1}%", savings_percent);

    let elapsed = start.elapsed();
    println!("✓ INT4 quantization test passed ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_fp16_precision(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 3: FP16 Precision Conversion");
    println!("-----------------------------------");

    let start = Instant::now();

    let tensor = Tensor::randn(0.0f32, 1.0f32, (256, 256), device)?;
    let original_size = tensor.dims().iter().product::<usize>() * 4;

    println!("Original tensor: {:?}, dtype: {:?}, size: {:.2} MB",
             tensor.dims(), tensor.dtype(), original_size as f64 / 1_048_576.0);

    let mut converter = PrecisionConverter::new(PrecisionType::Float16, device.clone());

    let converted = converter.to_float16(&tensor)?;

    assert_eq!(converted.dtype(), DType::F16);

    let converted_size = converted.dims().iter().product::<usize>() * 2;
    let savings_percent = (1.0 - (converted_size as f64 / original_size as f64)) * 100.0;

    println!("Converted dtype: {:?}, size: {:.2} MB",
             converted.dtype(), converted_size as f64 / 1_048_576.0);
    println!("Memory savings: {:.1}%", savings_percent);

    // Check statistics
    let stats = converter.get_stats();
    println!("Conversions: {}, Total saved: {:.2} MB",
             stats.conversions, stats.memory_saved_mb);

    let elapsed = start.elapsed();
    println!("✓ FP16 precision test passed ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_bf16_precision(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 4: BF16 Precision Conversion");
    println!("-----------------------------------");

    let start = Instant::now();

    let tensor = Tensor::randn(0.0f32, 1.0f32, (512, 512), device)?;
    let original_size = tensor.dims().iter().product::<usize>() * 4;

    println!("Original size: {:.2} MB", original_size as f64 / 1_048_576.0);

    let mut converter = PrecisionConverter::new(PrecisionType::BFloat16, device.clone());

    let converted = converter.to_bfloat16(&tensor)?;

    assert_eq!(converted.dtype(), DType::BF16);

    let converted_size = converted.dims().iter().product::<usize>() * 2;
    let savings_percent = (1.0 - (converted_size as f64 / original_size as f64)) * 100.0;

    println!("Converted size: {:.2} MB", converted_size as f64 / 1_048_576.0);
    println!("Memory savings: {:.1}%", savings_percent);

    let elapsed = start.elapsed();
    println!("✓ BF16 precision test passed ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_full_optimization_pipeline(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 5: Full Optimization Pipeline");
    println!("------------------------------------");

    let start = Instant::now();
    let mut stats = MemoryStats::new();

    // Step 1: Baseline F32 model
    let model_tensor = Tensor::randn(0.0f32, 1.0f32, (512, 512), device)?;
    let baseline_size = (model_tensor.dims().iter().product::<usize>() * 4) as f64 / 1_048_576.0;

    stats.add_component("baseline_f32", baseline_size);
    stats.update_peak(baseline_size);
    println!("Step 1: Baseline (F32): {:.2} MB", baseline_size);

    // Step 2: FP16 conversion
    let mut precision_converter = PrecisionConverter::new(PrecisionType::Float16, device.clone());
    let fp16_tensor = precision_converter.to_float16(&model_tensor)?;

    let fp16_size = (fp16_tensor.dims().iter().product::<usize>() * 2) as f64 / 1_048_576.0;
    let precision_savings = baseline_size - fp16_size;
    stats.add_component("fp16_model", fp16_size);
    stats.savings_mb += precision_savings;
    println!("Step 2: FP16 model: {:.2} MB (saved {:.2} MB)", fp16_size, precision_savings);

    // Step 3: INT8 quantization
    let fp32_for_quant = precision_converter.to_float32(&fp16_tensor)?;

    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };

    let mut quantizer = Quantizer::new(quant_config, device.clone());
    let quantized = quantizer.quantize_tensor(&fp32_for_quant, "optimized_model")?;

    let quantized_size = quantized.memory_bytes() as f64 / 1_048_576.0;
    let quant_savings = fp16_size - quantized_size;
    stats.add_component("int8_fp16", quantized_size);
    stats.savings_mb += quant_savings;
    println!("Step 3: INT8+FP16: {:.2} MB (saved {:.2} MB)", quantized_size, quant_savings);

    // Summary
    let total_savings = baseline_size - quantized_size;
    let savings_percent = (total_savings / baseline_size) * 100.0;

    println!("\n--- Pipeline Summary ---");
    println!("Baseline:     {:.2} MB", baseline_size);
    println!("Optimized:    {:.2} MB", quantized_size);
    println!("Total Saved:  {:.2} MB ({:.1}%)", total_savings, savings_percent);
    println!("Fits 4GB GPU: {}", if quantized_size < 3500.0 { "✓ YES" } else { "✗ NO" });

    let elapsed = start.elapsed();
    println!("\n✓ Full pipeline test passed ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_4gb_compatibility(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 6: 4GB GPU Compatibility Analysis");
    println!("----------------------------------------");

    // Simulate different model configurations
    let configs = vec![
        ("MAMBA-2 F32 Baseline", 500.0, QuantizationType::None, PrecisionType::Float32),
        ("MAMBA-2 INT8", 500.0, QuantizationType::Int8, PrecisionType::Float32),
        ("MAMBA-2 FP16", 500.0, QuantizationType::None, PrecisionType::Float16),
        ("MAMBA-2 INT8+FP16", 500.0, QuantizationType::Int8, PrecisionType::Float16),
        ("DQN F32", 150.0, QuantizationType::None, PrecisionType::Float32),
        ("DQN INT8+FP16", 150.0, QuantizationType::Int8, PrecisionType::Float16),
        ("PPO F32", 200.0, QuantizationType::None, PrecisionType::Float32),
        ("PPO INT8+FP16", 200.0, QuantizationType::Int8, PrecisionType::Float16),
    ];

    println!("Model configurations for 4GB GPU (3500MB usable):\n");

    for (name, base_size_mb, quant_type, precision) in configs {
        let memory_multiplier = precision.memory_multiplier();
        let quant_savings = match quant_type {
            QuantizationType::None => 1.0,
            QuantizationType::Int8 => 0.25,
            QuantizationType::Int4 => 0.125,
            QuantizationType::Dynamic => 0.25,
        };

        let final_size = base_size_mb * memory_multiplier * quant_savings;
        let fits = final_size <= 3500.0;

        println!(
            "{:20} {:>8.1} MB  [{:>5}] (quant={:?}, prec={:?})",
            name,
            final_size,
            if fits { "✓ FIT" } else { "✗ BIG" },
            quant_type,
            precision
        );
    }

    println!("\n✓ 4GB compatibility analysis complete\n");

    Ok(())
}