pub mod allocation;
pub mod management;
pub mod vendors;
use std::collections::HashMap;
use std::ffi::c_void;
use std::ptr::NonNull;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
pub use allocation::{
AllocationStrategy, AllocationStrategyManager, AllocatorType, ArenaAllocator, BuddyAllocator,
MemoryPool, SlabAllocator, UnifiedAllocator, UnifiedConfig,
};
use allocation::strategies::AllocationStats;
pub use management::{
AccessType, DefragmentationEngine, EvictionEngine, GarbageCollectionEngine,
IntegratedMemoryManager, ManagementStats, MemoryManagementConfig, MemoryManagementError,
MemoryRegion, PrefetchingEngine,
};
use management::eviction_policies::{CacheObject, ObjectPriority, ObjectType, RegionType};
pub use vendors::{
CudaConfig, CudaError, CudaMemoryBackend, CudaMemoryType, GpuBackendFactory, GpuVendor,
MetalConfig, MetalError, MetalMemoryBackend, MetalMemoryType, OneApiConfig, OneApiError,
OneApiMemoryBackend, OneApiMemoryType, RocmConfig, RocmError, RocmMemoryBackend,
RocmMemoryType, UnifiedGpuBackend, UnifiedGpuError, UnifiedMemoryStats, VendorConfig,
};
#[derive(Debug, Clone)]
pub struct GpuMemorySystemConfig {
pub vendor_config: VendorConfig,
pub allocation_config: UnifiedConfig,
pub management_config: MemoryManagementConfig,
pub system_config: SystemConfig,
}
#[derive(Debug, Clone)]
pub struct SystemConfig {
pub enable_unified_interface: bool,
pub enable_cross_vendor_sharing: bool,
pub enable_performance_monitoring: bool,
pub monitoring_interval: Duration,
pub memory_budget: f64,
pub enable_auto_optimization: bool,
pub optimization_interval: Duration,
pub enable_memory_compression: bool,
pub thread_pool_size: usize,
}
impl Default for SystemConfig {
fn default() -> Self {
Self {
enable_unified_interface: true,
enable_cross_vendor_sharing: false,
enable_performance_monitoring: true,
monitoring_interval: Duration::from_millis(500),
memory_budget: 0.9,
enable_auto_optimization: true,
optimization_interval: Duration::from_secs(60),
enable_memory_compression: false,
thread_pool_size: 4,
}
}
}
impl Default for GpuMemorySystemConfig {
fn default() -> Self {
let vendor = GpuBackendFactory::get_preferred_vendor();
Self {
vendor_config: GpuBackendFactory::create_default_config(vendor),
allocation_config: UnifiedConfig::default(),
management_config: MemoryManagementConfig::default(),
system_config: SystemConfig::default(),
}
}
}
pub struct GpuMemorySystem {
gpu_backend: UnifiedGpuBackend,
allocation_engine: UnifiedAllocator,
memory_manager: IntegratedMemoryManager,
config: GpuMemorySystemConfig,
stats: SystemStats,
memory_regions: HashMap<*mut c_void, MemoryAllocation>,
monitoring_enabled: bool,
last_optimization: Instant,
}
#[derive(Debug, Clone)]
pub struct MemoryAllocation {
pub ptr: *mut c_void,
pub size: usize,
pub allocator_type: AllocatorType,
pub vendor_memory_type: String,
pub allocated_at: Instant,
pub last_accessed: Option<Instant>,
pub access_count: u64,
pub ref_count: u32,
}
#[derive(Debug, Clone, Default)]
pub struct SystemStats {
pub total_allocations: u64,
pub total_deallocations: u64,
pub bytes_allocated: u64,
pub bytes_deallocated: u64,
pub active_allocations: u64,
pub peak_memory_usage: usize,
pub fragmentation_ratio: f64,
pub allocation_efficiency: f64,
pub vendor_stats: UnifiedMemoryStats,
pub allocation_stats: AllocationStats,
pub management_stats: ManagementStats,
pub uptime: Duration,
pub optimization_cycles: u64,
}
impl GpuMemorySystem {
pub fn new(config: GpuMemorySystemConfig) -> Result<Self, GpuMemorySystemError> {
let mut gpu_backend = UnifiedGpuBackend::new(config.vendor_config.clone())?;
let mut total_size =
(config.system_config.memory_budget * gpu_backend.get_total_memory() as f64) as usize;
if config.allocation_config.enable_buddy {
total_size = total_size.next_power_of_two();
}
let base_ptr = gpu_backend
.allocate(total_size)
.map_err(GpuMemorySystemError::BackendError)?;
let allocation_engine = UnifiedAllocator::new(
unsafe { NonNull::new_unchecked(base_ptr as *mut u8) },
total_size,
config.allocation_config.clone(),
)
.map_err(|e| GpuMemorySystemError::AllocationError(format!("{:?}", e)))?;
let memory_manager = IntegratedMemoryManager::new(config.management_config.clone());
Ok(Self {
gpu_backend,
allocation_engine,
memory_manager,
config,
stats: SystemStats::default(),
memory_regions: HashMap::new(),
monitoring_enabled: false,
last_optimization: Instant::now(),
})
}
pub fn auto_create() -> Result<Self, GpuMemorySystemError> {
let config = GpuMemorySystemConfig::default();
Self::new(config)
}
pub fn start(&mut self) -> Result<(), GpuMemorySystemError> {
if self.config.system_config.enable_performance_monitoring {
self.memory_manager
.start_background_management()
.map_err(|e| GpuMemorySystemError::ManagementError(format!("{}", e)))?;
self.monitoring_enabled = true;
}
Ok(())
}
pub fn allocate(
&mut self,
size: usize,
alignment: Option<usize>,
) -> Result<*mut c_void, GpuMemorySystemError> {
let start_time = Instant::now();
let allocator_type = self.choose_allocator(size);
let ptr_nonnull =
self.allocation_engine
.allocate(size, allocator_type.clone(), alignment)?;
let ptr = ptr_nonnull.as_ptr() as *mut c_void;
let allocation = MemoryAllocation {
ptr,
size,
allocator_type,
vendor_memory_type: self.get_vendor_memory_type(),
allocated_at: Instant::now(),
last_accessed: Some(Instant::now()),
access_count: 1,
ref_count: 1,
};
self.memory_regions.insert(ptr, allocation);
self.update_allocation_stats(size, start_time.elapsed());
self.handle_memory_pressure()?;
Ok(ptr)
}
pub fn free(&mut self, ptr: *mut c_void) -> Result<(), GpuMemorySystemError> {
let start_time = Instant::now();
let allocation = self
.memory_regions
.remove(&ptr)
.ok_or_else(|| GpuMemorySystemError::InvalidPointer("Pointer not found".to_string()))?;
self.allocation_engine
.free(ptr, allocation.allocator_type)?;
self.update_deallocation_stats(allocation.size, start_time.elapsed());
Ok(())
}
pub fn reallocate(
&mut self,
ptr: *mut c_void,
new_size: usize,
) -> Result<*mut c_void, GpuMemorySystemError> {
let allocation = self
.memory_regions
.get(&ptr)
.ok_or_else(|| GpuMemorySystemError::InvalidPointer("Pointer not found".to_string()))?;
let old_size = allocation.size;
let allocator_type = allocation.allocator_type.clone();
if let Ok(new_ptr) = self
.allocation_engine
.reallocate(ptr, new_size, allocator_type)
{
if new_ptr == ptr {
if let Some(allocation) = self.memory_regions.get_mut(&ptr) {
allocation.size = new_size;
allocation.last_accessed = Some(Instant::now());
allocation.access_count += 1;
}
return Ok(ptr);
}
}
let new_ptr = self.allocate(new_size, None)?;
unsafe {
std::ptr::copy_nonoverlapping(
ptr as *const u8,
new_ptr as *mut u8,
old_size.min(new_size),
);
}
self.free(ptr)?;
Ok(new_ptr)
}
pub fn record_access(
&mut self,
ptr: *mut c_void,
access_type: AccessType,
) -> Result<(), GpuMemorySystemError> {
if let Some(allocation) = self.memory_regions.get_mut(&ptr) {
allocation.last_accessed = Some(Instant::now());
allocation.access_count += 1;
self.memory_manager
.update_access_pattern(ptr, allocation.size, access_type)?;
}
Ok(())
}
pub fn get_memory_info(&self, ptr: *mut c_void) -> Option<&MemoryAllocation> {
self.memory_regions.get(&ptr)
}
pub fn get_stats(&mut self) -> SystemStats {
self.stats.vendor_stats = self.gpu_backend.get_memory_stats();
let unified_stats = self.allocation_engine.get_stats();
self.stats.allocation_stats = AllocationStats {
total_allocations: unified_stats.total_allocations,
total_deallocations: unified_stats.total_deallocations,
cache_hits: unified_stats.routing_cache_hits,
cache_misses: unified_stats.routing_decisions - unified_stats.routing_cache_hits,
fragmentation_events: 0,
total_allocated_bytes: unified_stats.bytes_allocated,
peak_allocated_bytes: unified_stats.peak_memory_usage as u64,
average_allocation_size: if unified_stats.total_allocations > 0 {
(unified_stats.bytes_allocated as f64) / (unified_stats.total_allocations as f64)
} else {
0.0
},
allocation_latency_ms: unified_stats.average_allocation_time_ns / 1_000_000.0,
};
self.stats.management_stats = self.memory_manager.get_stats().clone();
self.calculate_system_metrics();
self.stats.clone()
}
pub fn optimize(&mut self) -> Result<(), GpuMemorySystemError> {
if !self.config.system_config.enable_auto_optimization {
return Ok(());
}
let now = Instant::now();
if now.duration_since(self.last_optimization)
< self.config.system_config.optimization_interval
{
return Ok(());
}
let memory_regions: HashMap<usize, management::MemoryRegion> = self
.memory_regions
.iter()
.map(|(ptr, alloc)| {
let mut objects = HashMap::new();
objects.insert(
*ptr as usize,
CacheObject {
address: *ptr as usize,
size: alloc.size,
created_at: alloc.allocated_at,
last_access: alloc.last_accessed.unwrap_or(alloc.allocated_at),
access_count: alloc.access_count as u32,
access_frequency: alloc.access_count as f64,
priority: ObjectPriority::Normal,
kernel_context: None,
object_type: ObjectType::Data,
eviction_cost: 1.0,
replacement_cost: 1.0,
},
);
(
*ptr as usize,
management::MemoryRegion {
base_addr: *ptr as usize,
size: alloc.size,
objects,
region_type: RegionType::Buffer,
pressure: 0.0,
last_eviction: None,
},
)
})
.collect();
let _ = self
.memory_manager
.run_garbage_collection(&memory_regions)?;
self.allocation_engine.optimize_strategies()?;
self.memory_manager.optimize_policies()?;
if self.stats.fragmentation_ratio > 0.3 {
let _ = self.memory_manager.defragment(&memory_regions)?;
}
self.last_optimization = now;
self.stats.optimization_cycles += 1;
Ok(())
}
fn handle_memory_pressure(&mut self) -> Result<(), GpuMemorySystemError> {
let memory_usage_ratio = self.calculate_memory_usage_ratio();
if memory_usage_ratio > self.config.system_config.memory_budget {
let memory_regions: HashMap<usize, management::MemoryRegion> = self
.memory_regions
.iter()
.map(|(ptr, alloc)| {
let mut objects = HashMap::new();
objects.insert(
*ptr as usize,
CacheObject {
address: *ptr as usize,
size: alloc.size,
created_at: alloc.allocated_at,
last_access: alloc.last_accessed.unwrap_or(alloc.allocated_at),
access_count: alloc.access_count as u32,
access_frequency: alloc.access_count as f64,
priority: ObjectPriority::Normal,
kernel_context: None,
object_type: ObjectType::Data,
eviction_cost: 1.0,
replacement_cost: 1.0,
},
);
(
*ptr as usize,
management::MemoryRegion {
base_addr: *ptr as usize,
size: alloc.size,
objects,
region_type: RegionType::Buffer,
pressure: 0.0,
last_eviction: None,
},
)
})
.collect();
self.memory_manager
.handle_memory_pressure(memory_usage_ratio, &memory_regions)?;
}
Ok(())
}
fn choose_allocator(&self, size: usize) -> AllocatorType {
if size < 1024 {
AllocatorType::Slab } else if size < 1024 * 1024 {
AllocatorType::Buddy } else {
AllocatorType::Arena }
}
fn get_vendor_memory_type(&self) -> String {
match self.gpu_backend.get_vendor() {
GpuVendor::Nvidia => "Device".to_string(),
GpuVendor::Amd => "Device".to_string(),
GpuVendor::Intel => "Device".to_string(),
GpuVendor::Apple => "Private".to_string(),
GpuVendor::Unknown => "Unknown".to_string(),
}
}
fn update_allocation_stats(&mut self, size: usize, _duration: Duration) {
self.stats.total_allocations += 1;
self.stats.bytes_allocated += size as u64;
self.stats.active_allocations += 1;
if self.stats.bytes_allocated > self.stats.peak_memory_usage as u64 {
self.stats.peak_memory_usage = self.stats.bytes_allocated as usize;
}
}
fn update_deallocation_stats(&mut self, size: usize, _duration: Duration) {
self.stats.total_deallocations += 1;
self.stats.bytes_deallocated += size as u64;
self.stats.active_allocations = self.stats.active_allocations.saturating_sub(1);
}
fn calculate_memory_usage_ratio(&self) -> f64 {
let total_memory = self.get_total_gpu_memory();
let used_memory = self.stats.bytes_allocated - self.stats.bytes_deallocated;
used_memory as f64 / total_memory as f64
}
fn get_total_gpu_memory(&self) -> usize {
match self.gpu_backend.get_vendor() {
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, }
}
fn calculate_system_metrics(&mut self) {
let total_allocated = self
.memory_regions
.values()
.map(|alloc| alloc.size)
.sum::<usize>();
let total_managed = self.stats.vendor_stats.bytes_allocated;
self.stats.fragmentation_ratio = if total_managed > 0 {
1.0 - (total_allocated as f64 / total_managed as f64)
} else {
0.0
};
self.stats.allocation_efficiency = if self.stats.total_allocations > 0 {
let successful_allocations = self.stats.total_allocations;
successful_allocations as f64 / self.stats.total_allocations as f64
} else {
1.0
};
}
}
unsafe impl Send for GpuMemorySystem {}
unsafe impl Sync for GpuMemorySystem {}
#[derive(Debug)]
pub enum GpuMemorySystemError {
BackendError(UnifiedGpuError),
AllocationError(String),
ManagementError(String),
InvalidPointer(String),
SystemNotStarted,
ConfigurationError(String),
OptimizationFailed(String),
InternalError(String),
}
impl std::fmt::Display for GpuMemorySystemError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
GpuMemorySystemError::BackendError(err) => write!(f, "Backend error: {}", err),
GpuMemorySystemError::AllocationError(msg) => write!(f, "Allocation error: {}", msg),
GpuMemorySystemError::ManagementError(msg) => write!(f, "Management error: {}", msg),
GpuMemorySystemError::InvalidPointer(msg) => write!(f, "Invalid pointer: {}", msg),
GpuMemorySystemError::SystemNotStarted => write!(f, "System not started"),
GpuMemorySystemError::ConfigurationError(msg) => {
write!(f, "Configuration error: {}", msg)
}
GpuMemorySystemError::OptimizationFailed(msg) => {
write!(f, "Optimization failed: {}", msg)
}
GpuMemorySystemError::InternalError(msg) => write!(f, "Internal error: {}", msg),
}
}
}
impl std::error::Error for GpuMemorySystemError {}
impl From<UnifiedGpuError> for GpuMemorySystemError {
fn from(err: UnifiedGpuError) -> Self {
GpuMemorySystemError::BackendError(err)
}
}
impl From<allocation::AllocationError> for GpuMemorySystemError {
fn from(err: allocation::AllocationError) -> Self {
GpuMemorySystemError::AllocationError(format!("{}", err))
}
}
impl From<MemoryManagementError> for GpuMemorySystemError {
fn from(err: MemoryManagementError) -> Self {
GpuMemorySystemError::ManagementError(format!("{}", err))
}
}
pub struct ThreadSafeGpuMemorySystem {
system: Arc<Mutex<GpuMemorySystem>>,
}
impl ThreadSafeGpuMemorySystem {
pub fn new(config: GpuMemorySystemConfig) -> Result<Self, GpuMemorySystemError> {
let system = GpuMemorySystem::new(config)?;
Ok(Self {
system: Arc::new(Mutex::new(system)),
})
}
pub fn allocate(
&self,
size: usize,
alignment: Option<usize>,
) -> Result<*mut c_void, GpuMemorySystemError> {
let mut system = self.system.lock().map_err(|_| {
GpuMemorySystemError::InternalError("memory system lock poisoned".into())
})?;
system.allocate(size, alignment)
}
pub fn free(&self, ptr: *mut c_void) -> Result<(), GpuMemorySystemError> {
let mut system = self.system.lock().map_err(|_| {
GpuMemorySystemError::InternalError("memory system lock poisoned".into())
})?;
system.free(ptr)
}
pub fn get_stats(&self) -> SystemStats {
let mut system = self.system.lock().unwrap_or_else(|e| e.into_inner());
system.get_stats()
}
pub fn optimize(&self) -> Result<(), GpuMemorySystemError> {
let mut system = self.system.lock().map_err(|_| {
GpuMemorySystemError::InternalError("memory system lock poisoned".into())
})?;
system.optimize()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_config() -> GpuMemorySystemConfig {
let mut config = GpuMemorySystemConfig::default();
config.system_config.memory_budget = 0.001; if let VendorConfig::Cuda(ref mut cuda_config) = config.vendor_config {
cuda_config.enable_memory_pools = false;
}
config
}
#[test]
fn test_system_creation() {
let config = create_test_config();
let system = GpuMemorySystem::new(config);
assert!(system.is_ok() || system.is_err());
}
#[test]
fn test_auto_create() {
let system = GpuMemorySystem::auto_create();
assert!(system.is_ok() || system.is_err());
}
#[test]
fn test_thread_safe_wrapper() {
let config = create_test_config();
let system = ThreadSafeGpuMemorySystem::new(config);
assert!(system.is_ok() || system.is_err());
}
#[test]
fn test_allocator_selection() {
let config = create_test_config();
if let Ok(system) = GpuMemorySystem::new(config) {
assert_eq!(system.choose_allocator(512), AllocatorType::Slab);
assert_eq!(system.choose_allocator(64 * 1024), AllocatorType::Buddy);
assert_eq!(
system.choose_allocator(2 * 1024 * 1024),
AllocatorType::Arena
);
}
}
}