/home/noah/src/trueno/src/backends/gpu/mod.rs

Source
//! GPU backend using wgpu (Vulkan/Metal/DX12/WebGPU)
//!
//! This backend provides GPU-accelerated compute for large-scale operations.
//! It uses wgpu for cross-platform GPU access and WGSL compute shaders.
//!
//! # Performance
//!
//! GPU backend is optimal for very large workloads (>100K elements for reductions,
//! >1000×1000 for matrix operations) where transfer overhead is amortized.
//!
//! Expected speedups vs SIMD:
//! - Matrix multiplication (large): 10-50x
//! - Reductions (large): 5-20x
//!
//! # Architecture
//!
//! - Device initialization is lazy (first GPU operation)
//! - Compute shaders written in WGSL
//! - Asynchronous execution with pollster for blocking
//! - Automatic fallback to CPU if GPU unavailable
//!
//! # Memory Hierarchy Abstractions
//!
//! - [`TensorView`] - Structured view into GPU memory with shape/stride metadata
//! - [`PartitionView`] - Tiling strategy for efficient GPU work distribution
//!
//! Based on cuda-tile-behavior.md Section 3.2.

#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
mod batch;

#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
mod device;

#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
mod shaders;

#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
pub mod runtime;

// Memory hierarchy abstractions (always available, no GPU feature required)
mod partition_view;
mod tensor_view;
mod tiled_reduction;

pub use partition_view::{PartitionView, TileInfo};
pub use tensor_view::{MemoryLayout, TensorView};
pub use tiled_reduction::{
    tiled_max_2d, tiled_min_2d, tiled_reduce_2d, tiled_reduce_partial, tiled_sum_2d, MaxOp, MinOp,
    ReduceOp, SumOp, TILE_SIZE,
};

#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
pub use batch::{BufferId, GpuCommandBatch};

// Export GpuDevice for both native and WASM GPU features
#[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
pub use device::GpuDevice;

/// GPU backend for compute operations (native only, uses sync wrappers)
#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
#[derive(Clone)]
pub struct GpuBackend {
    device: Option<GpuDevice>,
}

#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
impl GpuBackend {
    /// Create a new GPU backend
    pub fn new() -> Self {
        Self { device: None }
    }

    /// Initialize GPU device (lazy)
    fn ensure_device(&mut self) -> Result<&GpuDevice, String> {
        if self.device.is_none() {
            self.device = Some(GpuDevice::new()?0);
        }
        Ok(self.device.as_ref().expect("device initialized above"))
    }

    /// Check if GPU is available
    pub fn is_available() -> bool {
        GpuDevice::is_available()
    }

    /// Vector addition on GPU: c = a + b
    ///
    /// # Arguments
    ///
    /// * `a` - Vector a
    /// * `b` - Vector b
    ///
    /// # Returns
    ///
    /// Vector c (element-wise sum)
    pub fn vec_add(&mut self, a: &[f32], b: &[f32]) -> Result<Vec<f32>, String> {
        if a.len() != b.len() {
            return Err(format!(
                "Vector length mismatch: {} != {}",
                a.len(),
                b.len()
            ));
        }

        // wgpu doesn't allow zero-sized buffers
        if a.is_empty() {
            return Err("Cannot perform GPU operation on empty vectors".to_string());
        }

        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; a.len()];

        // Execute GPU compute
        device.vec_add(a, b, &mut result)?;

