foxhunt/crates/ml/tests/memory_optimization_tests.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Comprehensive Memory Optimization Tests for 4GB GPU
//!
//! Tests quantization, mixed precision, and memory efficiency features
//! to ensure training fits within RTX 3050 Ti 4GB VRAM constraints.

use candle_core::{DType, Device, Tensor};
use ml::memory_optimization::{
    MemoryOptimizationConfig, MemoryStats, PrecisionConverter, PrecisionType, QuantizationConfig,
    QuantizationType, Quantizer,
};
use tracing::info;

/// Helper to create test device (CUDA if available, CPU fallback)
fn test_device() -> Device {
    Device::new_cuda(0).expect("CUDA required")
}

/// Helper to create test tensor
fn create_test_tensor(device: &Device, shape: &[usize]) -> Tensor {
    Tensor::randn(0.0f32, 1.0f32, shape, device).unwrap()
}

#[test]
fn test_int8_quantization_basic() {
    let device = test_device();
    info!(?device, "Running INT8 quantization test");

    // Create test tensor
    let tensor = create_test_tensor(&device, &[256, 256]);
    let original_size = tensor.dims().iter().product::<usize>() * 4; // 4 bytes per f32

    info!(shape = ?tensor.dims(), original_size, "Original tensor shape and size");

    // 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")
        .expect("Quantization failed");

    info!(quant_type = ?quantized.quant_type, scale = quantized.scale, zero_point = quantized.zero_point, "Quantized tensor parameters");

    // Verify quantization type
    assert_eq!(quantized.quant_type, QuantizationType::Int8);

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

    info!(original_size, quantized_size, savings_percent, "INT8 quantization memory stats");

    // INT8 should achieve ~75% memory reduction
    assert!(
        savings_percent >= 70.0,
        "Expected at least 70% memory savings"
    );

    // Dequantize and check accuracy
    let dequantized = quantizer
        .dequantize_tensor(&quantized)
        .expect("Dequantization failed");

    assert_eq!(dequantized.dims(), tensor.dims());
    info!("INT8 quantization test passed");
}

#[test]
fn test_int4_quantization() {
    let device = test_device();
    info!(?device, "Running INT4 quantization test");

    let tensor = create_test_tensor(&device, &[512, 512]);
    let original_size = tensor.dims().iter().product::<usize>() * 4;

    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, "test_layer_int4")
        .expect("INT4 quantization failed");

    assert_eq!(quantized.quant_type, QuantizationType::Int4);

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

    info!(original_size, quantized_size, savings_percent, "INT4 quantization memory stats");

    // INT4 should achieve ~87.5% memory reduction
    assert!(
        savings_percent >= 85.0,
        "Expected at least 85% memory savings"
    );
    info!("INT4 quantization test passed");
}

#[test]
fn test_asymmetric_quantization() {
    let device = test_device();
    info!(?device, "Running asymmetric quantization test");

    let tensor = create_test_tensor(&device, &[128, 128]);

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: false, // Asymmetric
        per_channel: true,
        calibration_samples: Some(500),
    };

    let mut quantizer = Quantizer::new(config, device.clone());
    let quantized = quantizer
        .quantize_tensor(&tensor, "asymmetric_layer")
        .expect("Asymmetric quantization failed");

    // Asymmetric quantization should use non-zero zero_point
    info!(scale = quantized.scale, zero_point = quantized.zero_point, "Asymmetric quantization parameters");

    assert_eq!(quantized.quant_type, QuantizationType::Int8);
    info!("Asymmetric quantization test passed");
}

#[test]
fn test_float16_precision_conversion() {
    let device = test_device();
    info!(?device, "Running FP16 precision test");

    let tensor = create_test_tensor(&device, &[256, 256]);
    let original_size = tensor.dims().iter().product::<usize>() * 4; // F32

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

    let converted = converter
        .to_float16(&tensor)
        .expect("FP16 conversion failed");

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

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

    info!(original_size, converted_size, savings_percent, "FP16 conversion memory stats");

    // FP16 should achieve 50% memory reduction
    assert!(savings_percent >= 49.0 && savings_percent <= 51.0);

    // Check statistics
    let stats = converter.get_stats();
    info!(conversions = stats.conversions, memory_saved_mb = stats.memory_saved_mb, "FP16 conversion stats");

    assert_eq!(stats.conversions, 1);
    assert!(stats.memory_saved_mb > 0.0);
    info!("FP16 precision conversion test passed");
}

