ruvector-onnx-embeddings 0.1.0

//! GPU Module Tests
//!
//! Comprehensive tests for GPU acceleration functionality.

use super::*;
use super::config::{GpuConfig, GpuMode, PowerPreference, GpuMemoryStats};
use super::backend::CpuBackend;
use super::shaders::ShaderModule;

// ==================== Configuration Tests ====================

#[test]
fn test_gpu_config_default() {
    let config = GpuConfig::default();

    assert_eq!(config.mode, GpuMode::Auto);
    assert_eq!(config.power_preference, PowerPreference::HighPerformance);
    assert_eq!(config.workgroup_size, 256);
    assert!(config.fallback_to_cpu);
    assert!(config.cache_shaders);
}

#[test]
fn test_gpu_config_builder() {
    let config = GpuConfig::auto()
        .with_mode(GpuMode::WebGpu)
        .with_power_preference(PowerPreference::LowPower)
        .with_workgroup_size(512)
        .with_min_batch_size(32)
        .with_min_dimension(256)
        .with_profiling(true);

    assert_eq!(config.mode, GpuMode::WebGpu);
    assert_eq!(config.power_preference, PowerPreference::LowPower);
    assert_eq!(config.workgroup_size, 512);
    assert_eq!(config.min_batch_size, 32);
    assert_eq!(config.min_dimension, 256);
    assert!(config.enable_profiling);
}

#[test]
fn test_should_use_gpu() {
    let config = GpuConfig::default()
        .with_min_batch_size(16)
        .with_min_dimension(128);

    // Below minimum batch size
    assert!(!config.should_use_gpu(8, 384));

    // Below minimum dimension
    assert!(!config.should_use_gpu(32, 64));

    // Both conditions met
    assert!(config.should_use_gpu(32, 384));

    // CPU only mode
    let cpu_config = GpuConfig::cpu_only();
    assert!(!cpu_config.should_use_gpu(1000, 1000));
}

#[test]
fn test_preset_configs() {
    let high_perf = GpuConfig::high_performance();
    assert_eq!(high_perf.workgroup_size, 512);
    assert_eq!(high_perf.min_batch_size, 8);

    let low_power = GpuConfig::low_power();
    assert_eq!(low_power.power_preference, PowerPreference::LowPower);
    assert_eq!(low_power.workgroup_size, 128);

    let cpu_only = GpuConfig::cpu_only();
    assert_eq!(cpu_only.mode, GpuMode::CpuOnly);
}

// ==================== Shader Tests ====================

#[test]
fn test_shader_registry_initialization() {
    let registry = ShaderRegistry::new();

    let expected_shaders = vec![
        "cosine_similarity",
        "batch_cosine_similarity",
        "dot_product",
        "euclidean_distance",
        "l2_normalize",
        "mean_pool",
        "max_pool",
        "cls_pool",
        "matmul",
        "vector_add",
        "vector_scale",
    ];

    for name in expected_shaders {
        assert!(registry.get(name).is_some(), "Missing shader: {}", name);
    }
}

#[test]
fn test_shader_module_content() {
    let registry = ShaderRegistry::new();

    // Check cosine similarity shader
    let cosine = registry.get("cosine_similarity").unwrap();
    assert!(cosine.source.contains("@compute"));
    assert!(cosine.source.contains("workgroup_size"));
    assert!(cosine.source.contains("cosine_similarity"));
    assert_eq!(cosine.entry_point, "cosine_similarity");
    assert_eq!(cosine.workgroup_size, [256, 1, 1]);

    // Check mean pool shader
    let mean_pool = registry.get("mean_pool").unwrap();
    assert!(mean_pool.source.contains("attention_mask"));
    assert!(mean_pool.source.contains("hidden_size"));
    assert_eq!(mean_pool.entry_point, "mean_pool");
}

#[test]
fn test_custom_shader_registration() {
    let mut registry = ShaderRegistry::new();

    let custom = ShaderModule {
        name: "custom_kernel".to_string(),
        source: "@compute @workgroup_size(64) fn custom() {}".to_string(),
        entry_point: "custom".to_string(),
        workgroup_size: [64, 1, 1],
    };

    registry.register(custom);

    assert!(registry.get("custom_kernel").is_some());
    let retrieved = registry.get("custom_kernel").unwrap();
    assert_eq!(retrieved.entry_point, "custom");
}

// ==================== Batch Operations Tests ====================

#[test]
fn test_batch_cosine_similarity() {
    let query = vec![1.0, 0.0, 0.0];
    let candidates: Vec<&[f32]> = vec![
        &[1.0, 0.0, 0.0][..],  // similarity = 1.0
        &[0.0, 1.0, 0.0][..],  // similarity = 0.0
        &[-1.0, 0.0, 0.0][..], // similarity = -1.0
    ];

    let results = batch_cosine_similarity_gpu(&query, &candidates);

    assert_eq!(results.len(), 3);
    assert!((results[0] - 1.0).abs() < 1e-6);
    assert!(results[1].abs() < 1e-6);
    assert!((results[2] - (-1.0)).abs() < 1e-6);
}

