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 .map(|(ptr, alloc)| {
427 let mut objects = HashMap::new();
428 objects.insert(
429 *ptr as usize,
430 CacheObject {
431 address: *ptr as usize,
432 size: alloc.size,
433 created_at: alloc.allocated_at,
434 last_access: alloc.last_accessed.unwrap_or(alloc.allocated_at),
435 access_count: alloc.access_count as u32,
436 access_frequency: alloc.access_count as f64,
437 priority: ObjectPriority::Normal,
438 kernel_context: None,
439 object_type: ObjectType::Data,
440 eviction_cost: 1.0,
441 replacement_cost: 1.0,
442 },
443 );
444
445 (
446 *ptr as usize,
447 management::MemoryRegion {
448 base_addr: *ptr as usize,
449 size: alloc.size,
450 objects,
451 region_type: RegionType::Buffer,
452 pressure: 0.0,
453 last_eviction: None,
454 },
455 )
456 })
457 .collect();
458
459 let _ = self
460 .memory_manager
461 .run_garbage_collection(&memory_regions)?;
462
463 self.allocation_engine.optimize_strategies()?;
465
466 self.memory_manager.optimize_policies()?;
468
469 if self.stats.fragmentation_ratio > 0.3 {
471 let _ = self.memory_manager.defragment(&memory_regions)?;
472 }
473
474 self.last_optimization = now;
475 self.stats.optimization_cycles += 1;
476
477 Ok(())
478 }
479
480 fn handle_memory_pressure(&mut self) -> Result<(), GpuMemorySystemError> {
482 let memory_usage_ratio = self.calculate_memory_usage_ratio();
483
484 if memory_usage_ratio > self.config.system_config.memory_budget {
485 let memory_regions: HashMap<usize, management::MemoryRegion> = self
486 .memory_regions
487 .iter()
488 .map(|(ptr, alloc)| {
489 let mut objects = HashMap::new();
490 objects.insert(
491 *ptr as usize,
492 CacheObject {
493 address: *ptr as usize,
494 size: alloc.size,
495 created_at: alloc.allocated_at,
496 last_access: alloc.last_accessed.unwrap_or(alloc.allocated_at),
497 access_count: alloc.access_count as u32,
498 access_frequency: alloc.access_count as f64,
499 priority: ObjectPriority::Normal,
500 kernel_context: None,
501 object_type: ObjectType::Data,
502 eviction_cost: 1.0,
503 replacement_cost: 1.0,
504 },
505 );
506
507 (
508 *ptr as usize,
509 management::MemoryRegion {
510 base_addr: *ptr as usize,
511 size: alloc.size,
512 objects,
513 region_type: RegionType::Buffer,
514 pressure: 0.0,
515 last_eviction: None,
516 },
517 )
518 })
519 .collect();
520
521 self.memory_manager
522 .handle_memory_pressure(memory_usage_ratio, &memory_regions)?;
523 }
524
525 Ok(())
526 }
527
528 fn choose_allocator(&self, size: usize) -> AllocatorType {
530 if size < 1024 {
532 AllocatorType::Slab } else if size < 1024 * 1024 {
534 AllocatorType::Buddy } else {
536 AllocatorType::Arena }
538 }
539
540 fn get_vendor_memory_type(&self) -> String {
542 match self.gpu_backend.get_vendor() {
543 GpuVendor::Nvidia => "Device".to_string(),
544 GpuVendor::Amd => "Device".to_string(),
545 GpuVendor::Intel => "Device".to_string(),
546 GpuVendor::Apple => "Private".to_string(),
547 GpuVendor::Unknown => "Unknown".to_string(),
548 }
549 }
550
551 fn update_allocation_stats(&mut self, size: usize, _duration: Duration) {
553 self.stats.total_allocations += 1;
554 self.stats.bytes_allocated += size as u64;
555 self.stats.active_allocations += 1;
556
557 if self.stats.bytes_allocated > self.stats.peak_memory_usage as u64 {
558 self.stats.peak_memory_usage = self.stats.bytes_allocated as usize;
559 }
560 }
561
562 fn update_deallocation_stats(&mut self, size: usize, _duration: Duration) {
564 self.stats.total_deallocations += 1;
565 self.stats.bytes_deallocated += size as u64;
566 self.stats.active_allocations = self.stats.active_allocations.saturating_sub(1);
567 }
568
569 fn calculate_memory_usage_ratio(&self) -> f64 {
571 let total_memory = self.get_total_gpu_memory();
572 let used_memory = self.stats.bytes_allocated - self.stats.bytes_deallocated;
573 used_memory as f64 / total_memory as f64
574 }
575
576 fn get_total_gpu_memory(&self) -> usize {
578 match self.gpu_backend.get_vendor() {
580 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, }
586 }
587
588 fn calculate_system_metrics(&mut self) {
590 let total_allocated = self
592 .memory_regions
593 .values()
594 .map(|alloc| alloc.size)
595 .sum::<usize>();
596 let total_managed = self.stats.vendor_stats.bytes_allocated;
597 self.stats.fragmentation_ratio = if total_managed > 0 {
598 1.0 - (total_allocated as f64 / total_managed as f64)
599 } else {
600 0.0
601 };
602
603 self.stats.allocation_efficiency = if self.stats.total_allocations > 0 {
605 let successful_allocations = self.stats.total_allocations;
606 successful_allocations as f64 / self.stats.total_allocations as f64
607 } else {
608 1.0
609 };
610 }
611}
612
613unsafe impl Send for GpuMemorySystem {}
621unsafe impl Sync for GpuMemorySystem {}
622
623#[derive(Debug)]
625pub enum GpuMemorySystemError {
626 BackendError(UnifiedGpuError),
627 AllocationError(String),
628 ManagementError(String),
629 InvalidPointer(String),
630 SystemNotStarted,
631 ConfigurationError(String),
632 OptimizationFailed(String),
633 InternalError(String),
634}
635
636impl std::fmt::Display for GpuMemorySystemError {
637 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
638 match self {
639 GpuMemorySystemError::BackendError(err) => write!(f, "Backend error: {}", err),
640 GpuMemorySystemError::AllocationError(msg) => write!(f, "Allocation error: {}", msg),
641 GpuMemorySystemError::ManagementError(msg) => write!(f, "Management error: {}", msg),
642 GpuMemorySystemError::InvalidPointer(msg) => write!(f, "Invalid pointer: {}", msg),
643 GpuMemorySystemError::SystemNotStarted => write!(f, "System not started"),
644 GpuMemorySystemError::ConfigurationError(msg) => {
645 write!(f, "Configuration error: {}", msg)
646 }
647 GpuMemorySystemError::OptimizationFailed(msg) => {
648 write!(f, "Optimization failed: {}", msg)
649 }
650 GpuMemorySystemError::InternalError(msg) => write!(f, "Internal error: {}", msg),
651 }
652 }
653}
654
655impl std::error::Error for GpuMemorySystemError {}
656
657impl From<UnifiedGpuError> for GpuMemorySystemError {
658 fn from(err: UnifiedGpuError) -> Self {
659 GpuMemorySystemError::BackendError(err)
660 }
661}
662
663impl From<allocation::AllocationError> for GpuMemorySystemError {
664 fn from(err: allocation::AllocationError) -> Self {
665 GpuMemorySystemError::AllocationError(format!("{}", err))
666 }
667}
668
669impl From<MemoryManagementError> for GpuMemorySystemError {
670 fn from(err: MemoryManagementError) -> Self {
671 GpuMemorySystemError::ManagementError(format!("{}", err))
672 }
673}
674
675pub struct ThreadSafeGpuMemorySystem {
677 system: Arc<Mutex<GpuMemorySystem>>,
678}
679
680impl ThreadSafeGpuMemorySystem {
681 pub fn new(config: GpuMemorySystemConfig) -> Result<Self, GpuMemorySystemError> {
682 let system = GpuMemorySystem::new(config)?;
683 Ok(Self {
684 system: Arc::new(Mutex::new(system)),
685 })
686 }
687
688 pub fn allocate(
689 &self,
690 size: usize,
691 alignment: Option<usize>,
692 ) -> Result<*mut c_void, GpuMemorySystemError> {
693 let mut system = self.system.lock().map_err(|_| {
694 GpuMemorySystemError::InternalError("memory system lock poisoned".into())
695 })?;
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().map_err(|_| {
701 GpuMemorySystemError::InternalError("memory system lock poisoned".into())
702 })?;
703 system.free(ptr)
704 }
705
706 pub fn get_stats(&self) -> SystemStats {
707 let mut system = self.system.lock().unwrap_or_else(|e| e.into_inner());
710 system.get_stats()
711 }
712
713 pub fn optimize(&self) -> Result<(), GpuMemorySystemError> {
714 let mut system = self.system.lock().map_err(|_| {
715 GpuMemorySystemError::InternalError("memory system lock poisoned".into())
716 })?;
717 system.optimize()
718 }
719}
720
721#[cfg(test)]
722mod tests {
723 use super::*;
724
725 fn create_test_config() -> GpuMemorySystemConfig {
727 let mut config = GpuMemorySystemConfig::default();
728 config.system_config.memory_budget = 0.001; if let VendorConfig::Cuda(ref mut cuda_config) = config.vendor_config {
732 cuda_config.enable_memory_pools = false;
733 }
734 config
735 }
736
737 #[test]
738 fn test_system_creation() {
739 let config = create_test_config();
740 let system = GpuMemorySystem::new(config);
741 assert!(system.is_ok() || system.is_err());
743 }
744
745 #[test]
746 fn test_auto_create() {
747 let system = GpuMemorySystem::auto_create();
749 assert!(system.is_ok() || system.is_err());
750 }
751
752 #[test]
753 fn test_thread_safe_wrapper() {
754 let config = create_test_config();
755 let system = ThreadSafeGpuMemorySystem::new(config);
756 assert!(system.is_ok() || system.is_err());
758 }
759
760 #[test]
761 fn test_allocator_selection() {
762 let config = create_test_config();
763 if let Ok(system) = GpuMemorySystem::new(config) {
765 assert_eq!(system.choose_allocator(512), AllocatorType::Slab);
766 assert_eq!(system.choose_allocator(64 * 1024), AllocatorType::Buddy);
767 assert_eq!(
768 system.choose_allocator(2 * 1024 * 1024),
769 AllocatorType::Arena
770 );
771 }
772 }
773}