torsh_tensor/
cache_optimization.rs

1// Cache optimization module for improving memory layout and access patterns
2
3#[cfg(feature = "simd")]
4use crate::storage::SimdStorage;
5use crate::{Tensor, TensorStorage};
6use std::collections::HashMap;
7use std::sync::{Arc, Mutex};
8use std::time::{Duration, Instant};
9use torsh_core::sync::MutexExt;
10use torsh_core::{
11    dtype::TensorElement,
12    error::{Result, TorshError},
13    shape::Shape,
14};
15
16#[cfg(feature = "simd")]
17use scirs2_core::simd_aligned::AlignedVec;
18
19/// Cache analysis report providing detailed performance metrics
20#[derive(Debug, Clone)]
21pub struct CacheAnalysisReport {
22    /// Overall cache efficiency score (0.0 to 1.0)
23    pub cache_efficiency: f64,
24    /// Estimated number of cache misses for typical access patterns
25    pub estimated_cache_misses: usize,
26    /// Spatial locality score (0.0 to 1.0)
27    pub spatial_locality_score: f64,
28    /// Temporal locality score (0.0 to 1.0)
29    pub temporal_locality_score: f64,
30    /// Whether current memory layout is optimal
31    pub memory_layout_optimal: bool,
32    /// List of recommended optimizations
33    pub recommended_optimizations: Vec<String>,
34}
35
36impl<T: TensorElement + Copy> Tensor<T> {
37    /// Memory layout optimization for cache efficiency
38    /// Analyzes and optimizes the tensor's memory layout to improve cache performance
39    pub fn optimize_cache_layout(&mut self) -> Result<()> {
40        // Check if tensor is large enough to benefit from optimization
41        if self.numel() < 1024 {
42            return Ok(()); // Skip small tensors
43        }
44
45        // Analyze current access pattern and stride layout
46        let current_strides = self.compute_strides();
47        let optimal_order = self.determine_optimal_dimension_order(&current_strides);
48
49        // If current layout is already optimal, return early
50        if optimal_order.iter().enumerate().all(|(i, &dim)| dim == i) {
51            return Ok(());
52        }
53
54        // Reorganize data for better cache locality
55        self.reorder_dimensions(&optimal_order)?;
56
57        // Add padding for cache line alignment if beneficial
58        self.add_cache_padding()?;
59
60        Ok(())
61    }
62
63    /// Determine optimal dimension order for cache efficiency
64    /// Prioritizes dimensions that are accessed more frequently together
65    fn determine_optimal_dimension_order(&self, strides: &[usize]) -> Vec<usize> {
66        let shape_binding = self.shape();
67        let dims = shape_binding.dims();
68        let mut dim_priorities: Vec<(usize, f64)> = (0..dims.len())
69            .map(|i| {
70                // Calculate priority based on dimension size and stride
71                let size_factor = dims[i] as f64;
72                let stride_factor = 1.0 / (strides[i] as f64 + 1.0);
73                let cache_friendliness = size_factor * stride_factor;
74                (i, cache_friendliness)
75            })
76            .collect();
77
78        // Sort by cache friendliness (higher is better)
79        dim_priorities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
80
81        dim_priorities.into_iter().map(|(dim, _)| dim).collect()
82    }
83
84    /// Reorder tensor dimensions for optimal cache access
85    fn reorder_dimensions(&mut self, optimal_order: &[usize]) -> Result<()> {
86        if optimal_order.len() != self.ndim() {
87            return Err(TorshError::InvalidOperation(
88                "Dimension order length mismatch".to_string(),
89            ));
90        }
91
92        // Create permutation for transpose operation
93        let data = self.to_vec()?;
94        let old_dims = self.shape().dims().to_vec();
95        let old_strides = self.compute_strides();
96
97        // Calculate new dimensions and create reordered data
98        let new_dims: Vec<usize> = optimal_order.iter().map(|&i| old_dims[i]).collect();
99        let new_numel = new_dims.iter().product::<usize>();
100        let mut new_data = vec![data[0]; new_numel]; // Initialize with first element
101
102        // Reorder data according to optimal dimension order
103        #[allow(clippy::needless_range_loop)]
104        for i in 0..new_numel {
105            let mut old_indices = vec![0; self.ndim()];
106            let mut remaining = i;
107
108            // Convert flat index to multi-dimensional indices in new layout
109            for (j, &dim_size) in new_dims.iter().enumerate().rev() {
110                old_indices[optimal_order[j]] = remaining % dim_size;
111                remaining /= dim_size;
112            }
113
114            // Calculate flat index in original layout
115            let old_flat_index: usize = old_indices
116                .iter()
117                .zip(old_strides.iter())
118                .map(|(&idx, &stride)| idx * stride)
119                .sum();
120
121            new_data[i] = data[old_flat_index];
122        }
123
124        // Update tensor with optimized layout
125        self.storage = TensorStorage::create_optimal(new_data)?;
126        self.shape = Shape::new(new_dims);
127
128        Ok(())
129    }
130
131    /// Add cache-line aligned padding for better memory access patterns
132    fn add_cache_padding(&mut self) -> Result<()> {
133        const CACHE_LINE_SIZE: usize = 64; // bytes
134        let element_size = std::mem::size_of::<T>();
135        let elements_per_cache_line = CACHE_LINE_SIZE / element_size;
136
137        // Only add padding if it would be beneficial
138        let shape_binding = self.shape();
139        let dims = shape_binding.dims();
140        if dims.is_empty() || dims[dims.len() - 1] % elements_per_cache_line == 0 {
141            return Ok(()); // Already aligned or no benefit
142        }
143
144        // Calculate padding needed for last dimension
145        let last_dim = dims[dims.len() - 1];
146        let padded_last_dim = last_dim.div_ceil(elements_per_cache_line) * elements_per_cache_line;
147        let padding_needed = padded_last_dim - last_dim;
148
149        // Only add padding if overhead is reasonable (< 25%)
150        if (padding_needed as f64 / last_dim as f64) > 0.25 {
151            return Ok(());
152        }
153
154        let data = self.to_vec()?;
155        let mut new_dims = dims.to_vec();
156        let last_idx = new_dims.len() - 1;
157        new_dims[last_idx] = padded_last_dim;
158
159        // Create padded data
160        let new_numel = new_dims.iter().product::<usize>();
161        let mut padded_data = Vec::with_capacity(new_numel);
162
163        let outer_size = new_numel / padded_last_dim;
164        for i in 0..outer_size {
165            let start_idx = i * last_dim;
166            let end_idx = (i + 1) * last_dim;
167
168            // Copy original data
169            padded_data.extend_from_slice(&data[start_idx..end_idx]);
170
171            // Add padding (zeros)
172            for _ in 0..padding_needed {
173                padded_data.push(data[0]); // Use first element as padding value
174            }
175        }
176
177        // Update tensor with padded layout
178        self.storage = TensorStorage::create_optimal(padded_data)?;
179        self.shape = Shape::new(new_dims);
180
181        Ok(())
182    }
183
184    /// Analyze memory access patterns and provide optimization recommendations
185    pub fn analyze_cache_performance(&self) -> CacheAnalysisReport {
186        let shape_binding = self.shape();
187        let dims = shape_binding.dims();
188        let strides = self.compute_strides();
189        let numel = self.numel();
190
191        // Calculate cache efficiency metrics
192        let mut cache_misses_estimate = 0f64;
193
194        // Estimate cache misses based on stride patterns
195        for (i, &stride) in strides.iter().enumerate() {
196            let dimension_accesses = dims[i] as f64;
197            let stride_penalty = if stride > 64 {
198                stride as f64 / 64.0
199            } else {
200                1.0
201            };
202            cache_misses_estimate += dimension_accesses * stride_penalty;
203        }
204
205        // Calculate spatial locality (how well adjacent elements are accessed together)
206        let spatial_locality_score = if strides.last().copied().unwrap_or(1) == 1usize {
207            1.0
208        } else {
209            1.0 / strides.last().copied().unwrap_or(1) as f64
210        };
211
212        // Calculate temporal locality (reuse of recently accessed data)
213        let temporal_locality_score = 1.0 / (numel as f64).log2().max(1.0);
214
215        CacheAnalysisReport {
216            cache_efficiency: (spatial_locality_score + temporal_locality_score) / 2.0,
217            estimated_cache_misses: cache_misses_estimate as usize,
218            spatial_locality_score,
219            temporal_locality_score,
220            memory_layout_optimal: strides.last().copied().unwrap_or(1) == 1usize,
221            recommended_optimizations: self.generate_optimization_recommendations(&strides),
222        }
223    }
224
225    /// Generate specific optimization recommendations based on current layout
226    fn generate_optimization_recommendations(&self, strides: &[usize]) -> Vec<String> {
227        let mut recommendations = Vec::new();
228        let shape_binding = self.shape();
229        let dims = shape_binding.dims();
230
231        // Check for non-contiguous memory layout
232        if strides.last().copied().unwrap_or(1) != 1 {
233            recommendations
234                .push("Consider using .contiguous() to ensure row-major layout".to_string());
235        }
236
237        // Check for small tensors that don't benefit from optimization
238        if self.numel() < 1024 {
239            recommendations.push("Tensor too small to benefit from cache optimization".to_string());
240        }
241
242        // Check for dimensions that could benefit from reordering
243        if dims.len() > 2 {
244            let largest_dim = dims.iter().enumerate().max_by_key(|(_, &size)| size);
245            if let Some((largest_idx, _)) = largest_dim {
246                if largest_idx != dims.len() - 1 {
247                    recommendations.push(format!(
248                        "Consider moving dimension {largest_idx} to the end for better cache locality"
249                    ));
250                }
251            }
252        }
253
254        // Check for padding opportunities
255        const CACHE_LINE_SIZE: usize = 64;
256        let element_size = std::mem::size_of::<T>();
257        let elements_per_cache_line = CACHE_LINE_SIZE / element_size;
258
259        if !dims.is_empty() {
260            let last_dim = dims[dims.len() - 1];
261            if last_dim % elements_per_cache_line != 0 {
262                recommendations
263                    .push("Consider adding cache-line padding for better alignment".to_string());
264            }
265        }
266
267        recommendations
268    }
269
270    /// Create a cache-optimized copy of the tensor
271    pub fn to_cache_optimized(&self) -> Result<Self> {
272        let mut optimized = self.clone();
273        optimized.optimize_cache_layout()?;
274        Ok(optimized)
275    }
276
277    /// Get memory usage statistics for the tensor
278    pub fn memory_stats(&self) -> MemoryStats {
279        let element_size = std::mem::size_of::<T>();
280        let total_elements = self.numel();
281        let total_bytes = total_elements * element_size;
282
283        // Estimate memory overhead based on storage type
284        let overhead_bytes = match &self.storage {
285            TensorStorage::InMemory(_) => {
286                // Arc + RwLock overhead
287                std::mem::size_of::<std::sync::Arc<std::sync::RwLock<Vec<T>>>>()
288            }
289            TensorStorage::MemoryMapped(_) => {
290                // Memory mapped storage overhead
291                1024 // Approximate overhead for file handles, cache, etc.
292            }
293            #[cfg(feature = "simd")]
294            TensorStorage::Aligned(_) => {
295                // Arc + RwLock + AlignedVec overhead
296                std::mem::size_of::<std::sync::Arc<std::sync::RwLock<AlignedVec<T>>>>()
297            }
298            #[cfg(feature = "simd")]
299            TensorStorage::SimdOptimized(_) => {
300                // Arc + SimdStorage overhead (no RwLock, so less overhead)
301                std::mem::size_of::<std::sync::Arc<SimdStorage<T>>>()
302            }
303            #[cfg(feature = "gpu")]
304            TensorStorage::Device { .. } => {
305                // Two Arc handles: the device allocation and the host cache.
306                std::mem::size_of::<std::sync::Arc<crate::storage::DeviceBuffer>>()
307                    + std::mem::size_of::<std::sync::Arc<std::sync::RwLock<Option<Vec<T>>>>>()
308            }
309        };
310
311        MemoryStats {
312            total_bytes,
313            element_size,
314            total_elements,
315            overhead_bytes,
316            is_memory_mapped: matches!(&self.storage, TensorStorage::MemoryMapped(_)),
317        }
318    }
319}
320
321/// Memory usage statistics for a tensor
322#[derive(Debug, Clone)]
323pub struct MemoryStats {
324    /// Total memory used by tensor data in bytes
325    pub total_bytes: usize,
326    /// Size of each element in bytes
327    pub element_size: usize,
328    /// Total number of elements
329    pub total_elements: usize,
330    /// Memory overhead from storage structures
331    pub overhead_bytes: usize,
332    /// Whether tensor uses memory-mapped storage
333    pub is_memory_mapped: bool,
334}
335
336impl MemoryStats {
337    /// Get effective memory usage (data + overhead)
338    pub fn effective_bytes(&self) -> usize {
339        self.total_bytes + self.overhead_bytes
340    }
341
342    /// Get memory efficiency (data bytes / total bytes)
343    pub fn efficiency(&self) -> f64 {
344        self.total_bytes as f64 / self.effective_bytes() as f64
345    }
346}
347
348/// Global memory pool for temporary tensor allocations
349pub struct TensorMemoryPool {
350    /// Pooled memory blocks organized by size
351    pool: Arc<Mutex<HashMap<usize, Vec<Vec<u8>>>>>,
352    /// Memory allocation statistics
353    stats: Arc<Mutex<PoolStatistics>>,
354    /// Maximum memory pool size in bytes
355    max_pool_size: usize,
356    /// Current pool size in bytes
357    current_pool_size: Arc<Mutex<usize>>,
358}
359
360#[derive(Debug, Clone, Default)]
361pub struct PoolStatistics {
362    pub allocations: usize,
363    pub deallocations: usize,
364    pub cache_hits: usize,
365    pub cache_misses: usize,
366    pub peak_memory_usage: usize,
367    pub total_memory_saved: usize,
368}
369
370impl TensorMemoryPool {
371    /// Create a new memory pool with specified maximum size
372    pub fn new(max_size_mb: usize) -> Self {
373        Self {
374            pool: Arc::new(Mutex::new(HashMap::new())),
375            stats: Arc::new(Mutex::new(PoolStatistics::default())),
376            max_pool_size: max_size_mb * 1024 * 1024,
377            current_pool_size: Arc::new(Mutex::new(0)),
378        }
379    }
380
381    /// Allocate memory from pool or create new
382    pub fn allocate(&self, size_bytes: usize) -> Vec<u8> {
383        let mut pool = self.pool.lock_or_recover();
384        let mut stats = self.stats.lock_or_recover();
385
386        stats.allocations += 1;
387
388        // Round up to next power of 2 for better pooling
389        let rounded_size = size_bytes.next_power_of_two();
390
391        if let Some(pool_vec) = pool.get_mut(&rounded_size) {
392            if let Some(memory) = pool_vec.pop() {
393                stats.cache_hits += 1;
394                let mut current_size = self.current_pool_size.lock_or_recover();
395                *current_size -= rounded_size;
396                return memory;
397            }
398        }
399
400        stats.cache_misses += 1;
401        vec![0u8; rounded_size]
402    }
403
404    /// Return memory to pool
405    pub fn deallocate(&self, mut memory: Vec<u8>) {
406        let size = memory.len();
407        let mut pool = self.pool.lock_or_recover();
408        let mut stats = self.stats.lock_or_recover();
409        let mut current_size = self.current_pool_size.lock_or_recover();
410
411        stats.deallocations += 1;
412
413        // Only pool if under size limit
414        if *current_size + size <= self.max_pool_size {
415            // Clear the memory before pooling for security
416            memory.fill(0);
417
418            pool.entry(size).or_default().push(memory);
419            *current_size += size;
420            stats.total_memory_saved += size;
421        }
422
423        stats.peak_memory_usage = stats.peak_memory_usage.max(*current_size);
424    }
425
426    /// Get pool statistics
427    pub fn get_statistics(&self) -> PoolStatistics {
428        self.stats.lock_or_recover().clone()
429    }
430
431    /// Clear the entire pool
432    pub fn clear(&self) {
433        let mut pool = self.pool.lock_or_recover();
434        let mut current_size = self.current_pool_size.lock_or_recover();
435
436        pool.clear();
437        *current_size = 0;
438    }
439}
440
441/// Memory pressure detection and adaptive allocation
442pub struct MemoryPressureMonitor {
443    /// Memory usage samples
444    samples: Arc<Mutex<Vec<(Instant, usize)>>>,
445    /// Current pressure level (0.0 to 1.0)
446    pressure_level: Arc<Mutex<f64>>,
447    /// System memory threshold for high pressure
448    high_pressure_threshold: usize,
449}
450
451impl MemoryPressureMonitor {
452    pub fn new(memory_limit_mb: usize) -> Self {
453        Self {
454            samples: Arc::new(Mutex::new(Vec::new())),
455            pressure_level: Arc::new(Mutex::new(0.0)),
456            high_pressure_threshold: memory_limit_mb * 1024 * 1024,
457        }
458    }
459
460    /// Record memory usage sample
461    pub fn record_usage(&self, bytes_used: usize) {
462        let mut samples = self.samples.lock_or_recover();
463        let mut pressure = self.pressure_level.lock_or_recover();
464
465        let now = Instant::now();
466        samples.push((now, bytes_used));
467
468        // Keep only recent samples (last 60 seconds)
469        samples.retain(|(time, _)| now.duration_since(*time) < Duration::from_secs(60));
470
471        // Calculate pressure based on recent usage
472        let avg_usage = if samples.is_empty() {
473            0.0
474        } else {
475            samples.iter().map(|(_, usage)| *usage as f64).sum::<f64>() / samples.len() as f64
476        };
477
478        *pressure = (avg_usage / self.high_pressure_threshold as f64).min(1.0);
479    }
480
481    /// Get current memory pressure level
482    pub fn get_pressure_level(&self) -> f64 {
483        *self.pressure_level.lock_or_recover()
484    }
485
486    /// Check if system is under high memory pressure
487    pub fn is_high_pressure(&self) -> bool {
488        self.get_pressure_level() > 0.8
489    }
490}
491
492/// NUMA-aware memory allocation hints
493#[derive(Debug, Clone, Copy)]
494pub enum NumaNode {
495    Local,
496    Node(u32),
497    Interleaved,
498}
499
500#[derive(Debug, Clone)]
501pub struct NumaAllocationHint {
502    pub preferred_node: NumaNode,
503    pub allow_fallback: bool,
504    pub bind_threads: bool,
505}
506
507impl<T: TensorElement + Copy + Default> Tensor<T> {
508    /// Advanced memory optimization with NUMA awareness
509    pub fn optimize_memory_layout(&mut self, numa_hint: Option<NumaAllocationHint>) -> Result<()> {
510        // Basic cache optimization
511        self.optimize_cache_layout()?;
512
513        // Apply NUMA optimization if hint provided
514        if let Some(hint) = numa_hint {
515            self.apply_numa_optimization(hint)?;
516        }
517
518        // Memory access pattern prediction
519        self.optimize_access_patterns()?;
520
521        Ok(())
522    }
523
524    /// Apply NUMA-specific optimizations
525    fn apply_numa_optimization(&mut self, _hint: NumaAllocationHint) -> Result<()> {
526        // NUMA optimization would require platform-specific implementation
527        // For now, we'll implement basic interleaving for large tensors
528        if self.numel() > 1_000_000 {
529            // Large tensors benefit from interleaved allocation
530            // This would require platform-specific NUMA API calls
531            // For now, just ensure contiguous layout
532            if !self.is_contiguous() {
533                let contiguous_tensor = self.contiguous()?;
534                *self = contiguous_tensor;
535            }
536        }
537        Ok(())
538    }
539
540    /// Optimize memory access patterns based on predicted usage
541    fn optimize_access_patterns(&mut self) -> Result<()> {
542        let shape_binding = self.shape();
543        let dims = shape_binding.dims();
544
545        // For matrices, optimize for row-major access
546        if dims.len() == 2 && dims[0] > 64 && dims[1] > 64 {
547            // Check if we should transpose for better cache behavior
548            let row_size = dims[1] * std::mem::size_of::<T>();
549            let cache_line_size = 64;
550
551            // If rows don't align well with cache lines, consider optimization
552            if row_size % cache_line_size != 0 && row_size < cache_line_size * 4 {
553                self.add_cache_padding()?;
554            }
555        }
556
557        // For 3D+ tensors, ensure innermost dimension is cache-friendly
558        if dims.len() >= 3 {
559            let innermost_size = dims[dims.len() - 1] * std::mem::size_of::<T>();
560            if !(32..=256).contains(&innermost_size) {
561                // Consider reshaping for better cache utilization
562                self.add_cache_padding()?;
563            }
564        }
565
566        Ok(())
567    }
568
569    /// Memory-mapped tensor creation with optimization hints
570    pub fn create_memory_mapped_optimized(
571        data: Vec<T>,
572        shape: Vec<usize>,
573        numa_hint: Option<NumaAllocationHint>,
574    ) -> Result<Self> {
575        let mut tensor = Self::from_data(data, shape, torsh_core::device::DeviceType::Cpu)?;
576        tensor.optimize_memory_layout(numa_hint)?;
577        Ok(tensor)
578    }
579
580    /// Prefetch memory pages for better performance
581    pub fn prefetch_data(&self) -> Result<()> {
582        // This would use madvise/PrefetchVirtualMemory on supported platforms
583        // For now, we'll implement a simple memory access pattern
584        if self.numel() > 10_000 {
585            let data = self.to_vec()?;
586            let stride = data.len() / 100; // Sample every 1% of data
587
588            // Touch memory at regular intervals to trigger prefetch
589            let mut _sum = T::default();
590            for i in (0..data.len()).step_by(stride.max(1)) {
591                _sum = data[i]; // Simple memory access to trigger prefetch
592            }
593        }
594        Ok(())
595    }
596}
597
598// Global memory pool instance
599static GLOBAL_MEMORY_POOL: std::sync::OnceLock<TensorMemoryPool> = std::sync::OnceLock::new();
600static MEMORY_PRESSURE_MONITOR: std::sync::OnceLock<MemoryPressureMonitor> =
601    std::sync::OnceLock::new();
602
603/// Get global memory pool
604pub fn get_memory_pool() -> &'static TensorMemoryPool {
605    GLOBAL_MEMORY_POOL.get_or_init(|| TensorMemoryPool::new(1024)) // 1GB default
606}
607
608/// Get memory pressure monitor
609pub fn get_memory_pressure_monitor() -> &'static MemoryPressureMonitor {
610    MEMORY_PRESSURE_MONITOR.get_or_init(|| MemoryPressureMonitor::new(8192)) // 8GB default
611}
612
613#[cfg(test)]
614mod tests {
615    use crate::creation::*;
616
617    #[test]
618    fn test_cache_optimization() {
619        let mut tensor = ones::<f32>(&[100, 100]).expect("ones creation should succeed");
620        assert!(tensor.optimize_cache_layout().is_ok());
621    }
622
623    #[test]
624    fn test_cache_analysis() {
625        let tensor = ones::<f32>(&[64, 64]).expect("ones creation should succeed");
626        let report = tensor.analyze_cache_performance();
627        assert!(report.cache_efficiency >= 0.0 && report.cache_efficiency <= 1.0);
628    }
629
630    #[test]
631    fn test_contiguous_layout() {
632        let tensor = ones::<f32>(&[10, 10]).expect("ones creation should succeed");
633        assert!(tensor.is_contiguous());
634
635        let contiguous = tensor
636            .contiguous()
637            .expect("contiguous conversion should succeed");
638        assert!(contiguous.is_contiguous());
639    }
640
641    #[test]
642    fn test_memory_stats() {
643        let tensor = ones::<f32>(&[100, 100]).expect("ones creation should succeed");
644        let stats = tensor.memory_stats();
645        assert_eq!(stats.total_elements, 10000);
646        assert_eq!(stats.element_size, 4); // f32 is 4 bytes
647        assert_eq!(stats.total_bytes, 40000);
648    }
649
650    #[test]
651    fn test_memory_pool() {
652        use super::*;
653
654        let pool = TensorMemoryPool::new(10); // 10 MB
655
656        // Test allocation
657        let memory1 = pool.allocate(1024);
658        assert_eq!(memory1.len(), 1024);
659
660        let memory2 = pool.allocate(2048);
661        assert_eq!(memory2.len(), 2048);
662
663        // Test deallocation and reuse
664        pool.deallocate(memory1);
665        let memory3 = pool.allocate(1024);
666        assert_eq!(memory3.len(), 1024);
667
668        // Check statistics
669        let stats = pool.get_statistics();
670        assert!(stats.allocations > 0);
671        assert!(stats.deallocations > 0);
672
673        pool.deallocate(memory2);
674        pool.deallocate(memory3);
675    }
676
677    #[test]
678    fn test_memory_pressure_monitor() {
679        use super::*;
680
681        let monitor = MemoryPressureMonitor::new(100); // 100 MB limit
682
683        // Test pressure calculation - monitor uses average of samples
684        monitor.record_usage(50 * 1024 * 1024); // 50 MB
685        assert!(monitor.get_pressure_level() < 0.6);
686
687        monitor.record_usage(90 * 1024 * 1024); // 90 MB
688                                                // Average of 50MB and 90MB = 70MB = 0.7 pressure
689        assert!(monitor.get_pressure_level() > 0.6);
690        assert!(monitor.get_pressure_level() < 0.8);
691        assert!(!monitor.is_high_pressure()); // 0.7 < 0.8, so not high pressure
692
693        // Add a higher pressure reading to trigger high pressure
694        monitor.record_usage(95 * 1024 * 1024); // 95 MB
695                                                // Average of 50MB, 90MB, and 95MB = ~78MB = 0.78 pressure (still < 0.8)
696        monitor.record_usage(100 * 1024 * 1024); // 100 MB
697                                                 // This should push the average above 0.8
698        assert!(monitor.is_high_pressure());
699    }
700
701    #[test]
702    fn test_advanced_memory_optimization() {
703        let mut tensor = ones::<f32>(&[64, 64]).expect("ones creation should succeed");
704
705        // Test with NUMA hint
706        let numa_hint = super::NumaAllocationHint {
707            preferred_node: super::NumaNode::Local,
708            allow_fallback: true,
709            bind_threads: false,
710        };
711
712        assert!(tensor.optimize_memory_layout(Some(numa_hint)).is_ok());
713        assert!(tensor.is_contiguous());
714    }
715
716    #[test]
717    fn test_cache_optimized_creation() {
718        let data: Vec<f32> = (0..10000).map(|i| i as f32).collect();
719        let shape = vec![100, 100];
720
721        let numa_hint = super::NumaAllocationHint {
722            preferred_node: super::NumaNode::Interleaved,
723            allow_fallback: true,
724            bind_threads: false,
725        };
726
727        let tensor = super::Tensor::create_memory_mapped_optimized(data, shape, Some(numa_hint));
728        assert!(tensor.is_ok());
729
730        let tensor = tensor.expect("operation should succeed");
731        // Shape may be optimized with padding for cache efficiency
732        let shape = tensor.shape();
733        let dims = shape.dims();
734        assert_eq!(dims[0], 100); // First dimension should be preserved
735        assert!(dims[1] >= 100); // Second dimension may have padding
736    }
737
738    #[test]
739    fn test_memory_prefetch() {
740        let tensor = ones::<f32>(&[200, 200]).expect("ones creation should succeed");
741        assert!(tensor.prefetch_data().is_ok());
742    }
743
744    #[test]
745    fn test_global_memory_pool_access() {
746        use super::*;
747
748        let pool = get_memory_pool();
749        let memory = pool.allocate(1024);
750        assert_eq!(memory.len(), 1024);
751        pool.deallocate(memory);
752
753        let monitor = get_memory_pressure_monitor();
754        monitor.record_usage(1024 * 1024); // 1 MB
755        assert!(monitor.get_pressure_level() >= 0.0);
756    }
757
758    #[test]
759    fn test_pool_statistics() {
760        use super::*;
761
762        let pool = TensorMemoryPool::new(5); // 5 MB
763
764        // Perform multiple allocations and deallocations
765        let mut memories = Vec::new();
766        for i in 0..10 {
767            let size = (i + 1) * 512;
768            memories.push(pool.allocate(size));
769        }
770
771        for memory in memories {
772            pool.deallocate(memory);
773        }
774
775        let stats = pool.get_statistics();
776        assert_eq!(stats.allocations, 10);
777        assert_eq!(stats.deallocations, 10);
778        assert!(stats.cache_hits + stats.cache_misses == 10);
779
780        pool.clear();
781    }
782
783    #[test]
784    fn test_memory_efficiency_calculation() {
785        let tensor = ones::<f32>(&[50, 50]).expect("ones creation should succeed");
786        let stats = tensor.memory_stats();
787
788        let efficiency = stats.efficiency();
789        assert!(efficiency > 0.0 && efficiency <= 1.0);
790
791        let effective = stats.effective_bytes();
792        assert!(effective >= stats.total_bytes);
793    }
794}
torsh_tensor/cache_optimization.rs

torsh_tensor/
cache_optimization.rs
