1#![allow(dead_code)]
3use crate::{Tensor, TensorStorage};
6use std::alloc::{handle_alloc_error, Layout};
7use std::collections::{HashMap, VecDeque};
8use std::marker::PhantomData;
9use std::mem::{ManuallyDrop, MaybeUninit};
10use std::ptr::NonNull;
11use std::sync::{Arc, Mutex, Weak};
12use torsh_core::{device::DeviceType, dtype::TensorElement, error::Result};
13
14use scirs2_core::memory::GlobalBufferPool;
16use scirs2_core::memory::LeakDetector;
17#[cfg(feature = "memory_efficient")]
20use scirs2_core::memory_efficient::{AccessMode, MemoryMappedArray};
21
22#[cfg(feature = "memory_efficient")]
25fn unique_mmap_path(tag: &str) -> std::path::PathBuf {
26 use std::sync::atomic::{AtomicU64, Ordering};
27 static COUNTER: AtomicU64 = AtomicU64::new(0);
28 let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
29 let nanos = std::time::SystemTime::now()
30 .duration_since(std::time::UNIX_EPOCH)
31 .unwrap_or_default()
32 .as_nanos();
33 std::env::temp_dir().join(format!(
34 "torsh_mmap_{tag}_{pid}_{nanos}_{seq}.bin",
35 pid = std::process::id()
36 ))
37}
38
39#[cfg(feature = "memory_efficient")]
46fn map_through_mmap_file<T: TensorElement>(
47 data: Vec<T>,
48 backing_path: &std::path::Path,
49) -> Result<Vec<T>> {
50 use scirs2_core::ndarray::Array1;
51
52 let array: Array1<T> = Array1::from(data);
54 let mmap = MemoryMappedArray::<T>::new(Some(&array), backing_path, AccessMode::Write, 0)
55 .map_err(|e| {
56 torsh_core::error::TorshError::IoError(format!(
57 "memory-mapped allocation failed at {path}: {e}",
58 path = backing_path.display()
59 ))
60 })?;
61
62 let mapped = mmap.as_slice().to_vec();
64
65 drop(mmap);
67 let _ = std::fs::remove_file(backing_path);
68
69 Ok(mapped)
70}
71
72static MEMORY_POOL: std::sync::OnceLock<Arc<Mutex<GlobalMemoryPool>>> = std::sync::OnceLock::new();
78
79pub fn init_memory_pool() -> Arc<Mutex<GlobalMemoryPool>> {
81 let arc = MEMORY_POOL
82 .get_or_init(|| {
83 let pool = Arc::new(Mutex::new(GlobalMemoryPool::new()));
84 if let Ok(mut guard) = pool.lock() {
86 guard.self_weak = Some(Arc::downgrade(&pool));
87 }
88 pool
89 })
90 .clone();
91 arc
92}
93
94pub fn get_memory_pool() -> Arc<Mutex<GlobalMemoryPool>> {
96 init_memory_pool()
97}
98
99struct RawEntry {
104 ptr: NonNull<u8>,
105 capacity_bytes: usize,
106 layout: Layout,
107}
108
109unsafe impl Send for RawEntry {}
111
112impl Drop for RawEntry {
113 fn drop(&mut self) {
114 unsafe { std::alloc::dealloc(self.ptr.as_ptr(), self.layout) };
116 }
117}
118
119pub struct ReusedBuffer<T: 'static> {
126 ptr: NonNull<T>,
127 capacity: usize,
128 layout: Layout,
129 pool: Weak<Mutex<GlobalMemoryPool>>,
130}
131
132unsafe impl<T: Send + 'static> Send for ReusedBuffer<T> {}
135
136impl<T: 'static> ReusedBuffer<T> {
137 pub fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<T>] {
139 unsafe {
141 std::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut MaybeUninit<T>, self.capacity)
142 }
143 }
144
145 pub fn capacity(&self) -> usize {
147 self.capacity
148 }
149
150 pub fn as_ptr_raw(&self) -> *mut T {
152 self.ptr.as_ptr()
153 }
154
155 pub fn into_vec(self, len: usize) -> Vec<T> {
163 debug_assert!(len <= self.capacity, "len must not exceed capacity");
164 let md = ManuallyDrop::new(self);
166 unsafe { Vec::from_raw_parts(md.ptr.as_ptr(), len, md.capacity) }
169 }
170
171 pub fn release_to_pool(self) {
175 let md = ManuallyDrop::new(self);
177 let raw_entry = RawEntry {
178 ptr: NonNull::new(md.ptr.as_ptr() as *mut u8)
179 .expect("ReusedBuffer pointer is non-null by construction"),
180 capacity_bytes: md.capacity * std::mem::size_of::<T>(),
181 layout: md.layout,
182 };
183 if let Some(pool_arc) = md.pool.upgrade() {
184 if let Ok(mut guard) = pool_arc.lock() {
185 let type_id = std::any::TypeId::of::<T>();
186 let size_class = guard.find_size_class(raw_entry.capacity_bytes);
187 let align = raw_entry.layout.align();
188 let pool_key = (type_id, size_class, align);
189 if let Some(bucket) = guard.pools.get_mut(&pool_key) {
190 if bucket.available_buffers.len() < bucket.max_buffers {
191 bucket.available_buffers.push_back(raw_entry);
192 bucket.deallocations += 1;
193 return;
195 }
196 }
197 }
198 }
199 }
201}
202
203impl<T: 'static> Drop for ReusedBuffer<T> {
204 fn drop(&mut self) {
205 let raw_entry = RawEntry {
208 ptr: NonNull::new(self.ptr.as_ptr() as *mut u8)
209 .expect("ReusedBuffer pointer is non-null by construction"),
210 capacity_bytes: self.capacity * std::mem::size_of::<T>(),
211 layout: self.layout,
212 };
213 if let Some(pool_arc) = self.pool.upgrade() {
214 if let Ok(mut guard) = pool_arc.lock() {
215 let type_id = std::any::TypeId::of::<T>();
216 let size_class = guard.find_size_class(raw_entry.capacity_bytes);
217 let align = raw_entry.layout.align();
218 let pool_key = (type_id, size_class, align);
219 if let Some(bucket) = guard.pools.get_mut(&pool_key) {
220 if bucket.available_buffers.len() < bucket.max_buffers {
221 let md_entry = ManuallyDrop::new(raw_entry);
224 bucket
227 .available_buffers
228 .push_back(unsafe { std::ptr::read(&*md_entry as *const RawEntry) });
229 bucket.deallocations += 1;
230 return;
231 }
232 }
233 }
234 }
235 }
237}
238
239pub struct GlobalMemoryPool {
243 pools: HashMap<(std::any::TypeId, usize, usize), MemoryPool>,
249 stats: PoolStatistics,
251 config: PoolConfig,
253 scirs2_pool: GlobalBufferPool,
255 leak_detector: LeakDetector,
257 self_weak: Option<Weak<Mutex<GlobalMemoryPool>>>,
259 }
264
265#[derive(Debug)]
267struct MemoryPool {
268 available_buffers: VecDeque<RawEntry>,
270 #[allow(dead_code)]
272 size_class: usize,
273 max_buffers: usize,
275 allocations: usize,
277 reuses: usize,
278 deallocations: usize,
279}
280
281#[derive(Debug, Clone)]
283pub struct PoolConfig {
284 pub max_buffers_per_class: usize,
286 pub max_total_memory: usize,
288 pub auto_cleanup: bool,
290 pub cleanup_threshold: f64,
292 pub size_classes: Vec<usize>,
294}
295
296#[derive(Debug, Default, Clone)]
298pub struct PoolStatistics {
299 pub total_allocations: usize,
301 pub pool_hits: usize,
303 pub pool_misses: usize,
305 pub total_bytes_allocated: usize,
307 pub bytes_in_pools: usize,
309 pub peak_memory_usage: usize,
311}
312
313#[derive(Debug)]
315pub struct PooledTensor<T: TensorElement + Default> {
316 tensor: Tensor<T>,
317 pool_key: Option<(std::any::TypeId, usize, usize)>,
318 _phantom: PhantomData<T>,
319}
320
321impl Default for PoolConfig {
322 fn default() -> Self {
323 let size_classes = (10..31) .map(|exp| 1 << exp)
326 .collect();
327
328 Self {
329 max_buffers_per_class: 16,
330 max_total_memory: 1024 * 1024 * 1024, auto_cleanup: true,
332 cleanup_threshold: 0.8,
333 size_classes,
334 }
335 }
336}
337
338impl Default for GlobalMemoryPool {
339 fn default() -> Self {
340 Self::new()
341 }
342}
343
344impl GlobalMemoryPool {
345 pub fn new() -> Self {
347 #[cfg(feature = "profiling")]
348 {
349 }
351 Self {
352 pools: HashMap::new(),
353 stats: PoolStatistics::default(),
354 config: PoolConfig::default(),
355 scirs2_pool: GlobalBufferPool::new(),
357 leak_detector: LeakDetector::new(Default::default())
358 .unwrap_or_else(|_| panic!("Failed to initialize leak detector")),
359 self_weak: None,
360 }
363 }
364
365 pub fn create_large_tensor<T: TensorElement>(
367 &mut self,
368 shape: &[usize],
369 device: DeviceType,
370 ) -> Result<Tensor<T>>
371 where
372 T: Clone + Default,
373 {
374 #[cfg(feature = "profiling")]
375 {
376 }
378 let total_elements: usize = shape.iter().product();
379 let total_bytes = total_elements * std::mem::size_of::<T>();
380
381 if total_bytes > 100 * 1024 * 1024 {
383 self.create_memory_mapped_tensor(shape, device)
385 } else if total_bytes > 10 * 1024 * 1024 {
386 self.create_chunked_tensor(shape, device)
388 } else if total_bytes > 1024 * 1024 {
389 self.create_pooled_tensor(shape, device)
391 } else {
392 Tensor::zeros(shape, device)
394 }
395 }
396
397 fn create_memory_mapped_tensor<T: TensorElement>(
404 &mut self,
405 shape: &[usize],
406 device: DeviceType,
407 ) -> Result<Tensor<T>>
408 where
409 T: Clone + Default,
410 {
411 let total_elements: usize = shape.iter().product();
412
413 let data = vec![T::default(); total_elements];
415
416 #[cfg(feature = "memory_efficient")]
417 {
418 let backing_path = unique_mmap_path("tensor");
421 let mapped = map_through_mmap_file::<T>(data, &backing_path)?;
422 Tensor::from_data(mapped, shape.to_vec(), device)
423 }
424
425 #[cfg(not(feature = "memory_efficient"))]
426 {
427 Tensor::from_data(data, shape.to_vec(), device)
429 }
430 }
431
432 fn create_chunked_tensor<T: TensorElement>(
434 &mut self,
435 shape: &[usize],
436 device: DeviceType,
437 ) -> Result<Tensor<T>>
438 where
439 T: Clone + Default,
440 {
441 let total_elements: usize = shape.iter().product();
442
443 let chunk_size = (1024 * 1024) / std::mem::size_of::<T>().max(1); let num_chunks = (total_elements + chunk_size - 1) / chunk_size;
446
447 let _ = (total_elements, num_chunks, chunk_size); let data = vec![T::default(); total_elements];
452
453 Tensor::from_data(data, shape.to_vec(), device)
458 }
459
460 fn create_pooled_tensor<T: TensorElement>(
462 &mut self,
463 shape: &[usize],
464 device: DeviceType,
465 ) -> Result<Tensor<T>>
466 where
467 T: Clone + Default,
468 {
469 let total_elements: usize = shape.iter().product();
470 let buffer_size = total_elements * std::mem::size_of::<T>();
471
472 let _ = (buffer_size, total_elements); let data = vec![T::default(); total_elements];
477
478 self.stats.pool_hits += 1;
480 Tensor::from_data(data, shape.to_vec(), device)
484 }
485
486 pub fn create_lazy_tensor<T: TensorElement>(
488 &mut self,
489 shape: &[usize],
490 device: DeviceType,
491 ) -> Result<Tensor<T>>
492 where
493 T: Clone + Default,
494 {
495 #[cfg(feature = "profiling")]
496 {
497 }
499 let total_elements: usize = shape.iter().product();
500
501 let data = vec![T::default(); total_elements];
503
504 Tensor::from_data(data, shape.to_vec(), device)
508 }
509
510 pub fn create_zero_copy_view<T: TensorElement>(
512 &self,
513 source: &Tensor<T>,
514 offset: usize,
515 shape: &[usize],
516 ) -> Result<Tensor<T>>
517 where
518 T: Clone,
519 {
520 #[cfg(feature = "profiling")]
521 {
522 }
524
525 let source_data = source.data()?;
527 let view_data = source_data[offset..offset + shape.iter().product::<usize>()].to_vec();
528
529 Tensor::from_data(view_data, shape.to_vec(), source.device())
530 }
531
532 pub fn get_enhanced_stats(&self) -> PoolStatistics {
534 self.stats.clone()
536 }
537
538 pub fn acquire_uninit<T: 'static>(&mut self, count: usize) -> ReusedBuffer<T> {
549 self.acquire_uninit_aligned::<T>(count, std::mem::align_of::<T>())
550 }
551
552 pub fn acquire_uninit_aligned<T: 'static>(
563 &mut self,
564 count: usize,
565 align: usize,
566 ) -> ReusedBuffer<T> {
567 let element_size = std::mem::size_of::<T>();
568 let element_align = std::mem::align_of::<T>();
569 assert!(
570 align.is_power_of_two(),
571 "alignment must be a power of two (got {align})"
572 );
573 assert!(
574 align >= element_align,
575 "alignment {align} must be >= align_of::<T>() ({element_align})"
576 );
577 let size_bytes = count * element_size;
578 let size_class = self.find_size_class(size_bytes);
579 let type_id = std::any::TypeId::of::<T>();
580 let pool_key = (type_id, size_class, align);
581
582 let layout =
583 Layout::from_size_align(size_bytes.max(1), align).expect("size and align are valid");
584
585 self.stats.total_allocations += 1;
587 self.stats.total_bytes_allocated += size_bytes;
588
589 if let Some(bucket) = self.pools.get_mut(&pool_key) {
591 let mut found_idx: Option<usize> = None;
593 for (i, entry) in bucket.available_buffers.iter().enumerate() {
594 if entry.capacity_bytes >= size_bytes && entry.layout.align() >= align {
595 found_idx = Some(i);
596 break;
597 }
598 }
599 if let Some(idx) = found_idx {
600 let raw_entry = bucket
601 .available_buffers
602 .remove(idx)
603 .expect("index was valid moments ago");
604 self.stats.pool_hits += 1;
605 bucket.reuses += 1;
606
607 let ptr = NonNull::new(raw_entry.ptr.as_ptr() as *mut T)
608 .expect("RawEntry pointer is non-null by construction");
609 let actual_capacity = raw_entry.capacity_bytes / element_size;
611 let entry_layout = raw_entry.layout;
612 std::mem::forget(raw_entry);
613
614 let weak = self.self_weak.clone().unwrap_or_else(Weak::new);
615 return ReusedBuffer {
616 ptr,
617 capacity: actual_capacity,
618 layout: entry_layout,
619 pool: weak,
620 };
621 }
622 }
623
624 self.stats.pool_misses += 1;
626
627 self.pools.entry(pool_key).or_insert_with(|| MemoryPool {
629 available_buffers: VecDeque::new(),
630 size_class,
631 max_buffers: self.config.max_buffers_per_class,
632 allocations: 0,
633 reuses: 0,
634 deallocations: 0,
635 });
636
637 if let Some(bucket) = self.pools.get_mut(&pool_key) {
638 bucket.allocations += 1;
639 }
640
641 let raw_ptr = unsafe { std::alloc::alloc(layout) };
643 let ptr = NonNull::new(raw_ptr as *mut T).unwrap_or_else(|| handle_alloc_error(layout));
644
645 let weak = self.self_weak.clone().unwrap_or_else(Weak::new);
646 ReusedBuffer {
647 ptr,
648 capacity: count,
649 layout,
650 pool: weak,
651 }
652 }
653
654 #[deprecated = "Use global_acquire_uninit instead for zero-copy buffer reuse"]
661 pub fn allocate<T: TensorElement + Default + 'static>(&mut self, count: usize) -> Vec<T> {
662 let mut buf = self.acquire_uninit::<T>(count);
663 for slot in buf.as_uninit_slice_mut() {
665 slot.write(T::default());
666 }
667 buf.into_vec(count)
668 }
669
670 pub fn find_size_class(&self, size_bytes: usize) -> usize {
672 self.config
673 .size_classes
674 .iter()
675 .position(|&class_size| size_bytes <= class_size)
676 .unwrap_or(self.config.size_classes.len() - 1)
677 }
678
679 pub fn deallocate<T: 'static>(&mut self, data: Vec<T>) {
686 drop(data);
688 }
689
690 pub fn clear(&mut self) {
692 self.pools.clear();
693 self.stats = PoolStatistics::default();
694 }
695
696 pub fn get_statistics(&self) -> &PoolStatistics {
698 &self.stats
699 }
700
701 pub fn hit_rate(&self) -> f64 {
703 if self.stats.total_allocations == 0 {
704 0.0
705 } else {
706 self.stats.pool_hits as f64 / self.stats.total_allocations as f64
707 }
708 }
709
710 pub fn cleanup(&mut self) {
712 if self.config.auto_cleanup {
713 let threshold_bytes =
714 (self.config.max_total_memory as f64 * self.config.cleanup_threshold) as usize;
715 if self.stats.total_bytes_allocated > threshold_bytes {
716 self.pools
717 .retain(|_, pool| !pool.available_buffers.is_empty());
718 }
719 }
720 }
721}
722
723impl std::fmt::Debug for GlobalMemoryPool {
724 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
725 f.debug_struct("GlobalMemoryPool")
726 .field("pools", &self.pools)
727 .field("stats", &self.stats)
728 .field("config", &self.config)
729 .field("scirs2_pool", &"<GlobalBufferPool>")
730 .field("leak_detector", &"<LeakDetector>")
731 .finish()
732 }
733}
734
735impl std::fmt::Debug for RawEntry {
738 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
739 f.debug_struct("RawEntry")
740 .field("capacity_bytes", &self.capacity_bytes)
741 .finish()
742 }
743}
744
745pub fn global_acquire_uninit<T: 'static>(count: usize) -> ReusedBuffer<T> {
758 let pool_arc = get_memory_pool();
759 let mut guard = pool_arc
760 .lock()
761 .expect("global memory pool lock should not be poisoned");
762 guard.acquire_uninit::<T>(count)
763}
764
765pub fn global_acquire_uninit_aligned<T: 'static>(count: usize, align: usize) -> ReusedBuffer<T> {
777 let pool_arc = get_memory_pool();
778 let mut guard = pool_arc
779 .lock()
780 .expect("global memory pool lock should not be poisoned");
781 guard.acquire_uninit_aligned::<T>(count, align)
782}
783
784pub type EnhancedMemoryStats = PoolStatistics;
788
789impl<T: TensorElement> Tensor<T> {
791 pub fn create_efficient(shape: &[usize], device: DeviceType) -> Result<Self>
793 where
794 T: Clone + Default,
795 {
796 let binding = get_memory_pool();
797 let mut pool = binding.lock().expect("lock should not be poisoned");
798 pool.create_large_tensor::<T>(shape, device)
799 }
800
801 pub fn lazy(shape: &[usize], device: DeviceType) -> Result<Self>
803 where
804 T: Clone + Default,
805 {
806 let binding = get_memory_pool();
807 let mut pool = binding.lock().expect("lock should not be poisoned");
808 pool.create_lazy_tensor::<T>(shape, device)
809 }
810
811 pub fn memory_mapped(shape: &[usize], device: DeviceType) -> Result<Self>
822 where
823 T: Clone + Default,
824 {
825 #[cfg(feature = "profiling")]
826 {
827 }
829
830 let total_elements: usize = shape.iter().product();
832 let data = vec![T::default(); total_elements];
833 Self::from_data(data, shape.to_vec(), device)
834 }
835
836 pub fn chunked(shape: &[usize], chunk_size: usize, device: DeviceType) -> Result<Self>
846 where
847 T: Clone + Default,
848 {
849 #[cfg(feature = "profiling")]
850 {
851 }
853 let total_elements: usize = shape.iter().product();
854
855 let effective_chunk_size = if chunk_size == 0 {
857 let default_chunk_bytes = 64 * 1024;
859 let element_size = std::mem::size_of::<T>();
860 (default_chunk_bytes / element_size.max(1)).max(1)
861 } else {
862 chunk_size
863 };
864
865 let cache_line_elements = 64 / std::mem::size_of::<T>().max(1);
867 let aligned_chunk_size = ((effective_chunk_size + cache_line_elements - 1)
868 / cache_line_elements)
869 * cache_line_elements;
870
871 let _ = (total_elements, effective_chunk_size, aligned_chunk_size); let data = vec![T::default(); total_elements];
876
877 Self::from_data(data, shape.to_vec(), device)
880 }
881
882 pub fn disk_backed(shape: &[usize], device: DeviceType, file_path: Option<&str>) -> Result<Self>
896 where
897 T: Clone + Default,
898 {
899 #[cfg(feature = "profiling")]
900 {
901 }
903 let total_elements: usize = shape.iter().product();
904
905 let backing_path = if let Some(path) = file_path {
907 std::path::PathBuf::from(path)
909 } else {
910 let temp_dir = std::env::temp_dir();
912 let timestamp = std::time::SystemTime::now()
913 .duration_since(std::time::UNIX_EPOCH)
914 .unwrap_or_default()
915 .as_secs();
916 temp_dir.join(format!(
917 "torsh_tensor_{}_{}.bin",
918 timestamp,
919 std::process::id()
920 ))
921 };
922
923 let _ = (total_elements, &backing_path); let data = vec![T::default(); total_elements];
929
930 let tensor = Self::from_data(data, shape.to_vec(), device)?;
933
934 Ok(tensor)
935 }
936
937 pub fn process_chunked<F, R>(&self, chunk_size: usize, mut processor: F) -> Result<Vec<R>>
939 where
940 F: FnMut(&[T]) -> Result<R>,
941 T: Clone,
942 {
943 #[cfg(feature = "profiling")]
944 {
945 }
947 let data = self.data()?;
948 let mut results = Vec::new();
949
950 let effective_chunk_size = chunk_size;
952
953 for chunk in data.chunks(effective_chunk_size) {
954 results.push(processor(chunk)?);
955 }
956
957 Ok(results)
958 }
959}
960
961impl MemoryPool {
962 fn new(size_class: usize, max_buffers: usize) -> Self {
963 Self {
964 available_buffers: VecDeque::new(),
965 size_class,
966 max_buffers,
967 allocations: 0,
968 reuses: 0,
969 deallocations: 0,
970 }
971 }
972}
973
974impl<T: TensorElement + Copy + Default> PooledTensor<T> {
975 pub fn new(shape: &[usize], device: DeviceType) -> Result<Self> {
977 let numel = shape.iter().product::<usize>();
978
979 let pool = get_memory_pool();
981 let data = {
982 let mut pool_guard = pool.lock().expect("lock should not be poisoned");
983 #[allow(deprecated)]
984 pool_guard.allocate::<T>(numel)
985 };
986
987 let tensor = Tensor::from_data(data, shape.to_vec(), device)?;
988 let type_id = std::any::TypeId::of::<T>();
989 let size_class = {
990 let pool_guard = pool.lock().expect("lock should not be poisoned");
991 pool_guard.find_size_class(numel * std::mem::size_of::<T>())
992 };
993 let align = std::mem::align_of::<T>();
994
995 Ok(Self {
996 tensor,
997 pool_key: Some((type_id, size_class, align)),
998 _phantom: PhantomData,
999 })
1000 }
1001
1002 pub fn zeros(shape: &[usize], device: DeviceType) -> Result<Self> {
1004 let mut pooled = Self::new(shape, device)?;
1005 let numel = shape.iter().product::<usize>();
1007 let data = vec![T::default(); numel];
1008 pooled.tensor.storage = TensorStorage::create_optimal(data)?;
1009 Ok(pooled)
1010 }
1011
1012 pub fn ones(shape: &[usize], device: DeviceType) -> Result<Self>
1014 where
1015 T: std::ops::Add<Output = T> + From<f32>,
1016 {
1017 let mut pooled = Self::new(shape, device)?;
1018 let numel = shape.iter().product::<usize>();
1020 let data = vec![T::from(1.0f32); numel];
1021 pooled.tensor.storage = TensorStorage::create_optimal(data)?;
1022 Ok(pooled)
1023 }
1024
1025 pub fn tensor(&self) -> &Tensor<T> {
1027 &self.tensor
1028 }
1029
1030 pub fn tensor_mut(&mut self) -> &mut Tensor<T> {
1032 &mut self.tensor
1033 }
1034
1035 pub fn into_tensor(mut self) -> Tensor<T> {
1037 self.pool_key = None; self.tensor.clone()
1039 }
1040}
1041
1042impl<T: TensorElement + std::default::Default> Drop for PooledTensor<T> {
1043 fn drop(&mut self) {
1044 if let Some((_type_id, _size_class, _align)) = self.pool_key {
1045 if let Ok(data) = self.tensor.to_vec() {
1047 let pool = get_memory_pool();
1048 let mut pool_guard = pool.lock().expect("lock should not be poisoned");
1049 pool_guard.deallocate(data);
1050 }
1051 }
1052 }
1053}
1054
1055impl<T: TensorElement + Copy + Default> Tensor<T> {
1057 pub fn pooled(shape: &[usize], device: DeviceType) -> Result<PooledTensor<T>> {
1059 PooledTensor::new(shape, device)
1060 }
1061
1062 pub fn temporary(shape: &[usize], device: DeviceType) -> Result<PooledTensor<T>> {
1064 PooledTensor::new(shape, device)
1065 }
1066}
1067
1068pub fn clear_memory_pool() {
1070 if let Some(pool) = MEMORY_POOL.get() {
1071 pool.lock().expect("lock should not be poisoned").clear();
1072 }
1073}
1074
1075pub fn get_pool_statistics() -> PoolStatistics {
1076 get_memory_pool()
1077 .lock()
1078 .expect("lock should not be poisoned")
1079 .get_statistics()
1080 .clone()
1081}
1082
1083pub fn get_pool_hit_rate() -> f64 {
1084 get_memory_pool()
1085 .lock()
1086 .expect("lock should not be poisoned")
1087 .hit_rate()
1088}
1089
1090pub fn cleanup_memory_pool() {
1091 get_memory_pool()
1092 .lock()
1093 .expect("lock should not be poisoned")
1094 .cleanup();
1095}
1096
1097#[cfg(test)]
1098mod tests {
1099 use super::*;
1100
1101 static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
1103
1104 #[test]
1105 fn test_memory_pool_basic() {
1106 clear_memory_pool();
1107
1108 let pooled = PooledTensor::<f32>::zeros(&[100, 100], DeviceType::Cpu)
1110 .expect("zeros creation should succeed");
1111 assert_eq!(pooled.tensor().numel(), 10000);
1112
1113 drop(pooled);
1115
1116 let _pooled2 = PooledTensor::<f32>::zeros(&[100, 100], DeviceType::Cpu)
1118 .expect("zeros creation should succeed");
1119
1120 let stats = get_pool_statistics();
1121 assert!(stats.pool_hits > 0 || stats.pool_misses > 0);
1122 }
1123
1124 #[test]
1125 fn test_pool_statistics() {
1126 clear_memory_pool();
1127
1128 let _pooled1 = PooledTensor::<f32>::zeros(&[50, 50], DeviceType::Cpu)
1129 .expect("zeros creation should succeed");
1130 let _pooled2 = PooledTensor::<f32>::ones(&[50, 50], DeviceType::Cpu)
1131 .expect("ones creation should succeed");
1132
1133 let stats = get_pool_statistics();
1134 assert!(stats.total_allocations >= 2);
1135 assert!(stats.total_bytes_allocated > 0);
1136 }
1137
1138 #[test]
1139 fn test_pool_cleanup() {
1140 clear_memory_pool();
1141
1142 for _ in 0..10 {
1144 let _temp = PooledTensor::<f32>::zeros(&[100, 100], DeviceType::Cpu)
1145 .expect("zeros creation should succeed");
1146 }
1147
1148 cleanup_memory_pool();
1149 let _stats = get_pool_statistics();
1150 }
1152
1153 #[test]
1154 fn test_pooled_tensor_conversion() {
1155 let pooled = PooledTensor::<f32>::ones(&[10, 10], DeviceType::Cpu)
1156 .expect("ones creation should succeed");
1157 let tensor = pooled.into_tensor();
1158 assert_eq!(tensor.numel(), 100);
1159 }
1160
1161 #[test]
1164 fn test_acquire_truly_reuses_allocation() {
1165 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1166 clear_memory_pool();
1167
1168 let buf1: ReusedBuffer<f32> = global_acquire_uninit::<f32>(1024);
1169 let ptr1 = buf1.as_ptr_raw();
1170 buf1.release_to_pool();
1171
1172 let buf2: ReusedBuffer<f32> = global_acquire_uninit::<f32>(1024);
1173 let ptr2 = buf2.as_ptr_raw();
1174 buf2.release_to_pool();
1175
1176 assert_eq!(
1177 ptr1, ptr2,
1178 "pool should return the same allocation on second acquire"
1179 );
1180 }
1181
1182 #[test]
1183 fn test_into_vec_transfers_ownership() {
1184 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1185 clear_memory_pool();
1186
1187 let mut buf: ReusedBuffer<f32> = global_acquire_uninit::<f32>(64);
1188 for slot in buf.as_uninit_slice_mut() {
1190 slot.write(1.0_f32);
1191 }
1192 let vec = buf.into_vec(64);
1193 assert_eq!(vec.len(), 64);
1194 assert!(vec.iter().all(|&x| x == 1.0_f32));
1195 }
1196
1197 #[test]
1198 fn test_drop_returns_to_pool() {
1199 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1200 clear_memory_pool();
1201
1202 {
1203 let buf: ReusedBuffer<f32> = global_acquire_uninit::<f32>(256);
1204 drop(buf);
1206 }
1207
1208 let buf2: ReusedBuffer<f32> = global_acquire_uninit::<f32>(256);
1210 buf2.release_to_pool();
1211
1212 let stats = get_pool_statistics();
1213 assert!(
1214 stats.pool_hits >= 1,
1215 "expected at least one pool hit after drop-return"
1216 );
1217 }
1218
1219 #[test]
1220 fn test_acquire_capacity_and_uninit_slice() {
1221 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1222 clear_memory_pool();
1223
1224 let buf: ReusedBuffer<u64> = global_acquire_uninit::<u64>(32);
1225 assert_eq!(buf.capacity(), 32);
1226 buf.release_to_pool();
1227 }
1228
1229 #[test]
1232 fn test_acquire_aligned_returns_simd_aligned_pointer() {
1233 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1234 clear_memory_pool();
1235
1236 let buf: ReusedBuffer<f32> = global_acquire_uninit_aligned::<f32>(1024, 32);
1238 assert_eq!(buf.capacity(), 1024);
1239 let addr = buf.as_ptr_raw() as usize;
1240 assert_eq!(
1241 addr % 32,
1242 0,
1243 "buffer pointer {addr:#x} must be 32-byte aligned"
1244 );
1245 buf.release_to_pool();
1246 }
1247
1248 #[test]
1249 fn test_acquire_aligned_pool_hit_on_release() {
1250 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1251 clear_memory_pool();
1252
1253 let buf1: ReusedBuffer<f32> = global_acquire_uninit_aligned::<f32>(2048, 32);
1254 let ptr1 = buf1.as_ptr_raw();
1255 let cap1 = buf1.capacity();
1256 buf1.release_to_pool();
1257
1258 let buf2: ReusedBuffer<f32> = global_acquire_uninit_aligned::<f32>(2048, 32);
1259 let ptr2 = buf2.as_ptr_raw();
1260 let cap2 = buf2.capacity();
1261 assert_eq!(
1262 ptr1, ptr2,
1263 "aligned bucket should return the same allocation on second acquire"
1264 );
1265 assert_eq!(cap1, cap2, "capacity should match across reuse");
1266 assert_eq!(ptr2 as usize % 32, 0, "reused buffer must remain aligned");
1268 buf2.release_to_pool();
1269 }
1270
1271 #[test]
1272 fn test_aligned_and_natural_buckets_are_independent() {
1273 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1274 clear_memory_pool();
1275
1276 let buf_aligned: ReusedBuffer<f32> = global_acquire_uninit_aligned::<f32>(512, 32);
1278 let ptr_aligned = buf_aligned.as_ptr_raw();
1279 buf_aligned.release_to_pool();
1280
1281 let buf_natural: ReusedBuffer<f32> = global_acquire_uninit::<f32>(512);
1284 let ptr_natural = buf_natural.as_ptr_raw();
1285 assert_ne!(
1286 ptr_aligned, ptr_natural,
1287 "naturally-aligned bucket must be distinct from the 32-byte bucket"
1288 );
1289 buf_natural.release_to_pool();
1290 }
1291
1292 #[test]
1293 #[should_panic(expected = "alignment must be a power of two")]
1294 fn test_acquire_aligned_rejects_non_power_of_two() {
1295 let _guard = TEST_LOCK.lock().expect("test mutex should not be poisoned");
1296 clear_memory_pool();
1297 let _buf: ReusedBuffer<f32> = global_acquire_uninit_aligned::<f32>(16, 6);
1298 }
1299
1300 #[cfg(feature = "memory_efficient")]
1309 #[test]
1310 fn test_map_through_mmap_file_roundtrips_known_data() {
1311 let known: Vec<f32> = (0..48).map(|i| (i as f32) * 1.5 - 7.25).collect();
1313
1314 let backing_path = unique_mmap_path("test_helper");
1315 assert!(
1316 backing_path.starts_with(std::env::temp_dir()),
1317 "backing file must live under the system temp directory"
1318 );
1319
1320 let mapped = map_through_mmap_file::<f32>(known.clone(), &backing_path)
1321 .expect("memory-mapped round-trip should succeed");
1322
1323 assert_eq!(
1324 mapped, known,
1325 "as_slice() must return exactly the data written to the memory-mapped file"
1326 );
1327
1328 let _ = std::fs::remove_file(&backing_path);
1330 }
1331
1332 #[cfg(feature = "memory_efficient")]
1335 #[test]
1336 fn test_memory_mapped_array_as_slice_direct() {
1337 use scirs2_core::memory_efficient::{AccessMode, MemoryMappedArray};
1338 use scirs2_core::ndarray::Array1;
1339
1340 let known: Vec<f64> = vec![3.5, -1.25, 42.0, 7.0, 0.5, 100.0, -8.0, 256.0];
1341 let backing_path = unique_mmap_path("test_direct");
1342
1343 let array = Array1::from(known.clone());
1344 let mmap = MemoryMappedArray::<f64>::new(Some(&array), &backing_path, AccessMode::Write, 0)
1345 .expect("memory-mapped array creation should succeed");
1346
1347 let read_back = mmap.as_slice().to_vec();
1348 drop(mmap);
1349 let _ = std::fs::remove_file(&backing_path);
1350
1351 assert_eq!(
1352 read_back, known,
1353 "as_slice() over a Write-mode memory map must return the written data"
1354 );
1355 }
1356
1357 #[cfg(feature = "memory_efficient")]
1360 #[test]
1361 fn test_create_memory_mapped_tensor_uses_mmap_path() {
1362 let mut pool = GlobalMemoryPool::new();
1363 let shape = [4usize, 5];
1364 let tensor = pool
1365 .create_memory_mapped_tensor::<f32>(&shape, DeviceType::Cpu)
1366 .expect("memory-mapped tensor creation should succeed");
1367
1368 assert_eq!(tensor.numel(), 20);
1369 let dims = tensor.shape();
1370 assert_eq!(dims.dims(), &[4, 5]);
1371
1372 let data = tensor.data().expect("tensor data should be readable");
1375 assert_eq!(data.len(), 20);
1376 assert!(data.iter().all(|&x| x == 0.0_f32));
1377 }
1378}