        Ok(result)
    }

    /// Dot product on GPU: result = sum(a[i] * b[i])
    ///
    /// # Arguments
    ///
    /// * `a` - Vector a
    /// * `b` - Vector b
    ///
    /// # Returns
    ///
    /// Scalar dot product result
    pub fn dot(&mut self, a: &[f32], b: &[f32]) -> Result<f32, String> {
        if a.len() != b.len() {
            return Err(format!(
                "Vector length mismatch: {} != {}",
                a.len(),
                b.len()
            ));
        }

        let device = self.ensure_device()?;

        // Execute GPU compute
        device.dot(a, b)
    }

    /// ReLU activation on GPU: result[i] = max(0, input[i])
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with ReLU applied element-wise
    pub fn relu(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.relu(input, &mut result)?;

        Ok(result)
    }

    /// Leaky ReLU activation on GPU: result[i] = max(negative_slope * input[i], input[i])
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    /// * `negative_slope` - Slope for negative values (typically 0.01)
    ///
    /// # Returns
    ///
    /// Vector with leaky ReLU applied element-wise
    pub fn leaky_relu(&mut self, input: &[f32], negative_slope: f32) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.leaky_relu(input, &mut result, negative_slope)?;

        Ok(result)
    }

    /// ELU activation on GPU: result[i] = x if x > 0, else alpha * (exp(x) - 1)
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    /// * `alpha` - Scaling factor for negative values (typically 1.0)
    ///
    /// # Returns
    ///
    /// Vector with ELU applied element-wise
    pub fn elu(&mut self, input: &[f32], alpha: f32) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.elu(input, &mut result, alpha)?;

        Ok(result)
    }

    /// Clip (clamp) operation on GPU: result[i] = clamp(input[i], min_val, max_val)
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    /// * `min_val` - Minimum value
    /// * `max_val` - Maximum value
    ///
    /// # Returns
    ///
    /// Vector with clip applied element-wise
    pub fn clip(&mut self, input: &[f32], min_val: f32, max_val: f32) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.clip(input, &mut result, min_val, max_val)?;

        Ok(result)
    }

    /// Sigmoid activation on GPU: result[i] = 1 / (1 + exp(-input[i]))
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with sigmoid applied element-wise
    pub fn sigmoid(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.sigmoid(input, &mut result)?;

        Ok(result)
    }

    /// Tanh activation on GPU: result[i] = tanh(input[i])
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with tanh applied element-wise
    pub fn tanh(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.tanh(input, &mut result)?;

        Ok(result)
    }

    /// Swish activation on GPU: result[i] = input[i] / (1 + exp(-input[i]))
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with swish applied element-wise
    pub fn swish(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.swish(input, &mut result)?;

        Ok(result)
    }

    /// GELU activation on GPU: result[i] = 0.5 * input[i] * (1 + tanh(...))
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with GELU applied element-wise
    pub fn gelu(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.gelu(input, &mut result)?;

        Ok(result)
    }

    /// Softmax activation on GPU: result[i] = exp(input[i] - max) / sum(exp(input - max))
    ///
    /// Uses multi-pass reduction for numerical stability:
    /// - Pass 1: Max reduction (parallel)
    /// - Pass 2: Exp-subtract (element-wise)
    /// - Pass 3: Sum reduction (parallel)
    /// - Pass 4: Normalize (element-wise)
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with softmax applied element-wise
    pub fn softmax(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.softmax(input, &mut result)?;

        Ok(result)
    }

    /// Log-softmax activation on GPU: result[i] = log(softmax(input)[i])
    ///
    /// Uses multi-pass reduction for numerical stability:
    /// - Pass 1: Max reduction (parallel)
    /// - Pass 2: Exp-subtract (element-wise)
    /// - Pass 3: Sum reduction (parallel)
    /// - Pass 4: Log-normalize (element-wise)
    ///
    /// # Arguments
    ///
    /// * `input` - Input vector
    ///
    /// # Returns
    ///
    /// Vector with log-softmax applied element-wise
    pub fn log_softmax(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Create output buffer
        let mut result = vec![0.0f32; input.len()];

        // Execute GPU compute
        device.log_softmax(input, &mut result)?;

        Ok(result)
    }

    /// 2D Convolution on GPU: output = input ⊗ kernel
    ///
    /// # Arguments
    ///
    /// * `input` - Input matrix (flattened row-major)
    /// * `kernel` - Convolution kernel (flattened row-major)
    /// * `input_rows` - Number of rows in input
    /// * `input_cols` - Number of columns in input
    /// * `kernel_rows` - Number of rows in kernel
    /// * `kernel_cols` - Number of columns in kernel
    ///
    /// # Returns
    ///
    /// Output matrix (flattened row-major, "valid" convolution)
    /// - output_rows = input_rows - kernel_rows + 1
    /// - output_cols = input_cols - kernel_cols + 1
    pub fn convolve2d(
        &mut self,
        input: &[f32],
        kernel: &[f32],
        input_rows: usize,
        input_cols: usize,
        kernel_rows: usize,
        kernel_cols: usize,
    ) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?;

        // Calculate output dimensions
        let output_rows = input_rows.saturating_sub(kernel_rows).saturating_add(1);
        let output_cols = input_cols.saturating_sub(kernel_cols).saturating_add(1);

        // Create output buffer
        let mut result = vec![0.0f32; output_rows * output_cols];

        // Execute GPU compute
        device.convolve2d(
            input,
            kernel,
            &mut result,
            input_rows,
            input_cols,
            kernel_rows,
            kernel_cols,
        )?;

        Ok(result)
    }

    /// Matrix multiplication on GPU: C = A × B
    ///
    /// # Arguments
    ///
    /// * `a` - Matrix A (m×k) in row-major order
    /// * `b` - Matrix B (k×n) in row-major order
    /// * `m` - Rows of A and C
    /// * `k` - Cols of A, rows of B
    /// * `n` - Cols of B and C
    ///
    /// # Returns
    ///
    /// Matrix C (m×n) in row-major order
    pub fn matmul(
        &mut self,
        a: &[f32],
        b: &[f32],
        m: usize,
        k: usize,
        n: usize,
    ) -> Result<Vec<f32>, String> {
        let device = self.ensure_device()?0;

        // Create output buffer
        let mut result = vec![0.0f32; m * n];

        // Execute GPU compute
        device.matmul(a, b, &mut result, m, k, n)?0;

        Ok(result)
    }

    /// Symmetric eigendecomposition on GPU
    ///
    /// Computes eigenvalues and eigenvectors using Jacobi algorithm with
    /// GPU-accelerated Givens rotations.
    ///
    /// # Arguments
    ///
    /// * `matrix` - Symmetric matrix data (row-major, n×n)
    /// * `n` - Matrix dimension
    ///
    /// # Returns
    ///
    /// Tuple of (eigenvalues, eigenvector_data) where eigenvector_data is row-major
    pub fn symmetric_eigen(
        &mut self,
        matrix: &[f32],
        n: usize,
    ) -> Result<(Vec<f32>, Vec<f32>), String> {
        let device = self.ensure_device()?;
        device.symmetric_eigen(matrix, n)
    }

    /// 2D Tiled Sum Reduction on GPU
    ///
    /// Uses 16×16 workgroups for efficient parallel reduction with
    /// optimal memory coalescing.
    ///
    /// # Arguments
    ///
    /// * `data` - Input 2D data in row-major order
    /// * `width` - Number of columns
    /// * `height` - Number of rows
    ///
    /// # Returns
    ///
    /// Sum of all elements
    pub fn tiled_sum_2d_gpu(
        &mut self,
        data: &[f32],
        width: usize,
        height: usize,
    ) -> Result<f32, String> {
        let device = self.ensure_device()?;
        device.tiled_sum_2d(data, width, height)
    }

    /// 2D Tiled Max Reduction on GPU
    ///
    /// Uses 16×16 workgroups for efficient parallel max reduction.
    pub fn tiled_max_2d_gpu(
        &mut self,
        data: &[f32],
        width: usize,
        height: usize,
    ) -> Result<f32, String> {
        let device = self.ensure_device()?;
        device.tiled_max_2d(data, width, height)
    }

    /// 2D Tiled Min Reduction on GPU
    ///
    /// Uses 16×16 workgroups for efficient parallel min reduction.
    pub fn tiled_min_2d_gpu(
        &mut self,
        data: &[f32],
        width: usize,
        height: usize,
    ) -> Result<f32, String> {
        let device = self.ensure_device()?;
        device.tiled_min_2d(data, width, height)
    }
}

