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            .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        // Optimize allocation strategies
465        self.allocation_engine.optimize_strategies()?;
466
467        // Optimize memory management policies
468        self.memory_manager.optimize_policies()?;
469
470        // Run defragmentation if needed
471        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    /// Check and handle memory pressure
482    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    /// Choose optimal allocator based on allocation size and patterns
531    fn choose_allocator(&self, size: usize) -> AllocatorType {
532        // Simple heuristics - could be enhanced with ML
533        if size < 1024 {
534            AllocatorType::Slab // Small allocations
535        } else if size < 1024 * 1024 {
536            AllocatorType::Buddy // Medium allocations
537        } else {
538            AllocatorType::Arena // Large allocations
539        }
540    }
541
542    /// Get vendor-specific memory type string
543    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    /// Update allocation statistics
554    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    /// Update deallocation statistics
565    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    /// Calculate memory usage ratio
572    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    /// Get total GPU memory (vendor-specific)
579    fn get_total_gpu_memory(&self) -> usize {
580        // This would be implemented based on vendor-specific device queries
581        match self.gpu_backend.get_vendor() {
582            GpuVendor::Nvidia => 8 * 1024 * 1024 * 1024, // 8GB typical
583            GpuVendor::Amd => 16 * 1024 * 1024 * 1024,   // 16GB typical
584            GpuVendor::Intel => 12 * 1024 * 1024 * 1024, // 12GB typical
585            GpuVendor::Apple => 32 * 1024 * 1024 * 1024, // 32GB unified memory
586            GpuVendor::Unknown => 4 * 1024 * 1024 * 1024, // 4GB fallback
587        }
588    }
589
590    /// Calculate system-wide metrics
591    fn calculate_system_metrics(&mut self) {
592        // Calculate fragmentation ratio
593        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        // Calculate allocation efficiency
606        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
615// Safety: GpuMemorySystem manages GPU memory through multiple components (backend, allocator, manager).
616// While it contains raw pointers via memory_regions HashMap<*mut c_void, MemoryAllocation>,
617// it's safe to share across threads when protected by Arc<Mutex<>> because:
618// 1. All raw pointers point to GPU memory managed by the GPU driver through the backend
619// 2. The Mutex provides exclusive access for all mutable operations
620// 3. All contained components (UnifiedGpuBackend, UnifiedAllocator, IntegratedMemoryManager) are already Send+Sync
621// 4. No thread-local state is maintained
622unsafe impl Send for GpuMemorySystem {}
623unsafe impl Sync for GpuMemorySystem {}
624
625/// GPU memory system errors
626#[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
677/// Thread-safe wrapper for GPU memory system
678pub 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().expect("lock poisoned");
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().expect("lock poisoned");
701        system.free(ptr)
702    }
703
704    pub fn get_stats(&self) -> SystemStats {
705        let mut system = self.system.lock().expect("lock poisoned");
706        system.get_stats()
707    }
708
709    pub fn optimize(&self) -> Result<(), GpuMemorySystemError> {
710        let mut system = self.system.lock().expect("lock poisoned");
711        system.optimize()
712    }
713}
714
715#[cfg(test)]
716mod tests {
717    use super::*;
718
719    // Helper function to create test configuration with small memory limits
720    fn create_test_config() -> GpuMemorySystemConfig {
721        let mut config = GpuMemorySystemConfig::default();
722        // Use very small memory budget for testing
723        config.system_config.memory_budget = 0.001; // 0.1% of total memory
724                                                    // Disable memory pools to avoid large allocations
725        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        // Accept both Ok and Err as valid outcomes since GPU might not be available
736        assert!(system.is_ok() || system.is_err());
737    }
738
739    #[test]
740    fn test_auto_create() {
741        // Auto-create may fail if no GPU is available, which is fine for testing
742        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        // Accept both Ok and Err as valid outcomes since GPU might not be available
751        assert!(system.is_ok() || system.is_err());
752    }
753
754    #[test]
755    fn test_allocator_selection() {
756        let config = create_test_config();
757        // Only run this test if we can actually create a system
758        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}