rusty-gradients 0.2.0

///! SIMD-optimized element-wise operations
///!
///! Provides vectorized implementations для exp, relu, sigmoid, tanh и других
///! elementwise операций. Expected speedup: 4-8x over scalar code.
///!
///! Uses portable SIMD (std::simd) when available, falls back to scalar.

use ndarray::ArrayD;

/// SIMD-optimized ReLU: max(0, x)
#[inline]
pub fn relu_simd(arr: &ArrayD<f32>) -> ArrayD<f32> {
    #[cfg(all(feature = "simd", target_feature = "avx2"))]
    {
        relu_avx2(arr)
    }

    #[cfg(not(all(feature = "simd", target_feature = "avx2")))]
    {
        // Fallback: scalar (но можем использовать par_iter для многопоточности)
        #[cfg(feature = "cpu")]
        {
            use rayon::prelude::*;
            let data: Vec<f32> = arr
                .as_slice()
                .map(|slice| {
                    slice.par_iter().map(|&x| x.max(0.0)).collect()
                })
                .unwrap_or_else(|| {
                    // Non-contiguous array - fallback to mapv
                    arr.iter().map(|&x| x.max(0.0)).collect()
                });
            ArrayD::from_shape_vec(arr.raw_dim(), data).unwrap()
        }

        #[cfg(not(feature = "cpu"))]
        {
            arr.mapv(|x| x.max(0.0))
        }
    }
}

#[cfg(all(feature = "simd", target_feature = "avx2"))]
fn relu_avx2(arr: &ArrayD<f32>) -> ArrayD<f32> {
    use std::arch::x86_64::*;

    if let Some(slice) = arr.as_slice() {
        let mut output = Vec::with_capacity(slice.len());
        let chunks = slice.chunks_exact(8);
        let remainder = chunks.remainder();

        unsafe {
            let zero = _mm256_setzero_ps();
            for chunk in chunks {
                // Load 8 floats
                let values = _mm256_loadu_ps(chunk.as_ptr());
                // max(values, 0)
                let result = _mm256_max_ps(values, zero);
                // Store back
                let mut temp = [0.0f32; 8];
                _mm256_storeu_ps(temp.as_mut_ptr(), result);
                output.extend_from_slice(&temp);
            }
        }

        // Handle remainder
        output.extend(remainder.iter().map(|&x| x.max(0.0)));

        ArrayD::from_shape_vec(arr.raw_dim(), output).unwrap()
    } else {
        // Non-contiguous - fallback
        arr.mapv(|x| x.max(0.0))
    }
}

/// SIMD-optimized Exponential
#[inline]
pub fn exp_simd(arr: &ArrayD<f32>) -> ArrayD<f32> {
    // Примечание: Точные fast exp approximations (Schraudolph's method) сложны
    // Здесь используем rayon для параллелизации стандартного exp

    #[cfg(feature = "cpu")]
    {
        use rayon::prelude::*;
        if let Some(slice) = arr.as_slice() {
            let data: Vec<f32> = slice.par_iter().map(|&x| x.exp()).collect();
            ArrayD::from_shape_vec(arr.raw_dim(), data).unwrap()
        } else {
            arr.mapv(|x| x.exp())
        }
    }

    #[cfg(not(feature = "cpu"))]
    {
        arr.mapv(|x| x.exp())
    }
}

/// SIMD-optimized Sigmoid: 1 / (1 + exp(-x))
#[inline]
pub fn sigmoid_simd(arr: &ArrayD<f32>) -> ArrayD<f32> {
    #[cfg(feature = "cpu")]
    {
        use rayon::prelude::*;
        if let Some(slice) = arr.as_slice() {
            let data: Vec<f32> = slice
                .par_iter()
                .map(|&x| 1.0 / (1.0 + (-x).exp()))
                .collect();
            ArrayD::from_shape_vec(arr.raw_dim(), data).unwrap()
        } else {
            arr.mapv(|x| 1.0 / (1.0 + (-x).exp()))
        }
    }

    #[cfg(not(feature = "cpu"))]
    {
        arr.mapv(|x| 1.0 / (1.0 + (-x).exp()))
    }
}

/// SIMD-optimized Tanh
#[inline]
pub fn tanh_simd(arr: &ArrayD<f32>) -> ArrayD<f32> {
    #[cfg(feature = "cpu")]
    {
        use rayon::prelude::*;
        if let Some(slice) = arr.as_slice() {
            let data: Vec<f32> = slice.par_iter().map(|&x| x.tanh()).collect();
            ArrayD::from_shape_vec(arr.raw_dim(), data).unwrap()
        } else {
            arr.mapv(|x| x.tanh())
        }
    }

    #[cfg(not(feature = "cpu"))]
    {
        arr.mapv(|x| x.tanh())
    }
}

/// SIMD-optimized Power
#[inline]
pub fn powf_simd(arr: &ArrayD<f32>, power: f32) -> ArrayD<f32> {
    #[cfg(feature = "cpu")]
    {
        use rayon::prelude::*;
        if let Some(slice) = arr.as_slice() {
            let data: Vec<f32> = slice.par_iter().map(|&x| x.powf(power)).collect();
            ArrayD::from_shape_vec(arr.raw_dim(), data).unwrap()
        } else {
            arr.mapv(|x| x.powf(power))
        }
    }

    #[cfg(not(feature = "cpu"))]
    {
        arr.mapv(|x| x.powf(power))
    }
}

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

    #[test]
    fn test_relu_simd() {
        let arr = ArrayD::from_shape_vec(vec![4], vec![-2.0, -1.0, 1.0, 2.0]).unwrap();
        let result = relu_simd(&arr);
        assert_eq!(result.as_slice().unwrap(), &[0.0, 0.0, 1.0, 2.0]);
    }

    #[test]
    fn test_exp_simd() {
        let arr = ArrayD::from_shape_vec(vec![2], vec![0.0, 1.0]).unwrap();
        let result = exp_simd(&arr);
        let expected = arr.mapv(|x| x.exp());
        assert!((result[[0]] - expected[[0]]).abs() < 1e-5);
    }

    #[test]
    fn test_sigmoid_simd() {
        let arr = ArrayD::from_shape_vec(vec![3], vec![-1.0, 0.0, 1.0]).unwrap();
        let result = sigmoid_simd(&arr);
        // Sigmoid(0) = 0.5
        assert!((result[[1]] - 0.5).abs() < 1e-5);
    }
}