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

//! Comprehensive GPU kernel tests for rtx-cfd
//!
//! This module tests all CUDA kernels for correctness and performance.
//! Tests compare GPU results with CPU reference implementations.

use approx::assert_relative_eq;
use rtx_cfd::kernels::*;
use rtx_cfd::*;

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

    const TEST_TOLERANCE: f32 = 1e-5;

    fn create_test_config() -> CfdConfig {
        CfdConfig {
            nx: 64,
            ny: 64,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 1.0,
            dt: 0.001,
            viscosity: 0.01,
            density: 1.0,
            device_id: 0,
            ..CfdConfig::default()
        }
    }

    #[test]
    fn test_cuda_kernel_manager_creation() -> CfdResult<()> {
        let config = create_test_config();
        let _manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        Ok(())
    }

    #[test]
    fn test_memory_operations() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test allocation
        let size = 1000;
        let device_array = manager.allocate_f32(size)?;
        assert_eq!(device_array.len(), size);

        // Test host to device copy
        let host_data: Vec<f32> = (0..size).map(|i| i as f32).collect();
        let device_data = manager.copy_to_device(&host_data)?;
        assert_eq!(device_data.len(), size);

        // Test device to host copy
        let result = manager.copy_from_device(&device_data)?;
        assert_eq!(result.len(), size);
        for (i, &val) in result.iter().enumerate() {
            assert_relative_eq!(val, i as f32, epsilon = TEST_TOLERANCE);
        }

        Ok(())
    }

    #[test]
    fn test_advection_kernel_1d() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = AdvectionKernel::new(&manager, AdvectionScheme::Upwind)?;

        let n = 100;
        let dx = 0.01;
        let dt = 0.001;
        let velocity = 1.0;

        // Create initial sine wave
        let phi_host: Vec<f32> = (0..n)
            .map(|i| (2.0 * std::f32::consts::PI * i as f32 / n as f32).sin())
            .collect();

        let phi = manager.copy_to_device(&phi_host)?;
        let mut phi_new = manager.allocate_f32(n)?;

        // Apply advection
        kernel.apply(&phi, &mut phi_new, velocity, dt, dx)?;

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

        // Verify result is reasonable (sine wave should be shifted)
        assert_eq!(result.len(), n);

        // Check that the solution has advected (values should be different but similar magnitude)
        let max_val = result.iter().fold(0.0f32, |acc, &x| acc.max(x.abs()));
        assert!(max_val > 0.5); // Should still have significant magnitude
        assert!(max_val < 1.5); // But not too large

        Ok(())
    }

    #[test]
    fn test_advection_kernel_2d() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = AdvectionKernel::new(&manager, AdvectionScheme::Upwind)?;

        let nx = 32;
        let ny = 32;
        let dx = 0.03125; // 1.0 / 32
        let dy = 0.03125;
        let dt = 0.001;

        // Create 2D Gaussian blob
        let mut phi_host = vec![0.0f32; nx * ny];
        let center_x = nx / 2;
        let center_y = ny / 2;
        for j in 0..ny {
            for i in 0..nx {
                let x = (i as f32 - center_x as f32) * dx;
                let y = (j as f32 - center_y as f32) * dy;
                let r2 = x * x + y * y;
                phi_host[j * nx + i] = (-10.0 * r2).exp();
            }
        }

        // Constant velocity field
        let u_host = vec![1.0f32; nx * ny];
        let v_host = vec![0.5f32; nx * ny];

        let phi = manager.copy_to_device(&phi_host)?;
        let u = manager.copy_to_device(&u_host)?;
        let v = manager.copy_to_device(&v_host)?;
        let mut phi_new = manager.allocate_f32(nx * ny)?;

        // Apply 2D advection
        kernel.apply_2d(&phi, &mut phi_new, &u, &v, dt, dx, dy, nx, ny)?;

        let result = manager.copy_from_device(&phi_new)?;

        // Verify result
        assert_eq!(result.len(), nx * ny);
        let max_val = result.iter().fold(0.0f32, |acc, &x| acc.max(x.abs()));
        assert!(max_val > 0.5); // Gaussian should still have significant magnitude
        assert!(max_val < 1.1); // But shouldn't grow

        Ok(())
    }

    #[test]
    fn test_diffusion_kernel_explicit() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = DiffusionKernel::new(&manager, DiffusionScheme::Explicit)?;

        let n = 100;
        let dx = 0.01;
        let dt = 0.00001; // Small time step for stability
        let alpha = 0.01;

        // Create initial step function
        let mut temp_host = vec![0.0f32; n];
        for i in n / 3..2 * n / 3 {
            temp_host[i] = 1.0;
        }

        let temp = manager.copy_to_device(&temp_host)?;
        let mut temp_new = manager.allocate_f32(n)?;

        // Apply diffusion
        kernel.apply(&temp, &mut temp_new, alpha, dt, dx)?;

        let result = manager.copy_from_device(&temp_new)?;

        // Verify diffusion smoothed the step function
        assert_eq!(result.len(), n);

        // Check that edges are smoothed
        let edge_smoothing = result[n / 3] > temp_host[n / 3 - 1]; // Should have diffused outward
        assert!(edge_smoothing);

        Ok(())
    }

    #[test]
    fn test_diffusion_kernel_2d() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = DiffusionKernel::new(&manager, DiffusionScheme::Explicit)?;

        let nx = 32;
        let ny = 32;
        let dx = 0.03125;
        let dy = 0.03125;
        let dt = 0.0001; // Small time step for 2D stability
        let alpha = 0.01;

        // Create 2D step function (hot center)
        let mut temp_host = vec![0.0f32; nx * ny];
        for j in ny / 3..2 * ny / 3 {
            for i in nx / 3..2 * nx / 3 {
                temp_host[j * nx + i] = 1.0;
            }
        }

        let temp = manager.copy_to_device(&temp_host)?;
        let mut temp_new = manager.allocate_f32(nx * ny)?;

        // Apply 2D diffusion
        kernel.apply_2d(&temp, &mut temp_new, alpha, dt, dx, dy, nx, ny)?;

        let result = manager.copy_from_device(&temp_new)?;

        // Verify 2D diffusion
        assert_eq!(result.len(), nx * ny);

        // Check that heat has spread
        let center_idx = (ny / 2) * nx + (nx / 2);
        assert!(result[center_idx] > 0.8); // Center should still be hot

        // Edge should have some heat
        let edge_idx = (ny / 3 - 1) * nx + (nx / 2);
        assert!(result[edge_idx] > 0.0);

        Ok(())
    }

    #[test]
    fn test_poisson_kernel_jacobi() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = PoissonKernel::new(&manager)?;

        let nx = 32;
        let ny = 32;
        let dx = 1.0 / (nx - 1) as f32;
        let dy = 1.0 / (ny - 1) as f32;

        // Create source term (constant source in center)
        let mut source_host = vec![0.0f32; nx * ny];
        for j in ny / 4..3 * ny / 4 {
            for i in nx / 4..3 * nx / 4 {
                source_host[j * nx + i] = 1.0;
            }
        }

        let source = manager.copy_to_device(&source_host)?;
        let mut phi = manager.allocate_f32(nx * ny)?; // Start with zeros

        // Solve Poisson equation
        let iterations = kernel.solve_jacobi_2d(&mut phi, &source, nx, ny, dx, dy, 100, 1e-6)?;

        let result = manager.copy_from_device(&phi)?;

        // Verify convergence
        assert!(iterations <= 100);
        assert_eq!(result.len(), nx * ny);

        // Solution should be non-zero in the source region
        let center_idx = (ny / 2) * nx + (nx / 2);
        // ∇²φ = f with zero boundaries: a positive source makes a NEGATIVE bump.
        assert!(result[center_idx] < -0.01);

        // Boundaries should remain zero (Dirichlet BC)
        assert_relative_eq!(result[0], 0.0, epsilon = TEST_TOLERANCE);
        assert_relative_eq!(result[nx - 1], 0.0, epsilon = TEST_TOLERANCE);

        Ok(())
    }

    #[test]
    fn test_matrix_ops_tridiagonal() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = MatrixOpsKernel::new(&manager)?;

        let n = 100;

        // Create tridiagonal matrix (2 on diagonal, -1 on off-diagonals)
        let diagonal = vec![2.0f32; n];
        let off_diagonal = vec![-1.0f32; n - 1];
        let x = vec![1.0f32; n];

        let d_diagonal = manager.copy_to_device(&diagonal)?;
        let d_off_diagonal = manager.copy_to_device(&off_diagonal)?;
        let d_x = manager.copy_to_device(&x)?;
        let mut d_y = manager.allocate_f32(n)?;

        // Perform matrix-vector multiplication
        kernel.tridiagonal_matvec(&d_diagonal, &d_off_diagonal, &d_x, &mut d_y)?;

        let result = manager.copy_from_device(&d_y)?;

        // Verify result (should be [1, 0, 0, ..., 0, 1] for this matrix)
        assert_eq!(result.len(), n);
        assert_relative_eq!(result[0], 1.0, epsilon = TEST_TOLERANCE); // 2*1 + (-1)*1 = 1
        assert_relative_eq!(result[n - 1], 1.0, epsilon = TEST_TOLERANCE); // (-1)*1 + 2*1 = 1

        // Interior points should be 0
        for i in 1..n - 1 {
            assert_relative_eq!(result[i], 0.0, epsilon = TEST_TOLERANCE); // (-1)*1 + 2*1 + (-1)*1 = 0
        }

        Ok(())
    }

    #[test]
    fn test_matrix_ops_dot_product() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = MatrixOpsKernel::new(&manager)?;

        let n = 1000;
        let x = vec![2.0f32; n];
        let y = vec![3.0f32; n];

        let d_x = manager.copy_to_device(&x)?;
        let d_y = manager.copy_to_device(&y)?;

        let result = kernel.dot_product(&d_x, &d_y)?;

        // Expected: 2 * 3 * 1000 = 6000
        let expected = 6000.0f32;
        assert_relative_eq!(result, expected, epsilon = TEST_TOLERANCE);

        Ok(())
    }

    #[test]
    fn test_matrix_ops_vector_norm() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = MatrixOpsKernel::new(&manager)?;

        let n = 100;
        let x = vec![3.0f32; n]; // Each element is 3

        let d_x = manager.copy_to_device(&x)?;

        let result = kernel.vector_norm(&d_x)?;

        // Expected: sqrt(3^2 * 100) = sqrt(900) = 30
        let expected = 30.0f32;
        assert_relative_eq!(result, expected, epsilon = TEST_TOLERANCE);

        Ok(())
    }

    #[test]
    fn test_matrix_ops_axpy() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);
        let kernel = MatrixOpsKernel::new(&manager)?;

        let n = 100;
        let alpha = 2.5f32;
        let x = vec![2.0f32; n];
        let y_initial = vec![1.0f32; n];

        let d_x = manager.copy_to_device(&x)?;
        let mut d_y = manager.copy_to_device(&y_initial)?;

        // y = alpha * x + y = 2.5 * 2 + 1 = 6
        kernel.axpy(alpha, &d_x, &mut d_y)?;

        let result = manager.copy_from_device(&d_y)?;

        let expected = 6.0f32;
        for &val in result.iter() {
            assert_relative_eq!(val, expected, epsilon = TEST_TOLERANCE);
        }

        Ok(())
    }

    #[test]
    fn test_performance_comparison() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        let n = 10000;
        let host_data: Vec<f32> = (0..n).map(|i| i as f32 * 0.1).collect();

        // Measure GPU memory operations
        let start = std::time::Instant::now();
        let device_data = manager.copy_to_device(&host_data)?;
        let _result = manager.copy_from_device(&device_data)?;
        let gpu_time = start.elapsed();

        println!("GPU memory round-trip for {} elements: {:?}", n, gpu_time);

        // GPU should be reasonably fast (< 10ms for this size)
        assert!(gpu_time.as_millis() < 100);

        Ok(())
    }

    #[test]
    fn test_kernel_error_handling() -> CfdResult<()> {
        let config = create_test_config();
        let manager = std::sync::Arc::new(CudaKernelManager::new(&config)?);

        // Test with mismatched array sizes (should not crash)
        let small_array = manager.allocate_f32(10)?;
        let large_array = manager.allocate_f32(100)?;

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

        // This should handle the size mismatch gracefully
        let mut phi_new = manager.allocate_f32(10)?;
        let result = kernel.apply(&small_array, &mut phi_new, 1.0, 0.001, 0.01);

        // Should succeed because kernel uses the smaller array's size
        assert!(result.is_ok());

        Ok(())
    }
}

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

    #[test]
    fn test_cuda_not_available() {
        let config = CfdConfig {
            nx: 32,
            ny: 32,
            nz: 1,
            lx: 1.0,
            ly: 1.0,
            lz: 1.0,
            dt: 0.001,
            viscosity: 0.01,
            density: 1.0,
            device_id: 0,
            gpu_memory_pool_size: 1024 * 1024 * 1024,
            reference_length: 1.0,
            reference_velocity: 1.0,
            use_gpu: true,
            ..CfdConfig::default()
        };

        let result = CudaKernelManager::new(&config);
        assert!(result.is_err());

        if let Err(e) = result {
            assert!(e.to_string().contains("CUDA not available"));
        }
    }
}