#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
impl Default for GpuBackend {
    fn default() -> Self {
        Self::new()
    }
}

// Stub implementation when GPU feature is disabled or on WASM
#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
#[derive(Clone)]
pub struct GpuBackend;

#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
impl GpuBackend {
    pub fn new() -> Self {
        Self
    }

    pub fn is_available() -> bool {
        false
    }
}

#[cfg(any(not(feature = "gpu"), target_arch = "wasm32"))]
impl Default for GpuBackend {
    fn default() -> Self {
        Self
    }
}

// Tests for stub implementation (when GPU feature is NOT enabled)
#[cfg(test)]
#[cfg(not(feature = "gpu"))]
mod stub_tests {
    use super::*;

    #[test]
    fn test_gpu_backend_stub_new() {
        let _backend = GpuBackend::new();
    }

    #[test]
    fn test_gpu_backend_stub_is_available() {
        assert!(!GpuBackend::is_available());
    }

    #[test]
    fn test_gpu_backend_stub_default() {
        let _backend = GpuBackend::default();
    }

    #[test]
    fn test_gpu_backend_stub_clone() {
        let backend = GpuBackend::new();
        let _cloned = backend.clone();
    }
}

// ===== GPU Tests =====

#[cfg(test)]
#[cfg(feature = "gpu")]
mod tests {
    use super::*;
    use std::sync::OnceLock;

    /// Shared GPU backend for fast test execution (initialized once)
    static SHARED_GPU: OnceLock<Option<GpuBackend>> = OnceLock::new();

    /// Get shared GPU backend (fast) or None if unavailable
    fn get_shared_gpu() -> Option<GpuBackend> {
        SHARED_GPU
            .get_or_init(|| {
                if GpuBackend::is_available() {
                    Some(GpuBackend::new())
                } else {
                    None
                }
            })
            .clone()
    }

    #[test]
    fn test_gpu_vec_add_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let a = vec![1.0, 2.0, 3.0, 4.0];
        let b = vec![5.0, 6.0, 7.0, 8.0];

        let result = gpu.vec_add(&a, &b);

