foxhunt/crates/ml/src/cuda_compat.rs

//! CUDA-compatible operations
//!
//! Manual implementations of operations missing in Candle CUDA kernels.
//! This module provides workarounds for operations that have CPU implementations
//! but lack CUDA kernel support.

use crate::MLError;
use candle_core::Tensor;

/// Manual sigmoid implementation for CUDA compatibility
///
/// Candle version `671de1db` lacks CUDA sigmoid kernel.
/// This function provides a manual implementation using CUDA-supported operations:
/// sigmoid(x) = 1 / (1 + exp(-x))
///
/// # Arguments
/// * `x` - Input tensor
///
/// # Returns
/// Tensor with sigmoid applied element-wise
///
/// # Example
/// ```ignore
/// let input = Tensor::new(&[1.0, 2.0, 3.0], &device)?;
/// let output = manual_sigmoid(&input)?;
/// ```
pub fn manual_sigmoid(x: &Tensor) -> Result<Tensor, MLError> {
    // sigmoid(x) = 1 / (1 + exp(-x))
    let neg_x = x.neg()?;
    let exp_neg_x = neg_x.exp()?;
    let one = Tensor::ones_like(&exp_neg_x)?;
    let denominator = (&one + &exp_neg_x)?;
    one.div(&denominator)
        .map_err(|e| MLError::ModelError(format!("Sigmoid computation failed: {}", e)))
}

/// Alternative sigmoid using tanh (for reference)
///
/// sigmoid(x) = 0.5 * (tanh(x/2) + 1)
///
/// This is an alternative implementation that may be useful if tanh has better
/// CUDA support in some Candle versions.
pub fn sigmoid_via_tanh(x: &Tensor) -> Result<Tensor, MLError> {
    let two = Tensor::new(&[2.0_f32], x.device())?;
    let half_x = x.broadcast_div(&two)?;
    let tanh_half = half_x.tanh()?;
    let one = Tensor::ones_like(&tanh_half)?;
    let numerator = (&tanh_half + &one)?;
    let half = Tensor::new(&[0.5_f32], x.device())?;
    numerator
        .broadcast_mul(&half)
        .map_err(|e| MLError::ModelError(format!("Sigmoid (via tanh) computation failed: {}", e)))
}

