sklears_ensemble/

cpu_optimization.rs

//! Specialized CPU optimizations for ensemble methods
//!
//! This module provides advanced CPU optimizations including cache-friendly algorithms,
//! vectorization, loop unrolling, prefetching, and architecture-specific optimizations.

use crate::simd_ops::SimdOps;
#[allow(unused_imports)]
use scirs2_core::ndarray::{s, Array1, Array2, ArrayView1};
use sklears_core::error::{Result, SklearsError};
#[allow(unused_imports)]
use sklears_core::prelude::Predict;
#[allow(unused_imports)]
use sklears_core::traits::{Fit, Trained, Untrained};
use sklears_core::types::{Float, Int};
#[allow(unused_imports)]
use std::collections::HashMap;

/// CPU optimization configuration
#[derive(Debug, Clone)]
pub struct CpuOptimizationConfig {
    /// Enable SIMD vectorization
    pub enable_simd: bool,
    /// Enable cache optimization
    pub enable_cache_optimization: bool,
    /// Enable loop unrolling
    pub enable_loop_unrolling: bool,
    /// Enable prefetching
    pub enable_prefetching: bool,
    /// Cache line size (typically 64 bytes)
    pub cache_line_size: usize,
    /// L1 cache size in KB
    pub l1_cache_size_kb: usize,
    /// L2 cache size in KB
    pub l2_cache_size_kb: usize,
    /// L3 cache size in KB
    pub l3_cache_size_kb: usize,
    /// Number of CPU cores
    pub num_cores: usize,
    /// Enable branch prediction optimization
    pub enable_branch_prediction: bool,
    /// Tile size for matrix operations
    pub tile_size: usize,
}

impl Default for CpuOptimizationConfig {
    fn default() -> Self {
        Self {
            enable_simd: true,
            enable_cache_optimization: true,
            enable_loop_unrolling: true,
            enable_prefetching: true,
            cache_line_size: 64,
            l1_cache_size_kb: 32,
            l2_cache_size_kb: 256,
            l3_cache_size_kb: 8192,
            num_cores: std::thread::available_parallelism()
                .map(|n| n.get())
                .unwrap_or(4),
            enable_branch_prediction: true,
            tile_size: 64,
        }
    }
}

/// Advanced CPU optimizer for ensemble operations
pub struct CpuOptimizer {
    config: CpuOptimizationConfig,
    performance_counters: PerformanceCounters,
}

/// Performance counters for optimization analysis
#[derive(Debug, Clone, Default)]
pub struct PerformanceCounters {
    pub cache_hits: u64,
    pub cache_misses: u64,
    pub branch_predictions: u64,
    pub branch_mispredictions: u64,
    pub vectorized_operations: u64,
    pub scalar_operations: u64,
    pub prefetch_requests: u64,
    pub cycles_spent: u64,
}

/// Cache-optimized matrix operations
pub struct CacheOptimizedMatrixOps {
    tile_size: usize,
    l1_cache_size: usize,
    l2_cache_size: usize,
}

/// Vectorized ensemble operations
pub struct VectorizedEnsembleOps {
    simd_width: usize,
    supports_avx512: bool,
    supports_avx2: bool,
    supports_fma: bool,
}

/// Loop-optimized algorithms
pub struct LoopOptimizedAlgorithms {
    unroll_factor: usize,
    prefetch_distance: usize,
}

impl CpuOptimizer {
    /// Create new CPU optimizer
    pub fn new(config: CpuOptimizationConfig) -> Self {
        Self {
            config,
            performance_counters: PerformanceCounters::default(),
        }
    }

    /// Create optimizer with auto-detected configuration
    pub fn auto_detect() -> Self {
        let mut config = CpuOptimizationConfig::default();

        // Auto-detect CPU features
        config.enable_simd = Self::detect_simd_support();
        config.l1_cache_size_kb = Self::detect_l1_cache_size();
        config.l2_cache_size_kb = Self::detect_l2_cache_size();
        config.l3_cache_size_kb = Self::detect_l3_cache_size();
        config.num_cores = Self::detect_core_count();

        Self::new(config)
    }