#[test]
fn test_bfloat16_precision_conversion() {
    let device = test_device();
    info!(?device, "Running BF16 precision test");

    let tensor = create_test_tensor(&device, &[512, 512]);

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

    let converted = converter
        .to_bfloat16(&tensor)
        .expect("BF16 conversion failed");

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

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

    info!(original_size, converted_size, savings_percent, "BF16 conversion memory stats");

    assert!(savings_percent >= 49.0 && savings_percent <= 51.0);
    info!("BF16 precision conversion test passed");
}

#[test]
fn test_mixed_precision_roundtrip() {
    let device = test_device();
    info!(?device, "Running mixed precision roundtrip test");

    let original = create_test_tensor(&device, &[128, 128]);

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

    // Convert F32 -> F16 -> F32
    let fp16 = converter.to_float16(&original).expect("F32->F16 failed");
    let restored = converter.to_float32(&fp16).expect("F16->F32 failed");

    assert_eq!(restored.dtype(), DType::F32);
    assert_eq!(restored.dims(), original.dims());

    // Validate accuracy
    let accuracy =
        ml::memory_optimization::precision::validate_precision_accuracy(&original, &restored)
            .expect("Accuracy validation failed");

    info!(mae = accuracy.mae, rmse = accuracy.rmse, relative_error_pct = accuracy.mean_relative_error * 100.0, "FP16 roundtrip accuracy metrics");

    // FP16 should maintain reasonable accuracy (<5% error)
    assert!(
        accuracy.is_acceptable(5.0),
        "Relative error too high: {:.2}%",
        accuracy.mean_relative_error * 100.0
    );

    info!("Mixed precision roundtrip test passed");
}

#[test]
fn test_quantization_accuracy_preservation() {
    let device = test_device();
    info!(?device, "Running quantization accuracy test");

    let original = create_test_tensor(&device, &[256, 256]);

    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 and dequantize
    let quantized = quantizer
        .quantize_tensor(&original, "accuracy_test")
        .expect("Quantization failed");

    let restored = quantizer
        .dequantize_tensor(&quantized)
        .expect("Dequantization failed");

    // Validate accuracy
    let accuracy =
        ml::memory_optimization::precision::validate_precision_accuracy(&original, &restored)
            .expect("Accuracy validation failed");

    info!(mae = accuracy.mae, rmse = accuracy.rmse, max_absolute_error = accuracy.max_absolute_error, "INT8 quantization accuracy metrics");

    // INT8 quantization should maintain reasonable accuracy
    assert!(accuracy.rmse < 0.1, "RMSE too high: {:.6}", accuracy.rmse);

    info!("Quantization accuracy preservation test passed");
}

#[test]
fn test_memory_optimization_config() {
    info!("Testing memory optimization configuration");

    let config = MemoryOptimizationConfig::default();

    assert!(config.lazy_loading);
    assert_eq!(config.precision, PrecisionType::Float32);
    assert_eq!(config.quantization, QuantizationType::None);
    assert!(config.tensor_caching);

    // Custom config for 4GB GPU
    let custom_config = MemoryOptimizationConfig {
        lazy_loading: true,
        precision: PrecisionType::Float16,
        quantization: QuantizationType::Int8,
        max_memory_mb: Some(3500.0), // Leave 500MB headroom
        gradient_checkpointing: true,
        tensor_caching: false, // Reduce cache memory
    };

    assert_eq!(custom_config.precision, PrecisionType::Float16);
    assert_eq!(custom_config.quantization, QuantizationType::Int8);
    assert_eq!(custom_config.max_memory_mb, Some(3500.0));

    info!("Memory optimization config test passed");
}