/// CUDA-compatible layer normalization
///
/// Candle version `671de1db` lacks CUDA layer normalization kernel.
/// This function provides a manual implementation using CUDA-supported operations:
/// LayerNorm(x) = γ * (x - μ) / sqrt(σ² + ε) + β
///
/// Where:
/// - μ = mean(x) across normalized dimensions
/// - σ² = variance(x) across normalized dimensions
/// - γ = learnable scale parameter (weight)
/// - β = learnable shift parameter (bias)
/// - ε = small constant for numerical stability
///
/// # Arguments
/// * `x` - Input tensor
/// * `normalized_shape` - Shape to normalize over (typically last dimension)
/// * `weight` - Optional learnable scale parameter
/// * `bias` - Optional learnable shift parameter
/// * `eps` - Small constant for numerical stability (typically 1e-5)
///
/// # Returns
/// Normalized tensor with same shape as input
///
/// # Example
/// ```ignore
/// let input = Tensor::new(&[[1.0, 2.0], [3.0, 4.0]], &device)?;
/// let weight = Tensor::ones(2, DType::F32, &device)?;
/// let bias = Tensor::zeros(2, DType::F32, &device)?;
/// let output = cuda_layer_norm(&input, &[2], Some(&weight), Some(&bias), 1e-5)?;
/// ```
pub fn cuda_layer_norm(
    x: &Tensor,
    normalized_shape: &[usize],
    weight: Option<&Tensor>,
    bias: Option<&Tensor>,
    eps: f64,
) -> Result<Tensor, MLError> {
    // Get the dimensions to normalize over
    let rank = x.dims().len();
    let norm_dims_count = normalized_shape.len();

    // Calculate dims to reduce over (last norm_dims_count dimensions)
    let dims_to_reduce: Vec<usize> = (rank - norm_dims_count..rank).collect();

    // BF16/F16: cast to F32 for layer norm precision (numerical stability for
    // mean/variance), then cast result back to original dtype.
    let original_dtype = x.dtype();
    let (x_compute, compute_dtype) = match original_dtype {
        candle_core::DType::BF16 | candle_core::DType::F16 => {
            (x.to_dtype(candle_core::DType::F32)?, candle_core::DType::F32)
        },
        candle_core::DType::F32 => (x.clone(), candle_core::DType::F32),
        candle_core::DType::F64 => (x.clone(), candle_core::DType::F64),
        candle_core::DType::U8 | candle_core::DType::U32 | candle_core::DType::I64
        | candle_core::DType::F8E4M3 => {
            return Err(MLError::ModelError(format!(
                "Unsupported dtype for layer norm: {:?}",
                original_dtype
            )))
        },
    };

    // Recalculate mean and centered on the (possibly cast) compute tensor
    let mean_compute = x_compute.mean_keepdim(dims_to_reduce.as_slice())?;
    let centered_compute = x_compute.broadcast_sub(&mean_compute)?;
    let variance_compute = centered_compute
        .sqr()?
        .mean_keepdim(dims_to_reduce.as_slice())?;

    // CRITICAL: Use compute_dtype to match the computation tensor's dtype
    let eps_tensor = match compute_dtype {
        candle_core::DType::F32 => Tensor::new(&[eps as f32], x.device())?,
        candle_core::DType::F64 => Tensor::new(&[eps], x.device())?,
        candle_core::DType::BF16 | candle_core::DType::F16
        | candle_core::DType::U8 | candle_core::DType::U32 | candle_core::DType::I64
        | candle_core::DType::F8E4M3 => {
            return Err(MLError::ModelError(format!(
                "Unsupported compute dtype for layer norm eps: {:?}",
                compute_dtype
            )))
        },
    };
    let variance_eps = variance_compute.broadcast_add(&eps_tensor)?;

    // Calculate standard deviation: sqrt(σ² + ε)
    let std = variance_eps.sqrt()?;

    // Normalize: (x - μ) / sqrt(σ² + ε)
    let normalized = centered_compute.broadcast_div(&std)?;

    // Apply scale (γ) if provided
    let scaled = if let Some(w) = weight {
        // Reshape weight to broadcast correctly
        let mut weight_shape = vec![1; rank];
        for (i, &dim) in normalized_shape.iter().enumerate() {
            weight_shape[rank - norm_dims_count + i] = dim;
        }
        let weight_reshaped = w.reshape(weight_shape)?;

        // FIX: Convert weight dtype AND device to match computation dtype
        let weight_converted = if weight_reshaped.dtype() != compute_dtype
            || !weight_reshaped.device().same_device(x.device())
        {
            let w_dtype = if weight_reshaped.dtype() != compute_dtype {
                weight_reshaped.to_dtype(compute_dtype)?
            } else {
                weight_reshaped
            };
            if !w_dtype.device().same_device(x.device()) {
                w_dtype.to_device(x.device())?
            } else {
                w_dtype
            }
        } else {
            weight_reshaped
        };

        normalized.broadcast_mul(&weight_converted)?
    } else {
        normalized
    };

    // Apply shift (β) if provided
    let shifted = if let Some(b) = bias {
        // Reshape bias to broadcast correctly
        let mut bias_shape = vec![1; rank];
        for (i, &dim) in normalized_shape.iter().enumerate() {
            bias_shape[rank - norm_dims_count + i] = dim;
        }
        let bias_reshaped = b.reshape(bias_shape)?;

        // FIX: Convert bias dtype AND device to match computation dtype
        let bias_converted = if bias_reshaped.dtype() != compute_dtype
            || !bias_reshaped.device().same_device(x.device())
        {
            let b_dtype = if bias_reshaped.dtype() != compute_dtype {
                bias_reshaped.to_dtype(compute_dtype)?
            } else {
                bias_reshaped
            };
            if !b_dtype.device().same_device(x.device()) {
                b_dtype.to_device(x.device())?
            } else {
                b_dtype
            }
        } else {
            bias_reshaped
        };

        scaled.broadcast_add(&bias_converted)?
    } else {
        scaled
    };

    // Cast back to original dtype if we promoted from BF16/F16
    let result = if shifted.dtype() != original_dtype {
        shifted.to_dtype(original_dtype)?
    } else {
        shifted
    };

    Ok(result)
}

