torsh_ffi/python/tensor/
memory.rs

1use crate::error::FfiResult;
2use once_cell::sync::Lazy;
3use parking_lot::Mutex;
4use std::collections::VecDeque;
5use std::sync::atomic::{AtomicUsize, Ordering};
6
7/// Memory pool for efficient tensor allocation
8#[derive(Debug)]
9pub struct MemoryPool {
10    /// Available memory blocks organized by size
11    free_blocks: Mutex<std::collections::HashMap<usize, VecDeque<Vec<f32>>>>,
12    /// Statistics for monitoring (using atomics for lock-free counter updates)
13    allocations: AtomicUsize,
14    deallocations: AtomicUsize,
15    pool_hits: AtomicUsize,
16    pool_misses: AtomicUsize,
17    max_pool_size: usize,
18}
19
20impl MemoryPool {
21    /// Create a new memory pool with specified maximum size
22    pub fn new(max_pool_size: usize) -> Self {
23        Self {
24            free_blocks: Mutex::new(std::collections::HashMap::new()),
25            allocations: AtomicUsize::new(0),
26            deallocations: AtomicUsize::new(0),
27            pool_hits: AtomicUsize::new(0),
28            pool_misses: AtomicUsize::new(0),
29            max_pool_size,
30        }
31    }
32
33    /// Allocate a vector from the pool or create new
34    pub fn allocate(&self, size: usize) -> FfiResult<Vec<f32>> {
35        let mut free_blocks = self.free_blocks.lock();
36
37        if let Some(blocks) = free_blocks.get_mut(&size) {
38            if let Some(mut block) = blocks.pop_front() {
39                // Reuse existing block
40                block.clear();
41                block.resize(size, 0.0);
42                self.pool_hits.fetch_add(1, Ordering::Relaxed);
43                self.allocations.fetch_add(1, Ordering::Relaxed);
44                return Ok(block);
45            }
46        }
47
48        // Create new block
49        self.pool_misses.fetch_add(1, Ordering::Relaxed);
50        self.allocations.fetch_add(1, Ordering::Relaxed);
51        Ok(vec![0.0; size])
52    }
53
54    /// Return a vector to the pool for reuse
55    pub fn deallocate(&self, mut data: Vec<f32>) -> FfiResult<()> {
56        let size = data.capacity();
57
58        // Only pool blocks that are reasonably sized and within our limits
59        if size > 0 && size <= self.max_pool_size {
60            let mut free_blocks = self.free_blocks.lock();
61
62            let blocks = free_blocks.entry(size).or_insert_with(VecDeque::new);
63
64            // Limit the number of blocks per size to prevent unbounded growth
65            if blocks.len() < 10 {
66                data.clear();
67                blocks.push_back(data);
68            }
69        }
70
71        self.deallocations.fetch_add(1, Ordering::Relaxed);
72        Ok(())
73    }
74
75    /// Get memory pool statistics
76    pub fn stats(&self) -> FfiResult<MemoryPoolStats> {
77        Ok(MemoryPoolStats {
78            allocations: self.allocations.load(Ordering::Relaxed),
79            deallocations: self.deallocations.load(Ordering::Relaxed),
80            pool_hits: self.pool_hits.load(Ordering::Relaxed),
81            pool_misses: self.pool_misses.load(Ordering::Relaxed),
82            active_blocks: self.free_blocks.lock().values().map(|v| v.len()).sum(),
83        })
84    }
85
86    /// Clear all pooled memory
87    pub fn clear(&self) -> FfiResult<()> {
88        self.free_blocks.lock().clear();
89        Ok(())
90    }
91}
92
93/// Statistics for memory pool usage
94#[derive(Debug, Clone)]
95pub struct MemoryPoolStats {
96    pub allocations: usize,
97    pub deallocations: usize,
98    pub pool_hits: usize,
99    pub pool_misses: usize,
100    pub active_blocks: usize,
101}
102
103/// Global memory pool instance
104pub static MEMORY_POOL: Lazy<MemoryPool> = Lazy::new(|| {
105    MemoryPool::new(1024 * 1024) // 1MB max pool size
106});
torsh_ffi/python/tensor/memory.rs

torsh_ffi/python/tensor/
memory.rs