#[test]
fn test_batch_dot_product() {
    let query = vec![1.0, 1.0, 1.0];
    let candidates: Vec<&[f32]> = vec![
        &[1.0, 1.0, 1.0][..],  // dot = 3.0
        &[2.0, 2.0, 2.0][..],  // dot = 6.0
        &[0.0, 0.0, 0.0][..],  // dot = 0.0
    ];

    let results = batch_dot_product_gpu(&query, &candidates);

    assert_eq!(results.len(), 3);
    assert!((results[0] - 3.0).abs() < 1e-6);
    assert!((results[1] - 6.0).abs() < 1e-6);
    assert!(results[2].abs() < 1e-6);
}

#[test]
fn test_batch_euclidean() {
    let query = vec![0.0, 0.0, 0.0];
    let candidates: Vec<&[f32]> = vec![
        &[3.0, 4.0, 0.0][..],  // dist = 5.0
        &[1.0, 0.0, 0.0][..],  // dist = 1.0
        &[0.0, 0.0, 0.0][..],  // dist = 0.0
    ];

    let results = batch_euclidean_gpu(&query, &candidates);

    assert_eq!(results.len(), 3);
    assert!((results[0] - 5.0).abs() < 1e-6);
    assert!((results[1] - 1.0).abs() < 1e-6);
    assert!(results[2].abs() < 1e-6);
}

// ==================== Pooling Tests (using public API) ====================

#[test]
fn test_mean_pool_via_api() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let pooler = GpuPooler::new(&backend, &shaders).unwrap();

    // batch=2, seq=2, hidden=3
    let tokens = vec![
        1.0, 2.0, 3.0,  // batch 0, seq 0
        4.0, 5.0, 6.0,  // batch 0, seq 1
        7.0, 8.0, 9.0,  // batch 1, seq 0
        10.0, 11.0, 12.0, // batch 1, seq 1
    ];
    let mask = vec![1i64, 1, 1, 1];

    let result = pooler.mean_pool(&tokens, &mask, 2, 2, 3).unwrap();

    assert_eq!(result.len(), 6);
    // Batch 0: mean of [1,2,3] and [4,5,6] = [2.5, 3.5, 4.5]
    assert!((result[0] - 2.5).abs() < 1e-6);
    assert!((result[1] - 3.5).abs() < 1e-6);
    assert!((result[2] - 4.5).abs() < 1e-6);
}

#[test]
fn test_cls_pool_via_api() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let pooler = GpuPooler::new(&backend, &shaders).unwrap();

    // batch=2, seq=3, hidden=4
    let tokens = vec![
        // Batch 0
        1.0, 2.0, 3.0, 4.0,    // CLS token
        5.0, 6.0, 7.0, 8.0,
        9.0, 10.0, 11.0, 12.0,
        // Batch 1
        10.0, 20.0, 30.0, 40.0, // CLS token
        50.0, 60.0, 70.0, 80.0,
        90.0, 100.0, 110.0, 120.0,
    ];

    let result = pooler.cls_pool(&tokens, 2, 4).unwrap();

    assert_eq!(result.len(), 8);

    // Batch 0: first token
    assert!((result[0] - 1.0).abs() < 1e-6);
    assert!((result[1] - 2.0).abs() < 1e-6);
    assert!((result[2] - 3.0).abs() < 1e-6);
    assert!((result[3] - 4.0).abs() < 1e-6);

    // Batch 1: first token
    assert!((result[4] - 10.0).abs() < 1e-6);
    assert!((result[5] - 20.0).abs() < 1e-6);
    assert!((result[6] - 30.0).abs() < 1e-6);
    assert!((result[7] - 40.0).abs() < 1e-6);
}

#[test]
fn test_max_pool_via_api() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let pooler = GpuPooler::new(&backend, &shaders).unwrap();

    // batch=1, seq=3, hidden=4
    let tokens = vec![
        1.0, 10.0, 3.0, 4.0,   // seq 0
        5.0, 2.0, 7.0, 8.0,    // seq 1
        9.0, 6.0, 11.0, 0.0,   // seq 2
    ];

    let mask = vec![1i64, 1, 1];

    let result = pooler.max_pool(&tokens, &mask, 1, 3, 4).unwrap();

    assert_eq!(result.len(), 4);

    // Max across all sequences for each dimension
    assert!((result[0] - 9.0).abs() < 1e-6);  // max(1, 5, 9)
    assert!((result[1] - 10.0).abs() < 1e-6); // max(10, 2, 6)
    assert!((result[2] - 11.0).abs() < 1e-6); // max(3, 7, 11)
    assert!((result[3] - 8.0).abs() < 1e-6);  // max(4, 8, 0)
}

