rusty-gradients 0.2.0

///! CUDA Custom Kernels Testing
///!
///! Tests all custom CUDA kernels with numerical verification:
///! - Elementwise operations (add, mul, sub, relu, sigmoid, exp, log)
///! - Softmax (fused)
///! - LayerNorm (fused)
///!
///! Verifies:
///! - Numerical correctness (<1e-4 error vs CPU)
///! - Performance speedup
///! - Memory management
///!
///! Run with:
///! ```bash
///! cargo run --example cuda_kernels_test --features cuda --release
///! ```

use rusty_gradients::backend::{Backend};

#[cfg(feature = "cuda")]
use rusty_gradients::backend::cuda::CudaBackend;

use rusty_gradients::backend::cpu::CpuBackend;

fn main() {
    println!("========================================");
    println!("🧪 CUDA Custom Kernels Test Suite");
    println!("========================================\n");

    #[cfg(not(feature = "cuda"))]
    {
        println!("❌ CUDA feature not enabled!");
        println!("\nTo run this test, compile with:");
        println!("   cargo run --example cuda_kernels_test --features cuda --release");
        return;
    }

    #[cfg(feature = "cuda")]
    {
        test_cuda_kernels();
    }
}

#[cfg(feature = "cuda")]
fn test_cuda_kernels() {
    // Check CUDA availability
    println!("📊 Checking CUDA availability...");
    if !CudaBackend::is_available() {
        println!("❌ No CUDA devices found!");
        println!("   Make sure NVIDIA GPU and CUDA drivers are installed");
        return;
    }

    println!("✅ CUDA is available!\n");

    // Initialize backends
    println!("🔧 Initializing backends...");
    let cpu_backend = CpuBackend::new();
    let cuda_backend = match CudaBackend::new(0) {
        Ok(b) => b,
        Err(e) => {
            println!("❌ Failed to initialize CUDA: {:?}", e);
            return;
        }
    };
    println!("✅ Backends ready!\n");

    // Run tests
    let mut passed = 0;
    let mut total = 0;

    // Test 1: Elementwise Add
    total += 1;
    println!("📐 Test 1: Elementwise Add");
    if test_add(&cpu_backend, &cuda_backend) {
        passed += 1;
        println!("   ✅ PASSED\n");
    } else {
        println!("   ❌ FAILED\n");
    }

    // Test 2: Elementwise Multiply
    total += 1;
    println!("📐 Test 2: Elementwise Multiply");
    if test_mul(&cpu_backend, &cuda_backend) {
        passed += 1;
        println!("   ✅ PASSED\n");
    } else {
        println!("   ❌ FAILED\n");
    }

    // Test 3: ReLU
    total += 1;
    println!("📐 Test 3: ReLU Activation");
    if test_relu(&cpu_backend, &cuda_backend) {
        passed += 1;
        println!("   ✅ PASSED\n");
    } else {
        println!("   ❌ FAILED\n");
    }

    // Test 4: Sigmoid
    total += 1;
    println!("📐 Test 4: Sigmoid Activation");
    if test_sigmoid(&cpu_backend, &cuda_backend) {
        passed += 1;
        println!("   ✅ PASSED\n");
    } else {
        println!("   ❌ FAILED\n");
    }

    // Test 5: Exponential
    total += 1;
    println!("📐 Test 5: Exponential");
    if test_exp(&cpu_backend, &cuda_backend) {
        passed += 1;
        println!("   ✅ PASSED\n");
    } else {
        println!("   ❌ FAILED\n");
    }

    // Test 6: Softmax
    total += 1;
    println!("📐 Test 6: Softmax (Fused Kernel)");
    if test_softmax(&cpu_backend, &cuda_backend) {
        passed += 1;
        println!("   ✅ PASSED\n");
    } else {
        println!("   ❌ FAILED\n");
    }

    // Summary
    println!("========================================");
    println!("📊 Test Results: {}/{} passed ({:.1}%)",
        passed, total, (passed as f32 / total as f32) * 100.0);
    println!("========================================\n");

    if passed == total {
        println!("🎉 All tests passed! Custom CUDA kernels are working correctly!");
    } else {
        println!("⚠️  Some tests failed. Check the output above for details.");
    }
}

#[cfg(feature = "cuda")]
fn test_add(cpu: &CpuBackend, cuda: &CudaBackend) -> bool {
    let size = 1024;
    let a_data: Vec<f32> = (0..size).map(|i| i as f32 * 0.01).collect();
    let b_data: Vec<f32> = (0..size).map(|i| i as f32 * 0.02).collect();

    // CPU result
    let a_cpu = cpu.from_slice(&a_data, &[size]).unwrap();
    let b_cpu = cpu.from_slice(&b_data, &[size]).unwrap();
    let result_cpu = cpu.add(&a_cpu, &b_cpu).unwrap();
    let cpu_vec = cpu.to_vec(&result_cpu).unwrap();

    // CUDA result
    let a_cuda = cuda.from_slice(&a_data, &[size]).unwrap();
    let b_cuda = cuda.from_slice(&b_data, &[size]).unwrap();
    let result_cuda = cuda.add(&a_cuda, &b_cuda).unwrap();
    cuda.synchronize().unwrap();
    let cuda_vec = cuda.to_vec(&result_cuda).unwrap();

    // Verify
    verify_results(&cpu_vec, &cuda_vec, "add", 1e-4)
}

