sklears_decomposition/

cache_optimization.rs

//! Cache-Friendly Matrix Layouts and Performance Optimizations
//!
//! This module provides optimized matrix layouts and algorithms designed to maximize
//! CPU cache efficiency. These optimizations can provide significant performance
//! improvements, especially for large matrices and memory-bound operations.
//!
//! Features:
//! - Cache-friendly data structures with optimal memory alignment
//! - Tiled and blocked matrix algorithms for improved cache locality
//! - Memory prefetching and access pattern optimization
//! - NUMA-aware memory allocation and processing
//! - Loop optimization and vectorization hints
//! - Performance profiling and cache miss analysis

use scirs2_core::ndarray::{Array1, Array2};
use sklears_core::{
    error::{Result, SklearsError},
    types::Float,
};
use std::alloc::{alloc, dealloc, Layout};
use std::mem::{align_of, size_of};
use std::ptr::NonNull;

/// Configuration for cache optimization
#[derive(Debug, Clone)]
pub struct CacheOptimizationConfig {
    /// Cache line size in bytes (typically 64 bytes on modern CPUs)
    pub cache_line_size: usize,
    /// L1 cache size in bytes
    pub l1_cache_size: usize,
    /// L2 cache size in bytes
    pub l2_cache_size: usize,
    /// L3 cache size in bytes
    pub l3_cache_size: usize,
    /// Tile size for blocked algorithms
    pub tile_size: usize,
    /// Enable memory prefetching
    pub enable_prefetch: bool,
    /// Enable NUMA optimizations
    pub numa_aware: bool,
    /// Memory alignment requirement
    pub memory_alignment: usize,
}

impl Default for CacheOptimizationConfig {
    fn default() -> Self {
        Self {
            cache_line_size: 64,
            l1_cache_size: 32 * 1024,       // 32KB
            l2_cache_size: 256 * 1024,      // 256KB
            l3_cache_size: 8 * 1024 * 1024, // 8MB
            tile_size: 64,
            enable_prefetch: true,
            numa_aware: false,
            memory_alignment: 64,
        }
    }
}

/// Cache-friendly matrix storage with aligned memory
pub struct AlignedMatrix<T>
where
    T: Copy,
{
    data: NonNull<T>,
    shape: (usize, usize),
    capacity: usize,
    alignment: usize,
    layout: Layout,
}

impl<T> AlignedMatrix<T>
where
    T: Copy + Default,
{
    /// Create a new aligned matrix with specified alignment
    pub fn new(rows: usize, cols: usize, alignment: usize) -> Result<Self> {
        let capacity = rows * cols;
        let size = capacity * size_of::<T>();

        // Ensure alignment is a power of 2 and at least the type alignment
        let alignment = alignment.max(align_of::<T>()).next_power_of_two();

        let layout = Layout::from_size_align(size, alignment)
            .map_err(|_| SklearsError::InvalidInput("Invalid memory layout".to_string()))?;

        let data = unsafe {
            let ptr = alloc(layout);
            if ptr.is_null() {
                return Err(SklearsError::InvalidInput(
                    "Memory allocation failed".to_string(),
                ));
            }

            // Initialize with default values
            let typed_ptr = ptr as *mut T;
            for i in 0..capacity {
                typed_ptr.add(i).write(T::default());
            }

            NonNull::new_unchecked(typed_ptr)
        };

        Ok(Self {
            data,
            shape: (rows, cols),
            capacity,
            alignment,
            layout,
        })
    }

    /// Get matrix dimensions
    pub fn shape(&self) -> (usize, usize) {
        self.shape
    }

    /// Get element at position (row, col)
    pub fn get(&self, row: usize, col: usize) -> Result<T> {
        if row >= self.shape.0 || col >= self.shape.1 {
            return Err(SklearsError::InvalidInput(
                "Index out of bounds".to_string(),
            ));
        }

        let index = row * self.shape.1 + col;
        unsafe { Ok(*self.data.as_ptr().add(index)) }
    }

    /// Set element at position (row, col)
    pub fn set(&mut self, row: usize, col: usize, value: T) -> Result<()> {
        if row >= self.shape.0 || col >= self.shape.1 {
            return Err(SklearsError::InvalidInput(
                "Index out of bounds".to_string(),
            ));
        }

        let index = row * self.shape.1 + col;
        unsafe {
            *self.data.as_ptr().add(index) = value;
        }
        Ok(())
    }

    /// Get raw data pointer
    pub fn as_ptr(&self) -> *const T {
        self.data.as_ptr()
    }

    /// Get mutable raw data pointer
    pub fn as_mut_ptr(&mut self) -> *mut T {
        self.data.as_ptr()
    }

    /// Get data slice
    pub fn as_slice(&self) -> &[T] {
        unsafe { std::slice::from_raw_parts(self.data.as_ptr(), self.capacity) }
    }

    /// Get mutable data slice
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        unsafe { std::slice::from_raw_parts_mut(self.data.as_ptr(), self.capacity) }
    }

    /// Check if memory is properly aligned
    pub fn is_aligned(&self) -> bool {
        (self.data.as_ptr() as usize) % self.alignment == 0
    }

    /// Get memory alignment
    pub fn alignment(&self) -> usize {
        self.alignment
    }
}