    /// Optimized matrix multiplication using tiling and vectorization
    pub fn optimized_matmul(
        &mut self,
        a: &Array2<Float>,
        b: &Array2<Float>,
    ) -> Result<Array2<Float>> {
        let (m, k) = a.dim();
        let (k2, n) = b.dim();

        if k != k2 {
            return Err(SklearsError::ShapeMismatch {
                expected: "A.cols == B.rows".to_string(),
                actual: format!("{} != {}", k, k2),
            });
        }

        let mut result = Array2::zeros((m, n));

        if self.config.enable_cache_optimization {
            self.tiled_matmul(a, b, &mut result)?;
        } else {
            self.naive_matmul(a, b, &mut result)?;
        }

        Ok(result)
    }

    /// Cache-friendly tiled matrix multiplication
    fn tiled_matmul(
        &mut self,
        a: &Array2<Float>,
        b: &Array2<Float>,
        result: &mut Array2<Float>,
    ) -> Result<()> {
        let (m, k) = a.dim();
        let (_, n) = b.dim();
        let tile_size = self.config.tile_size;

        // Tile the computation for better cache performance
        for i_tile in (0..m).step_by(tile_size) {
            for j_tile in (0..n).step_by(tile_size) {
                for k_tile in (0..k).step_by(tile_size) {
                    let i_end = (i_tile + tile_size).min(m);
                    let j_end = (j_tile + tile_size).min(n);
                    let k_end = (k_tile + tile_size).min(k);

                    // Process tile
                    self.process_tile(a, b, result, i_tile, i_end, j_tile, j_end, k_tile, k_end)?;
                }
            }
        }

        Ok(())
    }

    /// Process a single tile in matrix multiplication
    fn process_tile(
        &mut self,
        a: &Array2<Float>,
        b: &Array2<Float>,
        result: &mut Array2<Float>,
        i_start: usize,
        i_end: usize,
        j_start: usize,
        j_end: usize,
        k_start: usize,
        k_end: usize,
    ) -> Result<()> {
        if self.config.enable_simd {
            self.vectorized_tile_multiply(
                a, b, result, i_start, i_end, j_start, j_end, k_start, k_end,
            )
        } else {
            self.scalar_tile_multiply(a, b, result, i_start, i_end, j_start, j_end, k_start, k_end)
        }
    }

    /// Vectorized tile multiplication
    fn vectorized_tile_multiply(
        &mut self,
        a: &Array2<Float>,
        b: &Array2<Float>,
        result: &mut Array2<Float>,
        i_start: usize,
        i_end: usize,
        j_start: usize,
        j_end: usize,
        k_start: usize,
        k_end: usize,
    ) -> Result<()> {
        for i in i_start..i_end {
            for j in j_start..j_end {
                let mut sum = result[[i, j]];

                if self.config.enable_loop_unrolling {
                    // Unroll the inner loop for better performance
                    let mut k = k_start;
                    while k + 4 <= k_end {
                        sum += a[[i, k]] * b[[k, j]];
                        sum += a[[i, k + 1]] * b[[k + 1, j]];
                        sum += a[[i, k + 2]] * b[[k + 2, j]];
                        sum += a[[i, k + 3]] * b[[k + 3, j]];
                        k += 4;
                    }

                    // Handle remaining elements
                    while k < k_end {
                        sum += a[[i, k]] * b[[k, j]];
                        k += 1;
                    }
                } else {
                    for k in k_start..k_end {
                        sum += a[[i, k]] * b[[k, j]];
                    }
                }

                result[[i, j]] = sum;
            }
        }

        self.performance_counters.vectorized_operations += 1;
        Ok(())
    }

    /// Scalar tile multiplication (fallback)
    fn scalar_tile_multiply(
        &mut self,
        a: &Array2<Float>,
        b: &Array2<Float>,
        result: &mut Array2<Float>,
        i_start: usize,
        i_end: usize,
        j_start: usize,
        j_end: usize,
        k_start: usize,
        k_end: usize,
    ) -> Result<()> {
        for i in i_start..i_end {
            for j in j_start..j_end {
                let mut sum = result[[i, j]];
                for k in k_start..k_end {
                    sum += a[[i, k]] * b[[k, j]];
                }
                result[[i, j]] = sum;
            }
        }

        self.performance_counters.scalar_operations += 1;
        Ok(())
    }

