use super::error::{ComputeError, ComputeResult};
use super::pressure_monitor::MonitoringStatistics;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct BufferHandle(u64);
impl BufferHandle {
pub fn new(id: u64) -> Self {
Self(id)
}
pub fn id(&self) -> u64 {
self.0
}
}
#[derive(Debug, Clone)]
pub struct MemoryStats {
pub total_allocated: usize,
pub available: usize,
pub buffer_count: usize,
pub fragmentation_ratio: f32,
}
impl Default for MemoryStats {
fn default() -> Self {
Self {
total_allocated: 0,
available: 1024 * 1024 * 1024, buffer_count: 0,
fragmentation_ratio: 0.0,
}
}
}
#[derive(Debug, Hash, PartialEq, Eq, Clone)]
pub enum BufferSize {
Small, Medium, Large, XLarge, }
#[derive(Debug, Clone)]
pub struct BufferInfo {
pub size: usize,
pub size_category: BufferSize,
pub in_use: bool,
pub last_used: std::time::Instant,
pub creation_time: std::time::Instant,
pub usage_count: usize,
}
impl BufferInfo {
pub fn new(size: usize) -> Self {
let size_category = match size {
0..=1024 => BufferSize::Small,
1025..=1_048_576 => BufferSize::Medium,
1_048_577..=104_857_600 => BufferSize::Large,
_ => BufferSize::XLarge,
};
Self {
size,
size_category,
in_use: false,
last_used: std::time::Instant::now(),
creation_time: std::time::Instant::now(),
usage_count: 0,
}
}
}
pub struct CpuMemoryManager {
buffers: Arc<Mutex<HashMap<BufferHandle, Vec<u8>>>>,
buffer_info: Arc<Mutex<HashMap<BufferHandle, BufferInfo>>>,
next_id: Arc<Mutex<u64>>,
stats: Arc<Mutex<MemoryStats>>,
}
impl CpuMemoryManager {
pub fn new() -> Self {
Self {
buffers: Arc::new(Mutex::new(HashMap::new())),
buffer_info: Arc::new(Mutex::new(HashMap::new())),
next_id: Arc::new(Mutex::new(1)),
stats: Arc::new(Mutex::new(MemoryStats::default())),
}
}
pub fn allocate(&self, size: usize) -> Result<BufferHandle, ComputeError> {
let mut buffers = self.buffers.lock().unwrap();
let mut buffer_info = self.buffer_info.lock().unwrap();
let mut next_id = self.next_id.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
let id = *next_id;
*next_id += 1;
let handle = BufferHandle(id);
let buffer = vec![0u8; size];
buffers.insert(handle, buffer);
let mut info = BufferInfo::new(size);
info.in_use = true;
buffer_info.insert(handle, info);
stats.total_allocated += size;
stats.buffer_count += 1;
Ok(handle)
}
pub fn deallocate(&self, handle: BufferHandle) -> Result<(), ComputeError> {
let mut buffers = self.buffers.lock().unwrap();
let mut buffer_info = self.buffer_info.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
if let Some(info) = buffer_info.get(&handle) {
stats.total_allocated -= info.size;
stats.buffer_count -= 1;
}
buffers.remove(&handle);
buffer_info.remove(&handle);
Ok(())
}
pub fn read(&self, handle: BufferHandle) -> Result<Vec<u8>, ComputeError> {
let buffers = self.buffers.lock().unwrap();
buffers
.get(&handle)
.cloned()
.ok_or_else(|| ComputeError::General("Buffer not found".to_string()))
}
pub fn write(&self, handle: BufferHandle, data: &[u8]) -> Result<(), ComputeError> {
let mut buffers = self.buffers.lock().unwrap();
let buffer = buffers
.get_mut(&handle)
.ok_or_else(|| ComputeError::General("Buffer not found".to_string()))?;
if buffer.len() != data.len() {
return Err(ComputeError::General("Buffer size mismatch".to_string()));
}
buffer.copy_from_slice(data);
Ok(())
}
pub fn get_stats(&self) -> MemoryStats {
self.stats.lock().unwrap().clone()
}
}
impl Default for CpuMemoryManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "gpu")]
pub mod webgpu_memory {
use super::*;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
#[derive(Debug)]
pub struct GpuBuffer {
pub id: u64,
pub size: u64,
pub usage: wgpu::BufferUsages,
_marker: std::marker::PhantomData<()>,
}
#[derive(Debug)]
pub struct GpuMemoryPool {
buffers: HashMap<BufferSize, Vec<GpuBuffer>>,
active_buffers: HashMap<u64, GpuBuffer>,
next_id: u64,
}
impl Default for GpuMemoryPool {
fn default() -> Self {
Self::new()
}
}
impl GpuMemoryPool {
pub fn new() -> Self {
Self {
buffers: HashMap::new(),
active_buffers: HashMap::new(),
next_id: 1,
}
}
pub fn allocate(&mut self, size: u64, usage: wgpu::BufferUsages) -> u64 {
let id = self.next_id;
self.next_id += 1;
let buffer = GpuBuffer {
id,
size,
usage,
_marker: std::marker::PhantomData,
};
self.active_buffers.insert(id, buffer);
id
}
pub fn deallocate(&mut self, id: u64) -> Option<GpuBuffer> {
self.active_buffers.remove(&id)
}
}
pub struct WebGPUMemoryManager {
pool: Arc<Mutex<GpuMemoryPool>>,
stats: Arc<Mutex<super::MemoryStats>>,
}
impl Default for WebGPUMemoryManager {
fn default() -> Self {
Self::new()
}
}
impl WebGPUMemoryManager {
pub fn new() -> Self {
Self {
pool: Arc::new(Mutex::new(GpuMemoryPool::new())),
stats: Arc::new(Mutex::new(super::MemoryStats::default())),
}
}
pub fn allocate_buffer(&self, size: usize) -> Result<super::BufferHandle, ComputeError> {
let mut pool = self.pool.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
let id = pool.allocate(size as u64, wgpu::BufferUsages::STORAGE);
stats.total_allocated += size;
stats.buffer_count += 1;
Ok(super::BufferHandle::new(id))
}
pub fn deallocate_buffer(&self, handle: super::BufferHandle) -> Result<(), ComputeError> {
let mut pool = self.pool.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
if let Some(buffer) = pool.deallocate(handle.id()) {
stats.total_allocated = stats.total_allocated.saturating_sub(buffer.size as usize);
stats.buffer_count = stats.buffer_count.saturating_sub(1);
}
Ok(())
}
pub fn get_stats(&self) -> super::MemoryStats {
self.stats.lock().unwrap().clone()
}
}
}
pub struct GpuMemoryManager {
#[cfg(feature = "gpu")]
webgpu_manager: Option<webgpu_memory::WebGPUMemoryManager>,
cpu_manager: CpuMemoryManager,
buffer_pools: Arc<Mutex<HashMap<BufferSize, VecDeque<BufferHandle>>>>,
allocated_buffers: Arc<Mutex<HashMap<BufferHandle, BufferInfo>>>,
stats: Arc<Mutex<MemoryManagerStats>>,
}
#[derive(Debug, Clone)]
pub struct MemoryManagerStats {
allocations_by_size: HashMap<BufferSize, usize>,
deallocations_by_size: HashMap<BufferSize, usize>,
pool_hits: HashMap<BufferSize, usize>,
pool_misses: HashMap<BufferSize, usize>,
total_allocations: usize,
total_deallocations: usize,
peak_memory_usage: usize,
current_pool_sizes: HashMap<BufferSize, usize>,
}
impl GpuMemoryManager {
pub fn new() -> Self {
Self {
#[cfg(feature = "gpu")]
webgpu_manager: None, cpu_manager: CpuMemoryManager::new(),
buffer_pools: Arc::new(Mutex::new(HashMap::new())),
allocated_buffers: Arc::new(Mutex::new(HashMap::new())),
stats: Arc::new(Mutex::new(MemoryManagerStats::new())),
}
}
pub fn allocate(&self, size: usize) -> Result<BufferHandle, ComputeError> {
let size_category = match size {
0..=1024 => BufferSize::Small,
1025..=1_048_576 => BufferSize::Medium,
1_048_577..=104_857_600 => BufferSize::Large,
_ => BufferSize::XLarge,
};
{
let mut pools = self.buffer_pools.lock().unwrap();
let mut allocated = self.allocated_buffers.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
if let Some(pool) = pools.get_mut(&size_category) {
while let Some(handle) = pool.pop_front() {
if let Some(info) = allocated.get_mut(&handle) {
if !info.in_use && info.size >= size {
info.in_use = true;
info.last_used = std::time::Instant::now();
info.usage_count += 1;
*stats.pool_hits.entry(size_category.clone()).or_insert(0) += 1;
return Ok(handle);
}
}
}
}
*stats.pool_misses.entry(size_category.clone()).or_insert(0) += 1;
}
#[cfg(feature = "gpu")]
if let Some(ref webgpu) = self.webgpu_manager {
let handle = webgpu.allocate_buffer(size)?;
self.track_allocation(handle, size);
return Ok(handle);
}
let handle = self.cpu_manager.allocate(size)?;
self.track_allocation(handle, size);
Ok(handle)
}
pub fn deallocate(&self, handle: BufferHandle) -> Result<(), ComputeError> {
let mut allocated = self.allocated_buffers.lock().unwrap();
let mut pools = self.buffer_pools.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
if let Some(mut info) = allocated.remove(&handle) {
info.in_use = false;
let pool = pools.entry(info.size_category.clone()).or_default();
const MAX_POOL_SIZE: usize = 100;
if pool.len() < MAX_POOL_SIZE {
allocated.insert(handle, info.clone());
pool.push_back(handle);
*stats
.current_pool_sizes
.entry(info.size_category.clone())
.or_insert(0) = pool.len();
} else {
#[cfg(feature = "gpu")]
if let Some(ref webgpu) = self.webgpu_manager {
return webgpu.deallocate_buffer(handle);
}
self.cpu_manager.deallocate(handle)?;
}
*stats
.deallocations_by_size
.entry(info.size_category)
.or_insert(0) += 1;
stats.total_deallocations += 1;
}
Ok(())
}
pub fn read(&self, handle: BufferHandle) -> Result<Vec<u8>, ComputeError> {
self.cpu_manager.read(handle)
}
pub fn write(&self, handle: BufferHandle, data: &[u8]) -> Result<(), ComputeError> {
self.cpu_manager.write(handle, data)
}
pub fn get_stats(&self) -> MemoryStats {
#[cfg(feature = "gpu")]
if let Some(ref webgpu) = self.webgpu_manager {
return webgpu.get_stats();
}
self.cpu_manager.get_stats()
}
pub fn get_detailed_stats(&self) -> MemoryManagerStats {
self.stats.lock().unwrap().clone()
}
pub fn cleanup_unused_buffers(&self, age_threshold: std::time::Duration) {
let mut allocated = self.allocated_buffers.lock().unwrap();
let mut pools = self.buffer_pools.lock().unwrap();
let now = std::time::Instant::now();
let mut to_remove = Vec::new();
for (handle, info) in allocated.iter() {
if !info.in_use && now.duration_since(info.last_used) > age_threshold {
to_remove.push(*handle);
}
}
for handle in to_remove {
for pool in pools.values_mut() {
pool.retain(|&h| h != handle);
}
allocated.remove(&handle);
#[cfg(feature = "gpu")]
if let Some(ref webgpu) = self.webgpu_manager {
let _ = webgpu.deallocate_buffer(handle);
continue;
}
let _ = self.cpu_manager.deallocate(handle);
}
}
fn track_allocation(&self, handle: BufferHandle, size: usize) {
let mut allocated = self.allocated_buffers.lock().unwrap();
let mut stats = self.stats.lock().unwrap();
let mut info = BufferInfo::new(size);
info.in_use = true;
info.usage_count = 1;
allocated.insert(handle, info.clone());
*stats
.allocations_by_size
.entry(info.size_category)
.or_insert(0) += 1;
stats.total_allocations += 1;
let current_usage: usize = allocated
.values()
.filter(|info| info.in_use)
.map(|info| info.size)
.sum();
if current_usage > stats.peak_memory_usage {
stats.peak_memory_usage = current_usage;
}
}
}
impl Default for GpuMemoryManager {
fn default() -> Self {
Self::new()
}
}
impl MemoryManagerStats {
fn new() -> Self {
Self {
allocations_by_size: HashMap::new(),
deallocations_by_size: HashMap::new(),
pool_hits: HashMap::new(),
pool_misses: HashMap::new(),
total_allocations: 0,
total_deallocations: 0,
peak_memory_usage: 0,
current_pool_sizes: HashMap::new(),
}
}
}
use super::buffer_pool::BufferCategory;
#[cfg(feature = "gpu")]
use super::buffer_pool::PoolStatisticsSnapshot;
#[cfg(feature = "gpu")]
use super::buffer_pool::{AdvancedBufferPool, MemoryPressure};
#[derive(Debug, Clone)]
pub struct GpuMemoryConfig {
pub enable_advanced_features: bool,
pub enable_monitoring: bool,
pub auto_start_monitoring: bool,
pub pressure_threshold: f32,
pub enable_circuit_breaker: bool,
pub max_allocation_latency: Duration,
pub enable_optimization: bool,
}
impl Default for GpuMemoryConfig {
fn default() -> Self {
Self {
enable_advanced_features: true,
enable_monitoring: true,
auto_start_monitoring: false, pressure_threshold: 0.8,
enable_circuit_breaker: true,
max_allocation_latency: Duration::from_millis(100),
enable_optimization: true,
}
}
}
#[derive(Debug, Clone)]
pub struct EnhancedMemoryStats {
#[cfg(feature = "gpu")]
pub pool_stats: Option<PoolStatisticsSnapshot>,
#[cfg(feature = "gpu")]
pub monitoring_stats: Option<MonitoringStatistics>,
pub legacy_stats: Option<MemoryStats>,
#[cfg(feature = "gpu")]
pub monitoring_enabled: bool,
#[cfg(feature = "gpu")]
pub current_pressure: MemoryPressure,
pub enhanced_features_available: bool,
}
impl EnhancedMemoryStats {
pub fn cache_hit_ratio(&self) -> f32 {
#[cfg(feature = "gpu")]
if let Some(ref stats) = self.pool_stats {
return stats.cache_hit_ratio();
}
0.0 }
pub fn total_allocated(&self) -> u64 {
#[cfg(feature = "gpu")]
if let Some(ref stats) = self.pool_stats {
return stats.global.total_memory_allocated;
}
if let Some(ref legacy) = self.legacy_stats {
return legacy.total_allocated as u64;
}
0
}
pub fn pressure_ratio(&self) -> f32 {
#[cfg(feature = "gpu")]
if self.current_pressure != MemoryPressure::None {
match self.current_pressure {
MemoryPressure::None => 0.0,
MemoryPressure::Low => 0.2,
MemoryPressure::Medium => 0.5,
MemoryPressure::High => 0.8,
MemoryPressure::Critical => 1.0,
}
} else if let Some(ref legacy) = self.legacy_stats {
if legacy.available > 0 {
legacy.total_allocated as f32 / (legacy.total_allocated + legacy.available) as f32
} else {
0.0
}
} else {
0.0
}
#[cfg(not(feature = "gpu"))]
if let Some(ref legacy) = self.legacy_stats {
if legacy.available > 0 {
legacy.total_allocated as f32 / (legacy.total_allocated + legacy.available) as f32
} else {
0.0
}
} else {
0.0
}
}
pub fn avg_allocation_latency_ns(&self) -> u64 {
#[cfg(feature = "gpu")]
if let Some(ref stats) = self.pool_stats {
return stats.global.avg_allocation_latency_ns;
}
1_000_000 }
pub fn performance_summary(&self) -> String {
#[cfg(feature = "gpu")]
if let Some(ref stats) = self.pool_stats {
return stats.performance_summary();
}
format!(
"Legacy Memory: {:.1}MB allocated, {:.1}% pressure",
self.total_allocated() as f64 / (1024.0 * 1024.0),
self.pressure_ratio() * 100.0
)
}
}
pub struct EnhancedGpuMemoryManager {
#[cfg(feature = "gpu")]
advanced_pool: Option<Arc<AdvancedBufferPool>>,
config: GpuMemoryConfig,
#[cfg(feature = "gpu")]
device: Arc<super::device::GpuDevice>,
#[cfg(feature = "gpu")]
statistics_cache: Arc<Mutex<Option<(Instant, EnhancedMemoryStats)>>>,
#[cfg(feature = "gpu")]
last_optimization: Arc<Mutex<Instant>>,
initialization_time: Instant,
}
impl EnhancedGpuMemoryManager {
#[cfg(feature = "gpu")]
pub fn new(device: Arc<super::device::GpuDevice>) -> ComputeResult<Self> {
Self::with_config(device, GpuMemoryConfig::default())
}
#[cfg(not(feature = "gpu"))]
pub fn new() -> ComputeResult<Self> {
Self::with_config(GpuMemoryConfig::default())
}
#[cfg(feature = "gpu")]
pub fn with_config(
device: Arc<super::device::GpuDevice>,
config: GpuMemoryConfig,
) -> ComputeResult<Self> {
let mut manager = Self {
advanced_pool: None,
config: config.clone(),
device: device.clone(),
statistics_cache: Arc::new(Mutex::new(None)),
last_optimization: Arc::new(Mutex::new(Instant::now())),
initialization_time: Instant::now(),
};
if config.enable_advanced_features {
manager.initialize_advanced_features()?;
}
Ok(manager)
}
#[cfg(not(feature = "gpu"))]
pub fn with_config(config: GpuMemoryConfig) -> ComputeResult<Self> {
let manager = Self {
config: config.clone(),
initialization_time: Instant::now(),
};
Ok(manager)
}
#[cfg(feature = "gpu")]
fn initialize_advanced_features(&mut self) -> ComputeResult<()> {
log::info!("Enhanced GPU memory manager initialized with advanced features");
Ok(())
}
pub fn allocate_buffer(&self, size: usize) -> ComputeResult<BufferHandle> {
#[cfg(feature = "gpu")]
if let Some(ref pool) = self.advanced_pool {
let buffer = pool.get_buffer(
size as u64,
::wgpu::BufferUsages::STORAGE
| ::wgpu::BufferUsages::COPY_DST
| ::wgpu::BufferUsages::COPY_SRC,
Some("enhanced_buffer"),
)?;
return Ok(BufferHandle::new(buffer.allocation_id()));
}
Err(ComputeError::General(
"No memory manager available".to_string(),
))
}
pub fn create_storage_buffer(
&self,
size: u64,
_label: Option<&str>,
) -> ComputeResult<BufferHandle> {
self.allocate_buffer(size as usize)
}
pub fn create_uniform_buffer(
&self,
size: u64,
_label: Option<&str>,
) -> ComputeResult<BufferHandle> {
#[cfg(feature = "gpu")]
if let Some(ref pool) = self.advanced_pool {
let buffer = pool.get_buffer(
size,
::wgpu::BufferUsages::UNIFORM | ::wgpu::BufferUsages::COPY_DST,
_label,
)?;
return Ok(BufferHandle::new(buffer.allocation_id()));
}
self.allocate_buffer(size as usize)
}
pub fn create_readback_buffer(
&self,
size: u64,
_label: Option<&str>,
) -> ComputeResult<BufferHandle> {
#[cfg(feature = "gpu")]
if let Some(ref pool) = self.advanced_pool {
let buffer = pool.get_buffer(
size,
::wgpu::BufferUsages::COPY_DST | ::wgpu::BufferUsages::MAP_READ,
_label,
)?;
return Ok(BufferHandle::new(buffer.allocation_id()));
}
self.allocate_buffer(size as usize)
}
pub fn deallocate_buffer(&self, handle: BufferHandle) -> ComputeResult<()> {
#[cfg(feature = "gpu")]
if let Some(ref _pool) = self.advanced_pool {
println!("Deallocating buffer {}", handle.id());
return Ok(());
}
Ok(())
}
pub fn get_stats(&self) -> ComputeResult<EnhancedMemoryStats> {
#[cfg(feature = "gpu")]
{
if let Ok(cache) = self.statistics_cache.lock() {
if let Some((timestamp, stats)) = cache.as_ref() {
if timestamp.elapsed() < Duration::from_millis(100) {
return Ok(stats.clone());
}
}
}
}
#[cfg_attr(not(feature = "gpu"), allow(unused_mut))]
let mut stats = EnhancedMemoryStats {
#[cfg(feature = "gpu")]
pool_stats: None,
#[cfg(feature = "gpu")]
monitoring_stats: None,
legacy_stats: None,
#[cfg(feature = "gpu")]
monitoring_enabled: self.config.enable_monitoring,
#[cfg(feature = "gpu")]
current_pressure: MemoryPressure::None,
enhanced_features_available: self.config.enable_advanced_features,
};
#[cfg(feature = "gpu")]
{
if let Some(ref pool) = self.advanced_pool {
stats.pool_stats = Some(pool.get_statistics());
}
if let Ok(mut cache) = self.statistics_cache.lock() {
*cache = Some((Instant::now(), stats.clone()));
}
}
Ok(stats)
}
pub fn cleanup(&self, _aggressiveness: f32) -> ComputeResult<()> {
#[cfg(feature = "gpu")]
if let Some(ref _pool) = self.advanced_pool {
println!("Cleaning up advanced buffer pool with aggressiveness: {_aggressiveness}");
}
Ok(())
}
pub fn uptime(&self) -> Duration {
self.initialization_time.elapsed()
}
pub fn is_enhanced(&self) -> bool {
#[cfg(feature = "gpu")]
return self.advanced_pool.is_some();
#[cfg(not(feature = "gpu"))]
false
}
}
#[derive(Debug, Clone, Default)]
pub struct OptimizationResult {
pub memory_reclaimed: u64,
pub fragmentation_reduced: f32,
pub performance_improvement: f32,
pub operations_optimized: u64,
}
impl From<BufferSize> for BufferCategory {
fn from(size: BufferSize) -> Self {
match size {
BufferSize::Small => BufferCategory::Micro,
BufferSize::Medium => BufferCategory::Small,
BufferSize::Large => BufferCategory::Medium,
BufferSize::XLarge => BufferCategory::Large,
}
}
}
pub use self::EnhancedGpuMemoryManager as GpuMemoryManagerEnhanced;
pub use self::GpuMemoryManager as WebGPUMemoryManager;