// ==================== Vector Operations Tests ====================

#[test]
fn test_normalize_batch() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let ops = GpuVectorOps::new(&backend, &shaders).unwrap();

    let mut vectors = vec![
        3.0, 4.0, 0.0,  // norm = 5, normalized = [0.6, 0.8, 0]
        0.0, 0.0, 5.0,  // norm = 5, normalized = [0, 0, 1]
    ];

    ops.normalize_batch(&mut vectors, 3).unwrap();

    // Check first vector
    assert!((vectors[0] - 0.6).abs() < 1e-6);
    assert!((vectors[1] - 0.8).abs() < 1e-6);
    assert!(vectors[2].abs() < 1e-6);

    // Check second vector
    assert!(vectors[3].abs() < 1e-6);
    assert!(vectors[4].abs() < 1e-6);
    assert!((vectors[5] - 1.0).abs() < 1e-6);
}

#[test]
fn test_matmul() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let ops = GpuVectorOps::new(&backend, &shaders).unwrap();

    // 2x3 matrix
    let matrix = vec![
        1.0, 2.0, 3.0,
        4.0, 5.0, 6.0,
    ];

    // 3x1 vector
    let vector = vec![1.0, 1.0, 1.0];

    let result = ops.matmul(&matrix, &vector, 2, 3).unwrap();

    assert_eq!(result.len(), 2);
    assert!((result[0] - 6.0).abs() < 1e-6);  // 1+2+3
    assert!((result[1] - 15.0).abs() < 1e-6); // 4+5+6
}

#[test]
fn test_batch_add() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let ops = GpuVectorOps::new(&backend, &shaders).unwrap();

    let a = vec![1.0, 2.0, 3.0, 4.0];
    let b = vec![5.0, 6.0, 7.0, 8.0];

    let result = ops.batch_add(&a, &b).unwrap();

    assert_eq!(result, vec![6.0, 8.0, 10.0, 12.0]);
}

#[test]
fn test_batch_scale() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let ops = GpuVectorOps::new(&backend, &shaders).unwrap();

    let mut vectors = vec![1.0, 2.0, 3.0, 4.0];

    ops.batch_scale(&mut vectors, 2.0).unwrap();

    assert_eq!(vectors, vec![2.0, 4.0, 6.0, 8.0]);
}

// ==================== Integration Tests ====================

#[test]
fn test_gpu_similarity_with_backend() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let similarity = GpuSimilarity::new(&backend, &shaders).unwrap();

    let query = vec![1.0, 0.0, 0.0];
    let candidates: Vec<&[f32]> = vec![
        &[1.0, 0.0, 0.0][..],
        &[0.0, 1.0, 0.0][..],
    ];

    let results = similarity.batch_cosine(&query, &candidates).unwrap();

    assert_eq!(results.len(), 2);
    assert!((results[0] - 1.0).abs() < 1e-6);
}

#[test]
fn test_top_k_similar() {
    let backend = CpuBackend;
    let shaders = ShaderRegistry::new();
    let similarity = GpuSimilarity::new(&backend, &shaders).unwrap();

    let query = vec![1.0, 0.0, 0.0];
    let candidates: Vec<&[f32]> = vec![
        &[0.0, 1.0, 0.0][..],   // sim = 0
        &[1.0, 0.0, 0.0][..],   // sim = 1 (best)
        &[0.5, 0.5, 0.0][..],   // sim ≈ 0.707
        &[-1.0, 0.0, 0.0][..],  // sim = -1 (worst)
    ];

    let top2 = similarity.top_k(&query, &candidates, 2).unwrap();

    assert_eq!(top2.len(), 2);
    assert_eq!(top2[0].0, 1); // Index of [1,0,0]
    assert_eq!(top2[1].0, 2); // Index of [0.5,0.5,0]
}

// ==================== Memory Stats Tests ====================

#[test]
fn test_memory_stats() {
    let stats = GpuMemoryStats {
        total: 1024 * 1024 * 1024, // 1GB
        used: 512 * 1024 * 1024,   // 512MB
        free: 512 * 1024 * 1024,
        peak: 768 * 1024 * 1024,
    };

    assert!((stats.usage_percent() - 50.0).abs() < 0.1);
}

#[test]
fn test_empty_memory_stats() {
    let stats = GpuMemoryStats::default();
    assert_eq!(stats.usage_percent(), 0.0);
}

// ==================== Backend Tests ====================

#[test]
fn test_cpu_backend_info() {
    let backend = CpuBackend;

    assert!(backend.is_available());

    let info = backend.device_info();
    assert_eq!(info.backend, "CPU");
    assert!(!info.supports_compute);
}