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

//! Tests for GPU reduction operations
//!
//! Following strict TDD methodology - verifying our GPU reduction
//! implementations work correctly with cudarc 0.17.3

#[cfg(feature = "cuda")]
mod gpu_reduction_tests {
    use rtx_cfd::kernels::CudaKernelManager;
    use rtx_cfd::{CfdConfig, CfdResult};
    use std::sync::Arc;

    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_gpu_sum_reduction() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

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

        // Create test data
        let test_data: Vec<f32> = (1..=100).map(|i| i as f32).collect();
        let expected_sum = test_data.iter().sum::<f32>();

        // Copy to GPU
        let gpu_data = manager.copy_to_device(&test_data)?;

        // Perform GPU reduction
        let gpu_sum = manager.reduce_sum(&gpu_data)?;

        // Verify result
        assert!(
            (gpu_sum - expected_sum).abs() < 1e-4,
            "GPU sum {} != expected {}",
            gpu_sum,
            expected_sum
        );

        Ok(())
    }

    #[test]
    fn test_gpu_abs_max_reduction() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

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

        // Create test data with known max absolute value
        let mut test_data: Vec<f32> = vec![1.0, -5.0, 3.0, -7.5, 2.0, 6.0];
        let expected_max = 7.5;

        // Copy to GPU
        let gpu_data = manager.copy_to_device(&test_data)?;

        // Perform GPU reduction
        let gpu_max = manager.reduce_abs_max(&gpu_data)?;

        // Verify result
        assert!(
            (gpu_max - expected_max).abs() < 1e-6,
            "GPU max {} != expected {}",
            gpu_max,
            expected_max
        );

        Ok(())
    }

    #[test]
    fn test_gpu_reduction_large_array() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

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

        // Create large test array
        let n = 1_000_000;
        let test_data: Vec<f32> = (0..n).map(|i| (i % 100) as f32).collect();
        let expected_sum = test_data.iter().sum::<f32>();

        // Copy to GPU
        let gpu_data = manager.copy_to_device(&test_data)?;

        // Perform GPU reduction
        let gpu_sum = manager.reduce_sum(&gpu_data)?;

        // Verify result (allow for floating point error accumulation)
        let relative_error = (gpu_sum - expected_sum).abs() / expected_sum;
        assert!(
            relative_error < 1e-5,
            "GPU sum {} != expected {}, relative error: {}",
            gpu_sum,
            expected_sum,
            relative_error
        );

        Ok(())
    }

    #[test]
    fn test_poisson_solver_with_gpu_reduction() -> CfdResult<()> {
        // Skip test if no GPU available
        if std::env::var("SKIP_GPU_TESTS").is_ok() {
            return Ok(());
        }

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

        // Create Poisson solver that uses GPU reduction
        use rtx_cfd::kernels::PoissonKernel;
        let poisson = PoissonKernel::new(&manager)?;

        // Create test problem: Laplace equation with boundary conditions
        let nx = 32;
        let ny = 32;
        let n = nx * ny;

        // Initialize fields
        let mut phi_data = vec![0.0f32; n];
        let source_data = vec![0.0f32; n]; // Laplace equation (no source)

        // Set boundary conditions (phi = 1 on top, 0 elsewhere)
        for i in 0..nx {
            phi_data[i + (ny - 1) * nx] = 1.0;
        }

        // Copy to GPU
        let mut phi = manager.copy_to_device(&phi_data)?;
        let source = manager.copy_to_device(&source_data)?;

        // Solve using GPU with GPU reduction for convergence check
        let iterations = poisson.solve_2d(
            &mut phi,
            &source,
            nx,
            ny,
            1.0 / nx as f32,
            1.0 / ny as f32,
            100,  // max iterations
            1e-4, // tolerance
        )?;

        // Verify convergence
        assert!(
            iterations > 0 && iterations <= 100,
            "Solver took {} iterations",
            iterations
        );

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

        // Verify boundary conditions are preserved
        for i in 0..nx {
            assert!(
                (result[i + (ny - 1) * nx] - 1.0).abs() < 1e-3,
                "Top boundary not preserved"
            );
        }

        Ok(())
    }
}