osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
1021 lines
31 KiB
Rust
1021 lines
31 KiB
Rust
//! GPU Memory Pool
|
|
//!
|
|
//! This module provides efficient GPU memory management through pooling:
|
|
//! - Reduces allocation overhead by reusing memory blocks
|
|
//! - Supports multiple allocation strategies (best-fit, first-fit)
|
|
//! - Provides memory statistics and fragmentation tracking
|
|
//! - Thread-safe for concurrent allocations
|
|
//! - Automatic defragmentation and compaction
|
|
//!
|
|
//! Memory pooling can reduce allocation overhead by 10-100x for
|
|
//! frequently allocated tensor sizes during training.
|
|
|
|
use crate::error::{DistributedError, Result};
|
|
use parking_lot::{Mutex, RwLock};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::{BTreeMap, HashMap, VecDeque};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::time::Instant;
|
|
|
|
// =============================================================================
|
|
// Configuration
|
|
// =============================================================================
|
|
|
|
/// Memory allocation strategy
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum AllocationStrategy {
|
|
/// Best-fit: Find smallest block that fits (less fragmentation)
|
|
BestFit,
|
|
/// First-fit: Use first block that fits (faster)
|
|
FirstFit,
|
|
/// Buddy allocator: Power-of-2 block sizes
|
|
Buddy,
|
|
/// Slab allocator: Fixed-size blocks for common sizes
|
|
Slab,
|
|
}
|
|
|
|
/// Configuration for memory pool
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct MemoryPoolConfig {
|
|
/// Initial pool size in bytes
|
|
pub initial_size_bytes: usize,
|
|
/// Maximum pool size in bytes (0 = unlimited)
|
|
pub max_size_bytes: usize,
|
|
/// Allocation strategy
|
|
pub strategy: AllocationStrategy,
|
|
/// Minimum block size in bytes
|
|
pub min_block_size: usize,
|
|
/// Alignment requirement in bytes
|
|
pub alignment: usize,
|
|
/// Enable automatic defragmentation
|
|
pub auto_defrag: bool,
|
|
/// Fragmentation threshold to trigger defrag (0.0-1.0)
|
|
pub defrag_threshold: f32,
|
|
/// Enable memory usage statistics
|
|
pub enable_stats: bool,
|
|
/// Cache freed blocks for reuse
|
|
pub cache_freed_blocks: bool,
|
|
/// Maximum cached blocks per size class
|
|
pub max_cached_per_size: usize,
|
|
}
|
|
|
|
impl Default for MemoryPoolConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
initial_size_bytes: 256 * 1024 * 1024, // 256MB
|
|
max_size_bytes: 0, // Unlimited
|
|
strategy: AllocationStrategy::BestFit,
|
|
min_block_size: 512,
|
|
alignment: 256, // GPU typically needs 256-byte alignment
|
|
auto_defrag: true,
|
|
defrag_threshold: 0.3, // Defrag when 30% fragmented
|
|
enable_stats: true,
|
|
cache_freed_blocks: true,
|
|
max_cached_per_size: 32,
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Memory Block
|
|
// =============================================================================
|
|
|
|
/// Unique identifier for a memory block
|
|
pub type BlockId = u64;
|
|
|
|
/// State of a memory block
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum BlockState {
|
|
/// Block is free and available
|
|
Free,
|
|
/// Block is allocated and in use
|
|
Allocated,
|
|
/// Block is cached (freed but held for reuse)
|
|
Cached,
|
|
}
|
|
|
|
/// A memory block in the pool
|
|
#[derive(Debug, Clone)]
|
|
pub struct MemoryBlock {
|
|
/// Unique block ID
|
|
pub id: BlockId,
|
|
/// Offset from pool base address
|
|
pub offset: usize,
|
|
/// Size in bytes
|
|
pub size: usize,
|
|
/// Current state
|
|
pub state: BlockState,
|
|
/// Allocation timestamp (if allocated)
|
|
pub allocated_at: Option<Instant>,
|
|
/// Tag for debugging (e.g., "gradient", "activation")
|
|
pub tag: Option<String>,
|
|
/// Device ID
|
|
pub device_id: i32,
|
|
}
|
|
|
|
impl MemoryBlock {
|
|
/// Create a new free block
|
|
pub fn new(id: BlockId, offset: usize, size: usize, device_id: i32) -> Self {
|
|
Self {
|
|
id,
|
|
offset,
|
|
size,
|
|
state: BlockState::Free,
|
|
allocated_at: None,
|
|
tag: None,
|
|
device_id,
|
|
}
|
|
}
|
|
|
|
/// Check if block can satisfy allocation request
|
|
pub fn can_satisfy(&self, size: usize, alignment: usize) -> bool {
|
|
if self.state != BlockState::Free && self.state != BlockState::Cached {
|
|
return false;
|
|
}
|
|
|
|
let aligned_offset = (self.offset + alignment - 1) & !(alignment - 1);
|
|
let padding = aligned_offset - self.offset;
|
|
self.size >= size + padding
|
|
}
|
|
|
|
/// Get aligned offset within this block
|
|
pub fn aligned_offset(&self, alignment: usize) -> usize {
|
|
(self.offset + alignment - 1) & !(alignment - 1)
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Size Class Cache
|
|
// =============================================================================
|
|
|
|
/// Cache for commonly-used allocation sizes
|
|
pub struct SizeClassCache {
|
|
/// Size class -> list of cached block IDs
|
|
cache: HashMap<usize, VecDeque<BlockId>>,
|
|
/// Maximum blocks per size class
|
|
max_per_size: usize,
|
|
/// Size class granularity
|
|
granularity: usize,
|
|
}
|
|
|
|
impl SizeClassCache {
|
|
/// Create a new size class cache
|
|
pub fn new(max_per_size: usize, granularity: usize) -> Self {
|
|
Self {
|
|
cache: HashMap::new(),
|
|
max_per_size,
|
|
granularity,
|
|
}
|
|
}
|
|
|
|
/// Get size class for a given size
|
|
fn size_class(&self, size: usize) -> usize {
|
|
// Round up to nearest granularity
|
|
((size + self.granularity - 1) / self.granularity) * self.granularity
|
|
}
|
|
|
|
/// Try to get a cached block for the given size
|
|
pub fn get(&mut self, size: usize) -> Option<BlockId> {
|
|
let size_class = self.size_class(size);
|
|
self.cache.get_mut(&size_class)?.pop_front()
|
|
}
|
|
|
|
/// Cache a freed block
|
|
pub fn put(&mut self, size: usize, block_id: BlockId) -> bool {
|
|
let size_class = self.size_class(size);
|
|
let queue = self.cache.entry(size_class).or_default();
|
|
|
|
if queue.len() < self.max_per_size {
|
|
queue.push_back(block_id);
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Clear all cached blocks
|
|
pub fn clear(&mut self) {
|
|
self.cache.clear();
|
|
}
|
|
|
|
/// Get total cached block count
|
|
pub fn cached_count(&self) -> usize {
|
|
self.cache
|
|
.values()
|
|
.map(std::collections::VecDeque::len)
|
|
.sum()
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Memory Pool Statistics
|
|
// =============================================================================
|
|
|
|
/// Statistics for memory pool usage
|
|
#[derive(Debug, Default, Clone)]
|
|
pub struct MemoryPoolStats {
|
|
/// Total pool size in bytes
|
|
pub total_size_bytes: usize,
|
|
/// Currently allocated bytes
|
|
pub allocated_bytes: usize,
|
|
/// Peak allocated bytes
|
|
pub peak_allocated_bytes: usize,
|
|
/// Free bytes (including fragmented)
|
|
pub free_bytes: usize,
|
|
/// Number of allocations
|
|
pub allocation_count: usize,
|
|
/// Number of deallocations
|
|
pub deallocation_count: usize,
|
|
/// Number of cache hits
|
|
pub cache_hits: usize,
|
|
/// Number of cache misses
|
|
pub cache_misses: usize,
|
|
/// Fragmentation ratio (0.0 = none, 1.0 = fully fragmented)
|
|
pub fragmentation_ratio: f32,
|
|
/// Number of free blocks
|
|
pub free_block_count: usize,
|
|
/// Average allocation size
|
|
pub avg_allocation_size: usize,
|
|
/// Number of defragmentation runs
|
|
pub defrag_count: usize,
|
|
}
|
|
|
|
impl MemoryPoolStats {
|
|
/// Calculate cache hit ratio
|
|
pub fn cache_hit_ratio(&self) -> f32 {
|
|
let total = self.cache_hits + self.cache_misses;
|
|
if total == 0 {
|
|
0.0
|
|
} else {
|
|
self.cache_hits as f32 / total as f32
|
|
}
|
|
}
|
|
|
|
/// Calculate utilization ratio
|
|
pub fn utilization(&self) -> f32 {
|
|
if self.total_size_bytes == 0 {
|
|
0.0
|
|
} else {
|
|
self.allocated_bytes as f32 / self.total_size_bytes as f32
|
|
}
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Memory Pool
|
|
// =============================================================================
|
|
|
|
/// GPU memory pool for efficient allocation
|
|
pub struct MemoryPool {
|
|
/// Configuration
|
|
config: MemoryPoolConfig,
|
|
/// Device ID
|
|
device_id: i32,
|
|
/// All memory blocks
|
|
blocks: RwLock<HashMap<BlockId, MemoryBlock>>,
|
|
/// Free blocks sorted by offset (for coalescing)
|
|
free_blocks: RwLock<BTreeMap<usize, BlockId>>,
|
|
/// Size class cache
|
|
size_cache: Mutex<SizeClassCache>,
|
|
/// Next block ID
|
|
next_block_id: AtomicUsize,
|
|
/// Statistics
|
|
stats: RwLock<MemoryPoolStats>,
|
|
/// Base address (simulated - real impl would use CUDA)
|
|
base_address: usize,
|
|
/// Current pool size
|
|
current_size: AtomicUsize,
|
|
}
|
|
|
|
impl MemoryPool {
|
|
/// Create a new memory pool
|
|
pub fn new(config: MemoryPoolConfig, device_id: i32) -> Result<Self> {
|
|
let initial_size = config.initial_size_bytes;
|
|
let size_cache = SizeClassCache::new(config.max_cached_per_size, config.min_block_size);
|
|
|
|
let pool = Self {
|
|
config,
|
|
device_id,
|
|
blocks: RwLock::new(HashMap::new()),
|
|
free_blocks: RwLock::new(BTreeMap::new()),
|
|
size_cache: Mutex::new(size_cache),
|
|
next_block_id: AtomicUsize::new(1),
|
|
stats: RwLock::new(MemoryPoolStats::default()),
|
|
base_address: 0x1000_0000, // Simulated base address
|
|
current_size: AtomicUsize::new(initial_size),
|
|
};
|
|
|
|
// Initialize with a single free block
|
|
let initial_block = MemoryBlock::new(0, 0, initial_size, device_id);
|
|
{
|
|
let mut blocks = pool.blocks.write();
|
|
let mut free_blocks = pool.free_blocks.write();
|
|
blocks.insert(0, initial_block);
|
|
free_blocks.insert(0, 0);
|
|
}
|
|
|
|
// Initialize stats
|
|
{
|
|
let mut stats = pool.stats.write();
|
|
stats.total_size_bytes = initial_size;
|
|
stats.free_bytes = initial_size;
|
|
stats.free_block_count = 1;
|
|
}
|
|
|
|
Ok(pool)
|
|
}
|
|
|
|
/// Allocate memory from the pool
|
|
pub fn allocate(&self, size: usize) -> Result<MemoryAllocation> {
|
|
self.allocate_with_tag(size, None)
|
|
}
|
|
|
|
/// Allocate memory with a debug tag
|
|
pub fn allocate_with_tag(&self, size: usize, tag: Option<String>) -> Result<MemoryAllocation> {
|
|
let aligned_size = self.align_size(size);
|
|
|
|
// Try cache first
|
|
if self.config.cache_freed_blocks {
|
|
let mut cache = self.size_cache.lock();
|
|
if let Some(block_id) = cache.get(aligned_size) {
|
|
let mut blocks = self.blocks.write();
|
|
if let Some(block) = blocks.get_mut(&block_id) {
|
|
if block.state == BlockState::Cached && block.size >= aligned_size {
|
|
block.state = BlockState::Allocated;
|
|
block.allocated_at = Some(Instant::now());
|
|
block.tag = tag;
|
|
|
|
// Update stats
|
|
let mut stats = self.stats.write();
|
|
stats.cache_hits += 1;
|
|
stats.allocation_count += 1;
|
|
stats.allocated_bytes += block.size;
|
|
stats.free_bytes -= block.size;
|
|
stats.peak_allocated_bytes =
|
|
stats.peak_allocated_bytes.max(stats.allocated_bytes);
|
|
|
|
return Ok(MemoryAllocation {
|
|
block_id,
|
|
ptr: self.base_address + block.offset,
|
|
size: block.size,
|
|
device_id: self.device_id,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
// Update cache miss stat
|
|
let mut stats = self.stats.write();
|
|
stats.cache_misses += 1;
|
|
}
|
|
|
|
// Find a suitable block using the configured strategy
|
|
let block_id = self.find_block(aligned_size)?;
|
|
|
|
// Allocate from the block
|
|
self.allocate_from_block(block_id, aligned_size, tag)
|
|
}
|
|
|
|
/// Find a suitable block for allocation
|
|
fn find_block(&self, size: usize) -> Result<BlockId> {
|
|
match self.config.strategy {
|
|
AllocationStrategy::BestFit => self.find_best_fit(size),
|
|
AllocationStrategy::FirstFit => self.find_first_fit(size),
|
|
AllocationStrategy::Buddy => self.find_buddy_block(size),
|
|
AllocationStrategy::Slab => self.find_slab_block(size),
|
|
}
|
|
}
|
|
|
|
/// Best-fit allocation strategy
|
|
fn find_best_fit(&self, size: usize) -> Result<BlockId> {
|
|
let blocks = self.blocks.read();
|
|
let free_blocks = self.free_blocks.read();
|
|
|
|
let mut best_id = None;
|
|
let mut best_size = usize::MAX;
|
|
|
|
for &block_id in free_blocks.values() {
|
|
if let Some(block) = blocks.get(&block_id) {
|
|
if block.can_satisfy(size, self.config.alignment) && block.size < best_size {
|
|
best_id = Some(block_id);
|
|
best_size = block.size;
|
|
}
|
|
}
|
|
}
|
|
|
|
best_id.ok_or_else(|| {
|
|
DistributedError::runtime(format!(
|
|
"No suitable block found for {} bytes allocation",
|
|
size
|
|
))
|
|
})
|
|
}
|
|
|
|
/// First-fit allocation strategy
|
|
fn find_first_fit(&self, size: usize) -> Result<BlockId> {
|
|
let blocks = self.blocks.read();
|
|
let free_blocks = self.free_blocks.read();
|
|
|
|
for &block_id in free_blocks.values() {
|
|
if let Some(block) = blocks.get(&block_id) {
|
|
if block.can_satisfy(size, self.config.alignment) {
|
|
return Ok(block_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
Err(DistributedError::runtime(format!(
|
|
"No suitable block found for {} bytes allocation",
|
|
size
|
|
)))
|
|
}
|
|
|
|
/// Buddy allocator strategy
|
|
fn find_buddy_block(&self, size: usize) -> Result<BlockId> {
|
|
// Round up to power of 2
|
|
let buddy_size = size.next_power_of_two();
|
|
self.find_best_fit(buddy_size)
|
|
}
|
|
|
|
/// Slab allocator strategy
|
|
fn find_slab_block(&self, size: usize) -> Result<BlockId> {
|
|
// Round up to slab size class
|
|
let slab_size = self.align_size(size);
|
|
self.find_best_fit(slab_size)
|
|
}
|
|
|
|
/// Allocate from a specific block
|
|
fn allocate_from_block(
|
|
&self,
|
|
block_id: BlockId,
|
|
size: usize,
|
|
tag: Option<String>,
|
|
) -> Result<MemoryAllocation> {
|
|
let mut blocks = self.blocks.write();
|
|
let mut free_blocks = self.free_blocks.write();
|
|
|
|
// First, read the block info we need
|
|
let (aligned_offset, block_offset, block_orig_size) = {
|
|
let block = blocks
|
|
.get(&block_id)
|
|
.ok_or_else(|| DistributedError::runtime("Block not found"))?;
|
|
|
|
if block.state != BlockState::Free {
|
|
return Err(DistributedError::runtime("Block is not free"));
|
|
}
|
|
|
|
(
|
|
block.aligned_offset(self.config.alignment),
|
|
block.offset,
|
|
block.size,
|
|
)
|
|
};
|
|
|
|
let remaining = block_orig_size - size;
|
|
|
|
// Split block if there's enough remaining space
|
|
if remaining >= self.config.min_block_size {
|
|
let new_block_id = self.next_block_id.fetch_add(1, Ordering::SeqCst) as BlockId;
|
|
let new_block =
|
|
MemoryBlock::new(new_block_id, block_offset + size, remaining, self.device_id);
|
|
|
|
free_blocks.insert(new_block.offset, new_block_id);
|
|
blocks.insert(new_block_id, new_block);
|
|
}
|
|
|
|
// Now update the original block
|
|
let alloc_size = if remaining >= self.config.min_block_size {
|
|
size
|
|
} else {
|
|
block_orig_size
|
|
};
|
|
if let Some(block) = blocks.get_mut(&block_id) {
|
|
if remaining >= self.config.min_block_size {
|
|
block.size = size;
|
|
}
|
|
block.state = BlockState::Allocated;
|
|
block.allocated_at = Some(Instant::now());
|
|
block.tag = tag;
|
|
}
|
|
|
|
// Remove from free list
|
|
free_blocks.remove(&block_offset);
|
|
|
|
let ptr = self.base_address + aligned_offset;
|
|
|
|
// Update stats
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.allocation_count += 1;
|
|
stats.allocated_bytes += alloc_size;
|
|
stats.free_bytes -= alloc_size;
|
|
stats.free_block_count = free_blocks.len();
|
|
stats.peak_allocated_bytes = stats.peak_allocated_bytes.max(stats.allocated_bytes);
|
|
|
|
if stats.allocation_count > 0 {
|
|
stats.avg_allocation_size = stats.allocated_bytes / stats.allocation_count;
|
|
}
|
|
}
|
|
|
|
Ok(MemoryAllocation {
|
|
block_id,
|
|
ptr,
|
|
size: alloc_size,
|
|
device_id: self.device_id,
|
|
})
|
|
}
|
|
|
|
/// Free an allocation
|
|
pub fn free(&self, allocation: &MemoryAllocation) -> Result<()> {
|
|
let mut blocks = self.blocks.write();
|
|
let mut free_blocks = self.free_blocks.write();
|
|
|
|
let block = blocks
|
|
.get_mut(&allocation.block_id)
|
|
.ok_or_else(|| DistributedError::runtime("Block not found"))?;
|
|
|
|
if block.state != BlockState::Allocated {
|
|
return Err(DistributedError::runtime("Block is not allocated"));
|
|
}
|
|
|
|
let block_size = block.size;
|
|
let block_offset = block.offset;
|
|
|
|
// Try to cache the block
|
|
if self.config.cache_freed_blocks {
|
|
let mut cache = self.size_cache.lock();
|
|
if cache.put(block_size, allocation.block_id) {
|
|
block.state = BlockState::Cached;
|
|
block.allocated_at = None;
|
|
block.tag = None;
|
|
|
|
// Update stats
|
|
let mut stats = self.stats.write();
|
|
stats.deallocation_count += 1;
|
|
stats.allocated_bytes -= block_size;
|
|
stats.free_bytes += block_size;
|
|
|
|
return Ok(());
|
|
}
|
|
}
|
|
|
|
// Mark as free
|
|
block.state = BlockState::Free;
|
|
block.allocated_at = None;
|
|
block.tag = None;
|
|
|
|
// Add to free list
|
|
free_blocks.insert(block_offset, allocation.block_id);
|
|
|
|
// Update stats
|
|
{
|
|
let mut stats = self.stats.write();
|
|
stats.deallocation_count += 1;
|
|
stats.allocated_bytes -= block_size;
|
|
stats.free_bytes += block_size;
|
|
stats.free_block_count = free_blocks.len();
|
|
}
|
|
|
|
// Try to coalesce with neighbors
|
|
drop(blocks);
|
|
drop(free_blocks);
|
|
self.coalesce_neighbors(allocation.block_id)?;
|
|
|
|
// Check if defrag is needed
|
|
if self.config.auto_defrag {
|
|
let stats = self.stats.read();
|
|
if stats.fragmentation_ratio > self.config.defrag_threshold {
|
|
drop(stats);
|
|
self.defragment()?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Coalesce free blocks with neighbors
|
|
fn coalesce_neighbors(&self, block_id: BlockId) -> Result<()> {
|
|
let mut blocks = self.blocks.write();
|
|
let mut free_blocks = self.free_blocks.write();
|
|
|
|
let block = match blocks.get(&block_id) {
|
|
Some(b) if b.state == BlockState::Free => b.clone(),
|
|
_ => return Ok(()),
|
|
};
|
|
|
|
// Find and merge with next block
|
|
let next_offset = block.offset + block.size;
|
|
if let Some(&next_id) = free_blocks.get(&next_offset) {
|
|
if let Some(next_block) = blocks.get(&next_id) {
|
|
if next_block.state == BlockState::Free {
|
|
let merged_size = block.size + next_block.size;
|
|
|
|
// Remove next block
|
|
free_blocks.remove(&next_offset);
|
|
blocks.remove(&next_id);
|
|
|
|
// Update current block
|
|
if let Some(current) = blocks.get_mut(&block_id) {
|
|
current.size = merged_size;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Update fragmentation
|
|
self.update_fragmentation(&blocks, &free_blocks);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Update fragmentation ratio
|
|
fn update_fragmentation(
|
|
&self,
|
|
blocks: &HashMap<BlockId, MemoryBlock>,
|
|
free_blocks: &BTreeMap<usize, BlockId>,
|
|
) {
|
|
if free_blocks.is_empty() {
|
|
let mut stats = self.stats.write();
|
|
stats.fragmentation_ratio = 0.0;
|
|
return;
|
|
}
|
|
|
|
// Calculate fragmentation as ratio of free blocks to total free space
|
|
let mut total_free = 0;
|
|
let mut largest_free = 0;
|
|
|
|
for &block_id in free_blocks.values() {
|
|
if let Some(block) = blocks.get(&block_id) {
|
|
total_free += block.size;
|
|
largest_free = largest_free.max(block.size);
|
|
}
|
|
}
|
|
|
|
let mut stats = self.stats.write();
|
|
if total_free > 0 {
|
|
// Fragmentation = 1 - (largest_free / total_free)
|
|
stats.fragmentation_ratio = 1.0 - (largest_free as f32 / total_free as f32);
|
|
} else {
|
|
stats.fragmentation_ratio = 0.0;
|
|
}
|
|
}
|
|
|
|
/// Defragment the pool
|
|
pub fn defragment(&self) -> Result<()> {
|
|
// In a real implementation, this would:
|
|
// 1. Wait for all allocations to complete
|
|
// 2. Copy data to compact free space
|
|
// 3. Update block offsets
|
|
// For simulation, we just clear the cache and update stats
|
|
|
|
let mut cache = self.size_cache.lock();
|
|
cache.clear();
|
|
|
|
let mut stats = self.stats.write();
|
|
stats.defrag_count += 1;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Align size to minimum block size
|
|
fn align_size(&self, size: usize) -> usize {
|
|
let min = self.config.min_block_size;
|
|
((size + min - 1) / min) * min
|
|
}
|
|
|
|
/// Get pool statistics
|
|
pub fn stats(&self) -> MemoryPoolStats {
|
|
self.stats.read().clone()
|
|
}
|
|
|
|
/// Get device ID
|
|
pub fn device_id(&self) -> i32 {
|
|
self.device_id
|
|
}
|
|
|
|
/// Get current pool size
|
|
pub fn current_size(&self) -> usize {
|
|
self.current_size.load(Ordering::SeqCst)
|
|
}
|
|
|
|
/// Clear the pool
|
|
pub fn clear(&self) -> Result<()> {
|
|
let mut blocks = self.blocks.write();
|
|
let mut free_blocks = self.free_blocks.write();
|
|
let mut cache = self.size_cache.lock();
|
|
|
|
blocks.clear();
|
|
free_blocks.clear();
|
|
cache.clear();
|
|
|
|
// Reinitialize with single free block
|
|
let size = self.current_size.load(Ordering::SeqCst);
|
|
let block = MemoryBlock::new(0, 0, size, self.device_id);
|
|
blocks.insert(0, block);
|
|
free_blocks.insert(0, 0);
|
|
|
|
// Reset stats
|
|
let mut stats = self.stats.write();
|
|
*stats = MemoryPoolStats {
|
|
total_size_bytes: size,
|
|
free_bytes: size,
|
|
free_block_count: 1,
|
|
..Default::default()
|
|
};
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Memory Allocation Handle
|
|
// =============================================================================
|
|
|
|
/// Handle to an allocated memory region
|
|
#[derive(Debug, Clone)]
|
|
pub struct MemoryAllocation {
|
|
/// Block ID in the pool
|
|
pub block_id: BlockId,
|
|
/// Pointer to allocated memory
|
|
pub ptr: usize,
|
|
/// Size in bytes
|
|
pub size: usize,
|
|
/// Device ID
|
|
pub device_id: i32,
|
|
}
|
|
|
|
impl MemoryAllocation {
|
|
/// Get the memory address
|
|
pub fn address(&self) -> usize {
|
|
self.ptr
|
|
}
|
|
|
|
/// Get size in bytes
|
|
pub fn size_bytes(&self) -> usize {
|
|
self.size
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Thread-Safe Wrapper
|
|
// =============================================================================
|
|
|
|
/// Thread-safe shared memory pool
|
|
pub type SharedMemoryPool = Arc<MemoryPool>;
|
|
|
|
/// Create a shared memory pool
|
|
pub fn shared_memory_pool(config: MemoryPoolConfig, device_id: i32) -> Result<SharedMemoryPool> {
|
|
Ok(Arc::new(MemoryPool::new(config, device_id)?))
|
|
}
|
|
|
|
// =============================================================================
|
|
// Multi-Device Pool Manager
|
|
// =============================================================================
|
|
|
|
/// Manages memory pools across multiple devices
|
|
pub struct MultiDevicePoolManager {
|
|
/// Pools per device
|
|
pools: HashMap<i32, SharedMemoryPool>,
|
|
/// Configuration
|
|
config: MemoryPoolConfig,
|
|
}
|
|
|
|
impl MultiDevicePoolManager {
|
|
/// Create a new multi-device pool manager
|
|
pub fn new(config: MemoryPoolConfig) -> Self {
|
|
Self {
|
|
pools: HashMap::new(),
|
|
config,
|
|
}
|
|
}
|
|
|
|
/// Get or create pool for a device
|
|
pub fn get_pool(&mut self, device_id: i32) -> Result<SharedMemoryPool> {
|
|
if let Some(pool) = self.pools.get(&device_id) {
|
|
return Ok(pool.clone());
|
|
}
|
|
|
|
let pool = shared_memory_pool(self.config.clone(), device_id)?;
|
|
self.pools.insert(device_id, pool.clone());
|
|
Ok(pool)
|
|
}
|
|
|
|
/// Get statistics for all devices
|
|
pub fn all_stats(&self) -> HashMap<i32, MemoryPoolStats> {
|
|
self.pools
|
|
.iter()
|
|
.map(|(&id, pool)| (id, pool.stats()))
|
|
.collect()
|
|
}
|
|
|
|
/// Clear all pools
|
|
pub fn clear_all(&self) -> Result<()> {
|
|
for pool in self.pools.values() {
|
|
pool.clear()?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// Tests
|
|
// =============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_memory_pool_config_default() {
|
|
let config = MemoryPoolConfig::default();
|
|
assert_eq!(config.initial_size_bytes, 256 * 1024 * 1024);
|
|
assert_eq!(config.strategy, AllocationStrategy::BestFit);
|
|
assert_eq!(config.alignment, 256);
|
|
}
|
|
|
|
#[test]
|
|
fn test_memory_block_creation() {
|
|
let block = MemoryBlock::new(1, 1024, 4096, 0);
|
|
assert_eq!(block.id, 1);
|
|
assert_eq!(block.offset, 1024);
|
|
assert_eq!(block.size, 4096);
|
|
assert_eq!(block.state, BlockState::Free);
|
|
}
|
|
|
|
#[test]
|
|
fn test_block_can_satisfy() {
|
|
let block = MemoryBlock::new(1, 0, 4096, 0);
|
|
assert!(block.can_satisfy(1024, 256));
|
|
assert!(block.can_satisfy(4096, 256));
|
|
assert!(!block.can_satisfy(8192, 256));
|
|
}
|
|
|
|
#[test]
|
|
fn test_memory_pool_creation() {
|
|
let config = MemoryPoolConfig::default();
|
|
let pool = MemoryPool::new(config, 0).unwrap();
|
|
|
|
assert_eq!(pool.device_id(), 0);
|
|
assert_eq!(pool.current_size(), 256 * 1024 * 1024);
|
|
}
|
|
|
|
#[test]
|
|
fn test_simple_allocation() {
|
|
let config = MemoryPoolConfig {
|
|
initial_size_bytes: 1024 * 1024, // 1MB
|
|
..Default::default()
|
|
};
|
|
let pool = MemoryPool::new(config, 0).unwrap();
|
|
|
|
let alloc = pool.allocate(4096).unwrap();
|
|
assert!(alloc.ptr > 0);
|
|
assert!(alloc.size >= 4096);
|
|
|
|
let stats = pool.stats();
|
|
assert_eq!(stats.allocation_count, 1);
|
|
assert!(stats.allocated_bytes >= 4096);
|
|
}
|
|
|
|
#[test]
|
|
fn test_allocation_and_free() {
|
|
let config = MemoryPoolConfig {
|
|
initial_size_bytes: 1024 * 1024,
|
|
..Default::default()
|
|
};
|
|
let pool = MemoryPool::new(config, 0).unwrap();
|
|
|
|
let alloc = pool.allocate(4096).unwrap();
|
|
let stats_before = pool.stats();
|
|
|
|
pool.free(&alloc).unwrap();
|
|
let stats_after = pool.stats();
|
|
|
|
assert_eq!(stats_after.deallocation_count, 1);
|
|
assert!(stats_after.allocated_bytes < stats_before.allocated_bytes);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_allocations() {
|
|
let config = MemoryPoolConfig {
|
|
initial_size_bytes: 1024 * 1024,
|
|
..Default::default()
|
|
};
|
|
let pool = MemoryPool::new(config, 0).unwrap();
|
|
|
|
let alloc1 = pool.allocate(1024).unwrap();
|
|
let alloc2 = pool.allocate(2048).unwrap();
|
|
let alloc3 = pool.allocate(4096).unwrap();
|
|
|
|
assert_ne!(alloc1.ptr, alloc2.ptr);
|
|
assert_ne!(alloc2.ptr, alloc3.ptr);
|
|
|
|
let stats = pool.stats();
|
|
assert_eq!(stats.allocation_count, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_size_class_cache() {
|
|
let mut cache = SizeClassCache::new(10, 512);
|
|
|
|
cache.put(1024, 1);
|
|
cache.put(1024, 2);
|
|
cache.put(2048, 3);
|
|
|
|
assert_eq!(cache.get(1024), Some(1));
|
|
assert_eq!(cache.get(1024), Some(2));
|
|
assert_eq!(cache.get(1024), None);
|
|
assert_eq!(cache.get(2048), Some(3));
|
|
}
|
|
|
|
#[test]
|
|
fn test_cache_hit() {
|
|
let config = MemoryPoolConfig {
|
|
initial_size_bytes: 1024 * 1024,
|
|
cache_freed_blocks: true,
|
|
..Default::default()
|
|
};
|
|
let pool = MemoryPool::new(config, 0).unwrap();
|
|
|
|
// Allocate and free
|
|
let alloc1 = pool.allocate(4096).unwrap();
|
|
pool.free(&alloc1).unwrap();
|
|
|
|
// Allocate same size - should hit cache
|
|
let _alloc2 = pool.allocate(4096).unwrap();
|
|
|
|
let stats = pool.stats();
|
|
assert!(stats.cache_hits >= 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_allocation_with_tag() {
|
|
let config = MemoryPoolConfig {
|
|
initial_size_bytes: 1024 * 1024,
|
|
..Default::default()
|
|
};
|
|
let pool = MemoryPool::new(config, 0).unwrap();
|
|
|
|
let alloc = pool
|
|
.allocate_with_tag(4096, Some("gradient".to_string()))
|
|
.unwrap();
|
|
assert!(alloc.size >= 4096);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pool_clear() {
|
|
let config = MemoryPoolConfig {
|
|
initial_size_bytes: 1024 * 1024,
|
|
..Default::default()
|
|
};
|
|
let pool = MemoryPool::new(config, 0).unwrap();
|
|
|
|
let _ = pool.allocate(4096).unwrap();
|
|
let _ = pool.allocate(4096).unwrap();
|
|
|
|
pool.clear().unwrap();
|
|
|
|
let stats = pool.stats();
|
|
assert_eq!(stats.allocation_count, 0);
|
|
assert_eq!(stats.free_block_count, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multi_device_manager() {
|
|
let config = MemoryPoolConfig {
|
|
initial_size_bytes: 1024 * 1024,
|
|
..Default::default()
|
|
};
|
|
let mut manager = MultiDevicePoolManager::new(config);
|
|
|
|
let pool0 = manager.get_pool(0).unwrap();
|
|
let pool1 = manager.get_pool(1).unwrap();
|
|
|
|
assert_eq!(pool0.device_id(), 0);
|
|
assert_eq!(pool1.device_id(), 1);
|
|
|
|
let all_stats = manager.all_stats();
|
|
assert_eq!(all_stats.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_memory_pool_stats() {
|
|
let stats = MemoryPoolStats {
|
|
total_size_bytes: 1000,
|
|
allocated_bytes: 500,
|
|
cache_hits: 10,
|
|
cache_misses: 5,
|
|
..Default::default()
|
|
};
|
|
|
|
assert!((stats.utilization() - 0.5).abs() < 0.01);
|
|
assert!((stats.cache_hit_ratio() - 0.666).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_allocation_strategy() {
|
|
assert_eq!(AllocationStrategy::BestFit, AllocationStrategy::BestFit);
|
|
assert_ne!(AllocationStrategy::BestFit, AllocationStrategy::FirstFit);
|
|
}
|
|
}
|