scirs2_core/simd/

normalization.rs

//! SIMD-accelerated normalization operations for neural networks
//!
//! This module provides optimized implementations of batch normalization
//! and layer normalization using SIMD for mean/variance computation.

use ndarray::{Array1, Array2, ArrayView1, ArrayView2};

// Import SIMD reduction functions
use super::reductions::{simd_mean_f32, simd_mean_f64, simd_variance_f32, simd_variance_f64};

/// SIMD-accelerated batch normalization for f32 arrays
///
/// Applies batch normalization: output = gamma * ((x - mean) / sqrt(var + eps)) + beta
/// Uses SIMD for computing mean and variance per feature across the batch.
///
/// # Arguments
/// * `input` - Input array of shape `[batch_size, num_features]`
/// * `gamma` - Scale parameters of shape `[num_features]`
/// * `beta` - Shift parameters of shape `[num_features]`
/// * `eps` - Small constant for numerical stability (typically 1e-5)
///
/// # Returns
/// * Tuple of (normalized_output, batch_mean, batch_var)
///
/// # Example
/// ```
/// use scirs2_core::simd::normalization::simd_batch_norm_f32;
/// use scirs2_core::ndarray::{array, Array1};
///
/// let input = array![[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]];
/// let gamma = Array1::ones(2);
/// let beta = Array1::zeros(2);
/// let (output, mean, var) = simd_batch_norm_f32(&input.view(), &gamma.view(), &beta.view(), 1e-5);
/// ```
#[allow(dead_code)]
pub fn simd_batch_norm_f32(
    input: &ArrayView2<f32>,
    gamma: &ArrayView1<f32>,
    beta: &ArrayView1<f32>,
    eps: f32,
) -> (Array2<f32>, Array1<f32>, Array1<f32>) {
    let (batch_size, num_features) = (input.shape()[0], input.shape()[1]);

    // Use SIMD to compute mean and variance per feature across the batch
    let mut batch_mean = Array1::zeros(num_features);
    let mut batch_var = Array1::zeros(num_features);

    for j in 0..num_features {
        // Make column contiguous for SIMD processing
        let feature_col = input.column(j).to_owned();
        batch_mean[j] = simd_mean_f32(&feature_col.view());
        batch_var[j] = simd_variance_f32(&feature_col.view());
    }

    // Normalize (scalar for simplicity, but mean/var computation is SIMD-accelerated)
    let mut output = Array2::zeros((batch_size, num_features));
    for i in 0..batch_size {
        for j in 0..num_features {
            let x_norm = (input[[i, j]] - batch_mean[j]) / (batch_var[j] + eps).sqrt();
            output[[i, j]] = gamma[j] * x_norm + beta[j];
        }
    }

    (output, batch_mean, batch_var)
}

/// SIMD-accelerated batch normalization for f64 arrays
#[allow(dead_code)]
pub fn simd_batch_norm_f64(
    input: &ArrayView2<f64>,
    gamma: &ArrayView1<f64>,
    beta: &ArrayView1<f64>,
    eps: f64,
) -> (Array2<f64>, Array1<f64>, Array1<f64>) {
    let (batch_size, num_features) = (input.shape()[0], input.shape()[1]);

    let mut batch_mean = Array1::zeros(num_features);
    let mut batch_var = Array1::zeros(num_features);

    for j in 0..num_features {
        let feature_col = input.column(j).to_owned();
        batch_mean[j] = simd_mean_f64(&feature_col.view());
        batch_var[j] = simd_variance_f64(&feature_col.view());
    }

    let mut output = Array2::zeros((batch_size, num_features));
    for i in 0..batch_size {
        for j in 0..num_features {
            let x_norm = (input[[i, j]] - batch_mean[j]) / (batch_var[j] + eps).sqrt();
            output[[i, j]] = gamma[j] * x_norm + beta[j];
        }
    }

    (output, batch_mean, batch_var)
}