/// Wrapper for layer normalization that automatically falls back to CPU if CUDA fails
///
/// This function attempts to use the native candle layer_norm implementation.
/// If it fails on CUDA (due to missing CUDA kernel), it falls back to our
/// manual CUDA-compatible implementation.
///
/// # Arguments
/// * `x` - Input tensor
/// * `normalized_shape` - Shape to normalize over
/// * `weight` - Optional learnable scale parameter
/// * `bias` - Optional learnable shift parameter
/// * `eps` - Small constant for numerical stability
///
/// # Returns
/// Normalized tensor
pub fn layer_norm_with_fallback(
    x: &Tensor,
    normalized_shape: &[usize],
    weight: Option<&Tensor>,
    bias: Option<&Tensor>,
    eps: f64,
) -> Result<Tensor, MLError> {
    // Always use manual implementation for CUDA devices
    // This avoids the "no cuda implementation for layer-norm" error
    if x.device().is_cuda() {
        return cuda_layer_norm(x, normalized_shape, weight, bias, eps);
    }

    // BUG FIX #18 (Agent 231): candle_nn::ops::layer_norm only supports F32
    // For F64 tensors, use manual implementation instead of native layer_norm
    if x.dtype() == candle_core::DType::F64 {
        return cuda_layer_norm(x, normalized_shape, weight, bias, eps);
    }

    // Use native implementation for CPU with F32
    // candle_nn::ops::layer_norm requires weight and bias (not optional)
    match (weight, bias) {
        (Some(w), Some(b)) => candle_nn::ops::layer_norm(x, w, b, eps as f32)
            .map_err(|e| MLError::ModelError(format!("Layer normalization failed: {}", e))),
        _ => {
            // If weight/bias not provided, use manual implementation
            cuda_layer_norm(x, normalized_shape, weight, bias, eps)
        },
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::{DType, Device};

    #[test]
    fn test_manual_sigmoid_cpu() -> Result<(), MLError> {
        let device = Device::Cpu;

        // Test sigmoid(0) = 0.5
        let zero = Tensor::new(&[0.0_f32], &device)?;
        let result = manual_sigmoid(&zero)?;
        let result_vec = result.to_vec1::<f32>()?;
        assert!((result_vec[0] - 0.5).abs() < 1e-6);

        // Test sigmoid(large positive) ≈ 1
        let large_pos = Tensor::new(&[10.0_f32], &device)?;
        let result = manual_sigmoid(&large_pos)?;
        let result_vec = result.to_vec1::<f32>()?;
        assert!(result_vec[0] > 0.9999);

        // Test sigmoid(large negative) ≈ 0
        let large_neg = Tensor::new(&[-10.0_f32], &device)?;
        let result = manual_sigmoid(&large_neg)?;
        let result_vec = result.to_vec1::<f32>()?;
        assert!(result_vec[0] < 0.0001);

        Ok(())
    }

    #[test]
    fn test_manual_sigmoid_batch() -> Result<(), MLError> {
        let device = Device::Cpu;

        let input = Tensor::new(&[-2.0_f32, -1.0, 0.0, 1.0, 2.0], &device)?;
        let result = manual_sigmoid(&input)?;
        let result_vec = result.to_vec1::<f32>()?;

        // Check all values are in (0, 1)
        for &val in &result_vec {
            assert!(val > 0.0 && val < 1.0);
        }

        // Check sigmoid(0) = 0.5
        assert!((result_vec[2] - 0.5).abs() < 1e-6);

        // Check symmetry: sigmoid(-x) + sigmoid(x) = 1
        assert!((result_vec[0] + result_vec[4] - 1.0).abs() < 1e-6);
        assert!((result_vec[1] + result_vec[3] - 1.0).abs() < 1e-6);

        Ok(())
    }

    #[test]
    #[cfg(feature = "cuda")]
    #[ignore = "Only run when GPU available"]
    fn test_manual_sigmoid_cuda() -> Result<(), MLError> {
        use candle_core::Device;

        let device = Device::cuda_if_available(0)?;
        if !device.is_cuda() {
            println!("Skipping CUDA test - GPU not available");
            return Ok(());
        }

        let input = Tensor::new(&[-2.0_f32, -1.0, 0.0, 1.0, 2.0], &device)?;
        let result = manual_sigmoid(&input)?;

        // Move back to CPU for validation
        let result_cpu = result.to_device(&Device::Cpu)?;
        let result_vec = result_cpu.to_vec1::<f32>()?;

        // Check all values are in (0, 1)
        for &val in &result_vec {
            assert!(val > 0.0 && val < 1.0);
        }

        // Check sigmoid(0) = 0.5
        assert!((result_vec[2] - 0.5).abs() < 1e-6);

        Ok(())
    }

    #[test]
    fn test_cuda_layer_norm_cpu() -> Result<(), MLError> {
        let device = Device::Cpu;

        // Test simple 2D tensor [batch_size=2, features=4]
        let input = Tensor::new(&[[1.0_f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device)?;

        let weight = Tensor::ones(4, DType::F32, &device)?;
        let bias = Tensor::zeros(4, DType::F32, &device)?;

        let output = cuda_layer_norm(&input, &[4], Some(&weight), Some(&bias), 1e-5)?;

        // Check output shape
        assert_eq!(output.dims(), &[2, 4]);

        // Verify normalization (mean ≈ 0, std ≈ 1)
        let output_vec = output.to_vec2::<f32>()?;
        for row in &output_vec {
            let mean: f32 = row.iter().sum::<f32>() / row.len() as f32;
            let variance: f32 =
                row.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / row.len() as f32;
            let std = variance.sqrt();

            assert!(mean.abs() < 1e-5, "Mean should be close to 0, got {}", mean);
            assert!(
                (std - 1.0).abs() < 1e-3,
                "Std should be close to 1, got {}",
                std
            );
        }

        Ok(())
    }

    #[test]
    fn test_cuda_layer_norm_3d() -> Result<(), MLError> {
        let device = Device::Cpu;

        // Test 3D tensor [batch_size=2, seq_len=3, features=4]
        let input_data = vec![
            1.0_f32, 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, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0,
        ];
        let input = Tensor::from_slice(&input_data, (2, 3, 4), &device)?;

        let weight = Tensor::ones(4, DType::F32, &device)?;
        let bias = Tensor::zeros(4, DType::F32, &device)?;

        let output = cuda_layer_norm(&input, &[4], Some(&weight), Some(&bias), 1e-5)?;

        // Check output shape matches input
        assert_eq!(output.dims(), &[2, 3, 4]);

        Ok(())
    }

    #[test]
    fn test_layer_norm_with_fallback_cpu() -> Result<(), MLError> {
        let device = Device::Cpu;

        let input = Tensor::new(&[[1.0_f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device)?;

        let weight = Tensor::ones(4, DType::F32, &device)?;
        let bias = Tensor::zeros(4, DType::F32, &device)?;

        // Test fallback function
        let output = layer_norm_with_fallback(&input, &[4], Some(&weight), Some(&bias), 1e-5)?;

        // Check output shape
        assert_eq!(output.dims(), &[2, 4]);

        Ok(())
    }

    #[test]
    fn test_cuda_layer_norm_without_affine() -> Result<(), MLError> {
        let device = Device::Cpu;

        let input = Tensor::new(&[[1.0_f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device)?;

        // Test without weight and bias
        let output = cuda_layer_norm(&input, &[4], None, None, 1e-5)?;

        // Check output shape
        assert_eq!(output.dims(), &[2, 4]);

        // Verify normalization
        let output_vec = output.to_vec2::<f32>()?;
        for row in &output_vec {
            let mean: f32 = row.iter().sum::<f32>() / row.len() as f32;
            assert!(mean.abs() < 1e-5, "Mean should be close to 0, got {}", mean);
        }

        Ok(())
    }

    #[test]
    #[cfg(feature = "cuda")]
    #[ignore = "Only run when GPU available"]
    fn test_cuda_layer_norm_gpu() -> Result<(), MLError> {
        let device = Device::cuda_if_available(0)?;
        if !device.is_cuda() {
            println!("Skipping CUDA test - GPU not available");
            return Ok(());
        }

        let input = Tensor::new(&[[1.0_f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device)?;

        let weight = Tensor::ones(4, DType::F32, &device)?;
        let bias = Tensor::zeros(4, DType::F32, &device)?;

        // Test CUDA implementation directly
        let output = cuda_layer_norm(&input, &[4], Some(&weight), Some(&bias), 1e-5)?;

        // Move to CPU for validation
        let output_cpu = output.to_device(&Device::Cpu)?;
        let output_vec = output_cpu.to_vec2::<f32>()?;

        // Verify normalization
        for row in &output_vec {
            let mean: f32 = row.iter().sum::<f32>() / row.len() as f32;
            let variance: f32 =
                row.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / row.len() as f32;
            let std = variance.sqrt();

            assert!(mean.abs() < 1e-4, "Mean should be close to 0, got {}", mean);
            assert!(
                (std - 1.0).abs() < 1e-2,
                "Std should be close to 1, got {}",
                std
            );
        }

        Ok(())
    }

    #[test]
    #[cfg(feature = "cuda")]
    #[ignore = "Only run when GPU available"]
    fn test_layer_norm_fallback_gpu() -> Result<(), MLError> {
        let device = Device::cuda_if_available(0)?;
        if !device.is_cuda() {
            println!("Skipping CUDA test - GPU not available");
            return Ok(());
        }

        let input = Tensor::new(&[[1.0_f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device)?;

        let weight = Tensor::ones(4, DType::F32, &device)?;
        let bias = Tensor::zeros(4, DType::F32, &device)?;

        // Test fallback wrapper on GPU
        let output = layer_norm_with_fallback(&input, &[4], Some(&weight), Some(&bias), 1e-5)?;

        // Check output shape
        assert_eq!(output.dims(), &[2, 4]);

        // Move to CPU for validation
        let output_cpu = output.to_device(&Device::Cpu)?;
        assert_eq!(output_cpu.dims(), &[2, 4]);

        Ok(())
    }
}