        if let Ok(c) = result {
            assert_eq!(c.len(), 4);
            assert!((c[0] - 6.0).abs() < 1e-4);
            assert!((c[1] - 8.0).abs() < 1e-4);
            assert!((c[2] - 10.0).abs() < 1e-4);
            assert!((c[3] - 12.0).abs() < 1e-4);
        } else {
            eprintln!("GPU vec_add failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_vec_add_large() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let size = 10000;
        let a: Vec<f32> = (0..size).map(|i| i as f32).collect();
        let b: Vec<f32> = (0..size).map(|i| (i * 2) as f32).collect();

        let result = gpu.vec_add(&a, &b);

        if let Ok(c) = result {
            assert_eq!(c.len(), size);
            // Check first few elements
            assert!((c[0] - 0.0).abs() < 1e-4); // 0 + 0
            assert!((c[1] - 3.0).abs() < 1e-4); // 1 + 2
            assert!((c[100] - 300.0).abs() < 1e-4); // 100 + 200
        } else {
            eprintln!("GPU vec_add large failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_vec_add_length_mismatch() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let a = vec![1.0, 2.0, 3.0];
        let b = vec![4.0, 5.0]; // Different length

        let result = gpu.vec_add(&a, &b);
        assert!(result.is_err());
    }

    #[test]
    fn test_gpu_dot_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let a = vec![1.0, 2.0, 3.0, 4.0];
        let b = vec![5.0, 6.0, 7.0, 8.0];

        let result = gpu.dot(&a, &b);

        // Expected: 1*5 + 2*6 + 3*7 + 4*8 = 5 + 12 + 21 + 32 = 70
        if let Ok(dot_product) = result {
            assert!((dot_product - 70.0).abs() < 1e-4);
        } else {
            eprintln!("GPU dot failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_dot_large() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let size = 10000;
        let a: Vec<f32> = (0..size).map(|i| i as f32).collect();
        let b: Vec<f32> = (0..size).map(|_| 1.0).collect();

        let result = gpu.dot(&a, &b);

        // Expected: sum of 0 + 1 + 2 + ... + 9999 = 9999 * 10000 / 2 = 49995000
        if let Ok(dot_product) = result {
            let expected = (size * (size - 1) / 2) as f32;
            assert!((dot_product - expected).abs() < 1.0); // Allow small floating point error
        } else {
            eprintln!("GPU dot large failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_dot_length_mismatch() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let a = vec![1.0, 2.0, 3.0];
        let b = vec![4.0, 5.0]; // Different length

        let result = gpu.dot(&a, &b);
        assert!(result.is_err());
    }

    #[test]
    fn test_gpu_vec_add_matches_scalar() {
        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        use super::super::scalar::ScalarBackend;
        use crate::backends::VectorBackend;

        let mut gpu = GpuBackend::new();
        let a = vec![1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5];
        let b = vec![8.5, 7.5, 6.5, 5.5, 4.5, 3.5, 2.5, 1.5];

        let gpu_result = gpu.vec_add(&a, &b);

        let mut scalar_result = vec![0.0; 8];
        // SAFETY: Test code calling backend trait methods marked unsafe
        unsafe {
            ScalarBackend::add(&a, &b, &mut scalar_result);
        }

        if let Ok(gpu_r) = gpu_result {
            for (g, s) in gpu_r.iter().zip(scalar_result.iter()) {
                assert!(
                    (g - s).abs() < 1e-4,
                    "GPU vs Scalar mismatch: gpu={}, scalar={}",
                    g,
                    s
                );
            }
        } else {
            eprintln!("GPU vec_add failed: {:?}", gpu_result);
        }
    }

    #[test]
    fn test_gpu_dot_matches_scalar() {
        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        use super::super::scalar::ScalarBackend;
        use crate::backends::VectorBackend;

        let mut gpu = GpuBackend::new();
        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let b = vec![8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0];

        let gpu_result = gpu.dot(&a, &b);
        // SAFETY: Test code calling backend trait methods marked unsafe
        let scalar_result = unsafe { ScalarBackend::dot(&a, &b) };

        if let Ok(gpu_r) = gpu_result {
            assert!(
                (gpu_r - scalar_result).abs() < 1e-4,
                "GPU vs Scalar dot mismatch: gpu={}, scalar={}",
                gpu_r,
                scalar_result
            );
        } else {
            eprintln!("GPU dot failed: {:?}", gpu_result);
        }
    }

    #[test]
    fn test_gpu_vec_add_empty() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let a: Vec<f32> = vec![];
        let b: Vec<f32> = vec![];

        let result = gpu.vec_add(&a, &b);

        // GPU backend returns error for empty vectors (wgpu doesn't allow zero-sized buffers)
        assert!(
            result.is_err(),
            "Expected error for empty vectors, got: {:?}",
            result
        );
    }

    #[test]
    fn test_gpu_relu_matches_scalar() {
        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        use super::super::scalar::ScalarBackend;
        use crate::backends::VectorBackend;

        let mut gpu = GpuBackend::new();
        let input = vec![-3.0, -1.0, 0.0, 1.0, 3.0, -2.5, 2.5, 0.5];

        let gpu_result = gpu.relu(&input);
        let mut scalar_result = vec![0.0; input.len()];
        // SAFETY: Test code calling backend trait methods marked unsafe
        unsafe {
            ScalarBackend::relu(&input, &mut scalar_result);
        }

        if let Ok(gpu_r) = gpu_result {
            for (g, s) in gpu_r.iter().zip(scalar_result.iter()) {
                assert!(
                    (g - s).abs() < 1e-4,
                    "GPU vs Scalar relu mismatch: gpu={}, scalar={}",
                    g,
                    s
                );
            }
        } else {
            eprintln!("GPU relu failed: {:?}", gpu_result);
        }
    }

    #[test]
    fn test_gpu_sigmoid_matches_scalar() {
        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        use super::super::scalar::ScalarBackend;
        use crate::backends::VectorBackend;

        let mut gpu = GpuBackend::new();
        let input = vec![-3.0, -1.0, 0.0, 1.0, 3.0, -2.5, 2.5, 0.5];

        let gpu_result = gpu.sigmoid(&input);
        let mut scalar_result = vec![0.0; input.len()];
        // SAFETY: Test code calling backend trait methods marked unsafe
        unsafe {
            ScalarBackend::sigmoid(&input, &mut scalar_result);
        }

        if let Ok(gpu_r) = gpu_result {
            for (g, s) in gpu_r.iter().zip(scalar_result.iter()) {
                assert!(
                    (g - s).abs() < 1e-3,
                    "GPU vs Scalar sigmoid mismatch: gpu={}, scalar={}",
                    g,
                    s
                );
            }
        } else {
            eprintln!("GPU sigmoid failed: {:?}", gpu_result);
        }
    }

    #[test]
    fn test_gpu_gelu_matches_scalar() {
        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        use super::super::scalar::ScalarBackend;
        use crate::backends::VectorBackend;

        let mut gpu = GpuBackend::new();
        let input = vec![-2.0, -1.0, 0.0, 1.0, 2.0, -0.5, 0.5, 1.5];

        let gpu_result = gpu.gelu(&input);
        let mut scalar_result = vec![0.0; input.len()];
        // SAFETY: Test code calling backend trait methods marked unsafe
        unsafe {
            ScalarBackend::gelu(&input, &mut scalar_result);
        }

        if let Ok(gpu_r) = gpu_result {
            for (g, s) in gpu_r.iter().zip(scalar_result.iter()) {
                assert!(
                    (g - s).abs() < 1e-2,
                    "GPU vs Scalar gelu mismatch: gpu={}, scalar={}",
                    g,
                    s
                );
            }
        } else {
            eprintln!("GPU gelu failed: {:?}", gpu_result);
        }
    }

    #[test]
    fn test_gpu_swish_matches_scalar() {
        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        use super::super::scalar::ScalarBackend;
        use crate::backends::VectorBackend;

        let mut gpu = GpuBackend::new();
        let input = vec![-3.0, -1.0, 0.0, 1.0, 3.0, -2.5, 2.5, 0.5];

        let gpu_result = gpu.swish(&input);
        let mut scalar_result = vec![0.0; input.len()];
        // SAFETY: Test code calling backend trait methods marked unsafe
        unsafe {
            ScalarBackend::swish(&input, &mut scalar_result);
        }

        if let Ok(gpu_r) = gpu_result {
            for (g, s) in gpu_r.iter().zip(scalar_result.iter()) {
                assert!(
                    (g - s).abs() < 1e-3,
                    "GPU vs Scalar swish mismatch: gpu={}, scalar={}",
                    g,
                    s
                );
            }
        } else {
            eprintln!("GPU swish failed: {:?}", gpu_result);
        }
    }

    #[test]
    fn test_gpu_clip_matches_scalar() {
        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        use super::super::scalar::ScalarBackend;
        use crate::backends::VectorBackend;

        let mut gpu = GpuBackend::new();
        let input = vec![1.0, 5.0, 10.0, 15.0, -3.0, 20.0, 7.5, 0.0];
        let min_val = 3.0;
        let max_val = 12.0;

        let gpu_result = gpu.clip(&input, min_val, max_val);
        let mut scalar_result = vec![0.0; input.len()];
        // SAFETY: Test code calling backend trait methods marked unsafe
        unsafe {
            ScalarBackend::clamp(&input, min_val, max_val, &mut scalar_result);
        }

        if let Ok(gpu_r) = gpu_result {
            for (g, s) in gpu_r.iter().zip(scalar_result.iter()) {
                assert!(
                    (g - s).abs() < 1e-4,
                    "GPU vs Scalar clip mismatch: gpu={}, scalar={}",
                    g,
                    s
                );
            }
        } else {
            eprintln!("GPU clip failed: {:?}", gpu_result);
        }
    }

    #[test]
    fn test_gpu_leaky_relu_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let input = vec![-3.0, -1.0, 0.0, 1.0, 3.0];
        let negative_slope = 0.01;

        let result = gpu.leaky_relu(&input, negative_slope);

        if let Ok(output) = result {
            // Expected: max(negative_slope * x, x)
            let expected = [-0.03, -0.01, 0.0, 1.0, 3.0];
            for (r, e) in output.iter().zip(expected.iter()) {
                assert!(
                    (r - e).abs() < 1e-4,
                    "Leaky ReLU mismatch: got={}, expected={}",
                    r,
                    e
                );
            }
        } else {
            eprintln!("GPU leaky_relu failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_elu_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let input = vec![-2.0, -1.0, 0.0, 1.0, 2.0];
        let alpha = 1.0;

        let result = gpu.elu(&input, alpha);

        if let Ok(output) = result {
            // Expected: x if x > 0, else alpha * (exp(x) - 1)
            for (i, (r, &x)) in output.iter().zip(input.iter()).enumerate() {
                let expected = if x > 0.0 { x } else { alpha * (x.exp() - 1.0) };
                assert!(
                    (r - expected).abs() < 1e-3,
                    "ELU mismatch at {}: got={}, expected={}",
                    i,
                    r,
                    expected
                );
            }
        } else {
            eprintln!("GPU elu failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_tanh_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let input = vec![-2.0, -1.0, 0.0, 1.0, 2.0];

        let result = gpu.tanh(&input);

        if let Ok(output) = result {
            for (r, &x) in output.iter().zip(input.iter()) {
                let expected = x.tanh();
                assert!(
                    (r - expected).abs() < 1e-4,
                    "Tanh mismatch: got={}, expected={}",
                    r,
                    expected
                );
            }
        } else {
            eprintln!("GPU tanh failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_tanh_not_hardcoded() {
        // EXTREME TDD: Kill mutant that replaces return with Ok(vec![-1.0])
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let input = vec![1.0, 2.0, 3.0];

        let result = gpu.tanh(&input).expect("GPU tanh should succeed");

        // Kill mutant: verify result is NOT all -1.0 values
        assert_ne!(
            result,
            vec![-1.0, -1.0, -1.0],
            "GPU tanh returned hardcoded -1.0 values (mutant not killed)"
        );

        // Verify correct computation
        for (i, &x) in input.iter().enumerate() {
            let expected = x.tanh();
            assert!(
                (result[i] - expected).abs() < 1e-4,
                "tanh({}) = {} (expected {})",
                x,
                result[i],
                expected
            );
        }
    }

    #[test]
    fn test_gpu_softmax_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let input = vec![1.0, 2.0, 3.0, 4.0];

        let result = gpu.softmax(&input);

        if let Ok(output) = result {
            // Softmax should sum to 1
            let sum: f32 = output.iter().sum();
            assert!(
                (sum - 1.0).abs() < 1e-3,
                "Softmax sum should be 1, got {}",
                sum
            );

            // All values should be positive
            for &v in &output {
                assert!(v > 0.0, "Softmax values should be positive");
            }

            // Later values should be larger (input is increasing)
            for i in 1..output.len() {
                assert!(
                    output[i] > output[i - 1],
                    "Softmax should preserve order for increasing input"
                );
            }
        } else {
            eprintln!("GPU softmax failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_log_softmax_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };
        let input = vec![1.0, 2.0, 3.0, 4.0];

        let result = gpu.log_softmax(&input);

        if let Ok(output) = result {
            // log_softmax values should all be negative (log of probability < 1)
            for &v in &output {
                assert!(v <= 0.0, "Log softmax values should be <= 0, got {}", v);
            }

            // exp(log_softmax) should sum to 1
            let exp_sum: f32 = output.iter().map(|x| x.exp()).sum();
            assert!(
                (exp_sum - 1.0).abs() < 1e-3,
                "exp(log_softmax) should sum to 1, got {}",
                exp_sum
            );
        } else {
            eprintln!("GPU log_softmax failed: {:?}", result);
        }
    }

    #[test]
    fn test_gpu_matmul_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // Simple 2x2 matrix multiplication
        // A = [[1, 2], [3, 4]]
        // B = [[5, 6], [7, 8]]
        // C = A * B = [[19, 22], [43, 50]]
        let a = vec![1.0, 2.0, 3.0, 4.0];
        let b = vec![5.0, 6.0, 7.0, 8.0];

        let res = gpu.matmul(&a, &b, 2, 2, 2);

        if let Ok(result) = res {
            assert!(
                (result[0] - 19.0).abs() < 1e-3,
                "Expected 19.0, got {}",
                result[0]
            );
            assert!(
                (result[1] - 22.0).abs() < 1e-3,
                "Expected 22.0, got {}",
                result[1]
            );
            assert!(
                (result[2] - 43.0).abs() < 1e-3,
                "Expected 43.0, got {}",
                result[2]
            );
            assert!(
                (result[3] - 50.0).abs() < 1e-3,
                "Expected 50.0, got {}",
                result[3]
            );
        } else {
            eprintln!("GPU matmul failed: {:?}", res);
        }
    }

    #[test]
    fn test_gpu_matmul_identity() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // Multiply by identity matrix
        // A = [[1, 2], [3, 4]]
        // I = [[1, 0], [0, 1]]
        // A * I = A
        let a = vec![1.0, 2.0, 3.0, 4.0];
        let identity = vec![1.0, 0.0, 0.0, 1.0];

        let res = gpu.matmul(&a, &identity, 2, 2, 2);

        if let Ok(result) = res {
            for i in 0..4 {
                assert!(
                    (result[i] - a[i]).abs() < 1e-3,
                    "Expected {}, got {}",
                    a[i],
                    result[i]
                );
            }
        } else {
            eprintln!("GPU matmul identity failed: {:?}", res);
        }
    }

    #[test]
    fn test_gpu_matmul_non_square() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // 2x3 matrix * 3x2 matrix = 2x2 matrix
        // A = [[1, 2, 3], [4, 5, 6]]
        // B = [[7, 8], [9, 10], [11, 12]]
        // C = [[58, 64], [139, 154]]
        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        let b = vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0];

        let res = gpu.matmul(&a, &b, 2, 3, 2);

        if let Ok(result) = res {
            assert!(
                (result[0] - 58.0).abs() < 1e-3,
                "Expected 58.0, got {}",
                result[0]
            );
            assert!(
                (result[1] - 64.0).abs() < 1e-3,
                "Expected 64.0, got {}",
                result[1]
            );
            assert!(
                (result[2] - 139.0).abs() < 1e-3,
                "Expected 139.0, got {}",
                result[2]
            );
            assert!(
                (result[3] - 154.0).abs() < 1e-3,
                "Expected 154.0, got {}",
                result[3]
            );
        } else {
            eprintln!("GPU matmul non-square failed: {:?}", res);
        }
    }

    #[test]
    fn test_gpu_convolve2d_basic() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // 3x3 input, 2x2 kernel -> 2x2 output
        let input = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
        let kernel = vec![1.0, 0.0, 0.0, 1.0];

        let res = gpu.convolve2d(&input, &kernel, 3, 3, 2, 2);

        if let Ok(result) = res {
            // For kernel [[1, 0], [0, 1]], each output is sum of diagonal elements
            // Output[0,0] = input[0,0]*1 + input[1,1]*1 = 1 + 5 = 6
            assert!(
                (result[0] - 6.0).abs() < 1e-3,
                "Expected 6.0, got {}",
                result[0]
            );
        } else {
            eprintln!("GPU convolve2d basic failed: {:?}", res);
        }
    }

    #[test]
    fn test_gpu_convolve2d_identity() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // 3x3 input with center-only kernel should extract center values
        let input = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
        // 3x3 kernel with center = 1, rest = 0
        let kernel = vec![0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0];

        let res = gpu.convolve2d(&input, &kernel, 3, 3, 3, 3);

        if let Ok(result) = res {
            // Should extract center value
            assert!(
                (result[0] - 5.0).abs() < 1e-3,
                "Expected 5.0, got {}",
                result[0]
            );
        } else {
            eprintln!("GPU convolve2d identity failed: {:?}", res);
        }
    }

    #[test]
    fn test_gpu_convolve2d_averaging() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // 4x4 input with 2x2 averaging kernel
        let input = vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
        ];
        // 2x2 averaging kernel
        let kernel = vec![0.25, 0.25, 0.25, 0.25];

        let res = gpu.convolve2d(&input, &kernel, 4, 4, 2, 2);

        if let Ok(result) = res {
            // First output: average of top-left 2x2 = (1+2+5+6)/4 = 3.5
            assert!(
                (result[0] - 3.5).abs() < 1e-3,
                "Expected 3.5, got {}",
                result[0]
            );
        } else {
            eprintln!("GPU convolve2d averaging failed: {:?}", res);
        }
    }

    #[test]
    fn test_gpu_tiled_sum_matches_cpu() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // Test various sizes
        let test_cases = [
            (4, 4),   // Small (single partial tile)
            (16, 16), // Exact tile
            (32, 32), // Multiple tiles (2x2)
            (20, 20), // Non-aligned
            (100, 5), // Wide
            (5, 100), // Tall
        ];

        for (width, height) in test_cases {
            let data: Vec<f32> = (1..=(width * height) as i32).map(|x| x as f32).collect();
            let cpu_result = tiled_sum_2d(&data, width, height);

            if let Ok(gpu_result) = gpu.tiled_sum_2d_gpu(&data, width, height) {
                let rel_err = (gpu_result - cpu_result).abs() / cpu_result.abs().max(1.0);
                assert!(
                    rel_err < 1e-4,
                    "GPU vs CPU tiled_sum mismatch for {}x{}: gpu={}, cpu={}, rel_err={}",
                    width,
                    height,
                    gpu_result,
                    cpu_result,
                    rel_err
                );
            } else {
                eprintln!("GPU tiled_sum_2d failed for {}x{}", width, height);
            }
        }
    }

    #[test]
    fn test_gpu_tiled_max_matches_cpu() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // Test data with varying values
        let data: Vec<f32> = (1..=256).map(|x| x as f32).collect();
        let width = 16;
        let height = 16;

        let cpu_result = tiled_max_2d(&data, width, height);

        if let Ok(gpu_result) = gpu.tiled_max_2d_gpu(&data, width, height) {
            assert!(
                (gpu_result - cpu_result).abs() < 1e-5,
                "GPU vs CPU tiled_max mismatch: gpu={}, cpu={}",
                gpu_result,
                cpu_result
            );
        } else {
            eprintln!("GPU tiled_max_2d failed");
        }
    }

    #[test]
    fn test_gpu_tiled_min_matches_cpu() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // Test data with varying values including negatives
        let data: Vec<f32> = (-128..=127).map(|x| x as f32).collect();
        let width = 16;
        let height = 16;

        let cpu_result = tiled_min_2d(&data, width, height);

        if let Ok(gpu_result) = gpu.tiled_min_2d_gpu(&data, width, height) {
            assert!(
                (gpu_result - cpu_result).abs() < 1e-5,
                "GPU vs CPU tiled_min mismatch: gpu={}, cpu={}",
                gpu_result,
                cpu_result
            );
        } else {
            eprintln!("GPU tiled_min_2d failed");
        }
    }