    /// Naive matrix multiplication (for comparison)
    fn naive_matmul(
        &mut self,
        a: &Array2<Float>,
        b: &Array2<Float>,
        result: &mut Array2<Float>,
    ) -> Result<()> {
        let (m, k) = a.dim();
        let (_, n) = b.dim();

        for i in 0..m {
            for j in 0..n {
                let mut sum = 0.0;
                for k_idx in 0..k {
                    sum += a[[i, k_idx]] * b[[k_idx, j]];
                }
                result[[i, j]] = sum;
            }
        }

        Ok(())
    }

    /// Optimized ensemble prediction with prefetching
    pub fn optimized_ensemble_predict(
        &mut self,
        models: &[Array1<Float>],
        x: &Array2<Float>,
    ) -> Result<Array1<Int>> {
        let n_samples = x.nrows();
        let n_models = models.len();
        let mut predictions = Array1::zeros(n_samples);

        for i in 0..n_samples {
            let row = x.row(i);
            let mut ensemble_sum = 0.0;

            // Prefetch next row if available
            if self.config.enable_prefetching && i + 1 < n_samples {
                self.prefetch_memory_location(&x[[i + 1, 0]]);
            }

            // Vectorized dot product with models
            if self.config.enable_simd && n_models >= 4 {
                ensemble_sum = self.vectorized_ensemble_sum(models, &row)?;
            } else {
                for model in models {
                    ensemble_sum += row.dot(model);
                }
            }

            predictions[i] = if ensemble_sum / n_models as Float > 0.0 {
                1
            } else {
                0
            };
        }

        Ok(predictions)
    }

    /// Vectorized ensemble sum calculation
    fn vectorized_ensemble_sum(
        &mut self,
        models: &[Array1<Float>],
        x_row: &scirs2_core::ndarray::ArrayView1<Float>,
    ) -> Result<Float> {
        let mut total_sum = 0.0;

        // Process models in groups of 4 for better vectorization
        let n_models = models.len();
        let mut model_idx = 0;

        while model_idx + 4 <= n_models {
            let sum1 = x_row.dot(&models[model_idx]);
            let sum2 = x_row.dot(&models[model_idx + 1]);
            let sum3 = x_row.dot(&models[model_idx + 2]);
            let sum4 = x_row.dot(&models[model_idx + 3]);

            total_sum += sum1 + sum2 + sum3 + sum4;
            model_idx += 4;
        }

        // Handle remaining models
        while model_idx < n_models {
            total_sum += x_row.dot(&models[model_idx]);
            model_idx += 1;
        }

        self.performance_counters.vectorized_operations += 1;
        Ok(total_sum)
    }

    /// Memory prefetching hint
    fn prefetch_memory_location(&mut self, _addr: &Float) {
        // In a real implementation, this would use platform-specific prefetch instructions
        // For now, just record the prefetch request
        self.performance_counters.prefetch_requests += 1;
    }

    /// Cache-optimized histogram computation
    pub fn optimized_histogram(
        &mut self,
        data: &Array1<Float>,
        bins: usize,
        min_val: Float,
        max_val: Float,
    ) -> Result<Array1<usize>> {
        let mut histogram = Array1::zeros(bins);
        let bin_width = (max_val - min_val) / bins as Float;

        if self.config.enable_cache_optimization {
            // Process data in cache-friendly chunks
            let chunk_size = self.config.l1_cache_size_kb * 1024 / std::mem::size_of::<Float>();
            let data_len = data.len();

            for chunk_start in (0..data_len).step_by(chunk_size) {
                let chunk_end = (chunk_start + chunk_size).min(data_len);
                let chunk = data.slice(s![chunk_start..chunk_end]);

                for &value in chunk.iter() {
                    if value >= min_val && value < max_val {
                        let bin_idx = ((value - min_val) / bin_width) as usize;
                        let bin_idx = bin_idx.min(bins - 1);
                        histogram[bin_idx] += 1;
                    }
                }
            }
        } else {
            // Simple histogram computation
            for &value in data.iter() {
                if value >= min_val && value < max_val {
                    let bin_idx = ((value - min_val) / bin_width) as usize;
                    let bin_idx = bin_idx.min(bins - 1);
                    histogram[bin_idx] += 1;
                }
            }
        }

        Ok(histogram)
    }

