use super::config::{CacheEvictionStrategy, MemoryConfig, TransformerBasedOptimizerConfig};
use crate::error::Result;
use scirs2_core::ndarray::{Array1, Array2, Array3, Axis};
use scirs2_core::numeric::Float;
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::fmt::Debug;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum MemoryManagementStrategy {
FIFO,
LRU,
LFU,
ARC,
Compressed,
Hierarchical,
}
pub struct TransformerMemoryManager<T: Float + Debug + Send + Sync + 'static> {
strategy: MemoryManagementStrategy,
config: MemoryConfig,
primary_cache: MemoryCache<T>,
secondary_cache: Option<MemoryCache<T>>,
compression_manager: Option<CompressionManager<T>>,
statistics: MemoryStatistics,
access_tracker: AccessTracker,
pressure_monitor: MemoryPressureMonitor,
model_dimension: usize,
}
impl<T: Float + Debug + Send + Sync + 'static> TransformerMemoryManager<T> {
pub fn new(config: &TransformerBasedOptimizerConfig<T>) -> Result<Self> {
let memory_config = config.memory_config.clone();
let strategy = match memory_config.eviction_strategy {
CacheEvictionStrategy::LRU => MemoryManagementStrategy::LRU,
CacheEvictionStrategy::LFU => MemoryManagementStrategy::LFU,
CacheEvictionStrategy::FIFO => MemoryManagementStrategy::FIFO,
CacheEvictionStrategy::Random => MemoryManagementStrategy::LRU, };
let primary_cache = MemoryCache::new(
memory_config.max_cache_size / 2,
memory_config.eviction_strategy,
)?;
let secondary_cache = if memory_config.max_cache_size > 1024 * 1024 * 100 {
Some(MemoryCache::new(
memory_config.max_cache_size / 2,
CacheEvictionStrategy::FIFO,
)?)
} else {
None
};
let compression_manager = if memory_config.enable_compression {
Some(CompressionManager::new(0.5)?) } else {
None
};
let statistics = MemoryStatistics::new();
let access_tracker = AccessTracker::new(1000);
let pressure_monitor = MemoryPressureMonitor::new();
Ok(Self {
strategy,
config: memory_config,
primary_cache,
secondary_cache,
compression_manager,
statistics,
access_tracker,
pressure_monitor,
model_dimension: config.model_dimension,
})
}
pub fn store(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
let start_time = Instant::now();
self.pressure_monitor.update(self.get_memory_usage());
if self.pressure_monitor.is_high_pressure() {
self.evict_memory()?;
}
let storage_result = match self.strategy {
MemoryManagementStrategy::LRU => self.store_lru(key.clone(), tensor.clone()),
MemoryManagementStrategy::LFU => self.store_lfu(key.clone(), tensor.clone()),
MemoryManagementStrategy::FIFO => self.store_fifo(key.clone(), tensor.clone()),
MemoryManagementStrategy::ARC => self.store_arc(key.clone(), tensor.clone()),
MemoryManagementStrategy::Compressed => {
self.store_compressed(key.clone(), tensor.clone())
}
MemoryManagementStrategy::Hierarchical => {
self.store_hierarchical(key.clone(), tensor.clone())
}
};
let tensor_len = tensor.len();
if storage_result.is_err() && self.secondary_cache.is_some() {
if let Some(ref mut secondary) = self.secondary_cache {
secondary.store(key.clone(), tensor)?;
}
}
let storage_time = start_time.elapsed();
self.statistics.record_storage(tensor_len, storage_time);
self.access_tracker.record_write(key);
storage_result
}
pub fn retrieve(&mut self, key: &str) -> Result<Option<Array2<T>>> {
let start_time = Instant::now();
let result = self.primary_cache.retrieve(key)?;
if result.is_some() {
self.access_tracker.record_read(key.to_string());
let retrieval_time = start_time.elapsed();
self.statistics.record_retrieval(retrieval_time, true);
return Ok(result);
}
if let Some(ref mut secondary) = self.secondary_cache {
let result = secondary.retrieve(key)?;
if result.is_some() {
self.access_tracker.record_read(key.to_string());
let retrieval_time = start_time.elapsed();
self.statistics.record_retrieval(retrieval_time, true);
return Ok(result);
}
}
if let Some(ref mut compression) = self.compression_manager {
if let Some(compressed_data) = compression.retrieve(key)? {
let decompressed = compression.decompress(&compressed_data)?;
self.access_tracker.record_read(key.to_string());
let retrieval_time = start_time.elapsed();
self.statistics.record_retrieval(retrieval_time, true);
return Ok(Some(decompressed));
}
}
let retrieval_time = start_time.elapsed();
self.statistics.record_retrieval(retrieval_time, false);
Ok(None)
}
pub fn remove(&mut self, key: &str) -> Result<bool> {
let mut removed = false;
if self.primary_cache.remove(key)? {
removed = true;
}
if let Some(ref mut secondary) = self.secondary_cache {
if secondary.remove(key)? {
removed = true;
}
}
if let Some(ref mut compression) = self.compression_manager {
if compression.remove(key)? {
removed = true;
}
}
self.access_tracker.record_removal(key.to_string());
Ok(removed)
}
pub fn clear(&mut self) -> Result<()> {
self.primary_cache.clear()?;
if let Some(ref mut secondary) = self.secondary_cache {
secondary.clear()?;
}
if let Some(ref mut compression) = self.compression_manager {
compression.clear()?;
}
self.statistics.reset();
self.access_tracker.clear();
self.pressure_monitor.reset();
Ok(())
}
pub fn get_memory_usage(&self) -> usize {
let primary_usage = self.primary_cache.get_memory_usage();
let secondary_usage = self
.secondary_cache
.as_ref()
.map(|cache| cache.get_memory_usage())
.unwrap_or(0);
let compression_usage = self
.compression_manager
.as_ref()
.map(|comp| comp.get_memory_usage())
.unwrap_or(0);
primary_usage + secondary_usage + compression_usage
}
pub fn optimize_memory(&mut self) -> Result<OptimizationReport> {
let start_time = Instant::now();
let initial_usage = self.get_memory_usage();
let access_patterns = self.access_tracker.analyze_patterns();
self.reorganize_by_frequency(&access_patterns)?;
if let Some(ref mut compression) = self.compression_manager {
compression.optimize_compression_ratios(&access_patterns)?;
}
self.defragment_memory()?;
let final_usage = self.get_memory_usage();
let optimization_time = start_time.elapsed();
Ok(OptimizationReport {
initial_memory_usage: initial_usage,
final_memory_usage: final_usage,
memory_saved: initial_usage.saturating_sub(final_usage),
optimization_time,
operations_performed: access_patterns.total_accesses,
})
}
pub fn prefetch(&mut self, keys: Vec<String>) -> Result<usize> {
let mut prefetched_count = 0;
for key in keys {
if !self.primary_cache.contains(&key) {
if let Some(ref mut secondary) = self.secondary_cache {
if let Some(tensor) = secondary.retrieve(&key)? {
if self.primary_cache.store(key.clone(), tensor).is_ok() {
secondary.remove(&key)?;
prefetched_count += 1;
}
}
}
if let Some(ref mut compression) = self.compression_manager {
if let Some(compressed_data) = compression.retrieve(&key)? {
let decompressed = compression.decompress(&compressed_data)?;
if self.primary_cache.store(key.clone(), decompressed).is_ok() {
compression.remove(&key)?;
prefetched_count += 1;
}
}
}
}
}
Ok(prefetched_count)
}
fn store_lru(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
self.primary_cache.store(key, tensor)
}
fn store_lfu(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
self.primary_cache.store(key, tensor)
}
fn store_fifo(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
self.primary_cache.store(key, tensor)
}
fn store_arc(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
self.primary_cache.store(key, tensor)
}
fn store_compressed(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
if let Some(ref mut compression) = self.compression_manager {
let compressed_data = compression.compress(&tensor)?;
compression.store(key, compressed_data)?;
Ok(())
} else {
self.primary_cache.store(key, tensor)
}
}
fn store_hierarchical(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
let tensor_size = tensor.len() * std::mem::size_of::<T>();
if tensor_size < self.config.allocation_block_size {
self.primary_cache.store(key, tensor)
} else if let Some(ref mut secondary) = self.secondary_cache {
secondary.store(key, tensor)
} else {
self.primary_cache.store(key, tensor)
}
}
fn evict_memory(&mut self) -> Result<()> {
self.primary_cache.evict_lru()?;
if self.pressure_monitor.is_high_pressure() {
if let Some(ref mut secondary) = self.secondary_cache {
secondary.evict_lru()?;
}
}
Ok(())
}
fn reorganize_by_frequency(&mut self, patterns: &AccessPatterns) -> Result<()> {
let frequent_keys: Vec<String> = patterns
.frequency_map
.iter()
.filter(|(_, &count)| count as f64 > patterns.average_frequency)
.map(|(key, _)| key.clone())
.collect();
self.prefetch(frequent_keys)?;
Ok(())
}
fn defragment_memory(&mut self) -> Result<()> {
let primary_items = self.primary_cache.get_all_items()?;
self.primary_cache.clear()?;
for (key, tensor) in primary_items {
self.primary_cache.store(key, tensor)?;
}
Ok(())
}
pub fn get_statistics(&self) -> &MemoryStatistics {
&self.statistics
}
pub fn get_access_patterns(&self) -> AccessPatterns {
self.access_tracker.analyze_patterns()
}
pub fn set_strategy(&mut self, strategy: MemoryManagementStrategy) {
self.strategy = strategy;
}
pub fn get_memory_pressure(&self) -> f64 {
self.pressure_monitor.get_pressure_ratio()
}
}
pub struct MemoryCache<T: Float + Debug + Send + Sync + 'static> {
storage: HashMap<String, CacheEntry<T>>,
access_order: VecDeque<String>,
access_frequency: HashMap<String, usize>,
max_size: usize,
current_size: usize,
eviction_strategy: CacheEvictionStrategy,
}
impl<T: Float + Debug + Send + Sync + 'static> MemoryCache<T> {
pub fn new(max_size: usize, eviction_strategy: CacheEvictionStrategy) -> Result<Self> {
Ok(Self {
storage: HashMap::new(),
access_order: VecDeque::new(),
access_frequency: HashMap::new(),
max_size,
current_size: 0,
eviction_strategy,
})
}
pub fn store(&mut self, key: String, tensor: Array2<T>) -> Result<()> {
let tensor_size = tensor.len() * std::mem::size_of::<T>();
while self.current_size + tensor_size > self.max_size && !self.storage.is_empty() {
self.evict_one()?;
}
if tensor_size > self.max_size {
return Err(crate::error::OptimError::Other(
"Tensor too large for cache".to_string(),
));
}
if let Some(old_entry) = self.storage.remove(&key) {
self.current_size -= old_entry.size;
self.remove_from_access_order(&key);
}
let entry = CacheEntry {
tensor,
size: tensor_size,
access_time: Instant::now(),
access_count: 1,
};
self.storage.insert(key.clone(), entry);
self.current_size += tensor_size;
self.update_access_tracking(&key);
Ok(())
}
pub fn retrieve(&mut self, key: &str) -> Result<Option<Array2<T>>> {
let tensor_result = if let Some(entry) = self.storage.get_mut(key) {
entry.access_time = Instant::now();
entry.access_count += 1;
Some(entry.tensor.clone())
} else {
None
};
if tensor_result.is_some() {
self.update_access_tracking(key);
}
Ok(tensor_result)
}
pub fn remove(&mut self, key: &str) -> Result<bool> {
if let Some(entry) = self.storage.remove(key) {
self.current_size -= entry.size;
self.remove_from_access_order(key);
self.access_frequency.remove(key);
Ok(true)
} else {
Ok(false)
}
}
pub fn contains(&self, key: &str) -> bool {
self.storage.contains_key(key)
}
pub fn clear(&mut self) -> Result<()> {
self.storage.clear();
self.access_order.clear();
self.access_frequency.clear();
self.current_size = 0;
Ok(())
}
pub fn get_memory_usage(&self) -> usize {
self.current_size
}
pub fn evict_lru(&mut self) -> Result<()> {
if let Some(oldest_key) = self.access_order.front().cloned() {
self.remove(&oldest_key)?;
}
Ok(())
}
fn evict_one(&mut self) -> Result<()> {
match self.eviction_strategy {
CacheEvictionStrategy::LRU => self.evict_lru(),
CacheEvictionStrategy::LFU => self.evict_lfu(),
CacheEvictionStrategy::FIFO => self.evict_fifo(),
CacheEvictionStrategy::Random => self.evict_random(),
}
}
fn evict_lfu(&mut self) -> Result<()> {
if let Some((min_freq, lfu_key)) = self
.access_frequency
.iter()
.min_by_key(|(_, &freq)| freq)
.map(|(key, &freq)| (freq, key.clone()))
{
self.remove(&lfu_key)?;
}
Ok(())
}
fn evict_fifo(&mut self) -> Result<()> {
if let Some(first_key) = self.access_order.front().cloned() {
self.remove(&first_key)?;
}
Ok(())
}
fn evict_random(&mut self) -> Result<()> {
if let Some(random_key) = self.storage.keys().next().cloned() {
self.remove(&random_key)?;
}
Ok(())
}
fn update_access_tracking(&mut self, key: &str) {
self.remove_from_access_order(key);
self.access_order.push_back(key.to_string());
*self.access_frequency.entry(key.to_string()).or_insert(0) += 1;
}
fn remove_from_access_order(&mut self, key: &str) {
self.access_order.retain(|k| k != key);
}
pub fn get_all_items(&self) -> Result<Vec<(String, Array2<T>)>> {
let items = self
.storage
.iter()
.map(|(key, entry)| (key.clone(), entry.tensor.clone()))
.collect();
Ok(items)
}
}
#[derive(Debug, Clone)]
pub struct CacheEntry<T: Float + Debug + Send + Sync + 'static> {
pub tensor: Array2<T>,
pub size: usize,
pub access_time: Instant,
pub access_count: usize,
}
pub struct CompressionManager<T: Float + Debug + Send + Sync + 'static> {
compressed_storage: HashMap<String, CompressedData<T>>,
compression_ratio: f64,
memory_usage: usize,
_phantom: std::marker::PhantomData<T>,
}
impl<T: Float + Debug + Send + Sync + 'static> CompressionManager<T> {
pub fn new(compression_ratio: f64) -> Result<Self> {
Ok(Self {
compressed_storage: HashMap::new(),
compression_ratio,
memory_usage: 0,
_phantom: std::marker::PhantomData,
})
}
pub fn compress(&self, tensor: &Array2<T>) -> Result<CompressedData<T>> {
let shape = tensor.shape().to_vec();
let data: Vec<T> = tensor.iter().cloned().collect();
let data_len = data.len();
Ok(CompressedData::<T> {
shape,
data,
original_size: tensor.len() * std::mem::size_of::<T>(),
compressed_size: data_len * std::mem::size_of::<T>() / 2, })
}
pub fn decompress(&self, compressed: &CompressedData<T>) -> Result<Array2<T>> {
let array = Array2::from_shape_vec(
(compressed.shape[0], compressed.shape[1]),
compressed.data.clone(),
)
.map_err(|_| crate::error::OptimError::Other("Decompression failed".to_string()))?;
Ok(array)
}
pub fn store(&mut self, key: String, compressed: CompressedData<T>) -> Result<()> {
self.memory_usage += compressed.compressed_size;
self.compressed_storage.insert(key, compressed);
Ok(())
}
pub fn retrieve(&self, key: &str) -> Result<Option<CompressedData<T>>> {
Ok(self.compressed_storage.get(key).cloned())
}
pub fn remove(&mut self, key: &str) -> Result<bool> {
if let Some(compressed) = self.compressed_storage.remove(key) {
self.memory_usage -= compressed.compressed_size;
Ok(true)
} else {
Ok(false)
}
}
pub fn clear(&mut self) -> Result<()> {
self.compressed_storage.clear();
self.memory_usage = 0;
Ok(())
}
pub fn get_memory_usage(&self) -> usize {
self.memory_usage
}
pub fn optimize_compression_ratios(&mut self, _patterns: &AccessPatterns) -> Result<()> {
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct CompressedData<T: Float + Debug + Send + Sync + 'static> {
pub shape: Vec<usize>,
pub data: Vec<T>, pub original_size: usize,
pub compressed_size: usize,
}
#[derive(Debug, Clone)]
pub struct MemoryStatistics {
pub total_stores: usize,
pub total_retrievals: usize,
pub cache_hits: usize,
pub cache_misses: usize,
pub total_bytes_stored: usize,
pub average_storage_time: Duration,
pub average_retrieval_time: Duration,
pub pressure_events: usize,
}
impl Default for MemoryStatistics {
fn default() -> Self {
Self::new()
}
}
impl MemoryStatistics {
pub fn new() -> Self {
Self {
total_stores: 0,
total_retrievals: 0,
cache_hits: 0,
cache_misses: 0,
total_bytes_stored: 0,
average_storage_time: Duration::new(0, 0),
average_retrieval_time: Duration::new(0, 0),
pressure_events: 0,
}
}
pub fn record_storage(&mut self, bytes: usize, time: Duration) {
self.total_stores += 1;
self.total_bytes_stored += bytes;
self.average_storage_time = (self.average_storage_time * (self.total_stores - 1) as u32
+ time)
/ self.total_stores as u32;
}
pub fn record_retrieval(&mut self, time: Duration, hit: bool) {
self.total_retrievals += 1;
if hit {
self.cache_hits += 1;
} else {
self.cache_misses += 1;
}
self.average_retrieval_time =
(self.average_retrieval_time * (self.total_retrievals - 1) as u32 + time)
/ self.total_retrievals as u32;
}
pub fn record_pressure_event(&mut self) {
self.pressure_events += 1;
}
pub fn get_hit_ratio(&self) -> f64 {
if self.total_retrievals > 0 {
self.cache_hits as f64 / self.total_retrievals as f64
} else {
0.0
}
}
pub fn reset(&mut self) {
*self = Self::new();
}
}
pub struct AccessTracker {
read_log: VecDeque<AccessEvent>,
write_log: VecDeque<AccessEvent>,
max_log_size: usize,
}
impl AccessTracker {
pub fn new(max_log_size: usize) -> Self {
Self {
read_log: VecDeque::new(),
write_log: VecDeque::new(),
max_log_size,
}
}
pub fn record_read(&mut self, key: String) {
self.read_log.push_back(AccessEvent {
key,
timestamp: Instant::now(),
});
if self.read_log.len() > self.max_log_size {
self.read_log.pop_front();
}
}
pub fn record_write(&mut self, key: String) {
self.write_log.push_back(AccessEvent {
key,
timestamp: Instant::now(),
});
if self.write_log.len() > self.max_log_size {
self.write_log.pop_front();
}
}
pub fn record_removal(&mut self, _key: String) {
}
pub fn analyze_patterns(&self) -> AccessPatterns {
let mut frequency_map = HashMap::new();
for event in self.read_log.iter().chain(self.write_log.iter()) {
*frequency_map.entry(event.key.clone()).or_insert(0) += 1;
}
let total_accesses: usize = frequency_map.values().sum();
let average_frequency = if frequency_map.is_empty() {
0.0
} else {
total_accesses as f64 / frequency_map.len() as f64
};
AccessPatterns {
frequency_map,
average_frequency,
total_accesses,
}
}
pub fn clear(&mut self) {
self.read_log.clear();
self.write_log.clear();
}
}
#[derive(Debug, Clone)]
pub struct AccessEvent {
pub key: String,
pub timestamp: Instant,
}
#[derive(Debug, Clone)]
pub struct AccessPatterns {
pub frequency_map: HashMap<String, usize>,
pub average_frequency: f64,
pub total_accesses: usize,
}
pub struct MemoryPressureMonitor {
current_usage: usize,
max_memory: usize,
warning_threshold: f64,
critical_threshold: f64,
pressure_history: VecDeque<f64>,
}
impl Default for MemoryPressureMonitor {
fn default() -> Self {
Self::new()
}
}
impl MemoryPressureMonitor {
pub fn new() -> Self {
Self {
current_usage: 0,
max_memory: 1024 * 1024 * 1024, warning_threshold: 0.7,
critical_threshold: 0.9,
pressure_history: VecDeque::new(),
}
}
pub fn update(&mut self, current_usage: usize) {
self.current_usage = current_usage;
let pressure_ratio = self.get_pressure_ratio();
self.pressure_history.push_back(pressure_ratio);
if self.pressure_history.len() > 100 {
self.pressure_history.pop_front();
}
}
pub fn get_pressure_ratio(&self) -> f64 {
if self.max_memory > 0 {
self.current_usage as f64 / self.max_memory as f64
} else {
0.0
}
}
pub fn is_high_pressure(&self) -> bool {
self.get_pressure_ratio() > self.critical_threshold
}
pub fn is_warning_pressure(&self) -> bool {
self.get_pressure_ratio() > self.warning_threshold
}
pub fn reset(&mut self) {
self.current_usage = 0;
self.pressure_history.clear();
}
}
#[derive(Debug, Clone)]
pub struct OptimizationReport {
pub initial_memory_usage: usize,
pub final_memory_usage: usize,
pub memory_saved: usize,
pub optimization_time: Duration,
pub operations_performed: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_memory_manager_creation() {
let config = super::super::config::TransformerBasedOptimizerConfig::<f32>::default();
let manager = TransformerMemoryManager::new(&config);
assert!(manager.is_ok());
}
#[test]
fn test_memory_cache() {
let cache = MemoryCache::<f32>::new(1024 * 1024, CacheEvictionStrategy::LRU);
assert!(cache.is_ok());
let mut c = cache.expect("unwrap failed");
let tensor = Array2::<f32>::ones((10, 10));
assert!(c.store("test".to_string(), tensor).is_ok());
assert!(c.contains("test"));
}
#[test]
fn test_compression_manager() {
let compression = CompressionManager::<f32>::new(0.5);
assert!(compression.is_ok());
let comp = compression.expect("unwrap failed");
let tensor = Array2::<f32>::ones((5, 5));
let compressed = comp.compress(&tensor);
assert!(compressed.is_ok());
let decompressed = comp.decompress(&compressed.expect("unwrap failed"));
assert!(decompressed.is_ok());
}
#[test]
fn test_access_tracker() {
let mut tracker = AccessTracker::new(100);
tracker.record_read("key1".to_string());
tracker.record_write("key2".to_string());
let patterns = tracker.analyze_patterns();
assert!(patterns.total_accesses > 0);
}
#[test]
fn test_memory_pressure_monitor() {
let mut monitor = MemoryPressureMonitor::new();
monitor.update(500 * 1024 * 1024); assert!(!monitor.is_high_pressure());
monitor.update(950 * 1024 * 1024); assert!(monitor.is_high_pressure());
}
}