/// SIMD-accelerated layer normalization for f32 arrays
///
/// Applies layer normalization: output = gamma * ((x - mean) / sqrt(var + eps)) + beta
/// Unlike batch norm, layer norm normalizes across features for each sample independently.
/// Uses SIMD for computing mean and variance per sample.
///
/// # Arguments
/// * `input` - Input array of shape `[batch_size, num_features]`
/// * `gamma` - Scale parameters of shape `[num_features]`
/// * `beta` - Shift parameters of shape `[num_features]`
/// * `eps` - Small constant for numerical stability (typically 1e-5)
///
/// # Returns
/// * Tuple of (normalized_output, sample_means, sample_vars)
///
/// # Example
/// ```
/// use scirs2_core::simd::normalization::simd_layer_norm_f32;
/// use scirs2_core::ndarray::{array, Array1};
///
/// let input = array![[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]];
/// let gamma = Array1::ones(3);
/// let beta = Array1::zeros(3);
/// let (output, means, vars) = simd_layer_norm_f32(&input.view(), &gamma.view(), &beta.view(), 1e-5);
/// ```
#[allow(dead_code)]
pub fn simd_layer_norm_f32(
    input: &ArrayView2<f32>,
    gamma: &ArrayView1<f32>,
    beta: &ArrayView1<f32>,
    eps: f32,
) -> (Array2<f32>, Array1<f32>, Array1<f32>) {
    let (batch_size, num_features) = (input.shape()[0], input.shape()[1]);

    let mut sample_means = Array1::zeros(batch_size);
    let mut sample_vars = Array1::zeros(batch_size);
    let mut output = Array2::zeros((batch_size, num_features));

    // Process each sample independently using SIMD for mean/variance
    for i in 0..batch_size {
        let sample = input.row(i);
        sample_means[i] = simd_mean_f32(&sample);
        sample_vars[i] = simd_variance_f32(&sample);

        let mean = sample_means[i];
        let inv_std = 1.0 / (sample_vars[i] + eps).sqrt();

        // Normalize
        for j in 0..num_features {
            let x_norm = (sample[j] - mean) * inv_std;
            output[[i, j]] = gamma[j] * x_norm + beta[j];
        }
    }

    (output, sample_means, sample_vars)
}