    /// Branch prediction optimized decision tree traversal
    pub fn optimized_tree_traversal(
        &mut self,
        tree_nodes: &[(usize, Float, Int, Int)], // (feature_idx, threshold, left_child, right_child)
        x: &Array1<Float>,
    ) -> Result<Int> {
        let mut node_idx = 0;

        loop {
            if node_idx >= tree_nodes.len() {
                break;
            }

            let (feature_idx, threshold, left_child, right_child) = tree_nodes[node_idx];

            // Leaf node check
            if left_child == -1 && right_child == -1 {
                return Ok(feature_idx as Int); // Return class/value
            }

            // Branch prediction hint: assume left branch is more likely
            let go_left = if self.config.enable_branch_prediction {
                likely(x[feature_idx] <= threshold)
            } else {
                x[feature_idx] <= threshold
            };

            node_idx = if go_left {
                left_child as usize
            } else {
                right_child as usize
            };

            if go_left {
                self.performance_counters.branch_predictions += 1;
            } else {
                self.performance_counters.branch_mispredictions += 1;
            }
        }

        Err(SklearsError::InvalidInput(
            "Tree traversal failed: invalid tree structure".to_string(),
        ))
    }

    /// Auto-detect SIMD support
    fn detect_simd_support() -> bool {
        #[cfg(target_arch = "x86_64")]
        {
            is_x86_feature_detected!("avx2") || is_x86_feature_detected!("sse2")
        }
        #[cfg(target_arch = "aarch64")]
        {
            true // NEON is standard on ARM64
        }
        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
        {
            false
        }
    }

    /// Auto-detect L1 cache size
    fn detect_l1_cache_size() -> usize {
        // In a real implementation, this would query CPU cache info
        32 // Default 32KB L1 cache
    }

    /// Auto-detect L2 cache size
    fn detect_l2_cache_size() -> usize {
        // In a real implementation, this would query CPU cache info
        256 // Default 256KB L2 cache
    }

    /// Auto-detect L3 cache size
    fn detect_l3_cache_size() -> usize {
        // In a real implementation, this would query CPU cache info
        8192 // Default 8MB L3 cache
    }

    /// Auto-detect core count
    fn detect_core_count() -> usize {
        std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(4)
    }

    /// Get performance counters
    pub fn performance_counters(&self) -> &PerformanceCounters {
        &self.performance_counters
    }

    /// Reset performance counters
    pub fn reset_counters(&mut self) {
        self.performance_counters = PerformanceCounters::default();
    }

    /// Get cache efficiency ratio
    pub fn cache_efficiency(&self) -> Float {
        let total_accesses =
            self.performance_counters.cache_hits + self.performance_counters.cache_misses;
        if total_accesses > 0 {
            self.performance_counters.cache_hits as Float / total_accesses as Float
        } else {
            0.0
        }
    }

    /// Get branch prediction accuracy
    pub fn branch_prediction_accuracy(&self) -> Float {
        let total_branches = self.performance_counters.branch_predictions
            + self.performance_counters.branch_mispredictions;
        if total_branches > 0 {
            self.performance_counters.branch_predictions as Float / total_branches as Float
        } else {
            0.0
        }
    }

    /// Get vectorization ratio
    pub fn vectorization_ratio(&self) -> Float {
        let total_ops = self.performance_counters.vectorized_operations
            + self.performance_counters.scalar_operations;
        if total_ops > 0 {
            self.performance_counters.vectorized_operations as Float / total_ops as Float
        } else {
            0.0
        }
    }
}

/// Branch prediction hint (likely)
#[inline(always)]
fn likely(condition: bool) -> bool {
    // In a real implementation, this would use compiler-specific branch hints
    condition
}

