tenflowers_core/memory/
manager.rs

1//! High-level memory management coordination
2//!
3//! This module provides the main memory management interface that coordinates
4//! memory pools, performance tracking, cache optimization, and multi-stream operations.
5
6use super::{
7    cache::{AccessPattern, CacheOptimizer},
8    pools::{MemoryPool, MemoryPoolStats},
9    streams::MultiStreamMemoryManager,
10    tracking::{global_monitor_arc, PerformanceMonitor},
11    views::{MemoryAliasDetector, StridedView},
12};
13use crate::Device;
14#[cfg(feature = "gpu")]
15use crate::TensorError;
16use std::collections::HashMap;
17use std::sync::{Arc, RwLock};
18
19/// Central memory manager that coordinates all memory subsystems
20pub struct MemoryManager {
21    pools: Arc<RwLock<HashMap<Device, Arc<MemoryPool>>>>,
22    multi_stream_managers: Arc<RwLock<HashMap<usize, Arc<MultiStreamMemoryManager>>>>,
23    alias_detector: Arc<MemoryAliasDetector>,
24    cache_optimizer: Arc<CacheOptimizer>,
25    performance_monitor: Arc<PerformanceMonitor>,
26    default_pool_size: usize,
27}
28
29impl MemoryManager {
30    /// Create a new memory manager
31    pub fn new() -> Self {
32        Self {
33            pools: Arc::new(RwLock::new(HashMap::new())),
34            multi_stream_managers: Arc::new(RwLock::new(HashMap::new())),
35            alias_detector: Arc::new(MemoryAliasDetector::new()),
36            cache_optimizer: Arc::new(CacheOptimizer::new()),
37            performance_monitor: global_monitor_arc(),
38            default_pool_size: 512 * 1024 * 1024, // 512MB default
39        }
40    }
41
42    /// Create a new memory manager with custom pool size
43    pub fn with_pool_size(pool_size: usize) -> Self {
44        let mut manager = Self::new();
45        manager.default_pool_size = pool_size;
46        manager
47    }
48
49    /// Get or create a memory pool for a specific device
50    #[cfg(feature = "gpu")]
51    pub fn get_pool(&self, device: Device) -> crate::Result<Arc<MemoryPool>> {
52        let pools = self.pools.read().map_err(|_| {
53            TensorError::invalid_operation_simple("manager read lock poisoned".to_string())
54        })?;
55
56        if let Some(pool) = pools.get(&device) {
57            return Ok(Arc::clone(pool));
58        }
59
60        drop(pools);
61
62        // Create new pool
63        let pool = match device {
64            Device::Gpu(device_id) => MemoryPool::new(device_id, self.default_pool_size)?,
65            Device::Cpu => {
66                return Err(TensorError::unsupported_operation_simple(
67                    "CPU memory pools not yet implemented".to_string(),
68                ))
69            }
70            #[cfg(feature = "rocm")]
71            Device::Rocm(_) => {
72                return Err(TensorError::unsupported_operation_simple(
73                    "ROCm memory pools not yet implemented".to_string(),
74                ))
75            }
76        };
77        let pool = Arc::new(pool);
78
79        let mut pools = self.pools.write().map_err(|_| {
80            TensorError::invalid_operation_simple("manager write lock poisoned".to_string())
81        })?;
82        pools.insert(device, Arc::clone(&pool));
83
84        Ok(pool)
85    }
86
87    /// Get or create a multi-stream memory manager for a device
88    #[cfg(feature = "gpu")]
89    pub fn get_multi_stream_manager(
90        &self,
91        device_id: usize,
92        num_streams: usize,
93    ) -> crate::Result<Arc<MultiStreamMemoryManager>> {
94        let managers = self.multi_stream_managers.read().map_err(|_| {
95            TensorError::invalid_operation_simple("manager read lock poisoned".to_string())
96        })?;
97
98        if let Some(manager) = managers.get(&device_id) {
99            return Ok(Arc::clone(manager));
100        }
101
102        drop(managers);
103
104        // Create new multi-stream manager
105        let stream_pool_size = self.default_pool_size / num_streams;
106        let manager = MultiStreamMemoryManager::new(device_id, num_streams, stream_pool_size)?;
107        let manager = Arc::new(manager);
108
109        let mut managers = self.multi_stream_managers.write().map_err(|_| {
110            TensorError::invalid_operation_simple("manager write lock poisoned".to_string())
111        })?;
112        managers.insert(device_id, Arc::clone(&manager));
113
114        Ok(manager)
115    }
116
117    /// Check for memory aliasing between tensor views
118    pub fn check_memory_alias(&self, buffer_id: usize, offset: usize, size: usize) -> bool {
119        self.alias_detector.check_alias(buffer_id, offset, size)
120    }
121
122    /// Register a new memory view for alias detection
123    pub fn register_memory_view(&self, buffer_id: usize, offset: usize, size: usize) {
124        self.alias_detector.register_view(buffer_id, offset, size);
125    }
126
127    /// Unregister a memory view
128    pub fn unregister_memory_view(&self, buffer_id: usize, offset: usize, size: usize) {
129        self.alias_detector.unregister_view(buffer_id, offset, size);
130    }
131
132    /// Get optimal memory access pattern for given dimensions
133    pub fn get_optimal_access_pattern(&self, dims: &[usize], element_size: usize) -> AccessPattern {
134        self.cache_optimizer
135            .optimize_access_pattern(dims, element_size)
136    }
137
138    /// Get optimal alignment for data of given size
139    pub fn get_optimal_alignment(&self, data_size: usize) -> usize {
140        self.cache_optimizer.get_optimal_alignment(data_size)
141    }
142
143    /// Record a memory operation for performance tracking
144    pub fn record_memory_operation(&self, operation: &str, size: usize) {
145        self.performance_monitor.record_allocation(operation, size);
146    }
147
148    /// Get comprehensive memory statistics across all devices
149    pub fn get_memory_statistics(&self) -> MemoryStatistics {
150        let mut stats = MemoryStatistics::new();
151
152        // Aggregate pool statistics
153        let pools = self
154            .pools
155            .read()
156            .unwrap_or_else(|poisoned| poisoned.into_inner());
157        for (device, pool) in pools.iter() {
158            let pool_stats = pool.stats();
159            stats.add_device_stats(*device, pool_stats);
160        }
161
162        // Add performance monitor data
163        stats.total_allocations = self.performance_monitor.get_allocation_stats().0;
164        stats.total_deallocations = self.performance_monitor.get_allocation_stats().1;
165        stats.current_memory_tracked = self.performance_monitor.get_current_memory();
166        stats.peak_memory_tracked = self.performance_monitor.get_peak_memory();
167
168        // Add alias detector statistics
169        let (alias_buffers, alias_views) = self.alias_detector.get_alias_statistics();
170        stats.active_alias_buffers = alias_buffers;
171        stats.active_alias_views = alias_views;
172
173        stats
174    }
175
176    /// Generate comprehensive memory management report
177    pub fn generate_report(&self) -> String {
178        let mut report = String::new();
179        report.push_str("=== Memory Manager Report ===\n\n");
180
181        // Performance monitoring
182        report.push_str("Performance Monitoring:\n");
183        let perf_report = self.performance_monitor.generate_report();
184        report.push_str(&perf_report);
185        report.push('\n');
186
187        // Memory pools
188        report.push_str("Memory Pools:\n");
189        let pools = self
190            .pools
191            .read()
192            .unwrap_or_else(|poisoned| poisoned.into_inner());
193        for (device, pool) in pools.iter() {
194            let stats = pool.stats();
195            report.push_str(&format!("  Device {:?}:\n", device));
196            report.push_str(&format!("    Allocated: {} bytes\n", stats.total_allocated));
197            report.push_str(&format!("    Free: {} bytes\n", stats.total_free));
198            report.push_str(&format!(
199                "    Fragmentation: {:.2}%\n",
200                stats.fragmentation_ratio * 100.0
201            ));
202            report.push_str(&format!(
203                "    Memory Pressure: {:.2}%\n",
204                stats.memory_pressure * 100.0
205            ));
206        }
207        report.push('\n');
208
209        // Multi-stream managers
210        report.push_str("Multi-Stream Managers:\n");
211        let managers = self
212            .multi_stream_managers
213            .read()
214            .unwrap_or_else(|poisoned| poisoned.into_inner());
215        for (device_id, manager) in managers.iter() {
216            report.push_str(&format!("  Device {}:\n", device_id));
217            report.push_str(&format!("    Streams: {}\n", manager.num_streams()));
218
219            let (total_allocated, total_free) = manager.total_memory_usage();
220            report.push_str(&format!("    Total Allocated: {} bytes\n", total_allocated));
221            report.push_str(&format!("    Total Free: {} bytes\n", total_free));
222        }
223        report.push('\n');
224
225        // Alias detection
226        let (alias_buffers, alias_views) = self.alias_detector.get_alias_statistics();
227        report.push_str("Memory Aliasing:\n");
228        report.push_str(&format!("  Active Buffers: {}\n", alias_buffers));
229        report.push_str(&format!("  Active Views: {}\n", alias_views));
230
231        report
232    }
233
234    /// Optimize memory layout for tensor operations
235    pub fn optimize_tensor_layout(
236        &self,
237        shape: &[usize],
238        element_size: usize,
239    ) -> TensorLayoutOptimization {
240        let access_pattern = self
241            .cache_optimizer
242            .optimize_access_pattern(shape, element_size);
243        let optimal_alignment = self
244            .cache_optimizer
245            .get_optimal_alignment(shape.iter().product::<usize>() * element_size);
246
247        TensorLayoutOptimization {
248            access_pattern,
249            alignment: optimal_alignment,
250            block_size: self
251                .cache_optimizer
252                .get_optimal_block_size(element_size, shape.iter().product()),
253        }
254    }
255
256    /// Create an optimized strided view for zero-copy operations
257    pub fn create_strided_view(
258        &self,
259        offset: usize,
260        shape: Vec<usize>,
261        strides: Vec<usize>,
262        element_size: usize,
263    ) -> StridedView {
264        StridedView::new(offset, shape, strides, element_size)
265    }
266
267    /// Set the default pool size for new memory pools
268    pub fn set_default_pool_size(&mut self, size: usize) {
269        self.default_pool_size = size;
270    }
271
272    /// Get the current default pool size
273    pub fn default_pool_size(&self) -> usize {
274        self.default_pool_size
275    }
276
277    /// Clear all memory pools and reset the manager
278    pub fn clear(&self) {
279        let mut pools = self
280            .pools
281            .write()
282            .unwrap_or_else(|poisoned| poisoned.into_inner());
283        pools.clear();
284
285        let mut managers = self
286            .multi_stream_managers
287            .write()
288            .unwrap_or_else(|poisoned| poisoned.into_inner());
289        managers.clear();
290
291        self.performance_monitor.clear();
292    }
293
294    /// Get cache optimizer reference
295    pub fn cache_optimizer(&self) -> &CacheOptimizer {
296        &self.cache_optimizer
297    }
298
299    /// Get performance monitor reference
300    pub fn performance_monitor(&self) -> &PerformanceMonitor {
301        &self.performance_monitor
302    }
303
304    /// Get alias detector reference
305    pub fn alias_detector(&self) -> &MemoryAliasDetector {
306        &self.alias_detector
307    }
308}
309
310impl Default for MemoryManager {
311    fn default() -> Self {
312        Self::new()
313    }
314}
315
316/// Comprehensive memory statistics across all subsystems
317#[derive(Debug, Clone)]
318pub struct MemoryStatistics {
319    pub device_stats: HashMap<Device, MemoryPoolStats>,
320    pub total_allocations: usize,
321    pub total_deallocations: usize,
322    pub current_memory_tracked: usize,
323    pub peak_memory_tracked: usize,
324    pub active_alias_buffers: usize,
325    pub active_alias_views: usize,
326}
327
328impl MemoryStatistics {
329    fn new() -> Self {
330        Self {
331            device_stats: HashMap::new(),
332            total_allocations: 0,
333            total_deallocations: 0,
334            current_memory_tracked: 0,
335            peak_memory_tracked: 0,
336            active_alias_buffers: 0,
337            active_alias_views: 0,
338        }
339    }
340
341    fn add_device_stats(&mut self, device: Device, stats: MemoryPoolStats) {
342        self.device_stats.insert(device, stats);
343    }
344
345    /// Get total memory allocated across all devices
346    pub fn total_allocated(&self) -> usize {
347        self.device_stats.values().map(|s| s.total_allocated).sum()
348    }
349
350    /// Get total memory free across all devices
351    pub fn total_free(&self) -> usize {
352        self.device_stats.values().map(|s| s.total_free).sum()
353    }
354
355    /// Get average fragmentation ratio across all devices
356    pub fn average_fragmentation(&self) -> f32 {
357        if self.device_stats.is_empty() {
358            return 0.0;
359        }
360
361        let total_fragmentation: f32 = self
362            .device_stats
363            .values()
364            .map(|s| s.fragmentation_ratio)
365            .sum();
366
367        total_fragmentation / self.device_stats.len() as f32
368    }
369
370    /// Get maximum memory pressure across all devices
371    pub fn max_memory_pressure(&self) -> f32 {
372        self.device_stats
373            .values()
374            .map(|s| s.memory_pressure)
375            .fold(0.0, f32::max)
376    }
377}
378
379/// Tensor layout optimization recommendations
380#[derive(Debug, Clone)]
381pub struct TensorLayoutOptimization {
382    pub access_pattern: AccessPattern,
383    pub alignment: usize,
384    pub block_size: usize,
385}
386
387/// Global memory manager instance
388static GLOBAL_MEMORY_MANAGER: std::sync::OnceLock<Arc<MemoryManager>> = std::sync::OnceLock::new();
389
390/// Get the global memory manager
391pub fn global_memory_manager() -> &'static MemoryManager {
392    GLOBAL_MEMORY_MANAGER.get_or_init(|| Arc::new(MemoryManager::new()))
393}
394
395/// Get the global memory manager as Arc
396pub fn global_memory_manager_arc() -> Arc<MemoryManager> {
397    GLOBAL_MEMORY_MANAGER
398        .get_or_init(|| Arc::new(MemoryManager::new()))
399        .clone()
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405    use crate::Device;
406
407    #[test]
408    fn test_memory_manager_creation() {
409        let manager = MemoryManager::new();
410        assert_eq!(manager.default_pool_size(), 512 * 1024 * 1024);
411
412        let custom_manager = MemoryManager::with_pool_size(1024 * 1024);
413        assert_eq!(custom_manager.default_pool_size(), 1024 * 1024);
414    }
415
416    #[test]
417    fn test_memory_statistics() {
418        let mut stats = MemoryStatistics::new();
419        assert_eq!(stats.total_allocated(), 0);
420        assert_eq!(stats.total_free(), 0);
421        assert_eq!(stats.average_fragmentation(), 0.0);
422        assert_eq!(stats.max_memory_pressure(), 0.0);
423
424        // Add some mock device stats
425        let device_stats = crate::memory::pools::MemoryPoolStats {
426            total_allocated: 1000,
427            total_free: 2000,
428            blocks_allocated: 10,
429            blocks_free: 5,
430            fragmentation_ratio: 0.1,
431            peak_allocated: 1500,
432            allocation_count: 100,
433            deallocation_count: 90,
434            defragmentation_count: 2,
435            largest_free_block: 1000,
436            average_block_size: 200.0,
437            memory_pressure: 0.3,
438        };
439
440        stats.add_device_stats(Device::Cpu, device_stats.clone());
441
442        assert_eq!(stats.total_allocated(), 1000);
443        assert_eq!(stats.total_free(), 2000);
444        assert_eq!(stats.average_fragmentation(), 0.1);
445        assert_eq!(stats.max_memory_pressure(), 0.3);
446    }
447
448    #[test]
449    fn test_tensor_layout_optimization() {
450        let manager = MemoryManager::new();
451        let optimization = manager.optimize_tensor_layout(&[100, 100], 4);
452
453        // Should have reasonable values
454        assert!(optimization.alignment > 0);
455        assert!(optimization.block_size > 0);
456        // access_pattern should be one of the variants
457        matches!(
458            optimization.access_pattern,
459            AccessPattern::Sequential | AccessPattern::Blocked { .. } | AccessPattern::Tiled { .. }
460        );
461    }
462
463    #[test]
464    fn test_strided_view_creation() {
465        let manager = MemoryManager::new();
466        let view = manager.create_strided_view(0, vec![2, 3], vec![12, 4], 4);
467
468        assert_eq!(view.offset, 0);
469        assert_eq!(view.shape, vec![2, 3]);
470        assert_eq!(view.strides, vec![12, 4]);
471        assert_eq!(view.element_size, 4);
472    }
473
474    #[test]
475    fn test_memory_alias_operations() {
476        let manager = MemoryManager::new();
477
478        // Register a view
479        manager.register_memory_view(0, 0, 100);
480
481        // Check for alias
482        assert!(manager.check_memory_alias(0, 50, 100)); // Should overlap
483        assert!(!manager.check_memory_alias(0, 100, 50)); // Should not overlap
484
485        // Unregister
486        manager.unregister_memory_view(0, 0, 100);
487        assert!(!manager.check_memory_alias(0, 50, 100)); // No longer aliased
488    }
489
490    #[test]
491    fn test_memory_operation_recording() {
492        let manager = MemoryManager::new();
493
494        // Record initial state (global monitor may have previous operations)
495        let initial_stats = manager.get_memory_statistics();
496        let initial_allocations = initial_stats.total_allocations;
497        let initial_memory = initial_stats.current_memory_tracked;
498
499        manager.record_memory_operation("test_alloc", 1024);
500
501        // Check that exactly one operation was recorded
502        let stats = manager.get_memory_statistics();
503        assert_eq!(stats.total_allocations, initial_allocations + 1);
504        assert_eq!(stats.current_memory_tracked, initial_memory + 1024);
505    }
506
507    #[test]
508    fn test_optimal_access_pattern() {
509        let manager = MemoryManager::new();
510
511        // Small tensor should get sequential access
512        let pattern = manager.get_optimal_access_pattern(&[10, 10], 4);
513        matches!(pattern, AccessPattern::Sequential);
514
515        // Large tensor should get more complex access pattern
516        let pattern = manager.get_optimal_access_pattern(&[1000, 1000], 4);
517        matches!(
518            pattern,
519            AccessPattern::Blocked { .. } | AccessPattern::Tiled { .. }
520        );
521    }
522
523    #[test]
524    fn test_optimal_alignment() {
525        let manager = MemoryManager::new();
526
527        let small_alignment = manager.get_optimal_alignment(32);
528        let large_alignment = manager.get_optimal_alignment(8192);
529
530        assert!(small_alignment > 0);
531        assert!(large_alignment > 0);
532        assert!(large_alignment >= small_alignment); // Larger data should have larger or equal alignment
533    }
534
535    #[test]
536    fn test_manager_clear() {
537        let manager = MemoryManager::new();
538
539        manager.record_memory_operation("test", 1024);
540        manager.register_memory_view(0, 0, 100);
541
542        let stats_before = manager.get_memory_statistics();
543        assert!(stats_before.total_allocations > 0);
544
545        manager.clear();
546
547        let stats_after = manager.get_memory_statistics();
548        assert_eq!(stats_after.total_allocations, 0);
549        assert_eq!(stats_after.current_memory_tracked, 0);
550    }
551
552    #[test]
553    fn test_global_memory_manager() {
554        let manager1 = global_memory_manager();
555        let manager2 = global_memory_manager();
556
557        // Should be the same instance
558        assert!(std::ptr::eq(manager1, manager2));
559
560        // Test that we can use it - check relative change instead of absolute value
561        let initial_stats = manager1.get_memory_statistics();
562        let initial_tracked = initial_stats.current_memory_tracked;
563
564        manager1.record_memory_operation("global_test", 512);
565        let final_stats = manager2.get_memory_statistics();
566        let final_tracked = final_stats.current_memory_tracked;
567
568        // Check that the memory tracked increased by at least 512
569        // (may be more due to other concurrent operations)
570        assert!(final_tracked >= initial_tracked + 512);
571    }
572
573    #[test]
574    fn test_report_generation() {
575        let manager = MemoryManager::new();
576        manager.record_memory_operation("test_op", 1024);
577
578        let report = manager.generate_report();
579        assert!(report.contains("Memory Manager Report"));
580        assert!(report.contains("Performance Monitoring"));
581        assert!(report.contains("Memory Pools"));
582        assert!(report.contains("Memory Aliasing"));
583    }
584
585    #[test]
586    fn test_set_default_pool_size() {
587        let mut manager = MemoryManager::new();
588        let original_size = manager.default_pool_size();
589
590        manager.set_default_pool_size(1024 * 1024);
591        assert_eq!(manager.default_pool_size(), 1024 * 1024);
592        assert_ne!(manager.default_pool_size(), original_size);
593    }
594
595    // Verifies that `get_pool` / `get_multi_stream_manager` return shared
596    // `Arc` handles to the SAME underlying object on a cache hit, instead of
597    // the former "sharing not yet implemented" placeholder error. A real GPU
598    // adapter is not guaranteed to be present in every environment that
599    // builds with `--features gpu` (e.g. a headless CI runner); `MemoryPool`
600    // and `MultiStreamMemoryManager` construction surface adapter/device
601    // creation failures as an honest `Err` rather than panicking, so each
602    // test attempts the GPU operation and skips its assertions - without
603    // failing the suite - if no adapter is available. This mirrors the
604    // convention used by `ops::einsum::gpu::gpu_delegate_tests`.
605    //
606    // These tests use a small custom pool size (rather than
607    // `MemoryManager`'s 512MB default) because `MemoryPool::new` eagerly
608    // allocates a single `wgpu::Buffer` of the exact requested pool size,
609    // and some adapters cap `max_buffer_size` well below 512MB (e.g. the
610    // wgpu downlevel default of 256MB); a too-large request surfaces as a
611    // hard wgpu validation panic rather than a catchable `Result::Err`, so
612    // it isn't covered by the skip-guard above. A few MB is comfortably
613    // within limits on any real adapter these tests may find.
614    #[cfg(feature = "gpu")]
615    mod gpu_tests {
616        use super::*;
617
618        /// Small enough to stay under every real adapter's `max_buffer_size`.
619        const TEST_POOL_SIZE: usize = 4 * 1024 * 1024;
620
621        #[test]
622        fn get_pool_returns_same_arc_on_cache_hit() {
623            let manager = MemoryManager::with_pool_size(TEST_POOL_SIZE);
624
625            let pool1 = match manager.get_pool(Device::Gpu(0)) {
626                Ok(pool) => pool,
627                Err(_) => return, // No GPU adapter available in this environment; skip.
628            };
629            let pool2 = manager
630                .get_pool(Device::Gpu(0))
631                .expect("test: second get_pool call should hit the cache and succeed");
632
633            assert!(
634                Arc::ptr_eq(&pool1, &pool2),
635                "get_pool should return a handle to the same pool on a cache hit"
636            );
637        }
638
639        #[test]
640        fn get_multi_stream_manager_returns_same_arc_on_cache_hit() {
641            let manager = MemoryManager::with_pool_size(TEST_POOL_SIZE);
642
643            let stream_manager1 = match manager.get_multi_stream_manager(0, 2) {
644                Ok(stream_manager) => stream_manager,
645                Err(_) => return, // No GPU adapter available in this environment; skip.
646            };
647            let stream_manager2 = manager.get_multi_stream_manager(0, 2).expect(
648                "test: second get_multi_stream_manager call should hit the cache and succeed",
649            );
650
651            assert!(
652                Arc::ptr_eq(&stream_manager1, &stream_manager2),
653                "get_multi_stream_manager should return a handle to the same manager on a cache hit"
654            );
655        }
656    }
657}
tenflowers_core/memory/manager.rs

tenflowers_core/memory/
manager.rs
