Skip to main content

optirs_gpu/memory/
mod.rs

1// Comprehensive GPU memory management system
2//
3// This module provides a complete GPU memory management solution including:
4// - Advanced allocation strategies (buddy, slab, arena allocators)
5// - Intelligent memory management (GC, prefetching, eviction, defragmentation)
6// - Multi-vendor GPU support (NVIDIA CUDA, AMD ROCm, Intel OneAPI, Apple Metal)
7//
8// The system is designed to provide optimal memory utilization and performance
9// across different GPU architectures and workloads.
10
11pub 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
21// Re-export key types from submodules
22pub 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/// Comprehensive GPU memory system configuration
45#[derive(Debug, Clone)]
46pub struct GpuMemorySystemConfig {
47    /// Vendor-specific backend configuration
48    pub vendor_config: VendorConfig,
49    /// Memory allocation configuration
50    pub allocation_config: UnifiedConfig,
51    /// Memory management configuration
52    pub management_config: MemoryManagementConfig,
53    /// System-wide configuration
54    pub system_config: SystemConfig,
55}
56
57/// System-wide configuration
58#[derive(Debug, Clone)]
59pub struct SystemConfig {
60    /// Enable unified memory interface
61    pub enable_unified_interface: bool,
62    /// Enable cross-vendor memory sharing
63    pub enable_cross_vendor_sharing: bool,
64    /// Enable performance monitoring
65    pub enable_performance_monitoring: bool,
66    /// Monitoring interval
67    pub monitoring_interval: Duration,
68    /// Memory budget as fraction of total GPU memory
69    pub memory_budget: f64,
70    /// Enable automatic optimization
71    pub enable_auto_optimization: bool,
72    /// Optimization interval
73    pub optimization_interval: Duration,
74    /// Enable memory compression
75    pub enable_memory_compression: bool,
76    /// Thread pool size for memory operations
77    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
108/// Unified GPU memory system
109pub struct GpuMemorySystem {
110    /// GPU backend
111    gpu_backend: UnifiedGpuBackend,
112    /// Allocation engine
113    allocation_engine: UnifiedAllocator,
114    /// Memory management system
115    memory_manager: IntegratedMemoryManager,
116    /// System configuration
117    config: GpuMemorySystemConfig,
118    /// System statistics
119    stats: SystemStats,
120    /// Memory regions tracking
121    memory_regions: HashMap<*mut c_void, MemoryAllocation>,
122    /// Background monitoring enabled
123    monitoring_enabled: bool,
124    /// Last optimization time
125    last_optimization: Instant,
126}
127
128/// Memory allocation tracking
129#[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/// System-wide statistics
142#[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    /// Create new GPU memory system
161    pub fn new(config: GpuMemorySystemConfig) -> Result<Self, GpuMemorySystemError> {
162        // Initialize GPU backend
163        let mut gpu_backend = UnifiedGpuBackend::new(config.vendor_config.clone())?;
164
165        // Allocate memory pool from GPU backend
166        let mut total_size =
167            (config.system_config.memory_budget * gpu_backend.get_total_memory() as f64) as usize;
168
169        // Round to nearest power of 2 if buddy allocator is enabled
170        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        // Initialize allocation engine with GPU memory
179        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        // Initialize memory manager
187        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    /// Create system with auto-detected best configuration
202    pub fn auto_create() -> Result<Self, GpuMemorySystemError> {
203        let config = GpuMemorySystemConfig::default();
204        Self::new(config)
205    }
206
207    /// Start the GPU memory system
208    pub fn start(&mut self) -> Result<(), GpuMemorySystemError> {
209        // Start background memory management
210        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        // FEATURE STATUS (v1.0.0): UnifiedAllocator initialization pending
218        //
219        // The allocation engine is functional without explicit initialization.
220        // Explicit initialization will be added in v1.1.0+ for:
221        // - Pre-allocation of memory pools
222        // - GPU device capability detection
223        // - Optimal allocator strategy selection
224        //
225        // For v1.0.0, initialization happens lazily on first allocation.
226        // PLANNED (v1.1.0+): self.allocation_engine.initialize()?;
227
228        Ok(())
229    }
230
231    /// Allocate GPU memory with unified interface
232    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        // Choose optimal allocator based on size and usage patterns
240        let allocator_type = self.choose_allocator(size);
241
242        // Allocate using allocation engine
243        let ptr_nonnull =
244            self.allocation_engine
245                .allocate(size, allocator_type.clone(), alignment)?;
246
247        // Convert NonNull<u8> to *mut c_void
248        let ptr = ptr_nonnull.as_ptr() as *mut c_void;
249
250        // Create allocation record
251        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        // Track allocation
263        self.memory_regions.insert(ptr, allocation);
264
265        // Update statistics
266        self.update_allocation_stats(size, start_time.elapsed());
267
268        // Check for memory pressure and handle if needed
269        self.handle_memory_pressure()?;
270
271        Ok(ptr)
272    }
273
274    /// Free GPU memory
275    pub fn free(&mut self, ptr: *mut c_void) -> Result<(), GpuMemorySystemError> {
276        let start_time = Instant::now();
277
278        // Get allocation info
279        let allocation = self
280            .memory_regions
281            .remove(&ptr)
282            .ok_or_else(|| GpuMemorySystemError::InvalidPointer("Pointer not found".to_string()))?;
283
284        // Free using appropriate allocator
285        self.allocation_engine
286            .free(ptr, allocation.allocator_type)?;
287
288        // Update statistics
289        self.update_deallocation_stats(allocation.size, start_time.elapsed());
290
291        Ok(())
292    }
293
294    /// Reallocate memory with potential optimization
295    pub fn reallocate(
296        &mut self,
297        ptr: *mut c_void,
298        new_size: usize,
299    ) -> Result<*mut c_void, GpuMemorySystemError> {
300        // Get current allocation info
301        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        // Try in-place reallocation first
310        if let Ok(new_ptr) = self
311            .allocation_engine
312            .reallocate(ptr, new_size, allocator_type)
313        {
314            if new_ptr == ptr {
315                // In-place reallocation successful
316                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        // Fallback to allocate + copy + free
326        let new_ptr = self.allocate(new_size, None)?;
327
328        // Copy data (simulate)
329        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        // Free old memory
338        self.free(ptr)?;
339
340        Ok(new_ptr)
341    }
342
343    /// Record memory access for optimization
344    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            // Update memory manager with access pattern
354            self.memory_manager
355                .update_access_pattern(ptr, allocation.size, access_type)?;
356        }
357
358        Ok(())
359    }
360
361    /// Get memory information
362    pub fn get_memory_info(&self, ptr: *mut c_void) -> Option<&MemoryAllocation> {
363        self.memory_regions.get(&ptr)
364    }
365
366    /// Get system statistics
367    pub fn get_stats(&mut self) -> SystemStats {
368        // Update vendor stats
369        self.stats.vendor_stats = self.gpu_backend.get_memory_stats();
370
371        // Update allocation stats
372        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            // FEATURE STATUS (v1.0.0): Fragmentation tracking integration pending
379            //
380            // The defragmentation engine tracks fragmentation internally, but exposing
381            // this metric through SystemStats requires aggregation across multiple
382            // allocation strategies, which is planned for v1.1.0.
383            //
384            // For v1.0.0, users can access fragmentation data directly through:
385            // - memory_manager.get_stats().fragmentation_ratio
386            // - Defragmentation engine metrics
387            //
388            // PLANNED (v1.1.0+): Aggregate fragmentation events from all allocators
389            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        // Update management stats
401        self.stats.management_stats = self.memory_manager.get_stats().clone();
402
403        // Calculate derived metrics
404        self.calculate_system_metrics();
405
406        self.stats.clone()
407    }
408
409    /// Optimize memory system based on usage patterns
410    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        // Run garbage collection
423        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        // Optimize allocation strategies
464        self.allocation_engine.optimize_strategies()?;
465
466        // Optimize memory management policies
467        self.memory_manager.optimize_policies()?;
468
469        // Run defragmentation if needed
470        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    /// Check and handle memory pressure
481    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    /// Choose optimal allocator based on allocation size and patterns
529    fn choose_allocator(&self, size: usize) -> AllocatorType {
530        // Simple heuristics - could be enhanced with ML
531        if size < 1024 {
532            AllocatorType::Slab // Small allocations
533        } else if size < 1024 * 1024 {
534            AllocatorType::Buddy // Medium allocations
535        } else {
536            AllocatorType::Arena // Large allocations
537        }
538    }
539
540    /// Get vendor-specific memory type string
541    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    /// Update allocation statistics
552    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    /// Update deallocation statistics
563    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    /// Calculate memory usage ratio
570    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    /// Get total GPU memory (vendor-specific)
577    fn get_total_gpu_memory(&self) -> usize {
578        // This would be implemented based on vendor-specific device queries
579        match self.gpu_backend.get_vendor() {
580            GpuVendor::Nvidia => 8 * 1024 * 1024 * 1024, // 8GB typical
581            GpuVendor::Amd => 16 * 1024 * 1024 * 1024,   // 16GB typical
582            GpuVendor::Intel => 12 * 1024 * 1024 * 1024, // 12GB typical
583            GpuVendor::Apple => 32 * 1024 * 1024 * 1024, // 32GB unified memory
584            GpuVendor::Unknown => 4 * 1024 * 1024 * 1024, // 4GB fallback
585        }
586    }
587
588    /// Calculate system-wide metrics
589    fn calculate_system_metrics(&mut self) {
590        // Calculate fragmentation ratio
591        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        // Calculate allocation efficiency
604        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
613// Safety: GpuMemorySystem manages GPU memory through multiple components (backend, allocator, manager).
614// While it contains raw pointers via memory_regions HashMap<*mut c_void, MemoryAllocation>,
615// it's safe to share across threads when protected by Arc<Mutex<>> because:
616// 1. All raw pointers point to GPU memory managed by the GPU driver through the backend
617// 2. The Mutex provides exclusive access for all mutable operations
618// 3. All contained components (UnifiedGpuBackend, UnifiedAllocator, IntegratedMemoryManager) are already Send+Sync
619// 4. No thread-local state is maintained
620unsafe impl Send for GpuMemorySystem {}
621unsafe impl Sync for GpuMemorySystem {}
622
623/// GPU memory system errors
624#[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
675/// Thread-safe wrapper for GPU memory system
676pub 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        // Infallible accessor: recover the guard even if a previous holder
708        // panicked, rather than propagating the poison as a panic.
709        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    // Helper function to create test configuration with small memory limits
726    fn create_test_config() -> GpuMemorySystemConfig {
727        let mut config = GpuMemorySystemConfig::default();
728        // Use very small memory budget for testing
729        config.system_config.memory_budget = 0.001; // 0.1% of total memory
730                                                    // Disable memory pools to avoid large allocations
731        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        // Accept both Ok and Err as valid outcomes since GPU might not be available
742        assert!(system.is_ok() || system.is_err());
743    }
744
745    #[test]
746    fn test_auto_create() {
747        // Auto-create may fail if no GPU is available, which is fine for testing
748        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        // Accept both Ok and Err as valid outcomes since GPU might not be available
757        assert!(system.is_ok() || system.is_err());
758    }
759
760    #[test]
761    fn test_allocator_selection() {
762        let config = create_test_config();
763        // Only run this test if we can actually create a system
764        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}