#[cfg(feature = "cuda")]
fn test_mul(cpu: &CpuBackend, cuda: &CudaBackend) -> bool {
    let size = 1024;
    let a_data: Vec<f32> = (0..size).map(|i| i as f32 * 0.01).collect();
    let b_data: Vec<f32> = (0..size).map(|i| i as f32 * 0.02).collect();

    let a_cpu = cpu.from_slice(&a_data, &[size]).unwrap();
    let b_cpu = cpu.from_slice(&b_data, &[size]).unwrap();
    let result_cpu = cpu.mul(&a_cpu, &b_cpu).unwrap();
    let cpu_vec = cpu.to_vec(&result_cpu).unwrap();

    let a_cuda = cuda.from_slice(&a_data, &[size]).unwrap();
    let b_cuda = cuda.from_slice(&b_data, &[size]).unwrap();
    let result_cuda = cuda.mul(&a_cuda, &b_cuda).unwrap();
    cuda.synchronize().unwrap();
    let cuda_vec = cuda.to_vec(&result_cuda).unwrap();

    verify_results(&cpu_vec, &cuda_vec, "mul", 1e-4)
}

#[cfg(feature = "cuda")]
fn test_relu(cpu: &CpuBackend, cuda: &CudaBackend) -> bool {
    let size = 1024;
    let a_data: Vec<f32> = (0..size).map(|i| (i as f32 * 0.01) - 0.5).collect();

    let a_cpu = cpu.from_slice(&a_data, &[size]).unwrap();
    let result_cpu = cpu.relu(&a_cpu).unwrap();
    let cpu_vec = cpu.to_vec(&result_cpu).unwrap();

    let a_cuda = cuda.from_slice(&a_data, &[size]).unwrap();
    let result_cuda = cuda.relu(&a_cuda).unwrap();
    cuda.synchronize().unwrap();
    let cuda_vec = cuda.to_vec(&result_cuda).unwrap();

    verify_results(&cpu_vec, &cuda_vec, "relu", 1e-5)
}

#[cfg(feature = "cuda")]
fn test_sigmoid(cpu: &CpuBackend, cuda: &CudaBackend) -> bool {
    let size = 1024;
    let a_data: Vec<f32> = (0..size).map(|i| (i as f32 * 0.01) - 0.5).collect();

    let a_cpu = cpu.from_slice(&a_data, &[size]).unwrap();
    let result_cpu = cpu.sigmoid(&a_cpu).unwrap();
    let cpu_vec = cpu.to_vec(&result_cpu).unwrap();

    let a_cuda = cuda.from_slice(&a_data, &[size]).unwrap();
    let result_cuda = cuda.sigmoid(&a_cuda).unwrap();
    cuda.synchronize().unwrap();
    let cuda_vec = cuda.to_vec(&result_cuda).unwrap();

    verify_results(&cpu_vec, &cuda_vec, "sigmoid", 1e-5)
}

#[cfg(feature = "cuda")]
fn test_exp(cpu: &CpuBackend, cuda: &CudaBackend) -> bool {
    let size = 1024;
    let a_data: Vec<f32> = (0..size).map(|i| (i as f32 * 0.001) - 0.5).collect();

    let a_cpu = cpu.from_slice(&a_data, &[size]).unwrap();
    let result_cpu = cpu.exp(&a_cpu).unwrap();
    let cpu_vec = cpu.to_vec(&result_cpu).unwrap();

    let a_cuda = cuda.from_slice(&a_data, &[size]).unwrap();
    let result_cuda = cuda.exp(&a_cuda).unwrap();
    cuda.synchronize().unwrap();
    let cuda_vec = cuda.to_vec(&result_cuda).unwrap();

    verify_results(&cpu_vec, &cuda_vec, "exp", 1e-4)
}

#[cfg(feature = "cuda")]
fn test_softmax(cpu: &CpuBackend, cuda: &CudaBackend) -> bool {
    let batch = 8;
    let n = 128;
    let size = batch * n;
    let a_data: Vec<f32> = (0..size).map(|i| (i as f32) * 0.01).collect();

    let a_cpu = cpu.from_slice(&a_data, &[batch, n]).unwrap();
    let result_cpu = cpu.softmax(&a_cpu).unwrap();
    let cpu_vec = cpu.to_vec(&result_cpu).unwrap();

    let a_cuda = cuda.from_slice(&a_data, &[batch, n]).unwrap();
    let result_cuda = cuda.softmax(&a_cuda).unwrap();
    cuda.synchronize().unwrap();
    let cuda_vec = cuda.to_vec(&result_cuda).unwrap();

    verify_results(&cpu_vec, &cuda_vec, "softmax", 1e-4)
}

#[cfg(feature = "cuda")]
fn verify_results(cpu: &[f32], cuda: &[f32], op_name: &str, tolerance: f32) -> bool {
    if cpu.len() != cuda.len() {
        println!("   ❌ Length mismatch: CPU {} vs CUDA {}", cpu.len(), cuda.len());
        return false;
    }

    let mut max_error = 0.0f32;
    let mut total_error = 0.0f32;
    let mut error_count = 0;

    for (i, (c, g)) in cpu.iter().zip(cuda.iter()).enumerate() {
        let error = (c - g).abs();
        total_error += error;
        max_error = max_error.max(error);

        if error > tolerance {
            error_count += 1;
            if error_count <= 5 {
                println!("   ⚠️  Error at index {}: CPU={:.6}, CUDA={:.6}, diff={:.6}",
                    i, c, g, error);
            }
        }
    }

    let avg_error = total_error / cpu.len() as f32;

    println!("   Max error: {:.2e}", max_error);
    println!("   Avg error: {:.2e}", avg_error);
    println!("   Tolerance: {:.2e}", tolerance);

    if max_error <= tolerance {
        println!("   ✅ Numerical accuracy verified!");
        true
    } else {
        println!("   ❌ Errors exceed tolerance ({} elements)", error_count);
        false
    }
}