1pub mod allocation;
12pub mod management;
13pub mod vendors;
14
15use std::collections::HashMap;
16use std::ffi::c_void;
17use std::ptr::NonNull;
18use std::sync::{Arc, Mutex};
19use std::time::{Duration, Instant};
20
21pub use allocation::{
23 AllocationStrategy, AllocationStrategyManager, AllocatorType, ArenaAllocator, BuddyAllocator,
24 MemoryPool, SlabAllocator, UnifiedAllocator, UnifiedConfig,
25};
26
27use allocation::strategies::AllocationStats;
28
29pub use management::{
30 AccessType, DefragmentationEngine, EvictionEngine, GarbageCollectionEngine,
31 IntegratedMemoryManager, ManagementStats, MemoryManagementConfig, MemoryManagementError,
32 MemoryRegion, PrefetchingEngine,
33};
34
35use management::eviction_policies::{CacheObject, ObjectPriority, ObjectType, RegionType};
36
37pub use vendors::{
38 CudaConfig, CudaError, CudaMemoryBackend, CudaMemoryType, GpuBackendFactory, GpuVendor,
39 MetalConfig, MetalError, MetalMemoryBackend, MetalMemoryType, OneApiConfig, OneApiError,
40 OneApiMemoryBackend, OneApiMemoryType, RocmConfig, RocmError, RocmMemoryBackend,
41 RocmMemoryType, UnifiedGpuBackend, UnifiedGpuError, UnifiedMemoryStats, VendorConfig,
42};
43
44#[derive(Debug, Clone)]
46pub struct GpuMemorySystemConfig {
47 pub vendor_config: VendorConfig,
49 pub allocation_config: UnifiedConfig,
51 pub management_config: MemoryManagementConfig,
53 pub system_config: SystemConfig,
55}
56
57#[derive(Debug, Clone)]
59pub struct SystemConfig {
60 pub enable_unified_interface: bool,
62 pub enable_cross_vendor_sharing: bool,
64 pub enable_performance_monitoring: bool,
66 pub monitoring_interval: Duration,
68 pub memory_budget: f64,
70 pub enable_auto_optimization: bool,
72 pub optimization_interval: Duration,
74 pub enable_memory_compression: bool,
76 pub thread_pool_size: usize,
78}
79
80impl Default for SystemConfig {
81 fn default() -> Self {
82 Self {
83 enable_unified_interface: true,
84 enable_cross_vendor_sharing: false,
85 enable_performance_monitoring: true,
86 monitoring_interval: Duration::from_millis(500),
87 memory_budget: 0.9,
88 enable_auto_optimization: true,
89 optimization_interval: Duration::from_secs(60),
90 enable_memory_compression: false,
91 thread_pool_size: 4,
92 }
93 }
94}
95
96impl Default for GpuMemorySystemConfig {
97 fn default() -> Self {
98 let vendor = GpuBackendFactory::get_preferred_vendor();
99 Self {
100 vendor_config: GpuBackendFactory::create_default_config(vendor),
101 allocation_config: UnifiedConfig::default(),
102 management_config: MemoryManagementConfig::default(),
103 system_config: SystemConfig::default(),
104 }
105 }
106}
107
108pub struct GpuMemorySystem {
110 gpu_backend: UnifiedGpuBackend,
112 allocation_engine: UnifiedAllocator,
114 memory_manager: IntegratedMemoryManager,
116 config: GpuMemorySystemConfig,
118 stats: SystemStats,
120 memory_regions: HashMap<*mut c_void, MemoryAllocation>,
122 monitoring_enabled: bool,
124 last_optimization: Instant,
126}
127
128#[derive(Debug, Clone)]
130pub struct MemoryAllocation {
131 pub ptr: *mut c_void,
132 pub size: usize,
133 pub allocator_type: AllocatorType,
134 pub vendor_memory_type: String,
135 pub allocated_at: Instant,
136 pub last_accessed: Option<Instant>,
137 pub access_count: u64,
138 pub ref_count: u32,
139}
140
141#[derive(Debug, Clone, Default)]
143pub struct SystemStats {
144 pub total_allocations: u64,
145 pub total_deallocations: u64,
146 pub bytes_allocated: u64,
147 pub bytes_deallocated: u64,
148 pub active_allocations: u64,
149 pub peak_memory_usage: usize,
150 pub fragmentation_ratio: f64,
151 pub allocation_efficiency: f64,
152 pub vendor_stats: UnifiedMemoryStats,
153 pub allocation_stats: AllocationStats,
154 pub management_stats: ManagementStats,
155 pub uptime: Duration,
156 pub optimization_cycles: u64,
157}
158
159impl GpuMemorySystem {
160 pub fn new(config: GpuMemorySystemConfig) -> Result<Self, GpuMemorySystemError> {
162 let mut gpu_backend = UnifiedGpuBackend::new(config.vendor_config.clone())?;
164
165 let mut total_size =
167 (config.system_config.memory_budget * gpu_backend.get_total_memory() as f64) as usize;
168
169 if config.allocation_config.enable_buddy {
171 total_size = total_size.next_power_of_two();
172 }
173
174 let base_ptr = gpu_backend
175 .allocate(total_size)
176 .map_err(GpuMemorySystemError::BackendError)?;
177
178 let allocation_engine = UnifiedAllocator::new(
180 unsafe { NonNull::new_unchecked(base_ptr as *mut u8) },
181 total_size,
182 config.allocation_config.clone(),
183 )
184 .map_err(|e| GpuMemorySystemError::AllocationError(format!("{:?}", e)))?;
185
186 let memory_manager = IntegratedMemoryManager::new(config.management_config.clone());
188
189 Ok(Self {
190 gpu_backend,
191 allocation_engine,
192 memory_manager,
193 config,
194 stats: SystemStats::default(),
195 memory_regions: HashMap::new(),
196 monitoring_enabled: false,
197 last_optimization: Instant::now(),
198 })
199 }
200
201 pub fn auto_create() -> Result<Self, GpuMemorySystemError> {
203 let config = GpuMemorySystemConfig::default();
204 Self::new(config)
205 }
206
207 pub fn start(&mut self) -> Result<(), GpuMemorySystemError> {
209 if self.config.system_config.enable_performance_monitoring {
211 self.memory_manager
212 .start_background_management()
213 .map_err(|e| GpuMemorySystemError::ManagementError(format!("{}", e)))?;
214 self.monitoring_enabled = true;
215 }
216
217 Ok(())
229 }
230
231 pub fn allocate(
233 &mut self,
234 size: usize,
235 alignment: Option<usize>,
236 ) -> Result<*mut c_void, GpuMemorySystemError> {
237 let start_time = Instant::now();
238
239 let allocator_type = self.choose_allocator(size);
241
242 let ptr_nonnull =
244 self.allocation_engine
245 .allocate(size, allocator_type.clone(), alignment)?;
246
247 let ptr = ptr_nonnull.as_ptr() as *mut c_void;
249
250 let allocation = MemoryAllocation {
252 ptr,
253 size,
254 allocator_type,
255 vendor_memory_type: self.get_vendor_memory_type(),
256 allocated_at: Instant::now(),
257 last_accessed: Some(Instant::now()),
258 access_count: 1,
259 ref_count: 1,
260 };
261
262 self.memory_regions.insert(ptr, allocation);
264
265 self.update_allocation_stats(size, start_time.elapsed());
267
268 self.handle_memory_pressure()?;
270
271 Ok(ptr)
272 }
273
274 pub fn free(&mut self, ptr: *mut c_void) -> Result<(), GpuMemorySystemError> {
276 let start_time = Instant::now();
277
278 let allocation = self
280 .memory_regions
281 .remove(&ptr)
282 .ok_or_else(|| GpuMemorySystemError::InvalidPointer("Pointer not found".to_string()))?;
283
284 self.allocation_engine
286 .free(ptr, allocation.allocator_type)?;
287
288 self.update_deallocation_stats(allocation.size, start_time.elapsed());
290
291 Ok(())
292 }
293
294 pub fn reallocate(
296 &mut self,
297 ptr: *mut c_void,
298 new_size: usize,
299 ) -> Result<*mut c_void, GpuMemorySystemError> {
300 let allocation = self
302 .memory_regions
303 .get(&ptr)
304 .ok_or_else(|| GpuMemorySystemError::InvalidPointer("Pointer not found".to_string()))?;
305
306 let old_size = allocation.size;
307 let allocator_type = allocation.allocator_type.clone();
308
309 if let Ok(new_ptr) = self
311 .allocation_engine
312 .reallocate(ptr, new_size, allocator_type)
313 {
314 if new_ptr == ptr {
315 if let Some(allocation) = self.memory_regions.get_mut(&ptr) {
317 allocation.size = new_size;
318 allocation.last_accessed = Some(Instant::now());
319 allocation.access_count += 1;
320 }
321 return Ok(ptr);
322 }
323 }
324
325 let new_ptr = self.allocate(new_size, None)?;
327
328 unsafe {
330 std::ptr::copy_nonoverlapping(
331 ptr as *const u8,
332 new_ptr as *mut u8,
333 old_size.min(new_size),
334 );
335 }
336
337 self.free(ptr)?;
339
340 Ok(new_ptr)
341 }
342
343 pub fn record_access(
345 &mut self,
346 ptr: *mut c_void,
347 access_type: AccessType,
348 ) -> Result<(), GpuMemorySystemError> {
349 if let Some(allocation) = self.memory_regions.get_mut(&ptr) {
350 allocation.last_accessed = Some(Instant::now());
351 allocation.access_count += 1;
352
353 self.memory_manager
355 .update_access_pattern(ptr, allocation.size, access_type)?;
356 }
357
358 Ok(())
359 }
360
361 pub fn get_memory_info(&self, ptr: *mut c_void) -> Option<&MemoryAllocation> {
363 self.memory_regions.get(&ptr)
364 }
365
366 pub fn get_stats(&mut self) -> SystemStats {
368 self.stats.vendor_stats = self.gpu_backend.get_memory_stats();
370
371 let unified_stats = self.allocation_engine.get_stats();
373 self.stats.allocation_stats = AllocationStats {
374 total_allocations: unified_stats.total_allocations,
375 total_deallocations: unified_stats.total_deallocations,
376 cache_hits: unified_stats.routing_cache_hits,
377 cache_misses: unified_stats.routing_decisions - unified_stats.routing_cache_hits,
378 fragmentation_events: 0,
390 total_allocated_bytes: unified_stats.bytes_allocated,
391 peak_allocated_bytes: unified_stats.peak_memory_usage as u64,
392 average_allocation_size: if unified_stats.total_allocations > 0 {
393 (unified_stats.bytes_allocated as f64) / (unified_stats.total_allocations as f64)
394 } else {
395 0.0
396 },
397 allocation_latency_ms: unified_stats.average_allocation_time_ns / 1_000_000.0,
398 };
399
400 self.stats.management_stats = self.memory_manager.get_stats().clone();
402
403 self.calculate_system_metrics();
405
406 self.stats.clone()
407 }
408
409 pub fn optimize(&mut self) -> Result<(), GpuMemorySystemError> {
411 if !self.config.system_config.enable_auto_optimization {
412 return Ok(());
413 }
414
415 let now = Instant::now();
416 if now.duration_since(self.last_optimization)
417 < self.config.system_config.optimization_interval
418 {
419 return Ok(());
420 }
421
422 let memory_regions: HashMap<usize, management::MemoryRegion> = self
424 .memory_regions
425 .iter()
426 .enumerate()
427 .map(|(i, (ptr, alloc))| {
428 let mut objects = HashMap::new();
429 objects.insert(
430 *ptr as usize,
431 CacheObject {
432 address: *ptr as usize,
433 size: alloc.size,
434 created_at: alloc.allocated_at,
435 last_access: alloc.last_accessed.unwrap_or(alloc.allocated_at),
436 access_count: alloc.access_count as u32,
437 access_frequency: alloc.access_count as f64,
438 priority: ObjectPriority::Normal,
439 kernel_context: None,
440 object_type: ObjectType::Data,
441 eviction_cost: 1.0,
442 replacement_cost: 1.0,
443 },
444 );
445
446 (
447 *ptr as usize,
448 management::MemoryRegion {
449 base_addr: *ptr as usize,
450 size: alloc.size,
451 objects,
452 region_type: RegionType::Buffer,
453 pressure: 0.0,
454 last_eviction: None,
455 },
456 )
457 })
458 .collect();
459
460 let _ = self
461 .memory_manager
462 .run_garbage_collection(&memory_regions)?;
463
464 self.allocation_engine.optimize_strategies()?;
466
467 self.memory_manager.optimize_policies()?;
469
470 if self.stats.fragmentation_ratio > 0.3 {
472 let _ = self.memory_manager.defragment(&memory_regions)?;
473 }
474
475 self.last_optimization = now;
476 self.stats.optimization_cycles += 1;
477
478 Ok(())
479 }
480
481 fn handle_memory_pressure(&mut self) -> Result<(), GpuMemorySystemError> {
483 let memory_usage_ratio = self.calculate_memory_usage_ratio();
484
485 if memory_usage_ratio > self.config.system_config.memory_budget {
486 let memory_regions: HashMap<usize, management::MemoryRegion> = self
487 .memory_regions
488 .iter()
489 .enumerate()
490 .map(|(i, (ptr, alloc))| {
491 let mut objects = HashMap::new();
492 objects.insert(
493 *ptr as usize,
494 CacheObject {
495 address: *ptr as usize,
496 size: alloc.size,
497 created_at: alloc.allocated_at,
498 last_access: alloc.last_accessed.unwrap_or(alloc.allocated_at),
499 access_count: alloc.access_count as u32,
500 access_frequency: alloc.access_count as f64,
501 priority: ObjectPriority::Normal,
502 kernel_context: None,
503 object_type: ObjectType::Data,
504 eviction_cost: 1.0,
505 replacement_cost: 1.0,
506 },
507 );
508
509 (
510 *ptr as usize,
511 management::MemoryRegion {
512 base_addr: *ptr as usize,
513 size: alloc.size,
514 objects,
515 region_type: RegionType::Buffer,
516 pressure: 0.0,
517 last_eviction: None,
518 },
519 )
520 })
521 .collect();
522
523 self.memory_manager
524 .handle_memory_pressure(memory_usage_ratio, &memory_regions)?;
525 }
526
527 Ok(())
528 }
529
530 fn choose_allocator(&self, size: usize) -> AllocatorType {
532 if size < 1024 {
534 AllocatorType::Slab } else if size < 1024 * 1024 {
536 AllocatorType::Buddy } else {
538 AllocatorType::Arena }
540 }
541
542 fn get_vendor_memory_type(&self) -> String {
544 match self.gpu_backend.get_vendor() {
545 GpuVendor::Nvidia => "Device".to_string(),
546 GpuVendor::Amd => "Device".to_string(),
547 GpuVendor::Intel => "Device".to_string(),
548 GpuVendor::Apple => "Private".to_string(),
549 GpuVendor::Unknown => "Unknown".to_string(),
550 }
551 }
552
553 fn update_allocation_stats(&mut self, size: usize, duration: Duration) {
555 self.stats.total_allocations += 1;
556 self.stats.bytes_allocated += size as u64;
557 self.stats.active_allocations += 1;
558
559 if self.stats.bytes_allocated > self.stats.peak_memory_usage as u64 {
560 self.stats.peak_memory_usage = self.stats.bytes_allocated as usize;
561 }
562 }
563
564 fn update_deallocation_stats(&mut self, size: usize, _duration: Duration) {
566 self.stats.total_deallocations += 1;
567 self.stats.bytes_deallocated += size as u64;
568 self.stats.active_allocations = self.stats.active_allocations.saturating_sub(1);
569 }
570
571 fn calculate_memory_usage_ratio(&self) -> f64 {
573 let total_memory = self.get_total_gpu_memory();
574 let used_memory = self.stats.bytes_allocated - self.stats.bytes_deallocated;
575 used_memory as f64 / total_memory as f64
576 }
577
578 fn get_total_gpu_memory(&self) -> usize {
580 match self.gpu_backend.get_vendor() {
582 GpuVendor::Nvidia => 8 * 1024 * 1024 * 1024, GpuVendor::Amd => 16 * 1024 * 1024 * 1024, GpuVendor::Intel => 12 * 1024 * 1024 * 1024, GpuVendor::Apple => 32 * 1024 * 1024 * 1024, GpuVendor::Unknown => 4 * 1024 * 1024 * 1024, }
588 }
589
590 fn calculate_system_metrics(&mut self) {
592 let total_allocated = self
594 .memory_regions
595 .values()
596 .map(|alloc| alloc.size)
597 .sum::<usize>();
598 let total_managed = self.stats.vendor_stats.bytes_allocated;
599 self.stats.fragmentation_ratio = if total_managed > 0 {
600 1.0 - (total_allocated as f64 / total_managed as f64)
601 } else {
602 0.0
603 };
604
605 self.stats.allocation_efficiency = if self.stats.total_allocations > 0 {
607 let successful_allocations = self.stats.total_allocations;
608 successful_allocations as f64 / self.stats.total_allocations as f64
609 } else {
610 1.0
611 };
612 }
613}
614
615unsafe impl Send for GpuMemorySystem {}
623unsafe impl Sync for GpuMemorySystem {}
624
625#[derive(Debug)]
627pub enum GpuMemorySystemError {
628 BackendError(UnifiedGpuError),
629 AllocationError(String),
630 ManagementError(String),
631 InvalidPointer(String),
632 SystemNotStarted,
633 ConfigurationError(String),
634 OptimizationFailed(String),
635 InternalError(String),
636}
637
638impl std::fmt::Display for GpuMemorySystemError {
639 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
640 match self {
641 GpuMemorySystemError::BackendError(err) => write!(f, "Backend error: {}", err),
642 GpuMemorySystemError::AllocationError(msg) => write!(f, "Allocation error: {}", msg),
643 GpuMemorySystemError::ManagementError(msg) => write!(f, "Management error: {}", msg),
644 GpuMemorySystemError::InvalidPointer(msg) => write!(f, "Invalid pointer: {}", msg),
645 GpuMemorySystemError::SystemNotStarted => write!(f, "System not started"),
646 GpuMemorySystemError::ConfigurationError(msg) => {
647 write!(f, "Configuration error: {}", msg)
648 }
649 GpuMemorySystemError::OptimizationFailed(msg) => {
650 write!(f, "Optimization failed: {}", msg)
651 }
652 GpuMemorySystemError::InternalError(msg) => write!(f, "Internal error: {}", msg),
653 }
654 }
655}
656
657impl std::error::Error for GpuMemorySystemError {}
658
659impl From<UnifiedGpuError> for GpuMemorySystemError {
660 fn from(err: UnifiedGpuError) -> Self {
661 GpuMemorySystemError::BackendError(err)
662 }
663}
664
665impl From<allocation::AllocationError> for GpuMemorySystemError {
666 fn from(err: allocation::AllocationError) -> Self {
667 GpuMemorySystemError::AllocationError(format!("{}", err))
668 }
669}
670
671impl From<MemoryManagementError> for GpuMemorySystemError {
672 fn from(err: MemoryManagementError) -> Self {
673 GpuMemorySystemError::ManagementError(format!("{}", err))
674 }
675}
676
677pub struct ThreadSafeGpuMemorySystem {
679 system: Arc<Mutex<GpuMemorySystem>>,
680}
681
682impl ThreadSafeGpuMemorySystem {
683 pub fn new(config: GpuMemorySystemConfig) -> Result<Self, GpuMemorySystemError> {
684 let system = GpuMemorySystem::new(config)?;
685 Ok(Self {
686 system: Arc::new(Mutex::new(system)),
687 })
688 }
689
690 pub fn allocate(
691 &self,
692 size: usize,
693 alignment: Option<usize>,
694 ) -> Result<*mut c_void, GpuMemorySystemError> {
695 let mut system = self.system.lock().unwrap();
696 system.allocate(size, alignment)
697 }
698
699 pub fn free(&self, ptr: *mut c_void) -> Result<(), GpuMemorySystemError> {
700 let mut system = self.system.lock().unwrap();
701 system.free(ptr)
702 }
703
704 pub fn get_stats(&self) -> SystemStats {
705 let mut system = self.system.lock().unwrap();
706 system.get_stats()
707 }
708
709 pub fn optimize(&self) -> Result<(), GpuMemorySystemError> {
710 let mut system = self.system.lock().unwrap();
711 system.optimize()
712 }
713}
714
715#[cfg(test)]
716mod tests {
717 use super::*;
718
719 fn create_test_config() -> GpuMemorySystemConfig {
721 let mut config = GpuMemorySystemConfig::default();
722 config.system_config.memory_budget = 0.001; if let VendorConfig::Cuda(ref mut cuda_config) = config.vendor_config {
726 cuda_config.enable_memory_pools = false;
727 }
728 config
729 }
730
731 #[test]
732 fn test_system_creation() {
733 let config = create_test_config();
734 let system = GpuMemorySystem::new(config);
735 assert!(system.is_ok() || system.is_err());
737 }
738
739 #[test]
740 fn test_auto_create() {
741 let system = GpuMemorySystem::auto_create();
743 assert!(system.is_ok() || system.is_err());
744 }
745
746 #[test]
747 fn test_thread_safe_wrapper() {
748 let config = create_test_config();
749 let system = ThreadSafeGpuMemorySystem::new(config);
750 assert!(system.is_ok() || system.is_err());
752 }
753
754 #[test]
755 fn test_allocator_selection() {
756 let config = create_test_config();
757 if let Ok(system) = GpuMemorySystem::new(config) {
759 assert_eq!(system.choose_allocator(512), AllocatorType::Slab);
760 assert_eq!(system.choose_allocator(64 * 1024), AllocatorType::Buddy);
761 assert_eq!(
762 system.choose_allocator(2 * 1024 * 1024),
763 AllocatorType::Arena
764 );
765 }
766 }
767}