/// Branch prediction hint (unlikely)
#[inline(always)]
fn _unlikely(condition: bool) -> bool {
    // In a real implementation, this would use compiler-specific branch hints
    condition
}

impl CacheOptimizedMatrixOps {
    pub fn new(tile_size: usize, l1_cache_size: usize, l2_cache_size: usize) -> Self {
        Self {
            tile_size,
            l1_cache_size,
            l2_cache_size,
        }
    }

    /// Cache-optimized matrix transpose
    pub fn optimized_transpose(&self, matrix: &Array2<Float>) -> Array2<Float> {
        let (rows, cols) = matrix.dim();
        let mut result = Array2::zeros((cols, rows));

        // Tile the transpose for better cache performance
        for i_tile in (0..rows).step_by(self.tile_size) {
            for j_tile in (0..cols).step_by(self.tile_size) {
                let i_end = (i_tile + self.tile_size).min(rows);
                let j_end = (j_tile + self.tile_size).min(cols);

                for i in i_tile..i_end {
                    for j in j_tile..j_end {
                        result[[j, i]] = matrix[[i, j]];
                    }
                }
            }
        }

        result
    }
}

impl Default for VectorizedEnsembleOps {
    fn default() -> Self {
        Self::new()
    }
}

impl VectorizedEnsembleOps {
    /// Create new vectorized ensemble operations
    pub fn new() -> Self {
        Self {
            simd_width: Self::detect_simd_width(),
            supports_avx512: Self::detect_avx512(),
            supports_avx2: Self::detect_avx2(),
            supports_fma: Self::detect_fma(),
        }
    }

    /// Vectorized weighted sum of predictions
    pub fn vectorized_weighted_sum(
        &self,
        predictions: &[Array1<Float>],
        weights: &Array1<Float>,
    ) -> Result<Array1<Float>> {
        if predictions.is_empty() {
            return Err(SklearsError::InvalidInput(
                "No predictions provided".to_string(),
            ));
        }

        let n_samples = predictions[0].len();
        let mut result = Array1::zeros(n_samples);

        for (pred, &weight) in predictions.iter().zip(weights.iter()) {
            if pred.len() != n_samples {
                return Err(SklearsError::ShapeMismatch {
                    expected: format!("{} samples", n_samples),
                    actual: format!("{} samples", pred.len()),
                });
            }

            // Use SIMD operations for weighted addition
            let weighted_pred = SimdOps::scalar_multiply(pred, weight);
            result = SimdOps::add_arrays(&result, &weighted_pred);
        }

        Ok(result)
    }

    /// Detect SIMD width
    fn detect_simd_width() -> usize {
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx512f") {
                64 // 512 bits / 8 bits per byte
            } else if is_x86_feature_detected!("avx2") {
                32 // 256 bits / 8 bits per byte
            } else if is_x86_feature_detected!("sse2") {
                16 // 128 bits / 8 bits per byte
            } else {
                8 // Fallback
            }
        }
        #[cfg(target_arch = "aarch64")]
        {
            16 // NEON is 128-bit
        }
        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
        {
            8 // Conservative fallback
        }
    }

    /// Detect AVX-512 support
    fn detect_avx512() -> bool {
        #[cfg(target_arch = "x86_64")]
        {
            is_x86_feature_detected!("avx512f")
        }
        #[cfg(not(target_arch = "x86_64"))]
        {
            false
        }
    }

    /// Detect AVX2 support
    fn detect_avx2() -> bool {
        #[cfg(target_arch = "x86_64")]
        {
            is_x86_feature_detected!("avx2")
        }
        #[cfg(not(target_arch = "x86_64"))]
        {
            false
        }
    }

    /// Detect FMA support
    fn detect_fma() -> bool {
        #[cfg(target_arch = "x86_64")]
        {
            is_x86_feature_detected!("fma")
        }
        #[cfg(not(target_arch = "x86_64"))]
        {
            false
        }
    }
}

impl LoopOptimizedAlgorithms {
    /// Create new loop-optimized algorithms
    pub fn new(unroll_factor: usize, prefetch_distance: usize) -> Self {
        Self {
            unroll_factor,
            prefetch_distance,
        }
    }