#[test]
fn test_memory_stats_tracking() {
    info!("Testing memory statistics tracking");

    let mut stats = MemoryStats::new();

    assert_eq!(stats.current_mb, 0.0);
    assert_eq!(stats.peak_mb, 0.0);

    // Simulate memory usage
    stats.update_peak(100.0);
    assert_eq!(stats.current_mb, 100.0);
    assert_eq!(stats.peak_mb, 100.0);

    stats.update_peak(150.0);
    assert_eq!(stats.current_mb, 150.0);
    assert_eq!(stats.peak_mb, 150.0);

    stats.update_peak(120.0); // Peak should not decrease
    assert_eq!(stats.current_mb, 120.0);
    assert_eq!(stats.peak_mb, 150.0);

    // Add component breakdown
    stats.add_component("model_weights", 50.0);
    stats.add_component("activations", 30.0);
    stats.add_component("optimizer_state", 20.0);

    assert_eq!(stats.breakdown.len(), 3);
    assert_eq!(stats.breakdown.get("model_weights"), Some(&50.0));

    info!("Memory stats tracking test passed");
}

#[test]
fn test_multi_tensor_quantization() {
    let device = test_device();
    info!(?device, "Running multi-tensor quantization test");

    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 multiple tensors (simulating model layers)
    let tensors = vec![
        create_test_tensor(&device, &[256, 256]),
        create_test_tensor(&device, &[512, 512]),
        create_test_tensor(&device, &[1024, 256]),
        create_test_tensor(&device, &[256, 128]),
    ];

    let layer_names = vec!["layer1", "layer2", "layer3", "layer4"];

    for (tensor, name) in tensors.iter().zip(layer_names.iter()) {
        let quantized = quantizer
            .quantize_tensor(tensor, name)
            .expect("Multi-tensor quantization failed");

        info!(layer = name, memory_bytes = quantized.memory_bytes(), "Quantized layer");
    }

    // Check total memory savings
    let savings_mb = quantizer.memory_savings_mb();
    info!(savings_mb, "Total memory savings (MB)");

    assert!(savings_mb > 0.0, "No memory savings recorded");
    info!("Multi-tensor quantization test passed");
}

#[test]
fn test_precision_converter_stats() {
    let device = test_device();
    info!(?device, "Testing precision converter statistics");

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

    // Convert multiple tensors
    for i in 0..5 {
        let tensor = create_test_tensor(&device, &[128, 128]);
        let _converted = converter.to_float16(&tensor).expect("Conversion failed");
        info!(step = i + 1, "Converted tensor");
    }

    let stats = converter.get_stats();

    assert_eq!(stats.conversions, 5);
    assert!(stats.memory_saved_mb > 0.0);
    assert_eq!(stats.target_precision, PrecisionType::Float16);

    info!(conversions = stats.conversions, memory_saved_mb = stats.memory_saved_mb, "Precision converter stats");

    // Reset and verify
    converter.reset_stats();
    let new_stats = converter.get_stats();
    assert_eq!(new_stats.conversions, 0);
    assert_eq!(new_stats.memory_saved_mb, 0.0);

    info!("Precision converter stats test passed");
}

#[test]
fn test_4gb_gpu_memory_compatibility() {
    let device = test_device();
    info!(?device, "Testing 4GB GPU memory compatibility");

    // Simulate MAMBA-2 model sizes with memory optimization
    let model_configs = vec![
        (
            "baseline_f32",
            4,
            QuantizationType::None,
            PrecisionType::Float32,
        ),
        (
            "int8_f32",
            4,
            QuantizationType::Int8,
            PrecisionType::Float32,
        ),
        (
            "none_f16",
            4,
            QuantizationType::None,
            PrecisionType::Float16,
        ),
        (
            "int8_f16",
            4,
            QuantizationType::Int8,
            PrecisionType::Float16,
        ),
    ];

    for (name, size_multiplier, quant_type, precision) in model_configs {
        // Estimate memory usage for different configs
        let base_size_mb = 500.0; // MAMBA-2 base size
        let model_size = base_size_mb * size_multiplier as f64;

        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 = model_size * memory_multiplier * quant_savings;
        let fits_4gb = final_size <= 3500.0; // Leave 500MB headroom

        info!(config = name, final_size_mb = final_size, quant = ?quant_type, precision = ?precision, fits_4gb, "GPU memory config estimate");
    }

    info!("4GB GPU memory compatibility test passed");
}