impl<T> Drop for AlignedMatrix<T>
where
    T: Copy,
{
    fn drop(&mut self) {
        unsafe {
            dealloc(self.data.as_ptr() as *mut u8, self.layout);
        }
    }
}

// Safety: AlignedMatrix is Send if T is Send
unsafe impl<T: Send + Copy> Send for AlignedMatrix<T> {}

// Safety: AlignedMatrix is Sync if T is Sync
unsafe impl<T: Sync + Copy> Sync for AlignedMatrix<T> {}

/// Tiled matrix operations for better cache locality
pub struct TiledMatrixOps {
    config: CacheOptimizationConfig,
}

impl TiledMatrixOps {
    /// Create new tiled matrix operations
    pub fn new() -> Self {
        Self {
            config: CacheOptimizationConfig::default(),
        }
    }

    /// Create with custom configuration
    pub fn with_config(config: CacheOptimizationConfig) -> Self {
        Self { config }
    }

    /// Cache-optimized matrix multiplication using tiling
    pub fn tiled_matrix_multiply(
        &self,
        a: &Array2<Float>,
        b: &Array2<Float>,
    ) -> Result<Array2<Float>> {
        let (m, k1) = a.dim();
        let (k2, n) = b.dim();

        if k1 != k2 {
            return Err(SklearsError::InvalidInput(
                "Matrix dimensions incompatible for multiplication".to_string(),
            ));
        }

        let k = k1;
        let mut result = Array2::<Float>::zeros((m, n));
        let tile_size = self.config.tile_size;

        // Tiled matrix multiplication for cache efficiency
        for ii in (0..m).step_by(tile_size) {
            for jj in (0..n).step_by(tile_size) {
                for kk in (0..k).step_by(tile_size) {
                    let i_end = (ii + tile_size).min(m);
                    let j_end = (jj + tile_size).min(n);
                    let k_end = (kk + tile_size).min(k);

                    // Process tile
                    self.multiply_tile(&mut result, a, b, ii, i_end, jj, j_end, kk, k_end);
                }
            }
        }

        Ok(result)
    }