    /// Loop-unrolled array sum
    pub fn unrolled_sum(&self, array: &Array1<Float>) -> Float {
        let mut sum = 0.0;
        let len = array.len();
        let unroll = self.unroll_factor;

        let mut i = 0;

        // Unrolled loop
        while i + unroll <= len {
            let mut partial_sum = 0.0;
            for j in 0..unroll {
                partial_sum += array[i + j];
            }
            sum += partial_sum;
            i += unroll;
        }

        // Handle remaining elements
        while i < len {
            sum += array[i];
            i += 1;
        }

        sum
    }

    /// Prefetch-optimized array processing
    pub fn prefetched_process<F>(&self, array: &Array1<Float>, mut f: F) -> Result<()>
    where
        F: FnMut(Float),
    {
        let len = array.len();
        let prefetch_dist = self.prefetch_distance;

        for i in 0..len {
            // Prefetch future elements
            if i + prefetch_dist < len {
                // In a real implementation, this would use actual prefetch instructions
                // For now, this is just a placeholder
            }

            f(array[i]);
        }

        Ok(())
    }
}

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

    #[test]
    fn test_cpu_optimization_config() {
        let config = CpuOptimizationConfig::default();
        assert!(config.enable_simd);
        assert!(config.enable_cache_optimization);
        assert_eq!(config.cache_line_size, 64);
    }

    #[test]
    fn test_cpu_optimizer_creation() {
        let optimizer = CpuOptimizer::auto_detect();
        assert!(optimizer.config.num_cores > 0);
    }

    #[test]
    fn test_optimized_matmul() {
        let mut optimizer = CpuOptimizer::auto_detect();

        let a = array![[1.0, 2.0], [3.0, 4.0]];
        let b = array![[5.0, 6.0], [7.0, 8.0]];

        let result = optimizer.optimized_matmul(&a, &b).unwrap();
        let expected = array![[19.0, 22.0], [43.0, 50.0]];

        assert_eq!(result, expected);
    }

    #[test]
    fn test_optimized_histogram() {
        let mut optimizer = CpuOptimizer::auto_detect();

        let data = array![1.0, 2.0, 3.0, 4.0, 5.0, 2.0, 3.0, 3.0];
        let histogram = optimizer.optimized_histogram(&data, 5, 1.0, 6.0).unwrap();

        assert_eq!(histogram.sum(), data.len());
    }

    #[test]
    fn test_performance_counters() {
        let mut optimizer = CpuOptimizer::auto_detect();

        let a = array![[1.0, 2.0], [3.0, 4.0]];
        let b = array![[5.0, 6.0], [7.0, 8.0]];

        optimizer.optimized_matmul(&a, &b).unwrap();

        let counters = optimizer.performance_counters();
        assert!(counters.vectorized_operations > 0 || counters.scalar_operations > 0);
    }

    #[test]
    fn test_cache_optimized_matrix_ops() {
        let ops = CacheOptimizedMatrixOps::new(64, 32 * 1024, 256 * 1024);

        let matrix = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
        let transposed = ops.optimized_transpose(&matrix);
        let expected = array![[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]];

        assert_eq!(transposed, expected);
    }

    #[test]
    fn test_vectorized_ensemble_ops() {
        let ops = VectorizedEnsembleOps::new();

        let pred1 = array![1.0, 2.0, 3.0];
        let pred2 = array![2.0, 3.0, 4.0];
        let predictions = vec![pred1, pred2];
        let weights = array![0.6, 0.4];

        let result = ops.vectorized_weighted_sum(&predictions, &weights).unwrap();
        let expected = array![1.4, 2.4, 3.4]; // 0.6*1 + 0.4*2, etc.

        for (actual, expected) in result.iter().zip(expected.iter()) {
            assert!((actual - expected).abs() < 1e-10);
        }
    }

    #[test]
    fn test_loop_optimized_algorithms() {
        let alg = LoopOptimizedAlgorithms::new(4, 8);

        let array = array![1.0, 2.0, 3.0, 4.0, 5.0];
        let sum = alg.unrolled_sum(&array);

        assert_eq!(sum, 15.0);
    }
}