rustytorch/crates/specialized/rtx-cfd/tests/gpu_kernel_comprehensive_tests.rs

//! Comprehensive GPU kernel tests for CFD operations
//!
//! Tests all GPU kernels including advection, diffusion, and Poisson solvers
//! with various schemes and boundary conditions.

use approx::assert_relative_eq;
use rtx_cfd::kernels::{AdvectionKernel, CudaKernelManager, DiffusionKernel, PoissonKernel};
use rtx_cfd::kernels::{AdvectionScheme, DiffusionScheme};
use rtx_cfd::{CfdConfig, CfdResult};

#[cfg(feature = "cuda")]
mod gpu_kernel_tests {
    use super::*;

    /// Helper to check if CUDA is available
    fn cuda_available() -> bool {
        std::env::var("CUDA_VISIBLE_DEVICES")
            .map(|devices| !devices.is_empty() && devices != "-1")
            .unwrap_or(true)
    }

    /// Test advection kernels with all schemes
    #[test]
    fn test_advection_kernels_all_schemes() -> CfdResult<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping test");
            return Ok(());
        }

        let config = CfdConfig {
            device_id: 0,
            nx: 128,
            ny: 128,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 0.0,
            dt: 0.001,
            ..Default::default()
        };

        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test all advection schemes
        let schemes = vec![
            AdvectionScheme::Upwind,
            AdvectionScheme::Central,
            AdvectionScheme::Quick,
            AdvectionScheme::Weno,
        ];

        for scheme in schemes {
            println!("Testing advection scheme: {:?}", scheme);

            let kernel = AdvectionKernel::new(&manager, scheme)?;

            // Create test data - Gaussian pulse
            let n = 128;
            let mut phi_host = vec![0.0f32; n];
            for i in 0..n {
                let x = i as f32 / n as f32;
                phi_host[i] = (-50.0 * (x - 0.5) * (x - 0.5)).exp();
            }

            // Copy to GPU
            let phi_device = manager.copy_to_device(&phi_host)?;
            let mut phi_new_device = manager.allocate_f32(n)?;

            // Apply advection
            let velocity = 1.0;
            let dt = 0.001;
            let dx = 1.0 / n as f32;

            kernel.apply(&phi_device, &mut phi_new_device, velocity, dt, dx)?;

            // Copy result back
            let phi_new_host = manager.copy_from_device(&phi_new_device)?;

            // Verify: pulse should have moved
            let expected_shift = (velocity * dt / dx) as usize;
            for i in expected_shift..n - expected_shift {
                // Check that values are finite and reasonable
                assert!(
                    phi_new_host[i].is_finite(),
                    "Non-finite value at index {} for scheme {:?}",
                    i,
                    scheme
                );
                assert!(
                    phi_new_host[i] >= -0.1 && phi_new_host[i] <= 1.1,
                    "Value out of range at index {} for scheme {:?}: {}",
                    i,
                    scheme,
                    phi_new_host[i]
                );
            }
        }

        Ok(())
    }

    /// Test 2D advection kernels
    #[test]
    fn test_advection_2d_kernels() -> CfdResult<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping test");
            return Ok(());
        }

        let config = CfdConfig {
            device_id: 0,
            nx: 64,
            ny: 64,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 0.0,
            dt: 0.001,
            ..Default::default()
        };

        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = AdvectionKernel::new(&manager, AdvectionScheme::Upwind)?;

        let nx = 64;
        let ny = 64;
        let n = nx * ny;

        // Create 2D Gaussian pulse
        let mut phi_host = vec![0.0f32; n];
        let mut u_host = vec![1.0f32; n]; // Uniform velocity in x
        let mut v_host = vec![0.5f32; n]; // Uniform velocity in y

        for j in 0..ny {
            for i in 0..nx {
                let idx = j * nx + i;
                let x = i as f32 / nx as f32;
                let y = j as f32 / ny as f32;
                phi_host[idx] = (-50.0 * ((x - 0.5) * (x - 0.5) + (y - 0.5) * (y - 0.5))).exp();
            }
        }

        // Copy to GPU
        let phi_device = manager.copy_to_device(&phi_host)?;
        let u_device = manager.copy_to_device(&u_host)?;
        let v_device = manager.copy_to_device(&v_host)?;
        let mut phi_new_device = manager.allocate_f32(n)?;

        // Apply 2D advection
        let dt = 0.001;
        let dx = 1.0 / nx as f32;
        let dy = 1.0 / ny as f32;

        kernel.apply_2d(
            &phi_device,
            &mut phi_new_device,
            &u_device,
            &v_device,
            dt,
            dx,
            dy,
            nx,
            ny,
        )?;

        // Copy result back
        let phi_new_host = manager.copy_from_device(&phi_new_device)?;

        // Verify all values are finite and reasonable
        for idx in 0..n {
            assert!(
                phi_new_host[idx].is_finite(),
                "Non-finite value at index {}",
                idx
            );
            assert!(
                phi_new_host[idx] >= -0.1 && phi_new_host[idx] <= 1.1,
                "Value out of range at index {}: {}",
                idx,
                phi_new_host[idx]
            );
        }

        Ok(())
    }

    /// Test diffusion kernels with all schemes
    #[test]
    fn test_diffusion_kernels_all_schemes() -> CfdResult<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping test");
            return Ok(());
        }

        let config = CfdConfig {
            device_id: 0,
            nx: 128,
            ny: 1,
            nz: 1,
            lx: 1.0,
            ly: 0.0,
            lz: 0.0,
            dt: 0.0001,
            ..Default::default()
        };

        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test all diffusion schemes
        let schemes = vec![
            DiffusionScheme::Explicit,
            DiffusionScheme::Implicit,
            DiffusionScheme::CrankNicolson,
        ];

        for scheme in schemes {
            println!("Testing diffusion scheme: {:?}", scheme);

            let kernel = DiffusionKernel::new(&manager, scheme)?;

            // Create test data - step function
            let n = 128;
            let mut temp_host = vec![0.0f32; n];
            for i in n / 4..3 * n / 4 {
                temp_host[i] = 1.0;
            }

            // Copy to GPU
            let temp_device = manager.copy_to_device(&temp_host)?;
            let mut temp_new_device = manager.allocate_f32(n)?;

            // Apply diffusion
            let alpha = 0.01; // Thermal diffusivity
            let dt = 0.0001;
            let dx = 1.0 / n as f32;

            kernel.apply(&temp_device, &mut temp_new_device, alpha, dt, dx)?;

            // Copy result back
            let temp_new_host = manager.copy_from_device(&temp_new_device)?;

            // Verify: step function should be smoothed
            for i in 0..n {
                assert!(
                    temp_new_host[i].is_finite(),
                    "Non-finite value at index {} for scheme {:?}",
                    i,
                    scheme
                );
                assert!(
                    temp_new_host[i] >= -0.1 && temp_new_host[i] <= 1.1,
                    "Value out of range at index {} for scheme {:?}: {}",
                    i,
                    scheme,
                    temp_new_host[i]
                );
            }

            // Check that diffusion occurred (values should be less sharp)
            if n > 4 {
                let mid = n / 2;
                assert!(
                    temp_new_host[mid] > 0.0 && temp_new_host[mid] <= 1.0,
                    "Diffusion did not occur properly for scheme {:?}",
                    scheme
                );
            }
        }

        Ok(())
    }

    /// Test 2D diffusion kernels
    #[test]
    fn test_diffusion_2d_kernels() -> CfdResult<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping test");
            return Ok(());
        }

        let config = CfdConfig {
            device_id: 0,
            nx: 64,
            ny: 64,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 0.0,
            dt: 0.0001,
            ..Default::default()
        };

        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = DiffusionKernel::new(&manager, DiffusionScheme::Explicit)?;

        let nx = 64;
        let ny = 64;
        let n = nx * ny;

        // Create 2D hot spot in center
        let mut temp_host = vec![0.0f32; n];
        for j in ny / 4..3 * ny / 4 {
            for i in nx / 4..3 * nx / 4 {
                let idx = j * nx + i;
                temp_host[idx] = 1.0;
            }
        }

        // Copy to GPU
        let temp_device = manager.copy_to_device(&temp_host)?;
        let mut temp_new_device = manager.allocate_f32(n)?;

        // Apply 2D diffusion
        let alpha = 0.01;
        let dt = 0.0001;
        let dx = 1.0 / nx as f32;
        let dy = 1.0 / ny as f32;

        kernel.apply_2d(
            &temp_device,
            &mut temp_new_device,
            alpha,
            dt,
            dx,
            dy,
            nx,
            ny,
        )?;

        // Copy result back
        let temp_new_host = manager.copy_from_device(&temp_new_device)?;

        // Verify all values are finite and in range
        for idx in 0..n {
            assert!(
                temp_new_host[idx].is_finite(),
                "Non-finite value at index {}",
                idx
            );
            assert!(
                temp_new_host[idx] >= -0.1 && temp_new_host[idx] <= 1.1,
                "Value out of range at index {}: {}",
                idx,
                temp_new_host[idx]
            );
        }

        Ok(())
    }

    /// Test Poisson solver
    #[test]
    fn test_poisson_solver() -> CfdResult<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping test");
            return Ok(());
        }

        let config = CfdConfig {
            device_id: 0,
            nx: 32,
            ny: 32,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 0.0,
            dt: 0.001,
            ..Default::default()
        };

        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = PoissonKernel::new(&manager)?;

        let nx = 32;
        let ny = 32;
        let n = nx * ny;

        // Create source term - point source in center
        let mut source_host = vec![0.0f32; n];
        let center = ny / 2 * nx + nx / 2;
        source_host[center] = 100.0;

        // Initial guess
        let phi_host = vec![0.0f32; n];

        // Copy to GPU
        let mut phi_device = manager.copy_to_device(&phi_host)?;
        let source_device = manager.copy_to_device(&source_host)?;

        // Solve Poisson equation
        let dx = 1.0 / nx as f32;
        let dy = 1.0 / ny as f32;
        let max_iterations = 1000;
        let tolerance = 1e-4;

        let iterations = kernel.solve_2d(
            &mut phi_device,
            &source_device,
            nx,
            ny,
            dx,
            dy,
            max_iterations,
            tolerance,
        )?;

        println!("Poisson solver converged in {} iterations", iterations);

        // Copy result back
        let phi_result = manager.copy_from_device(&phi_device)?;

        // Verify solution
        for idx in 0..n {
            assert!(
                phi_result[idx].is_finite(),
                "Non-finite value at index {}",
                idx
            );
        }

        // Check that center has highest value
        let center_val = phi_result[center];
        assert!(center_val > 0.0, "Center value should be positive");

        Ok(())
    }

    /// Test memory management and large arrays
    #[test]
    fn test_gpu_memory_management() -> CfdResult<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping test");
            return Ok(());
        }

        let config = CfdConfig {
            device_id: 0,
            nx: 256,
            ny: 256,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 0.0,
            dt: 0.001,
            ..Default::default()
        };

        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test allocation and deallocation of large arrays
        let sizes = vec![1024, 65536, 262144];

        for size in sizes {
            println!("Testing allocation of size {}", size);

            // Allocate
            let _gpu_array = manager.allocate_f32(size)?;

            // Create test data
            let host_data: Vec<f32> = (0..size).map(|i| i as f32 / 1000.0).collect();

            // Copy to GPU and back
            let gpu_data = manager.copy_to_device(&host_data)?;
            let result = manager.copy_from_device(&gpu_data)?;

            // Verify
            for i in 0..size.min(100) {
                // Check first 100 elements
                assert_relative_eq!(result[i], host_data[i], epsilon = 1e-5);
            }
        }

        Ok(())
    }

    /// Test kernel performance and timing
    #[test]
    fn test_kernel_performance() -> CfdResult<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping test");
            return Ok(());
        }

        let config = CfdConfig {
            device_id: 0,
            nx: 512,
            ny: 512,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 0.0,
            dt: 0.001,
            ..Default::default()
        };

        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = AdvectionKernel::new(&manager, AdvectionScheme::Upwind)?;

        let n = 512 * 512;
        let phi_host = vec![1.0f32; n];
        let u_host = vec![1.0f32; n];
        let v_host = vec![0.5f32; n];

        let phi_device = manager.copy_to_device(&phi_host)?;
        let u_device = manager.copy_to_device(&u_host)?;
        let v_device = manager.copy_to_device(&v_host)?;
        let mut phi_new_device = manager.allocate_f32(n)?;

        // Warm up
        kernel.apply_2d(
            &phi_device,
            &mut phi_new_device,
            &u_device,
            &v_device,
            0.001,
            0.01,
            0.01,
            512,
            512,
        )?;

        // Time multiple iterations
        let start = std::time::Instant::now();
        let iterations = 10;

        for _ in 0..iterations {
            kernel.apply_2d(
                &phi_device,
                &mut phi_new_device,
                &u_device,
                &v_device,
                0.001,
                0.01,
                0.01,
                512,
                512,
            )?;
            manager.synchronize()?;
        }

        let elapsed = start.elapsed();
        let avg_time = elapsed.as_secs_f64() / iterations as f64;

        println!("Average kernel execution time: {:.3} ms", avg_time * 1000.0);
        println!("Throughput: {:.2} Mcells/s", n as f64 / avg_time / 1e6);

        // Performance should be reasonable
        assert!(avg_time < 1.0, "Kernel too slow: {:.3} s", avg_time);

        Ok(())
    }
}

#[cfg(not(feature = "cuda"))]
mod gpu_kernel_tests {
    use super::*;

    #[test]
    fn test_cuda_not_available() {
        println!("CUDA feature not enabled, GPU tests skipped");
        assert!(true);
    }
}