rustytorch/examples/pinn_mre_helmholtz/tests/validate_uber_kernel.rs

//! Validation tests for uber-kernel correctness
//!
//! These tests verify that the fused CUDA kernels produce output that matches
//! the reference cuBLAS-based forward pass within floating-point tolerance.

use pinn_mre_helmholtz::{Config, Mre1DPinnSolver, ForwardWorkspace};
use rtx_tensor::{Device, Tensor};

#[cfg(feature = "cuda")]
mod uber_kernel_tests {
    use super::*;
    use pinn_mre_helmholtz::uber_kernel::{
        forward_uber_value_only, forward_uber_value_only_out, forward_fused_lite,
    };

    /// Maximum allowed difference between fused kernel and cuBLAS output
    /// Using 1e-4 because different accumulation orders can cause small differences
    const MAX_ERROR: f32 = 1e-4;

    /// Compare two tensors element-wise, return max absolute difference
    fn max_abs_diff(a: &Tensor, b: &Tensor) -> f32 {
        let a_data = a.to_vec().expect("Failed to get a data");
        let b_data = b.to_vec().expect("Failed to get b data");

        assert_eq!(a_data.len(), b_data.len(), "Tensor sizes don't match");

        a_data.iter()
            .zip(b_data.iter())
            .map(|(x, y)| (x - y).abs())
            .fold(0.0f32, |acc, x| acc.max(x))
    }

    #[test]
    fn test_uber_value_only_matches_forward() {
        let cfg = Config {
            n_data: 200,
            n_pde: 200,
            epochs: 1,
            ..Config::default()
        };

        let device = Device::Cuda(0);
        let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
        let mut ws = ForwardWorkspace::new(200, &cfg, &device).expect("Failed to create workspace");

        // Reference: cuBLAS-based forward pass (writes to ws.output)
        solver.u_net
            .forward_with_workspace(&solver.x_data, &mut ws)
            .expect("Reference forward failed");

        // Test: uber_value_only kernel
        let weights = solver.u_net.to_uber_kernel_weights(&cfg);
        let scale = 2.0 * std::f64::consts::PI as f32;

        let uber_output = forward_uber_value_only(&solver.x_data, &weights, scale)
            .expect("Uber kernel forward failed");

        let max_diff = max_abs_diff(&ws.output, &uber_output);
        println!("uber_value_only vs cuBLAS max diff: {}", max_diff);

        assert!(
            max_diff < MAX_ERROR,
            "uber_value_only output differs from cuBLAS by {} (max allowed: {})",
            max_diff, MAX_ERROR
        );
    }

    #[test]
    fn test_uber_value_only_out_matches_forward() {
        let cfg = Config {
            n_data: 200,
            n_pde: 200,
            epochs: 1,
            ..Config::default()
        };

        let device = Device::Cuda(0);
        let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
        let mut ws = ForwardWorkspace::new(200, &cfg, &device).expect("Failed to create workspace");

        // Reference: cuBLAS-based forward pass (writes to ws.output)
        solver.u_net
            .forward_with_workspace(&solver.x_data, &mut ws)
            .expect("Reference forward failed");

        // Test: uber_value_only_out kernel with pre-allocated output
        let weights = solver.u_net.to_uber_kernel_weights(&cfg);
        let scale = 2.0 * std::f64::consts::PI as f32;

        let mut output = Tensor::zeros([200, 2], &device)
            .expect("Failed to allocate output");

        forward_uber_value_only_out(&solver.x_data, &weights, scale, &mut output)
            .expect("Uber kernel forward_out failed");

        let max_diff = max_abs_diff(&ws.output, &output);
        println!("uber_value_only_out vs cuBLAS max diff: {}", max_diff);

        assert!(
            max_diff < MAX_ERROR,
            "uber_value_only_out output differs from cuBLAS by {} (max allowed: {})",
            max_diff, MAX_ERROR
        );
    }

    #[test]
    fn test_fused_lite_matches_forward() {
        let cfg = Config {
            n_data: 200,
            n_pde: 200,
            epochs: 1,
            ..Config::default()
        };

        let device = Device::Cuda(0);
        let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
        let mut ws = ForwardWorkspace::new(200, &cfg, &device).expect("Failed to create workspace");

        // Reference: cuBLAS-based forward pass (writes to ws.output)
        solver.u_net
            .forward_with_workspace(&solver.x_data, &mut ws)
            .expect("Reference forward failed");

        // Test: fused_lite kernel with shared memory
        let weights = solver.u_net.to_uber_kernel_weights(&cfg);
        let scale = 2.0 * std::f64::consts::PI as f32;

        let mut output = Tensor::zeros([200, 2], &device)
            .expect("Failed to allocate output");

        forward_fused_lite(&solver.x_data, &weights, scale, &mut output)
            .expect("Fused-lite kernel failed");

        let max_diff = max_abs_diff(&ws.output, &output);
        println!("fused_lite vs cuBLAS max diff: {}", max_diff);

        assert!(
            max_diff < MAX_ERROR,
            "fused_lite output differs from cuBLAS by {} (max allowed: {})",
            max_diff, MAX_ERROR
        );
    }

    #[test]
    fn test_all_kernels_consistent_larger_batch() {
        let cfg = Config {
            n_data: 1000,
            n_pde: 1000,
            epochs: 1,
            ..Config::default()
        };

        let device = Device::Cuda(0);
        let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
        let mut ws = ForwardWorkspace::new(1000, &cfg, &device).expect("Failed to create workspace");

        // Reference: cuBLAS-based forward pass (writes to ws.output)
        solver.u_net
            .forward_with_workspace(&solver.x_data, &mut ws)
            .expect("Reference forward failed");

        let weights = solver.u_net.to_uber_kernel_weights(&cfg);
        let scale = 2.0 * std::f64::consts::PI as f32;

        // Test uber_value_only
        let uber_output = forward_uber_value_only(&solver.x_data, &weights, scale)
            .expect("Uber kernel forward failed");
        let uber_diff = max_abs_diff(&ws.output, &uber_output);
        println!("uber_value_only (1000 pts) vs cuBLAS max diff: {}", uber_diff);

        // Test fused_lite
        let mut fused_output = Tensor::zeros([1000, 2], &device)
            .expect("Failed to allocate output");
        forward_fused_lite(&solver.x_data, &weights, scale, &mut fused_output)
            .expect("Fused-lite kernel failed");
        let fused_diff = max_abs_diff(&ws.output, &fused_output);
        println!("fused_lite (1000 pts) vs cuBLAS max diff: {}", fused_diff);

        assert!(uber_diff < MAX_ERROR, "uber_value_only error too large: {}", uber_diff);
        assert!(fused_diff < MAX_ERROR, "fused_lite error too large: {}", fused_diff);
    }
}