1#![allow(dead_code)]
17use std::alloc::{GlobalAlloc, Layout, System};
18use std::collections::{BTreeMap, HashMap, VecDeque};
19use std::mem::{align_of, size_of};
20use std::ptr::NonNull;
21use std::sync::{Arc, Mutex, RwLock};
22use std::time::{Duration, Instant};
23use torsh_core::sync::{MutexExt, RwLockExt};
24
25use torsh_core::{
27 dtype::TensorElement,
28 error::{Result, TorshError},
29};
30
31#[derive(Debug, Clone)]
33pub struct MemoryConfig {
34 pub enable_pooling: bool,
36 pub pool_size: usize,
38 pub max_cached_per_size: usize,
40 pub enable_compression: bool,
42 pub compression_threshold: usize,
44 pub enable_numa_awareness: bool,
46 pub cache_line_size: usize,
48 pub enable_predictive_allocation: bool,
50 pub memory_pressure_threshold: f64,
52}
53
54impl Default for MemoryConfig {
55 fn default() -> Self {
56 Self {
57 enable_pooling: true,
58 pool_size: 1024 * 1024 * 1024, max_cached_per_size: 64,
60 enable_compression: true,
61 compression_threshold: 100 * 1024 * 1024, enable_numa_awareness: false, cache_line_size: 64,
64 enable_predictive_allocation: true,
65 memory_pressure_threshold: 0.8,
66 }
67 }
68}
69
70pub struct AdvancedMemoryPool<T: TensorElement> {
72 config: MemoryConfig,
73 size_class_pools: RwLock<BTreeMap<usize, VecDeque<NonNull<T>>>>,
75 stats: RwLock<MemoryStats>,
77 allocation_history: Mutex<VecDeque<AllocationRecord>>,
79 predictor: Mutex<Option<AllocationPredictor>>,
81 compression_manager: Arc<CompressionManager>,
83 numa_allocators: Vec<Arc<Mutex<NumaAllocator>>>,
85}
86
87impl<T: TensorElement> AdvancedMemoryPool<T> {
88 pub fn new() -> Self {
90 Self::with_config(MemoryConfig::default())
91 }
92
93 pub fn with_config(config: MemoryConfig) -> Self {
95 let numa_nodes = if config.enable_numa_awareness {
96 detect_numa_nodes()
97 } else {
98 1
99 };
100
101 let numa_allocators = (0..numa_nodes)
102 .map(|node_id| Arc::new(Mutex::new(NumaAllocator::new(node_id))))
103 .collect();
104
105 Self {
106 config,
107 size_class_pools: RwLock::new(BTreeMap::new()),
108 stats: RwLock::new(MemoryStats::default()),
109 allocation_history: Mutex::new(VecDeque::with_capacity(10000)),
110 predictor: Mutex::new(None),
111 compression_manager: Arc::new(CompressionManager::new()),
112 numa_allocators,
113 }
114 }
115
116 pub fn allocate(&self, size: usize) -> Result<NonNull<T>> {
118 #[cfg(feature = "profiling")]
119 {
120 }
122 let aligned_size = self.align_size(size);
123
124 if self.config.enable_compression && size > self.config.compression_threshold {
126 return self.allocate_compressed(aligned_size);
127 }
128
129 if let Some(ptr) = self.try_reuse_from_pool(aligned_size)? {
131 self.record_allocation(aligned_size, true);
132 return Ok(ptr);
133 }
134
135 if self.config.enable_predictive_allocation {
137 self.maybe_predictive_allocate(aligned_size)?;
138 }
139
140 let ptr = self.allocate_new(aligned_size)?;
142 self.record_allocation(aligned_size, false);
143
144 Ok(ptr)
145 }
146
147 pub fn deallocate(&self, ptr: NonNull<T>, size: usize) -> Result<()> {
149 #[cfg(feature = "profiling")]
150 {
151 }
153 let aligned_size = self.align_size(size);
154
155 if self.compression_manager.is_compressed(ptr) {
157 return self.compression_manager.deallocate(ptr);
158 }
159
160 if self.should_cache_allocation(aligned_size) {
162 let mut pools = self.size_class_pools.write_or_recover();
163 let pool = pools.entry(aligned_size).or_insert_with(VecDeque::new);
164
165 if pool.len() < self.config.max_cached_per_size {
166 pool.push_back(ptr);
167 self.update_stats(|stats| stats.pooled_allocations += 1);
168 return Ok(());
169 }
170 }
171
172 self.free_allocation(ptr, aligned_size)?;
174 Ok(())
175 }
176
177 fn try_reuse_from_pool(&self, size: usize) -> Result<Option<NonNull<T>>> {
179 if !self.config.enable_pooling {
180 return Ok(None);
181 }
182
183 let mut pools = self.size_class_pools.write_or_recover();
184
185 if let Some(pool) = pools.get_mut(&size) {
187 if let Some(ptr) = pool.pop_front() {
188 self.update_stats(|stats| stats.pool_hits += 1);
189 return Ok(Some(ptr));
190 }
191 }
192
193 let max_oversized = size * 2; for (&pool_size, pool) in pools.range_mut(size..).take(5) {
197 if pool_size > max_oversized {
198 break;
199 }
200
201 if let Some(ptr) = pool.pop_front() {
202 self.update_stats(|stats| {
203 stats.pool_hits += 1;
204 stats.oversized_reuse += 1;
205 });
206 return Ok(Some(ptr));
207 }
208 }
209
210 self.update_stats(|stats| stats.pool_misses += 1);
211 Ok(None)
212 }
213
214 fn allocate_new(&self, size: usize) -> Result<NonNull<T>> {
216 let layout = Layout::from_size_align(
217 size * size_of::<T>(),
218 align_of::<T>().max(self.config.cache_line_size),
219 )
220 .map_err(|_| TorshError::InvalidArgument("Invalid memory layout".to_string()))?;
221
222 if self.config.enable_numa_awareness && !self.numa_allocators.is_empty() {
224 let numa_node = self.select_numa_node();
225 let allocator = &self.numa_allocators[numa_node];
226 let mut allocator = allocator.lock_or_recover();
227 return allocator.allocate(layout);
228 }
229
230 unsafe {
232 let ptr = System.alloc(layout);
233 if ptr.is_null() {
234 return Err(TorshError::AllocationError(
235 "Failed to allocate memory".to_string(),
236 ));
237 }
238
239 self.prefault_pages(ptr, layout.size());
241
242 Ok(NonNull::new_unchecked(ptr as *mut T))
243 }
244 }
245
246 fn allocate_compressed(&self, size: usize) -> Result<NonNull<T>> {
248 self.compression_manager.allocate_compressed(size)
249 }
250
251 fn free_allocation(&self, ptr: NonNull<T>, size: usize) -> Result<()> {
253 let layout = Layout::from_size_align(
254 size * size_of::<T>(),
255 align_of::<T>().max(self.config.cache_line_size),
256 )
257 .map_err(|_| TorshError::InvalidArgument("Invalid memory layout".to_string()))?;
258
259 unsafe {
260 System.dealloc(ptr.as_ptr() as *mut u8, layout);
261 }
262
263 self.update_stats(|stats| stats.direct_deallocations += 1);
264 Ok(())
265 }
266
267 fn maybe_predictive_allocate(&self, size: usize) -> Result<()> {
269 let mut predictor_guard = self.predictor.lock_or_recover();
270
271 if predictor_guard.is_none() {
272 *predictor_guard = Some(AllocationPredictor::new());
273 }
274
275 if let Some(predictor) = predictor_guard.as_mut() {
276 if let Some(predicted_sizes) = predictor.predict_next_allocations(size) {
277 for predicted_size in predicted_sizes {
279 if predicted_size != size && predicted_size > 0 {
280 let _ = self.allocate_new(predicted_size);
282 }
283 }
284 }
285 }
286
287 Ok(())
288 }
289
290 fn align_size(&self, size: usize) -> usize {
292 let cache_line = self.config.cache_line_size;
293 ((size + cache_line - 1) / cache_line) * cache_line
294 }
295
296 fn should_cache_allocation(&self, size: usize) -> bool {
298 self.config.enable_pooling &&
299 size <= self.config.pool_size / 100 && !self.is_memory_pressure_high()
301 }
302
303 fn is_memory_pressure_high(&self) -> bool {
305 let stats = self.stats.read_or_recover();
307 let total_allocations = stats.pool_hits + stats.pool_misses + stats.direct_allocations;
308
309 if total_allocations == 0 {
310 return false;
311 }
312
313 let cache_hit_rate = stats.pool_hits as f64 / total_allocations as f64;
314 cache_hit_rate < (1.0 - self.config.memory_pressure_threshold)
315 }
316
317 fn prefault_pages(&self, ptr: *mut u8, size: usize) {
319 const PAGE_SIZE: usize = 4096;
320 let page_count = (size + PAGE_SIZE - 1) / PAGE_SIZE;
321
322 unsafe {
323 for i in 0..page_count {
324 let page_ptr = ptr.add(i * PAGE_SIZE);
325 std::ptr::write_volatile(page_ptr, 0);
326 }
327 }
328 }
329
330 fn select_numa_node(&self) -> usize {
332 let stats = self.stats.read_or_recover();
334 (stats.total_allocations % self.numa_allocators.len()) as usize
335 }
336
337 fn record_allocation(&self, size: usize, was_reused: bool) {
339 let record = AllocationRecord {
340 size,
341 timestamp: Instant::now(),
342 was_reused,
343 };
344
345 let mut history = self.allocation_history.lock_or_recover();
346 history.push_back(record);
347
348 if history.len() > 10000 {
350 history.pop_front();
351 }
352
353 self.update_stats(|stats| {
354 stats.total_allocations += 1;
355 if was_reused {
356 stats.reused_allocations += 1;
357 } else {
358 stats.direct_allocations += 1;
359 }
360 });
361 }
362
363 fn update_stats<F>(&self, f: F)
365 where
366 F: FnOnce(&mut MemoryStats),
367 {
368 let mut stats = self.stats.write_or_recover();
369 f(&mut *stats);
370 }
371
372 pub fn get_stats(&self) -> MemoryStats {
374 self.stats.read_or_recover().clone()
375 }
376
377 pub fn defragment(&self) -> Result<DefragmentationReport> {
379 #[cfg(feature = "profiling")]
380 {
381 }
383 let start_time = Instant::now();
384 let mut report = DefragmentationReport::default();
385
386 {
388 let mut pools = self.size_class_pools.write_or_recover();
389 let initial_pools = pools.len();
390 pools.retain(|_, pool| !pool.is_empty());
391 report.pools_cleaned = initial_pools - pools.len();
392 }
393
394 if self.config.enable_compression {
396 report.compression_stats = self.compression_manager.compress_fragmented()?;
397 }
398
399 report.duration = start_time.elapsed();
401 report.memory_freed = self.estimate_memory_freed();
402
403 Ok(report)
404 }
405
406 fn estimate_memory_freed(&self) -> usize {
408 let stats = self.stats.read_or_recover();
410 stats
411 .total_allocations
412 .saturating_sub(stats.reused_allocations)
413 * 1024 }
415}
416
417impl<T: TensorElement> Default for AdvancedMemoryPool<T> {
418 fn default() -> Self {
419 Self::new()
420 }
421}
422
423#[derive(Debug, Clone, Default)]
425pub struct MemoryStats {
426 pub total_allocations: usize,
427 pub direct_allocations: usize,
428 pub reused_allocations: usize,
429 pub pooled_allocations: usize,
430 pub pool_hits: usize,
431 pub pool_misses: usize,
432 pub oversized_reuse: usize,
433 pub direct_deallocations: usize,
434 pub compression_saves: usize,
435 pub numa_allocations: usize,
436}
437
438impl MemoryStats {
439 pub fn hit_rate(&self) -> f64 {
441 let total_pool_requests = self.pool_hits + self.pool_misses;
442 if total_pool_requests == 0 {
443 0.0
444 } else {
445 self.pool_hits as f64 / total_pool_requests as f64
446 }
447 }
448
449 pub fn reuse_rate(&self) -> f64 {
451 if self.total_allocations == 0 {
452 0.0
453 } else {
454 self.reused_allocations as f64 / self.total_allocations as f64
455 }
456 }
457}
458
459#[derive(Debug, Clone)]
461struct AllocationRecord {
462 size: usize,
463 timestamp: Instant,
464 was_reused: bool,
465}
466
467struct AllocationPredictor {
469 size_patterns: HashMap<usize, Vec<usize>>,
470 temporal_patterns: VecDeque<(Instant, usize)>,
471 max_history: usize,
472}
473
474impl AllocationPredictor {
475 fn new() -> Self {
476 Self {
477 size_patterns: HashMap::new(),
478 temporal_patterns: VecDeque::new(),
479 max_history: 1000,
480 }
481 }
482
483 fn predict_next_allocations(&mut self, size: usize) -> Option<Vec<usize>> {
485 self.temporal_patterns.push_back((Instant::now(), size));
487
488 if self.temporal_patterns.len() > self.max_history {
490 self.temporal_patterns.pop_front();
491 }
492
493 if let Some(following_sizes) = self.size_patterns.get(&size) {
495 let mut counts: HashMap<usize, usize> = HashMap::new();
497 for &following_size in following_sizes {
498 *counts.entry(following_size).or_insert(0) += 1;
499 }
500
501 let mut sorted: Vec<_> = counts.into_iter().collect();
502 sorted.sort_by(|a, b| b.1.cmp(&a.1));
503
504 Some(sorted.into_iter().take(3).map(|(size, _)| size).collect())
505 } else {
506 None
507 }
508 }
509}
510
511struct CompressionManager {
513 compressed_allocations: RwLock<HashMap<usize, CompressedAllocation>>,
514}
515
516impl CompressionManager {
517 fn new() -> Self {
518 Self {
519 compressed_allocations: RwLock::new(HashMap::new()),
520 }
521 }
522
523 fn allocate_compressed<T: TensorElement>(&self, size: usize) -> Result<NonNull<T>> {
524 let compressed_size = size / 2; let layout = Layout::from_size_align(compressed_size, align_of::<T>())
528 .map_err(|_| TorshError::InvalidArgument("Invalid layout".to_string()))?;
529
530 unsafe {
531 let ptr = System.alloc(layout);
532 if ptr.is_null() {
533 return Err(TorshError::AllocationError(
534 "Compression allocation failed".to_string(),
535 ));
536 }
537
538 let allocation = CompressedAllocation {
539 original_size: size,
540 compressed_size,
541 compression_ratio: 0.5,
542 };
543
544 self.compressed_allocations
545 .write_or_recover()
546 .insert(ptr as usize, allocation);
547 Ok(NonNull::new_unchecked(ptr as *mut T))
548 }
549 }
550
551 fn is_compressed<T: TensorElement>(&self, ptr: NonNull<T>) -> bool {
552 self.compressed_allocations
553 .read_or_recover()
554 .contains_key(&(ptr.as_ptr() as usize))
555 }
556
557 fn deallocate<T: TensorElement>(&self, ptr: NonNull<T>) -> Result<()> {
558 let ptr_key = ptr.as_ptr() as usize;
559 let mut allocations = self.compressed_allocations.write_or_recover();
560
561 if let Some(allocation) = allocations.remove(&ptr_key) {
562 let layout = Layout::from_size_align(allocation.compressed_size, align_of::<T>())
563 .map_err(|_| TorshError::InvalidArgument("Invalid layout".to_string()))?;
564
565 unsafe {
566 System.dealloc(ptr_key as *mut u8, layout);
567 }
568 Ok(())
569 } else {
570 Err(TorshError::InvalidArgument(
571 "Allocation not found".to_string(),
572 ))
573 }
574 }
575
576 fn compress_fragmented(&self) -> Result<CompressionStats> {
577 Ok(CompressionStats {
579 allocations_compressed: 0,
580 memory_saved: 0,
581 average_compression_ratio: 0.0,
582 })
583 }
584}
585
586#[derive(Debug, Clone)]
588struct CompressedAllocation {
589 original_size: usize,
590 compressed_size: usize,
591 compression_ratio: f64,
592}
593
594struct NumaAllocator {
596 node_id: usize,
597 allocations: usize,
598}
599
600impl NumaAllocator {
601 fn new(node_id: usize) -> Self {
602 Self {
603 node_id,
604 allocations: 0,
605 }
606 }
607
608 fn allocate<T: TensorElement>(&mut self, layout: Layout) -> Result<NonNull<T>> {
609 unsafe {
611 let ptr = System.alloc(layout);
612 if ptr.is_null() {
613 return Err(TorshError::AllocationError(
614 "NUMA allocation failed".to_string(),
615 ));
616 }
617 self.allocations += 1;
618 Ok(NonNull::new_unchecked(ptr as *mut T))
619 }
620 }
621}
622
623#[derive(Debug, Default)]
625pub struct DefragmentationReport {
626 pub duration: Duration,
627 pub pools_cleaned: usize,
628 pub memory_freed: usize,
629 pub compression_stats: CompressionStats,
630}
631
632#[derive(Debug, Default)]
634pub struct CompressionStats {
635 pub allocations_compressed: usize,
636 pub memory_saved: usize,
637 pub average_compression_ratio: f64,
638}
639
640fn detect_numa_nodes() -> usize {
642 1 }
645
646pub struct GlobalMemoryOptimizer {
648 f32_pool: AdvancedMemoryPool<f32>,
649 f64_pool: AdvancedMemoryPool<f64>,
650 i32_pool: AdvancedMemoryPool<i32>,
651 i64_pool: AdvancedMemoryPool<i64>,
652 config: MemoryConfig,
653}
654
655impl GlobalMemoryOptimizer {
656 pub fn new() -> Self {
658 let config = MemoryConfig::default();
659 Self::with_config(config)
660 }
661
662 pub fn with_config(config: MemoryConfig) -> Self {
664 Self {
665 f32_pool: AdvancedMemoryPool::with_config(config.clone()),
666 f64_pool: AdvancedMemoryPool::with_config(config.clone()),
667 i32_pool: AdvancedMemoryPool::with_config(config.clone()),
668 i64_pool: AdvancedMemoryPool::with_config(config.clone()),
669 config,
670 }
671 }
672
673 pub fn get_pool<T: TensorElement>(&self) -> Option<&AdvancedMemoryPool<T>> {
675 None }
678
679 pub fn global_defragmentation(&self) -> Result<Vec<DefragmentationReport>> {
681 let mut reports = Vec::new();
682
683 reports.push(self.f32_pool.defragment()?);
684 reports.push(self.f64_pool.defragment()?);
685 Ok(reports)
688 }
689
690 pub fn get_aggregate_stats(&self) -> AggregateMemoryStats {
692 AggregateMemoryStats {
693 f32_stats: self.f32_pool.get_stats(),
694 f64_stats: self.f64_pool.get_stats(),
695 i32_stats: self.i32_pool.get_stats(),
696 i64_stats: self.i64_pool.get_stats(),
697 }
698 }
699}
700
701impl Default for GlobalMemoryOptimizer {
702 fn default() -> Self {
703 Self::new()
704 }
705}
706
707#[derive(Debug)]
709pub struct AggregateMemoryStats {
710 pub f32_stats: MemoryStats,
711 pub f64_stats: MemoryStats,
712 pub i32_stats: MemoryStats,
713 pub i64_stats: MemoryStats,
714}
715
716impl AggregateMemoryStats {
717 pub fn overall_hit_rate(&self) -> f64 {
719 let total_hits = self.f32_stats.pool_hits
720 + self.f64_stats.pool_hits
721 + self.i32_stats.pool_hits
722 + self.i64_stats.pool_hits;
723 let total_misses = self.f32_stats.pool_misses
724 + self.f64_stats.pool_misses
725 + self.i32_stats.pool_misses
726 + self.i64_stats.pool_misses;
727
728 let total_requests = total_hits + total_misses;
729 if total_requests == 0 {
730 0.0
731 } else {
732 total_hits as f64 / total_requests as f64
733 }
734 }
735
736 pub fn total_allocations(&self) -> usize {
738 self.f32_stats.total_allocations
739 + self.f64_stats.total_allocations
740 + self.i32_stats.total_allocations
741 + self.i64_stats.total_allocations
742 }
743}
744
745#[cfg(test)]
746mod tests {
747 use super::*;
748 use std::ptr;
749
750 #[test]
751 fn test_memory_config_default() {
752 let config = MemoryConfig::default();
753 assert!(config.enable_pooling);
754 assert!(config.pool_size > 0);
755 assert!(config.cache_line_size > 0);
756 }
757
758 #[test]
759 fn test_advanced_memory_pool_creation() {
760 let pool: AdvancedMemoryPool<f32> = AdvancedMemoryPool::new();
761 let stats = pool.get_stats();
762
763 assert_eq!(stats.total_allocations, 0);
764 assert_eq!(stats.pool_hits, 0);
765 assert_eq!(stats.pool_misses, 0);
766 }
767
768 #[test]
769 fn test_memory_allocation_and_deallocation() {
770 let pool: AdvancedMemoryPool<f32> = AdvancedMemoryPool::new();
771
772 let ptr = pool.allocate(1024).expect("allocation should succeed");
774 pool.deallocate(ptr, 1024)
778 .expect("deallocation should succeed");
779
780 let stats = pool.get_stats();
781 assert_eq!(stats.total_allocations, 1);
782 }
783
784 #[test]
785 fn test_memory_pool_reuse() {
786 let pool: AdvancedMemoryPool<f32> = AdvancedMemoryPool::new();
787
788 let ptr1 = pool.allocate(1024).expect("allocation should succeed");
790 pool.deallocate(ptr1, 1024)
791 .expect("deallocation should succeed");
792
793 let ptr2 = pool.allocate(1024).expect("allocation should succeed");
795 pool.deallocate(ptr2, 1024)
796 .expect("deallocation should succeed");
797
798 let stats = pool.get_stats();
799 assert_eq!(stats.total_allocations, 2);
800 }
802
803 #[test]
804 fn test_memory_stats_calculations() {
805 let mut stats = MemoryStats::default();
806 stats.pool_hits = 80;
807 stats.pool_misses = 20;
808 stats.total_allocations = 100;
809 stats.reused_allocations = 80;
810
811 assert_eq!(stats.hit_rate(), 0.8);
812 assert_eq!(stats.reuse_rate(), 0.8);
813 }
814
815 #[test]
816 fn test_size_alignment() {
817 let pool: AdvancedMemoryPool<f32> = AdvancedMemoryPool::with_config(MemoryConfig {
818 cache_line_size: 64,
819 ..Default::default()
820 });
821
822 assert_eq!(pool.align_size(1), 64);
823 assert_eq!(pool.align_size(65), 128);
824 assert_eq!(pool.align_size(128), 128);
825 }
826
827 #[test]
828 fn test_defragmentation() {
829 let pool: AdvancedMemoryPool<f32> = AdvancedMemoryPool::new();
830
831 for i in 0..10 {
833 let ptr = pool
834 .allocate(1024 * (i + 1))
835 .expect("allocation should succeed");
836 pool.deallocate(ptr, 1024 * (i + 1))
837 .expect("deallocation should succeed");
838 }
839
840 let report = pool.defragment().expect("defragmentation should succeed");
841 let _ = report.duration; }
846
847 #[test]
848 fn test_global_memory_optimizer() {
849 let optimizer = GlobalMemoryOptimizer::new();
850 let stats = optimizer.get_aggregate_stats();
851
852 assert_eq!(stats.total_allocations(), 0);
853 assert_eq!(stats.overall_hit_rate(), 0.0);
854 }
855
856 #[test]
857 fn test_compression_manager() {
858 let manager = CompressionManager::new();
859 let ptr = NonNull::new(ptr::null_mut::<f32>().wrapping_add(0x1000))
860 .expect("pointer should be non-null");
861
862 assert!(!manager.is_compressed(ptr));
863 }
864
865 #[test]
866 fn test_allocation_predictor() {
867 let mut predictor = AllocationPredictor::new();
868
869 let predictions = predictor.predict_next_allocations(1024);
871 assert!(predictions.is_none());
872 }
873
874 #[test]
875 fn test_memory_pressure_detection() {
876 let pool: AdvancedMemoryPool<f32> = AdvancedMemoryPool::with_config(MemoryConfig {
877 memory_pressure_threshold: 0.5,
878 ..Default::default()
879 });
880
881 assert!(!pool.is_memory_pressure_high());
883 }
884}