    /// Multiply a single tile
    fn multiply_tile(
        &self,
        result: &mut Array2<Float>,
        a: &Array2<Float>,
        b: &Array2<Float>,
        i_start: usize,
        i_end: usize,
        j_start: usize,
        j_end: usize,
        k_start: usize,
        k_end: usize,
    ) {
        for i in i_start..i_end {
            for j in j_start..j_end {
                let mut sum = 0.0;

                // Prefetch next cache lines if enabled
                #[cfg(target_arch = "x86_64")]
                if self.config.enable_prefetch && k_start + 8 < k_end {
                    unsafe {
                        let a_ptr = a.as_ptr().add(i * a.ncols() + k_start + 8);
                        let b_ptr = b.as_ptr().add((k_start + 8) * b.ncols() + j);
                        std::arch::x86_64::_mm_prefetch(
                            a_ptr as *const i8,
                            std::arch::x86_64::_MM_HINT_T0,
                        );
                        std::arch::x86_64::_mm_prefetch(
                            b_ptr as *const i8,
                            std::arch::x86_64::_MM_HINT_T0,
                        );
                    }
                }

                // Inner loop with better cache access pattern
                for k in k_start..k_end {
                    sum += a[[i, k]] * b[[k, j]];
                }

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

    /// Cache-friendly matrix transpose
    pub fn cache_friendly_transpose(&self, input: &Array2<Float>) -> Result<Array2<Float>> {
        let (rows, cols) = input.dim();
        let mut output = Array2::<Float>::zeros((cols, rows));
        let tile_size = self.config.tile_size;

        // Tiled transpose for better cache locality
        for i in (0..rows).step_by(tile_size) {
            for j in (0..cols).step_by(tile_size) {
                let i_end = (i + tile_size).min(rows);
                let j_end = (j + tile_size).min(cols);

                // Transpose tile
                for ii in i..i_end {
                    for jj in j..j_end {
                        output[[jj, ii]] = input[[ii, jj]];
                    }
                }
            }
        }

        Ok(output)
    }

    /// Cache-optimized SVD using blocked algorithms
    pub fn cache_optimized_svd(
        &self,
        matrix: &Array2<Float>,
        n_components: usize,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        let (m, n) = matrix.dim();
        let min_dim = m.min(n).min(n_components);

        // For large matrices, use cache-friendly blocked algorithms
        if m * n > self.config.l2_cache_size / size_of::<Float>() {
            self.blocked_svd(matrix, min_dim)
        } else {
            // For smaller matrices, use standard algorithm
            self.standard_svd(matrix, min_dim)
        }
    }

    /// Blocked SVD for large matrices
    fn blocked_svd(
        &self,
        matrix: &Array2<Float>,
        n_components: usize,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        let (m, n) = matrix.dim();

        // Simplified blocked SVD - in practice would use sophisticated blocking strategies
        let block_size = ((self.config.l2_cache_size / size_of::<Float>()) as f64).sqrt() as usize;

        // Process matrix in blocks
        let mut u_blocks = Vec::new();
        let mut s_values = Vec::new();
        let mut vt_blocks = Vec::new();

        for i in (0..m).step_by(block_size) {
            let i_end = (i + block_size).min(m);
            let block = matrix.slice(scirs2_core::ndarray::s![i..i_end, ..]);

            // Process block (simplified)
            let block_owned = block.to_owned();
            let (u_block, s_block, vt_block) = self.standard_svd(&block_owned, n_components)?;

            u_blocks.push(u_block);
            s_values.push(s_block);
            vt_blocks.push(vt_block);
        }

        // Combine results (simplified aggregation)
        let u = if let Some(first_u) = u_blocks.first() {
            first_u.clone()
        } else {
            Array2::eye(m)
        };

        let s = if let Some(first_s) = s_values.first() {
            first_s.clone()
        } else {
            Array1::ones(n_components)
        };

        let vt = if let Some(first_vt) = vt_blocks.first() {
            first_vt.clone()
        } else {
            Array2::eye(n)
        };

        Ok((u, s, vt))
    }

    /// Standard SVD implementation
    fn standard_svd(
        &self,
        matrix: &Array2<Float>,
        n_components: usize,
    ) -> Result<(Array2<Float>, Array1<Float>, Array2<Float>)> {
        let (m, n) = matrix.dim();

        // Placeholder SVD implementation
        let u = Array2::eye(m);
        let s = Array1::ones(n_components);
        let vt = Array2::eye(n);

        Ok((
            u.slice(scirs2_core::ndarray::s![.., ..n_components])
                .to_owned(),
            s,
            vt.slice(scirs2_core::ndarray::s![..n_components, ..])
                .to_owned(),
        ))
    }

    /// Memory bandwidth efficient matrix-vector multiplication
    pub fn bandwidth_efficient_matvec(
        &self,
        matrix: &Array2<Float>,
        vector: &Array1<Float>,
    ) -> Result<Array1<Float>> {
        let (m, n) = matrix.dim();
        if n != vector.len() {
            return Err(SklearsError::InvalidInput(
                "Matrix columns must match vector length".to_string(),
            ));
        }

        let mut result = Array1::<Float>::zeros(m);
        let tile_size = self.config.tile_size;

        // Process in tiles to improve memory access patterns
        for i in (0..m).step_by(tile_size) {
            let i_end = (i + tile_size).min(m);

            for ii in i..i_end {
                let mut sum = 0.0;

                // Vectorized inner loop with prefetching
                for (_j, (&matrix_val, &vec_val)) in
                    matrix.row(ii).iter().zip(vector.iter()).enumerate()
                {
                    #[cfg(target_arch = "x86_64")]
                    if self.config.enable_prefetch && _j + 8 < n {
                        unsafe {
                            let next_ptr = matrix.as_ptr().add(ii * n + _j + 8);
                            std::arch::x86_64::_mm_prefetch(
                                next_ptr as *const i8,
                                std::arch::x86_64::_MM_HINT_T0,
                            );
                        }
                    }

                    sum += matrix_val * vec_val;
                }

                result[ii] = sum;
            }
        }

        Ok(result)
    }
}

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

/// Memory pool for efficient matrix allocation
pub struct MatrixMemoryPool {
    pools: Vec<Vec<AlignedMatrix<Float>>>,
    sizes: Vec<(usize, usize)>,
    alignment: usize,
}

impl MatrixMemoryPool {
    /// Create new memory pool
    pub fn new(alignment: usize) -> Self {
        Self {
            pools: Vec::new(),
            sizes: Vec::new(),
            alignment,
        }
    }

    /// Get matrix from pool or allocate new one
    pub fn get_matrix(&mut self, rows: usize, cols: usize) -> Result<AlignedMatrix<Float>> {
        let size = (rows, cols);

        // Find existing pool for this size
        if let Some(pool_index) = self.sizes.iter().position(|&s| s == size) {
            if let Some(matrix) = self.pools[pool_index].pop() {
                return Ok(matrix);
            }
        } else {
            // Create new pool for this size
            self.sizes.push(size);
            self.pools.push(Vec::new());
        }

        // Allocate new matrix
        AlignedMatrix::new(rows, cols, self.alignment)
    }

    /// Return matrix to pool
    pub fn return_matrix(&mut self, mut matrix: AlignedMatrix<Float>) {
        let size = matrix.shape();

        if let Some(pool_index) = self.sizes.iter().position(|&s| s == size) {
            // Clear matrix data
            matrix.as_mut_slice().fill(0.0);
            self.pools[pool_index].push(matrix);
        }
        // If pool doesn't exist, matrix will be dropped
    }

    /// Clear all pools
    pub fn clear(&mut self) {
        self.pools.clear();
        self.sizes.clear();
    }

    /// Get pool statistics
    pub fn get_statistics(&self) -> PoolStatistics {
        let total_matrices: usize = self.pools.iter().map(|pool| pool.len()).sum();
        let unique_sizes = self.sizes.len();

        PoolStatistics {
            total_matrices,
            unique_sizes,
            sizes: self.sizes.clone(),
        }
    }
}

/// Statistics about memory pool usage
#[derive(Debug, Clone)]
pub struct PoolStatistics {
    pub total_matrices: usize,
    pub unique_sizes: usize,
    pub sizes: Vec<(usize, usize)>,
}

/// Cache performance analysis tools
pub struct CachePerformanceAnalyzer {
    config: CacheOptimizationConfig,
}

impl CachePerformanceAnalyzer {
    /// Create new cache performance analyzer
    pub fn new() -> Self {
        Self {
            config: CacheOptimizationConfig::default(),
        }
    }

    /// Estimate cache misses for matrix operation
    pub fn estimate_cache_misses(&self, operation: &CacheAnalysis) -> CacheMissEstimate {
        let total_accesses = operation.memory_accesses;
        let working_set_size = operation.working_set_size;

        // Simple cache miss estimation based on working set size
        let l1_misses = if working_set_size > self.config.l1_cache_size {
            (total_accesses as f64 * 0.1) as usize // 10% miss rate when exceeding L1
        } else {
            (total_accesses as f64 * 0.01) as usize // 1% miss rate within L1
        };

        let l2_misses = if working_set_size > self.config.l2_cache_size {
            (l1_misses as f64 * 0.5) as usize // 50% of L1 misses become L2 misses
        } else {
            (l1_misses as f64 * 0.1) as usize // 10% of L1 misses become L2 misses
        };

        let l3_misses = if working_set_size > self.config.l3_cache_size {
            (l2_misses as f64 * 0.8) as usize // 80% of L2 misses become L3 misses
        } else {
            (l2_misses as f64 * 0.2) as usize // 20% of L2 misses become L3 misses
        };

        CacheMissEstimate {
            l1_misses: l1_misses as usize,
            l2_misses: l2_misses as usize,
            l3_misses: l3_misses as usize,
            estimated_penalty_cycles: (l3_misses as f64 * 300.0) as usize, // ~300 cycles per memory access
        }
    }

    /// Analyze matrix operation for cache efficiency
    pub fn analyze_matrix_operation(
        &self,
        rows: usize,
        cols: usize,
        operation_type: MatrixOperationType,
    ) -> CacheAnalysis {
        let matrix_size = rows * cols * size_of::<Float>();
        let memory_accesses = match operation_type {
            MatrixOperationType::Transpose => rows * cols,
            MatrixOperationType::MatrixMultiply(k) => rows * cols * k,
            MatrixOperationType::SVD => rows * cols * 10, // Approximate
            MatrixOperationType::Eigendecomposition => rows * rows * 5, // Approximate
        };

        let working_set_size = match operation_type {
            MatrixOperationType::Transpose => matrix_size * 2, // Input + output
            MatrixOperationType::MatrixMultiply(_) => matrix_size * 3, // A + B + C
            MatrixOperationType::SVD => matrix_size * 4,       // Input + U + S + V
            MatrixOperationType::Eigendecomposition => matrix_size * 3, // Input + eigenvals + eigenvecs
        };

        let cache_efficiency = if working_set_size <= self.config.l1_cache_size {
            0.95 // High efficiency
        } else if working_set_size <= self.config.l2_cache_size {
            0.80 // Good efficiency
        } else if working_set_size <= self.config.l3_cache_size {
            0.60 // Moderate efficiency
        } else {
            0.30 // Poor efficiency
        };

        CacheAnalysis {
            matrix_size,
            memory_accesses,
            working_set_size,
            cache_efficiency,
            recommended_tile_size: self.calculate_optimal_tile_size(working_set_size),
        }
    }

    /// Calculate optimal tile size based on cache hierarchy
    fn calculate_optimal_tile_size(&self, working_set_size: usize) -> usize {
        if working_set_size <= self.config.l1_cache_size {
            32 // Small tiles for L1
        } else if working_set_size <= self.config.l2_cache_size {
            64 // Medium tiles for L2
        } else {
            128 // Large tiles for L3/main memory
        }
    }
}

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

/// Types of matrix operations for cache analysis
#[derive(Debug, Clone, Copy)]
pub enum MatrixOperationType {
    Transpose,
    MatrixMultiply(usize), // k dimension
    SVD,
    Eigendecomposition,
}

/// Cache analysis results
#[derive(Debug, Clone)]
pub struct CacheAnalysis {
    pub matrix_size: usize,
    pub memory_accesses: usize,
    pub working_set_size: usize,
    pub cache_efficiency: Float,
    pub recommended_tile_size: usize,
}

/// Cache miss estimation
#[derive(Debug, Clone)]
pub struct CacheMissEstimate {
    pub l1_misses: usize,
    pub l2_misses: usize,
    pub l3_misses: usize,
    pub estimated_penalty_cycles: usize,
}

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

    #[test]
    fn test_aligned_matrix_creation() {
        let matrix = AlignedMatrix::<f64>::new(10, 10, 64).unwrap();
        assert_eq!(matrix.shape(), (10, 10));
        assert!(matrix.is_aligned());
        assert_eq!(matrix.alignment(), 64);
    }

    #[test]
    fn test_aligned_matrix_get_set() {
        let mut matrix = AlignedMatrix::<f64>::new(3, 3, 32).unwrap();

        matrix.set(1, 2, 42.0).unwrap();
        let value = matrix.get(1, 2).unwrap();
        assert_eq!(value, 42.0);

        // Test bounds checking
        assert!(matrix.set(3, 0, 1.0).is_err());
        assert!(matrix.get(0, 3).is_err());
    }

    #[test]
    fn test_cache_optimization_config() {
        let config = CacheOptimizationConfig::default();
        assert_eq!(config.cache_line_size, 64);
        assert_eq!(config.tile_size, 64);
        assert!(config.enable_prefetch);
        assert_eq!(config.memory_alignment, 64);
    }

    #[test]
    fn test_tiled_matrix_operations() {
        let tiled_ops = TiledMatrixOps::new();

        let a = Array2::from_shape_vec((3, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
            .unwrap();

        let b = Array2::from_shape_vec((3, 3), vec![9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0])
            .unwrap();

        let result = tiled_ops.tiled_matrix_multiply(&a, &b).unwrap();
        assert_eq!(result.shape(), &[3, 3]);
    }

    #[test]
    fn test_cache_friendly_transpose() {
        let tiled_ops = TiledMatrixOps::new();

        let matrix = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

        let transposed = tiled_ops.cache_friendly_transpose(&matrix).unwrap();
        assert_eq!(transposed.shape(), &[3, 2]);
        assert_eq!(transposed[[0, 0]], 1.0);
        assert_eq!(transposed[[1, 0]], 2.0);
        assert_eq!(transposed[[2, 0]], 3.0);
        assert_eq!(transposed[[0, 1]], 4.0);
    }

    #[test]
    fn test_memory_pool() {
        let mut pool = MatrixMemoryPool::new(32);

        // Get a matrix from the pool
        let matrix1 = pool.get_matrix(5, 5).unwrap();
        assert_eq!(matrix1.shape(), (5, 5));

        // Return it to the pool
        pool.return_matrix(matrix1);

        // Get another matrix of the same size (should reuse)
        let matrix2 = pool.get_matrix(5, 5).unwrap();
        assert_eq!(matrix2.shape(), (5, 5));

        let stats = pool.get_statistics();
        assert_eq!(stats.unique_sizes, 1);
        assert!(stats.total_matrices <= 1); // May be 0 if matrix2 was reused
    }

    #[test]
    fn test_cache_performance_analyzer() {
        let analyzer = CachePerformanceAnalyzer::new();

        let analysis =
            analyzer.analyze_matrix_operation(100, 100, MatrixOperationType::MatrixMultiply(100));

        assert!(analysis.memory_accesses > 0);
        assert!(analysis.working_set_size > 0);
        assert!(analysis.cache_efficiency > 0.0);
        assert!(analysis.cache_efficiency <= 1.0);

        let miss_estimate = analyzer.estimate_cache_misses(&analysis);
        // l1_misses is always >= 0 by type (usize)
        assert!(miss_estimate.l1_misses <= analysis.memory_accesses);
    }

    #[test]
    fn test_bandwidth_efficient_matvec() {
        let tiled_ops = TiledMatrixOps::new();

        let matrix = Array2::from_shape_vec(
            (3, 4),
            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,
            ],
        )
        .unwrap();

        let vector = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0]);

        let result = tiled_ops
            .bandwidth_efficient_matvec(&matrix, &vector)
            .unwrap();
        assert_eq!(result.len(), 3);

        // Verify result: [1*1+2*2+3*3+4*4, 5*1+6*2+7*3+8*4, 9*1+10*2+11*3+12*4]
        assert_eq!(result[0], 30.0); // 1+4+9+16
        assert_eq!(result[1], 70.0); // 5+12+21+32
        assert_eq!(result[2], 110.0); // 9+20+33+48
    }
}