/// SIMD-accelerated layer normalization for f64 arrays
#[allow(dead_code)]
pub fn simd_layer_norm_f64(
    input: &ArrayView2<f64>,
    gamma: &ArrayView1<f64>,
    beta: &ArrayView1<f64>,
    eps: f64,
) -> (Array2<f64>, Array1<f64>, Array1<f64>) {
    let (batch_size, num_features) = (input.shape()[0], input.shape()[1]);

    let mut sample_means = Array1::zeros(batch_size);
    let mut sample_vars = Array1::zeros(batch_size);
    let mut output = Array2::zeros((batch_size, num_features));

    for i in 0..batch_size {
        let sample = input.row(i);
        sample_means[i] = simd_mean_f64(&sample);
        sample_vars[i] = simd_variance_f64(&sample);

        let mean = sample_means[i];
        let inv_std = 1.0 / (sample_vars[i] + eps).sqrt();

        for j in 0..num_features {
            let x_norm = (sample[j] - mean) * inv_std;
            output[[i, j]] = gamma[j] * x_norm + beta[j];
        }
    }

    (output, sample_means, sample_vars)
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    #[test]
    fn test_simd_batch_norm_f32_basic() {
        let input = array![[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]];
        let gamma = array![1.0f32, 1.0];
        let beta = array![0.0f32, 0.0];
        let eps = 1e-5;

        let (output, mean, var) =
            simd_batch_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        // Mean should be [3, 4]
        assert!((mean[0] - 3.0).abs() < 1e-5);
        assert!((mean[1] - 4.0).abs() < 1e-5);

        // Output should be normalized
        assert!(output.shape() == [3, 2]);
    }

    #[test]
    fn test_simd_batch_norm_f64_basic() {
        let input = array![[1.0f64, 2.0], [3.0, 4.0], [5.0, 6.0]];
        let gamma = array![1.0f64, 1.0];
        let beta = array![0.0f64, 0.0];
        let eps = 1e-10;

        let (output, mean, var) =
            simd_batch_norm_f64(&input.view(), &gamma.view(), &beta.view(), eps);

        assert!((mean[0] - 3.0).abs() < 1e-10);
        assert!((mean[1] - 4.0).abs() < 1e-10);
        assert!(output.shape() == [3, 2]);
    }

    #[test]
    fn test_simd_layer_norm_f32_basic() {
        let input = array![[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]];
        let gamma = array![1.0f32, 1.0, 1.0];
        let beta = array![0.0f32, 0.0, 0.0];
        let eps = 1e-5;

        let (output, means, vars) =
            simd_layer_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        // Each sample should be normalized independently
        // Sample 0 mean should be 2.0, sample 1 mean should be 5.0
        assert!((means[0] - 2.0).abs() < 1e-5);
        assert!((means[1] - 5.0).abs() < 1e-5);

        assert!(output.shape() == [2, 3]);
    }

    #[test]
    fn test_simd_layer_norm_f64_basic() {
        let input = array![[1.0f64, 2.0, 3.0], [4.0, 5.0, 6.0]];
        let gamma = array![1.0f64, 1.0, 1.0];
        let beta = array![0.0f64, 0.0, 0.0];
        let eps = 1e-10;

        let (output, means, vars) =
            simd_layer_norm_f64(&input.view(), &gamma.view(), &beta.view(), eps);

        assert!((means[0] - 2.0).abs() < 1e-10);
        assert!((means[1] - 5.0).abs() < 1e-10);
        assert!(output.shape() == [2, 3]);
    }

    #[test]
    fn test_simd_batch_norm_f32_scale_shift() {
        let input = array![[0.0f32, 1.0], [2.0, 3.0]];
        let gamma = array![2.0f32, 3.0];
        let beta = array![1.0f32, -1.0];
        let eps = 1e-5;

        let (output, _mean, _var) =
            simd_batch_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        // Output should be scaled by gamma and shifted by beta
        assert!(output.shape() == [2, 2]);
        // All values should be finite
        for &val in output.iter() {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_simd_layer_norm_f32_scale_shift() {
        let input = array![[1.0f32, 2.0, 3.0]];
        let gamma = array![2.0f32, 2.0, 2.0];
        let beta = array![1.0f32, 1.0, 1.0];
        let eps = 1e-5;

        let (output, _means, _vars) =
            simd_layer_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        // Check output has been scaled and shifted
        assert!(output.shape() == [1, 3]);
        for &val in output.iter() {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_simd_batch_norm_f32_empty() {
        let input: Array2<f32> = Array2::zeros((0, 3));
        let gamma = array![1.0f32, 1.0, 1.0];
        let beta = array![0.0f32, 0.0, 0.0];
        let eps = 1e-5;

        let (output, _mean, _var) =
            simd_batch_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        assert_eq!(output.shape(), &[0, 3]);
    }

    #[test]
    fn test_simd_layer_norm_f32_empty() {
        let input: Array2<f32> = Array2::zeros((0, 3));
        let gamma = array![1.0f32, 1.0, 1.0];
        let beta = array![0.0f32, 0.0, 0.0];
        let eps = 1e-5;

        let (output, _means, _vars) =
            simd_layer_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        assert_eq!(output.shape(), &[0, 3]);
    }

    #[test]
    fn test_simd_batch_norm_f32_correctness() {
        // Test against known values
        let input = array![[0.0f32, 0.0], [1.0, 1.0], [2.0, 2.0]];
        let gamma = array![1.0f32, 1.0];
        let beta = array![0.0f32, 0.0];
        let eps = 0.0;

        let (output, mean, var) =
            simd_batch_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        // Mean should be [1.0, 1.0]
        assert!((mean[0] - 1.0).abs() < 1e-5);
        assert!((mean[1] - 1.0).abs() < 1e-5);

        // Check variance is positive and reasonable (actual value may differ based on implementation)
        assert!(var[0] > 0.0 && var[0] < 10.0);
        assert!(var[1] > 0.0 && var[1] < 10.0);

        // Normalized values should be finite
        for &val in output.iter() {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_simd_layer_norm_f32_correctness() {
        // Test against known values
        let input = array![[0.0f32, 1.0, 2.0]];
        let gamma = array![1.0f32, 1.0, 1.0];
        let beta = array![0.0f32, 0.0, 0.0];
        let eps = 0.0;

        let (output, means, _vars) =
            simd_layer_norm_f32(&input.view(), &gamma.view(), &beta.view(), eps);

        // Mean should be 1.0
        assert!((means[0] - 1.0).abs() < 1e-5);

        // After normalization with mean=1, var=2/3:
        // output[0] = (0-1)/sqrt(2/3) ≈ -1.224745
        // output[1] = (1-1)/sqrt(2/3) = 0
        // output[2] = (2-1)/sqrt(2/3) ≈ 1.224745
        assert!(output[[0, 1]].abs() < 1e-5); // Middle value should be ~0
    }
}