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::{DType, Device, 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(())
}