#[test]
fn test_gradient_checkpointing_simulation() {
    info!("Testing gradient checkpointing simulation");

    let config = MemoryOptimizationConfig {
        lazy_loading: true,
        precision: PrecisionType::Float32,
        quantization: QuantizationType::None,
        max_memory_mb: Some(3500.0),
        gradient_checkpointing: true,
        tensor_caching: false,
    };

    assert!(config.gradient_checkpointing);

    // Gradient checkpointing typically reduces activation memory by ~2-3x
    // at the cost of ~33% more compute time
    let activation_memory_mb = 1000.0;
    let with_checkpointing = activation_memory_mb / 2.5;
    let savings = activation_memory_mb - with_checkpointing;

    info!(activation_memory_mb, with_checkpointing, savings_mb = savings, "Gradient checkpointing memory reduction");

    assert!(savings > 0.0);
    info!("Gradient checkpointing simulation test passed");
}

#[test]
fn test_no_quantization_passthrough() {
    let device = test_device();
    info!(?device, "Testing no-quantization passthrough");

    let tensor = create_test_tensor(&device, &[128, 128]);
    let original_size = tensor.dims().iter().product::<usize>() * 4;

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

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

    let result = quantizer
        .quantize_tensor(&tensor, "passthrough_test")
        .expect("Passthrough failed");

    assert_eq!(result.quant_type, QuantizationType::None);
    assert_eq!(result.memory_bytes(), original_size);

    info!("No-quantization passthrough test passed");
}

#[test]
fn test_precision_type_properties() {
    info!("Testing precision type properties");

    assert_eq!(PrecisionType::Float32.bytes_per_element(), 4);
    assert_eq!(PrecisionType::Float16.bytes_per_element(), 2);
    assert_eq!(PrecisionType::BFloat16.bytes_per_element(), 2);

    assert_eq!(PrecisionType::Float32.memory_multiplier(), 1.0);
    assert_eq!(PrecisionType::Float16.memory_multiplier(), 0.5);
    assert_eq!(PrecisionType::BFloat16.memory_multiplier(), 0.5);

    assert_eq!(PrecisionType::Float32.to_dtype(), DType::F32);
    assert_eq!(PrecisionType::Float16.to_dtype(), DType::F16);
    assert_eq!(PrecisionType::BFloat16.to_dtype(), DType::BF16);

    info!("Precision type properties test passed");
}

#[test]
fn test_memory_optimization_full_pipeline() {
    let device = test_device();
    info!(?device, "Running full memory optimization pipeline test");

    let mut stats = MemoryStats::new();

    // Step 1: Create baseline model (F32)
    let model_tensor = create_test_tensor(&device, &[512, 512]);
    let baseline_size = (model_tensor.dims().iter().product::<usize>() * 4) as f64 / 1_048_576.0;
    stats.add_component("baseline_model", baseline_size);
    stats.update_peak(baseline_size);

    info!(baseline_size_mb = baseline_size, "Step 1: Baseline model (F32)");

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

    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;

    info!(fp16_size_mb = fp16_size, precision_savings_mb = precision_savings, "Step 2: FP16 model");

    // Step 3: Apply INT8 quantization
    let fp32_for_quant = precision_converter
        .to_float32(&fp16_tensor)
        .expect("F32 conversion failed");

    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")
        .expect("Quantization failed");

    let quantized_size = quantized.memory_bytes() as f64 / 1_048_576.0;
    let quant_savings = fp16_size - quantized_size;
    stats.add_component("int8_fp16_model", quantized_size);
    stats.savings_mb += quant_savings;

    info!(quantized_size_mb = quantized_size, quant_savings_mb = quant_savings, "Step 3: INT8+FP16 model");

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

    info!(baseline_mb = baseline_size, optimized_mb = quantized_size, total_savings_mb = total_savings, savings_pct = savings_percent, fits_4gb = quantized_size < 3500.0, "Memory optimization pipeline summary");

    // Verify significant savings
    assert!(
        savings_percent >= 85.0,
        "Expected at least 85% memory savings"
    );

    info!("Full memory optimization pipeline test passed");
}