foxhunt/tests/gpu/gpu_memory_management_test.rs

//! GPU Memory Management Tests
//! 
//! Tests GPU memory allocation, deallocation, and memory pool management

use super::utils::*;
use candle_core::{Device, DType, Tensor};
use log::{info, warn};
use std::collections::HashMap;

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_gpu_memory_allocation() {
        init_gpu_test_env();
        info!("💾 Testing GPU memory allocation");
        
        let device = get_test_device();
        
        // Test progressive memory allocation
        let sizes = vec![
            (10, 10),        // Small tensor
            (100, 100),      // Medium tensor  
            (1000, 100),     // Large tensor
            (10, 10, 10),    // 3D tensor
        ];
        
        let mut tensors = Vec::new();
        
        for (i, size) in sizes.into_iter().enumerate() {
            let tensor = match size.len() {
                2 => Tensor::zeros((size[0], size[1]), DType::F32, &device),
                3 => Tensor::zeros((size[0], size[1], size[2]), DType::F32, &device),
                _ => panic!("Unsupported tensor dimensions"),
            };
            
            match tensor {
                Ok(t) => {
                    let elem_count: usize = t.shape().dims().iter().product();
                    let memory_size = elem_count * 4; // f32 = 4 bytes
                    
                    info!("✅ Allocated tensor {}: {:?} ({} bytes)", 
                          i, t.shape(), memory_size);
                    tensors.push(t);
                }
                Err(e) => {
                    warn!("⚠️  Failed to allocate tensor {}: {}", i, e);
                    break;
                }
            }
        }
        
        info!("✅ GPU memory allocation test completed - {} tensors allocated", tensors.len());
        assert!(!tensors.is_empty(), "Should allocate at least one tensor");
    }
    
    #[test]
    fn test_gpu_memory_deallocation() {
        init_gpu_test_env();
        info!("🗑️  Testing GPU memory deallocation");
        
        let device = get_test_device();
        
        // Allocate and deallocate tensors in a loop
        for iteration in 0..5 {
            info!("Iteration {}: Allocating tensors", iteration);
            
            let mut temp_tensors = Vec::new();
            
            // Allocate multiple tensors
            for i in 0..10 {
                match Tensor::randn(0.0, 1.0, (100, 100), &device) {
                    Ok(tensor) => {
                        temp_tensors.push(tensor);
                    }
                    Err(e) => {
                        warn!("Allocation failed at tensor {}: {}", i, e);
                        break;
                    }
                }
            }
            
            info!("Iteration {}: Allocated {} tensors", iteration, temp_tensors.len());
            
            // Tensors automatically deallocated when temp_tensors goes out of scope
        }
        
        info!("✅ GPU memory deallocation test completed");
    }
    
    #[test]
    fn test_gpu_memory_limits() {
        init_gpu_test_env();
        
        if !cuda_available() {
            info!("⚠️  Skipping GPU memory limits test - CUDA not available");
            return;
        }
        
        info!("📏 Testing GPU memory limits");
        
        let device = get_test_device();
        
        // Try to allocate progressively larger tensors until we hit limits
        let mut max_successful_size = 0;
        let mut step_size = 1000;
        
        for size in (step_size..=100_000).step_by(step_size) {
            match Tensor::zeros((size, size), DType::F32, &device) {
                Ok(_tensor) => {
                    max_successful_size = size;
                    info!("✅ Successfully allocated {}x{} tensor", size, size);
                }
                Err(e) => {
                    warn!("❌ Failed to allocate {}x{} tensor: {}", size, size, e);
                    break;
                }
            }
        }
        
        info!("Maximum successful tensor size: {}x{}", max_successful_size, max_successful_size);
        
        if max_successful_size > 0 {
            let memory_estimate = (max_successful_size * max_successful_size * 4) as f64 / (1024.0 * 1024.0);
            info!("Estimated memory usage: {:.2} MB", memory_estimate);
        }
        
        assert!(max_successful_size > 0, "Should be able to allocate at least small tensors");
    }
    
    #[test]
    fn test_gpu_memory_fragmentation() {
        init_gpu_test_env();
        info!("🧩 Testing GPU memory fragmentation handling");
        
        let device = get_test_device();
        
        // Allocate tensors of different sizes to test fragmentation
        let mut tensors: HashMap<String, Tensor> = HashMap::new();
        
        // Phase 1: Allocate various sized tensors
        let allocations = vec![
            ("small_1", (50, 50)),
            ("large_1", (500, 500)),
            ("small_2", (75, 75)),
            ("medium_1", (200, 200)),
            ("small_3", (60, 60)),
            ("large_2", (400, 400)),
        ];
        
        for (name, (rows, cols)) in allocations {
            match Tensor::randn(0.0, 1.0, (rows, cols), &device) {
                Ok(tensor) => {
                    info!("✅ Allocated {}: {}x{}", name, rows, cols);
                    tensors.insert(name.to_string(), tensor);
                }
                Err(e) => {
                    warn!("❌ Failed to allocate {}: {}", name, e);
                }
            }
        }
        
        // Phase 2: Deallocate some tensors (simulate fragmentation)
        tensors.remove("small_2");
        tensors.remove("large_1");
        info!("🗑️  Deallocated some tensors to create fragmentation");
        
        // Phase 3: Try to allocate new tensors in fragmented space
        match Tensor::randn(0.0, 1.0, (300, 300), &device) {
            Ok(_tensor) => {
                info!("✅ Successfully allocated tensor in fragmented memory space");
            }
            Err(e) => {
                warn!("⚠️  Failed to allocate in fragmented space: {}", e);
            }
        }
        
        info!("✅ GPU memory fragmentation test completed");
    }
    
    #[test]
    fn test_gpu_memory_pool_behavior() {
        init_gpu_test_env();
        info!("🏊 Testing GPU memory pool behavior");
        
        let device = get_test_device();
        
        // Test memory pool efficiency by allocating/deallocating repeatedly
        let tensor_size = (100, 100);
        let num_iterations = 20;
        
        for i in 0..num_iterations {
            // Allocate tensor
            let tensor = Tensor::zeros(tensor_size, DType::F32, &device);
            
            match tensor {
                Ok(t) => {
                    // Do some basic operations to ensure tensor is actually used
                    let sum = t.sum_all().unwrap_or_else(|_| {
                        Tensor::new(0.0f32, &device).expect("Should create scalar")
                    });
                    
                    assert_eq!(sum.to_scalar::<f32>().unwrap_or(0.0), 0.0);
                    
                    if i % 5 == 0 {
                        info!("✅ Iteration {}: Memory pool operation successful", i);
                    }
                }
                Err(e) => {
                    warn!("❌ Memory pool operation failed at iteration {}: {}", i, e);
                    break;
                }
            }
            
            // Tensor automatically deallocated here
        }
        
        info!("✅ GPU memory pool behavior test completed");
    }
    
    #[cfg(feature = "cuda")]
    #[test]
    fn test_cuda_memory_info() {
        init_gpu_test_env();
        
        if !cuda_available() {
            info!("⚠️  Skipping CUDA memory info test - CUDA not available");
            return;
        }
        
        info!("📊 Testing CUDA memory information");
        
        match cudarc::driver::CudaDevice::new(0) {
            Ok(device) => {
                if let Ok(total_memory) = device.total_memory() {
                    info!("Total GPU memory: {} bytes ({:.2} GB)", 
                          total_memory, total_memory as f64 / (1024.0 * 1024.0 * 1024.0));
                    
                    // Test memory allocation tracking
                    let test_tensor_size = 1024 * 1024; // 1M floats = 4MB
                    let candle_device = Device::cuda_if_available(0).expect("CUDA should be available");
                    
                    match Tensor::zeros((test_tensor_size,), DType::F32, &candle_device) {
                        Ok(_tensor) => {
                            info!("✅ Allocated test tensor of {} elements", test_tensor_size);
                            
                            // In a real implementation, we could query actual memory usage here
                            // For now, we just verify the allocation succeeded
                        }
                        Err(e) => {
                            warn!("❌ Failed to allocate test tensor: {}", e);
                        }
                    }
                }
                
                info!("✅ CUDA memory info test completed");
            }
            Err(e) => {
                warn!("❌ Failed to initialize CUDA device: {}", e);
            }
        }
    }
}