Files
rustytorch/crates/specialized/rtx-cfd/tests/kernels_tests.rs
T
2026-03-04 00:08:42 +00:00

305 lines
10 KiB
Rust

//! Tests for GPU kernels module
//!
//! This module tests the real CUDA implementations of CFD kernels including:
//! - Advection schemes (upwind, central differencing)
//! - Diffusion schemes (implicit/explicit)
//! - Pressure Poisson solver
//! - Matrix operations for CFD
use approx::assert_relative_eq;
use rtx_cfd::{CfdConfig, CfdResult};
#[cfg(feature = "cuda")]
mod cuda_tests {
use super::*;
use rtx_cfd::kernels::{
AdvectionKernel, AdvectionScheme, CudaKernelManager, DiffusionKernel, DiffusionScheme,
MatrixOpsKernel, PoissonKernel,
};
#[tokio::test]
async fn test_advection_kernel_upwind() -> CfdResult<()> {
let config = CfdConfig::new().with_gpu(true);
let kernel_manager = CudaKernelManager::new(&config)?;
let advection_kernel = AdvectionKernel::new(&kernel_manager, AdvectionScheme::Upwind)?;
// Test data: 1D advection with known analytical solution
let nx = 128;
let dx = 1.0 / (nx as f64);
let dt = 0.001;
let velocity = 1.0;
// Initial condition: Gaussian pulse
let mut phi = vec![0.0f32; nx];
let center = nx / 4;
let sigma = 5.0;
for i in 0..nx {
let x = (i as f64 - center as f64) * dx;
phi[i] = (-(x * x) / (2.0 * sigma * sigma)).exp() as f32;
}
// Allocate GPU memory
let mut d_phi = kernel_manager.allocate_f32(nx)?;
let d_phi_new = kernel_manager.allocate_f32(nx)?;
// Copy to GPU
kernel_manager.copy_to_device(&phi, &mut d_phi)?;
// Run advection kernel
advection_kernel.apply(&d_phi, &d_phi_new, velocity as f32, dt as f32, dx as f32)?;
// Copy result back
let mut result = vec![0.0f32; nx];
kernel_manager.copy_from_device(&d_phi_new, &mut result)?;
// Verify that the pulse has moved (mass conservation)
let initial_mass: f32 = phi.iter().sum();
let final_mass: f32 = result.iter().sum();
assert_relative_eq!(initial_mass, final_mass, epsilon = 1e-6);
// Verify that maximum has shifted in the correct direction
let initial_max_idx = phi
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.unwrap()
.0;
let final_max_idx = result
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.unwrap()
.0;
assert!(
final_max_idx > initial_max_idx,
"Pulse should move in positive direction"
);
Ok(())
}
#[tokio::test]
async fn test_advection_kernel_central() -> CfdResult<()> {
let config = CfdConfig::new().with_gpu(true);
let kernel_manager = CudaKernelManager::new(&config)?;
let advection_kernel = AdvectionKernel::new(&kernel_manager, AdvectionScheme::Central)?;
// Test data: smooth sinusoidal wave
let nx = 256;
let dx = 2.0 * std::f64::consts::PI / (nx as f64);
let dt = 0.001;
let velocity = 1.0;
let mut phi = vec![0.0f32; nx];
for i in 0..nx {
let x = i as f64 * dx;
phi[i] = (2.0 * x).sin() as f32;
}
let mut d_phi = kernel_manager.allocate_f32(nx)?;
let d_phi_new = kernel_manager.allocate_f32(nx)?;
kernel_manager.copy_to_device(&phi, &mut d_phi)?;
advection_kernel.apply(&d_phi, &d_phi_new, velocity as f32, dt as f32, dx as f32)?;
let mut result = vec![0.0f32; nx];
kernel_manager.copy_from_device(&d_phi_new, &mut result)?;
// For central scheme, verify mass conservation and smoothness
let initial_mass: f32 = phi.iter().sum();
let final_mass: f32 = result.iter().sum();
assert_relative_eq!(initial_mass, final_mass, epsilon = 1e-5);
Ok(())
}
#[tokio::test]
async fn test_diffusion_kernel_explicit() -> CfdResult<()> {
let config = CfdConfig::new().with_gpu(true);
let kernel_manager = CudaKernelManager::new(&config)?;
let diffusion_kernel = DiffusionKernel::new(&kernel_manager, DiffusionScheme::Explicit)?;
// Test 1D heat equation with analytical solution
let nx = 128;
let dx = 1.0 / (nx as f64);
let dt = 0.0001; // Small dt for stability
let alpha = 0.1; // Thermal diffusivity
// Initial condition: step function
let mut temperature = vec![0.0f32; nx];
for i in nx / 4..3 * nx / 4 {
temperature[i] = 1.0;
}
let mut d_temp = kernel_manager.allocate_f32(nx)?;
let d_temp_new = kernel_manager.allocate_f32(nx)?;
kernel_manager.copy_to_device(&temperature, &mut d_temp)?;
diffusion_kernel.apply(&d_temp, &d_temp_new, alpha as f32, dt as f32, dx as f32)?;
let mut result = vec![0.0f32; nx];
kernel_manager.copy_from_device(&d_temp_new, &mut result)?;
// Verify diffusion: edges should be smoother, total heat conserved
let initial_total: f32 = temperature.iter().sum();
let final_total: f32 = result.iter().sum();
assert_relative_eq!(initial_total, final_total, epsilon = 1e-6);
// Verify smoothing: gradient at edges should be reduced
let initial_gradient = (temperature[nx / 4] - temperature[nx / 4 - 1]).abs();
let final_gradient = (result[nx / 4] - result[nx / 4 - 1]).abs();
assert!(
final_gradient < initial_gradient,
"Diffusion should smooth gradients"
);
Ok(())
}
#[tokio::test]
async fn test_poisson_kernel_2d() -> CfdResult<()> {
let config = CfdConfig::new().with_gpu(true);
let kernel_manager = CudaKernelManager::new(&config)?;
let poisson_kernel = PoissonKernel::new(&kernel_manager)?;
// Test 2D Poisson equation: ∇²φ = f
// Analytical solution: φ(x,y) = sin(πx)sin(πy), f = -2π²sin(πx)sin(πy)
let nx = 64;
let ny = 64;
let dx = 1.0 / (nx as f64);
let dy = 1.0 / (ny as f64);
// Setup source term f
let mut source = vec![0.0f32; nx * ny];
for j in 0..ny {
for i in 0..nx {
let x = i as f64 * dx;
let y = j as f64 * dy;
let idx = j * nx + i;
source[idx] = (-2.0
* std::f64::consts::PI.powi(2)
* (std::f64::consts::PI * x).sin()
* (std::f64::consts::PI * y).sin()) as f32;
}
}
// Initial guess (zeros)
let phi = vec![0.0f32; nx * ny];
let mut d_phi = kernel_manager.allocate_f32(nx * ny)?;
let mut d_source = kernel_manager.allocate_f32(nx * ny)?;
kernel_manager.copy_to_device(&phi, &mut d_phi)?;
kernel_manager.copy_to_device(&source, &mut d_source)?;
// Solve Poisson equation
let max_iterations = 1000;
let tolerance = 1e-6;
let iterations = poisson_kernel.solve_2d(
&d_phi,
&d_source,
nx,
ny,
dx as f32,
dy as f32,
max_iterations,
tolerance,
)?;
let mut result = vec![0.0f32; nx * ny];
kernel_manager.copy_from_device(&d_phi, &mut result)?;
// Verify against analytical solution
let mut max_error = 0.0f32;
for j in 1..ny - 1 {
for i in 1..nx - 1 {
let x = i as f64 * dx;
let y = j as f64 * dy;
let idx = j * nx + i;
let analytical =
(std::f64::consts::PI * x).sin() * (std::f64::consts::PI * y).sin();
let error = (result[idx] - analytical as f32).abs();
max_error = max_error.max(error);
}
}
assert!(iterations < max_iterations, "Solver should converge");
assert!(max_error < 0.1, "Solution should be reasonably accurate");
Ok(())
}
#[tokio::test]
async fn test_matrix_ops_kernel() -> CfdResult<()> {
let config = CfdConfig::new().with_gpu(true);
let kernel_manager = CudaKernelManager::new(&config)?;
let matrix_ops = MatrixOpsKernel::new(&kernel_manager)?;
// Test sparse matrix-vector multiplication (typical in CFD)
let n = 1000;
let mut diagonal = vec![2.0f32; n];
let mut off_diagonal = vec![-1.0f32; n - 1];
let mut x = vec![1.0f32; n];
// Create tridiagonal matrix (common in CFD discretizations)
let mut d_diag = kernel_manager.allocate_f32(n)?;
let mut d_off_diag = kernel_manager.allocate_f32(n - 1)?;
let mut d_x = kernel_manager.allocate_f32(n)?;
let d_y = kernel_manager.allocate_f32(n)?;
kernel_manager.copy_to_device(&diagonal, &mut d_diag)?;
kernel_manager.copy_to_device(&off_diagonal, &mut d_off_diag)?;
kernel_manager.copy_to_device(&x, &mut d_x)?;
// Perform A*x = y operation
matrix_ops.tridiagonal_matvec(&d_diag, &d_off_diag, &d_x, &d_y, n)?;
let mut result = vec![0.0f32; n];
kernel_manager.copy_from_device(&d_y, &mut result)?;
// Verify result manually for first few elements
assert_relative_eq!(result[0], 2.0 * x[0] - x[1], epsilon = 1e-6);
assert_relative_eq!(result[1], -x[0] + 2.0 * x[1] - x[2], epsilon = 1e-6);
assert_relative_eq!(result[n - 1], -x[n - 2] + 2.0 * x[n - 1], epsilon = 1e-6);
Ok(())
}
}
#[cfg(not(feature = "cuda"))]
mod cpu_fallback_tests {
use super::*;
#[test]
fn test_cpu_fallback_advection() {
// Basic test to ensure CPU fallback works
let config = CfdConfig::new().with_gpu(false);
// Implementation will provide CPU fallback
assert!(config.validate().is_ok());
}
}
#[tokio::test]
async fn test_kernel_manager_initialization() -> CfdResult<()> {
let config = CfdConfig::new();
#[cfg(feature = "cuda")]
{
// Should successfully create kernel manager
let result = rtx_cfd::kernels::CudaKernelManager::new(&config);
// May fail if no CUDA device available - that's expected in CI
match result {
Ok(_) => println!("CUDA kernel manager created successfully"),
Err(e) => println!("CUDA not available (expected in CI): {}", e),
}
}
#[cfg(not(feature = "cuda"))]
{
println!("CUDA feature not enabled - using CPU fallback");
}
Ok(())
}