    #[test]
    fn test_gpu_tiled_sum_large_matrix() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // Large 64×64 matrix (16 tiles)
        let width = 64;
        let height = 64;
        let data: Vec<f32> = vec![1.0; width * height];

        let cpu_result = tiled_sum_2d(&data, width, height);
        let expected = (width * height) as f32;

        // Verify CPU is correct
        assert!((cpu_result - expected).abs() < 1e-3);

        if let Ok(gpu_result) = gpu.tiled_sum_2d_gpu(&data, width, height) {
            let rel_err = (gpu_result - expected).abs() / expected;
            assert!(
                rel_err < 1e-4,
                "GPU tiled_sum large matrix: got {}, expected {}, rel_err={}",
                gpu_result,
                expected,
                rel_err
            );
        } else {
            eprintln!("GPU tiled_sum_2d large matrix failed");
        }
    }

    #[test]
    fn test_gpu_tiled_max_with_negatives() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        // All negative values - max should be -1.0
        let data: Vec<f32> = (-100..-1).map(|x| x as f32).collect();
        let width = 9;
        let height = 11;

        let cpu_result = tiled_max_2d(&data, width, height);

        if let Ok(gpu_result) = gpu.tiled_max_2d_gpu(&data, width, height) {
            assert!(
                (gpu_result - cpu_result).abs() < 1e-5,
                "GPU vs CPU tiled_max (negatives): gpu={}, cpu={}",
                gpu_result,
                cpu_result
            );
        } else {
            eprintln!("GPU tiled_max_2d with negatives failed");
        }
    }

    #[test]
    fn test_gpu_tiled_min_single_element() {
        let Some(mut gpu) = get_shared_gpu() else {
            eprintln!("GPU not available, skipping test");
            return;
        };

        let data = vec![42.0];

        if let Ok(gpu_result) = gpu.tiled_min_2d_gpu(&data, 1, 1) {
            assert!(
                (gpu_result - 42.0).abs() < 1e-5,
                "GPU tiled_min single element: got {}, expected 42.0",
                gpu_result
            );
        } else {
            eprintln!("GPU tiled_min_2d single element failed");
        }
    }
}

Coverage Report

Created: 2026-01-25 15:05