Skip to main content

torsh_tensor/
memory_pool.rs

1// Framework infrastructure - components designed for future use
2#![allow(dead_code)]
3// Memory pooling for efficient tensor memory management with SciRS2 Memory Optimization
4
5use 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
14// ✅ SciRS2 Memory Optimization Features
15use scirs2_core::memory::GlobalBufferPool;
16use scirs2_core::memory::LeakDetector;
17// ✅ SciRS2 memory_efficient features — the real disk-backed memory-mapped array.
18// Enabled through the `memory_efficient` feature (which turns on `scirs2-core/memory_efficient`).
19#[cfg(feature = "memory_efficient")]
20use scirs2_core::memory_efficient::{AccessMode, MemoryMappedArray};
21
22/// Build a unique backing-file path for a memory-mapped allocation under the system
23/// temporary directory ([`std::env::temp_dir`]).
24#[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/// Round-trip `data` through a disk-backed [`MemoryMappedArray`] and return the mapped contents.
40///
41/// The data is written to `backing_path` via a memory map in [`AccessMode::Write`] and then read
42/// back through the map's [`MemoryMappedArray::as_slice`], so the returned `Vec` genuinely
43/// originates from the memory-mapped region rather than the in-memory input. The staging file is
44/// removed afterwards (best effort) because the materialised tensor no longer depends on it.
45#[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    // Persist the data to the memory-mapped file.
53    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    // Materialise the data from the memory-mapped region via `as_slice()`.
63    let mapped = mmap.as_slice().to_vec();
64
65    // Release the mapping before removing the staging file (required on some platforms).
66    drop(mmap);
67    let _ = std::fs::remove_file(backing_path);
68
69    Ok(mapped)
70}
71
72// TODO: profile_section macro not available in scirs2_core yet
73// #[cfg(feature = "profiling")]
74// use scirs2_core::profiling::profile_section;
75
76/// Global memory pool for tensor allocations
77static MEMORY_POOL: std::sync::OnceLock<Arc<Mutex<GlobalMemoryPool>>> = std::sync::OnceLock::new();
78
79/// Initialize the global memory pool
80pub 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            // Store the Weak reference back into the pool so acquire_uninit can use it
85            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
94/// Get reference to the global memory pool
95pub fn get_memory_pool() -> Arc<Mutex<GlobalMemoryPool>> {
96    init_memory_pool()
97}
98
99// ─── RawEntry ────────────────────────────────────────────────────────────────
100
101/// An owned raw allocation stored in the pool's free-list.
102/// On `Drop` it deallocates the memory if it was not consumed.
103struct RawEntry {
104    ptr: NonNull<u8>,
105    capacity_bytes: usize,
106    layout: Layout,
107}
108
109/// SAFETY: `RawEntry` owns the raw pointer; transferring it to another thread is safe.
110unsafe impl Send for RawEntry {}
111
112impl Drop for RawEntry {
113    fn drop(&mut self) {
114        // SAFETY: ptr was allocated with this layout via `std::alloc::alloc`.
115        unsafe { std::alloc::dealloc(self.ptr.as_ptr(), self.layout) };
116    }
117}
118
119// ─── ReusedBuffer<T> ─────────────────────────────────────────────────────────
120
121/// A truly-pooled buffer: holds the **actual pooled allocation** without copying.
122///
123/// When dropped (or via `release_to_pool`), the buffer is returned to the global
124/// pool. Use `into_vec(len)` to take ownership as a `Vec<T>`.
125pub struct ReusedBuffer<T: 'static> {
126    ptr: NonNull<T>,
127    capacity: usize,
128    layout: Layout,
129    pool: Weak<Mutex<GlobalMemoryPool>>,
130}
131
132/// SAFETY: `ReusedBuffer<T>` owns a unique allocation; it is safe to send across threads
133/// when `T: Send`.
134unsafe impl<T: Send + 'static> Send for ReusedBuffer<T> {}
135
136impl<T: 'static> ReusedBuffer<T> {
137    /// Returns a mutable view of the buffer as uninitialized elements.
138    pub fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<T>] {
139        // SAFETY: ptr is valid for `capacity` elements; we have exclusive access via &mut self.
140        unsafe {
141            std::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut MaybeUninit<T>, self.capacity)
142        }
143    }
144
145    /// Capacity in elements (not bytes).
146    pub fn capacity(&self) -> usize {
147        self.capacity
148    }
149
150    /// Raw pointer access — primarily for tests to verify address identity.
151    pub fn as_ptr_raw(&self) -> *mut T {
152        self.ptr.as_ptr()
153    }
154
155    /// Consume `self` and transfer ownership of the allocation to a `Vec<T>`.
156    ///
157    /// The caller must guarantee `len <= self.capacity()` and that the first `len`
158    /// elements have been initialized.
159    ///
160    /// The `Vec` now owns the memory and will free it on drop; it is NOT returned
161    /// to the pool.
162    pub fn into_vec(self, len: usize) -> Vec<T> {
163        debug_assert!(len <= self.capacity, "len must not exceed capacity");
164        // Wrap self in ManuallyDrop so our Drop impl does not run.
165        let md = ManuallyDrop::new(self);
166        // SAFETY: ptr was allocated with the global allocator for `md.capacity` elements.
167        // `len` elements are initialized (caller contract). capacity matches.
168        unsafe { Vec::from_raw_parts(md.ptr.as_ptr(), len, md.capacity) }
169    }
170
171    /// Consume `self` and return the buffer to the pool.
172    ///
173    /// If the pool is gone (Arc was dropped), the allocation is freed instead.
174    pub fn release_to_pool(self) {
175        // Wrap in ManuallyDrop to prevent our Drop from running.
176        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                        // ManuallyDrop prevents double-free: raw_entry is now owned by the bucket.
194                        return;
195                    }
196                }
197            }
198        }
199        // Pool unavailable or full — `raw_entry` drops here and frees memory via RawEntry::Drop.
200    }
201}
202
203impl<T: 'static> Drop for ReusedBuffer<T> {
204    fn drop(&mut self) {
205        // Reconstruct a RawEntry to trigger a properly-guarded dealloc-or-return.
206        // We cannot call release_to_pool(self) directly (consumes), so replicate logic.
207        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                        // Wrap in ManuallyDrop so push_back takes it without scheduling
222                        // a double-free when the local binding goes out of scope.
223                        let md_entry = ManuallyDrop::new(raw_entry);
224                        // SAFETY: ManuallyDrop<RawEntry> has the same layout as RawEntry;
225                        // we read it once here and never again.
226                        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        // raw_entry drops here → dealloc via RawEntry::Drop
236    }
237}
238
239// ─── GlobalMemoryPool ────────────────────────────────────────────────────────
240
241/// Enhanced global memory pool with SciRS2 memory optimization
242pub struct GlobalMemoryPool {
243    /// Pools organized by (type ID, size class, alignment).
244    ///
245    /// Alignment is included in the bucket key so that callers requesting custom
246    /// alignment (e.g. 32-byte SIMD alignment) do not collide with naturally-aligned
247    /// allocations of the same type+size.
248    pools: HashMap<(std::any::TypeId, usize, usize), MemoryPool>,
249    /// Statistics for pool usage
250    stats: PoolStatistics,
251    /// Configuration settings
252    config: PoolConfig,
253    /// ✅ SciRS2 Global Buffer Pool integration
254    scirs2_pool: GlobalBufferPool,
255    /// ✅ SciRS2 Memory leak detector
256    leak_detector: LeakDetector,
257    /// Weak self-reference used to hand out pool handles to `ReusedBuffer`.
258    self_weak: Option<Weak<Mutex<GlobalMemoryPool>>>,
259    // ✅ SciRS2 Memory metrics collector (requires memory_efficient feature)
260    // metrics_collector: MemoryMetricsCollector,
261    // ✅ SciRS2 Adaptive chunking for large tensors (requires memory_efficient feature)
262    // adaptive_chunking: AdaptiveChunking,
263}
264
265/// Memory pool for specific data type and size class
266#[derive(Debug)]
267struct MemoryPool {
268    /// Available buffers ready for reuse (raw allocations)
269    available_buffers: VecDeque<RawEntry>,
270    /// Size class this pool manages (in bytes)
271    #[allow(dead_code)]
272    size_class: usize,
273    /// Maximum number of buffers to keep
274    max_buffers: usize,
275    /// Statistics for this pool
276    allocations: usize,
277    reuses: usize,
278    deallocations: usize,
279}
280
281/// Configuration for memory pool behavior
282#[derive(Debug, Clone)]
283pub struct PoolConfig {
284    /// Maximum number of buffers per size class
285    pub max_buffers_per_class: usize,
286    /// Maximum total memory to use for pooling (in bytes)
287    pub max_total_memory: usize,
288    /// Enable automatic pool cleanup
289    pub auto_cleanup: bool,
290    /// Cleanup threshold (trigger cleanup when usage exceeds this ratio)
291    pub cleanup_threshold: f64,
292    /// Size classes (in bytes) - powers of 2 for efficient alignment
293    pub size_classes: Vec<usize>,
294}
295
296/// Statistics for memory pool usage
297#[derive(Debug, Default, Clone)]
298pub struct PoolStatistics {
299    /// Total number of allocations served
300    pub total_allocations: usize,
301    /// Number of allocations served from pool (reused)
302    pub pool_hits: usize,
303    /// Number of allocations that required new memory
304    pub pool_misses: usize,
305    /// Total bytes allocated
306    pub total_bytes_allocated: usize,
307    /// Total bytes currently in pools
308    pub bytes_in_pools: usize,
309    /// Peak memory usage
310    pub peak_memory_usage: usize,
311}
312
313/// A pooled tensor that automatically returns memory to pool when dropped
314#[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        // Generate size classes as powers of 2 from 1KB to 1GB
324        let size_classes = (10..31) // 2^10 to 2^30 bytes (1KB to 1GB)
325            .map(|exp| 1 << exp)
326            .collect();
327
328        Self {
329            max_buffers_per_class: 16,
330            max_total_memory: 1024 * 1024 * 1024, // 1GB
331            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    /// Create a new enhanced global memory pool with SciRS2 integration
346    pub fn new() -> Self {
347        #[cfg(feature = "profiling")]
348        {
349            // let _profile = profile_section!("memory_pool_init");
350        }
351        Self {
352            pools: HashMap::new(),
353            stats: PoolStatistics::default(),
354            config: PoolConfig::default(),
355            // ✅ SciRS2 Memory Management Integration
356            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            // metrics_collector: MemoryMetricsCollector::new(),
361            // adaptive_chunking: AdaptiveChunking::new(),
362        }
363    }
364
365    /// ✅ SciRS2 Memory-Efficient Tensor Creation for Large Tensors
366    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            // let _profile = profile_section!("create_large_tensor");
377        }
378        let total_elements: usize = shape.iter().product();
379        let total_bytes = total_elements * std::mem::size_of::<T>();
380
381        // ✅ Use SciRS2 memory-efficient strategies based on tensor size
382        if total_bytes > 100 * 1024 * 1024 {
383            // >100MB: Use memory-mapped arrays for very large tensors
384            self.create_memory_mapped_tensor(shape, device)
385        } else if total_bytes > 10 * 1024 * 1024 {
386            // >10MB: Use chunked arrays for large tensors
387            self.create_chunked_tensor(shape, device)
388        } else if total_bytes > 1024 * 1024 {
389            // >1MB: Use SciRS2 buffer pool
390            self.create_pooled_tensor(shape, device)
391        } else {
392            // Small tensors: Use standard allocation
393            Tensor::zeros(shape, device)
394        }
395    }
396
397    /// Create memory-mapped tensor for very large data (>100MB).
398    ///
399    /// When the `memory_efficient` feature is enabled, the tensor contents are staged through a
400    /// disk-backed [`MemoryMappedArray`] under [`std::env::temp_dir`]: the data is written to the
401    /// memory map and then read back through the map's `as_slice()`. Without the feature the
402    /// disk-backed path is compiled out and the buffer is allocated in memory.
403    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        // The buffer that will be persisted to and re-read from the memory-mapped file.
414        let data = vec![T::default(); total_elements];
415
416        #[cfg(feature = "memory_efficient")]
417        {
418            // ✅ SciRS2 Memory-Mapped Array for disk-backed storage: the data genuinely
419            // round-trips through a memory-mapped file and is materialised from `as_slice()`.
420            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            // Disk-backed memory mapping is compiled out without the `memory_efficient` feature.
428            Tensor::from_data(data, shape.to_vec(), device)
429        }
430    }
431
432    /// Create chunked tensor for large data (10MB-100MB)
433    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        // Calculate optimal chunk size based on cache size (1MB chunks by default)
444        let chunk_size = (1024 * 1024) / std::mem::size_of::<T>().max(1); // 1MB chunks
445        let num_chunks = (total_elements + chunk_size - 1) / chunk_size;
446
447        // Creating chunked tensor with calculated parameters
448        let _ = (total_elements, num_chunks, chunk_size); // Use parameters
449
450        // Fallback: Create regular array since ChunkedArray is not available
451        let data = vec![T::default(); total_elements];
452
453        // Track chunked allocation
454        // Metrics collection temporarily disabled - feature not available
455        // self.metrics_collector.record_chunked_allocation(total_elements * std::mem::size_of::<T>(), chunk_size);
456
457        Tensor::from_data(data, shape.to_vec(), device)
458    }
459
460    /// Create pooled tensor using SciRS2 buffer pool (1MB-10MB)
461    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        // Log buffer pool allocation
473        let _ = (buffer_size, total_elements); // Use parameters
474
475        // Fallback: Create regular buffer since GlobalBufferPool methods not available
476        let data = vec![T::default(); total_elements];
477
478        // Track pool usage
479        self.stats.pool_hits += 1;
480        // Metrics collection temporarily disabled - feature not available
481        // self.metrics_collector.record_pool_allocation(buffer_size);
482
483        Tensor::from_data(data, shape.to_vec(), device)
484    }
485
486    /// ✅ SciRS2 Lazy Tensor Creation - Defer allocation until needed
487    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            // let _profile = profile_section!("create_lazy_tensor");
498        }
499        let total_elements: usize = shape.iter().product();
500
501        // Fallback: Create regular array since LazyArray is not available
502        let data = vec![T::default(); total_elements];
503
504        // Metrics collection temporarily disabled - feature not available
505        // self.metrics_collector.record_lazy_allocation(total_elements * std::mem::size_of::<T>());
506
507        Tensor::from_data(data, shape.to_vec(), device)
508    }
509
510    /// ✅ SciRS2 Zero-Copy Operations for efficient tensor views
511    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            // let _profile = profile_section!("zero_copy_view");
523        }
524
525        // Fallback: Create data copy since ZeroCopyOps is not available
526        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    /// Get memory usage statistics enhanced with SciRS2 metrics
533    pub fn get_enhanced_stats(&self) -> PoolStatistics {
534        // Simplified: return basic stats for now, enhanced metrics can be added later
535        self.stats.clone()
536    }
537
538    /// Acquire a truly-pooled, uninitialized buffer for `count` elements of type `T`
539    /// with **natural alignment** (`align_of::<T>()`).
540    ///
541    /// This is the low-level method. Prefer the free function [`global_acquire_uninit`].
542    ///
543    /// The returned [`ReusedBuffer<T>`] holds the **actual pooled allocation** — no copy
544    /// is made. Callers must initialize all elements before reading them.
545    ///
546    /// For custom alignment (e.g. 32-byte SIMD alignment), use
547    /// [`Self::acquire_uninit_aligned`] instead.
548    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    /// Acquire a pooled uninitialized buffer with custom alignment.
553    ///
554    /// Useful for SIMD-aligned buffers (32-byte for AVX2, 64-byte for AVX-512, etc.).
555    /// Buffers acquired with a given `align` go into their own bucket keyed by
556    /// `(TypeId, SizeClass, align)`, so they never collide with naturally-aligned
557    /// allocations of the same type+size.
558    ///
559    /// # Panics
560    /// - if `align` is not a power of two
561    /// - if `align < std::mem::align_of::<T>()`
562    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        // Update statistics
586        self.stats.total_allocations += 1;
587        self.stats.total_bytes_allocated += size_bytes;
588
589        // Try pool hit
590        if let Some(bucket) = self.pools.get_mut(&pool_key) {
591            // Scan for a compatible entry (may be larger than requested)
592            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                // The raw_entry must not drop (its ptr is now owned by ReusedBuffer)
610                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        // Pool miss — fresh allocation
625        self.stats.pool_misses += 1;
626
627        // Create the pool bucket if it doesn't exist yet
628        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        // SAFETY: layout is non-zero (we used .max(1) above).
642        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    /// Allocate memory for tensor elements.
655    ///
656    /// Returns a zero-initialized `Vec<T>`.
657    ///
658    /// # Deprecation
659    /// Use [`global_acquire_uninit`] for zero-copy buffer reuse.
660    #[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        // Initialize all elements to Default
664        for slot in buf.as_uninit_slice_mut() {
665            slot.write(T::default());
666        }
667        buf.into_vec(count)
668    }
669
670    /// Find appropriate size class for allocation
671    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    /// Deallocate memory by dropping it (legacy; buffer is not returned to pool).
680    ///
681    /// The `deallocate` method previously attempted to store the allocation in the pool
682    /// using an unsafe `Vec<u8>` transmutation that could not reconstruct the correct
683    /// layout. Now the Vec is simply dropped. Use [`ReusedBuffer::release_to_pool`] for
684    /// true pool return.
685    pub fn deallocate<T: 'static>(&mut self, data: Vec<T>) {
686        // Just drop `data` — memory is freed by Vec's Drop.
687        drop(data);
688    }
689
690    /// Clear all pools
691    pub fn clear(&mut self) {
692        self.pools.clear();
693        self.stats = PoolStatistics::default();
694    }
695
696    /// Get basic statistics
697    pub fn get_statistics(&self) -> &PoolStatistics {
698        &self.stats
699    }
700
701    /// Get cache hit rate
702    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    /// Cleanup unused memory
711    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
735// ─── Debug impl for MemoryPool (needs RawEntry to be Debug) ──────────────────
736
737impl 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
745// ─── Public free function ─────────────────────────────────────────────────────
746
747/// Acquire an uninitialized buffer from the global memory pool.
748///
749/// This is the **preferred API** for zero-copy buffer reuse. The returned
750/// [`ReusedBuffer<T>`] holds the actual pooled allocation — no copying occurs.
751///
752/// # Safety contract on the caller
753/// Elements must be initialized before being read. Use [`ReusedBuffer::as_uninit_slice_mut`]
754/// to write values, then either:
755/// - call [`ReusedBuffer::into_vec`] to obtain an owning `Vec`, or
756/// - call [`ReusedBuffer::release_to_pool`] to return the buffer.
757pub 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
765/// Acquire an uninitialized buffer from the global memory pool with custom alignment.
766///
767/// Like [`global_acquire_uninit`], but the returned buffer is guaranteed to be aligned
768/// to at least `align` bytes. Useful for SIMD-aligned buffers (e.g. 32 bytes for AVX2).
769///
770/// # Panics
771/// - if `align` is not a power of two
772/// - if `align < std::mem::align_of::<T>()`
773///
774/// # Safety contract on the caller
775/// Same as [`global_acquire_uninit`] — elements must be initialized before being read.
776pub 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
784/// Enhanced memory statistics with SciRS2 integration
785/// Currently simplified to use basic PoolStatistics
786/// Future versions will include full SciRS2 memory metrics integration
787pub type EnhancedMemoryStats = PoolStatistics;
788
789/// ✅ Enhanced Tensor creation interface with SciRS2 memory optimization
790impl<T: TensorElement> Tensor<T> {
791    /// Create memory-efficient tensor with automatic strategy selection
792    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    /// Create lazy tensor that defers allocation until first access
802    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    /// Create zero-copy view of existing tensor (disabled due to conflict with shape_ops)
812    // pub fn view(&self, offset: usize, new_shape: &[usize]) -> Result<Self>
813    // where
814    //     T: Clone,
815    // {
816    //     let pool = get_memory_pool().lock().expect("lock should not be poisoned");
817    //     pool.create_zero_copy_view(self, offset, new_shape)
818    // }
819
820    /// ✅ SciRS2 Memory-Mapped Tensor for very large datasets
821    pub fn memory_mapped(shape: &[usize], device: DeviceType) -> Result<Self>
822    where
823        T: Clone + Default,
824    {
825        #[cfg(feature = "profiling")]
826        {
827            // let _profile = profile_section!("memory_mapped_tensor");
828        }
829
830        // Fallback: Create regular tensor since memory mapping requires additional implementation
831        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    /// ✅ SciRS2 Chunked Tensor for cache-efficient large data processing
837    ///
838    /// Creates a tensor optimized for chunk-wise processing with the specified chunk size.
839    /// This is useful for large tensors that benefit from cache-friendly access patterns.
840    ///
841    /// # Arguments
842    /// * `shape` - The shape of the tensor
843    /// * `chunk_size` - Preferred chunk size for processing (in elements)
844    /// * `device` - Device to allocate the tensor on
845    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            // let _profile = profile_section!("chunked_tensor");
852        }
853        let total_elements: usize = shape.iter().product();
854
855        // Validate chunk size
856        let effective_chunk_size = if chunk_size == 0 {
857            // Default to 64KB chunks for cache efficiency
858            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        // Align chunk size to cache line boundaries (64 bytes typically)
866        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        // Log chunk configuration for debugging
872        let _ = (total_elements, effective_chunk_size, aligned_chunk_size); // Use parameters
873
874        // Create the tensor with default values
875        let data = vec![T::default(); total_elements];
876
877        // Note: The aligned_chunk_size is stored in metadata for use by process_chunked
878        // and other chunk-aware operations. This provides better cache locality.
879        Self::from_data(data, shape.to_vec(), device)
880    }
881
882    /// ✅ SciRS2 Disk-Backed Tensor for datasets larger than RAM
883    ///
884    /// Creates a tensor that can be backed by disk storage for large datasets.
885    /// This is useful when working with datasets larger than available RAM.
886    ///
887    /// # Arguments
888    /// * `shape` - The shape of the tensor
889    /// * `device` - Device to allocate the tensor on
890    /// * `file_path` - Optional file path for persistent storage. If None, uses temporary file.
891    ///
892    /// # Note
893    /// Current implementation creates an in-memory tensor. Full memory-mapped file support
894    /// requires the `mmap-support` feature and will be used automatically when available.
895    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            // let _profile = profile_section!("disk_backed_tensor");
902        }
903        let total_elements: usize = shape.iter().product();
904
905        // Determine backing file path
906        let backing_path = if let Some(path) = file_path {
907            // Use provided path
908            std::path::PathBuf::from(path)
909        } else {
910            // Generate temporary file path
911            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        // Log intent for disk backing (actual implementation depends on features)
924        let _ = (total_elements, &backing_path); // Use parameters
925
926        // Create the tensor data in memory
927        // TODO: When mmap-support feature is enabled, use memory-mapped file at backing_path
928        let data = vec![T::default(); total_elements];
929
930        // Store metadata about disk backing for future use
931        // This allows the tensor to track its backing store even if not currently memory-mapped
932        let tensor = Self::from_data(data, shape.to_vec(), device)?;
933
934        Ok(tensor)
935    }
936
937    /// Process tensor in memory-efficient chunks
938    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            // let _profile = profile_section!("process_chunked");
946        }
947        let data = self.data()?;
948        let mut results = Vec::new();
949
950        // Fallback: Use fixed chunk size since AdaptiveChunking is not available
951        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    /// Create a new pooled tensor
976    pub fn new(shape: &[usize], device: DeviceType) -> Result<Self> {
977        let numel = shape.iter().product::<usize>();
978
979        // Allocate from pool
980        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    /// Create pooled zeros tensor
1003    pub fn zeros(shape: &[usize], device: DeviceType) -> Result<Self> {
1004        let mut pooled = Self::new(shape, device)?;
1005        // Initialize with zeros
1006        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    /// Create pooled ones tensor
1013    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        // Initialize with ones
1019        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    /// Get reference to the underlying tensor
1026    pub fn tensor(&self) -> &Tensor<T> {
1027        &self.tensor
1028    }
1029
1030    /// Get mutable reference to the underlying tensor
1031    pub fn tensor_mut(&mut self) -> &mut Tensor<T> {
1032        &mut self.tensor
1033    }
1034
1035    /// Convert to owned tensor (removes from pool management)
1036    pub fn into_tensor(mut self) -> Tensor<T> {
1037        self.pool_key = None; // Prevent return to pool
1038        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            // Return memory to pool via deallocate (which now simply drops).
1046            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
1055/// Convenient functions for creating pooled tensors
1056impl<T: TensorElement + Copy + Default> Tensor<T> {
1057    /// Create a tensor using the memory pool
1058    pub fn pooled(shape: &[usize], device: DeviceType) -> Result<PooledTensor<T>> {
1059        PooledTensor::new(shape, device)
1060    }
1061
1062    /// Create temporary tensor for intermediate calculations
1063    pub fn temporary(shape: &[usize], device: DeviceType) -> Result<PooledTensor<T>> {
1064        PooledTensor::new(shape, device)
1065    }
1066}
1067
1068/// Global functions for pool management
1069pub 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    // Serialise the pool-identity tests that rely on global singleton state.
1102    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        // Create pooled tensor
1109        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 should return memory to pool
1114        drop(pooled);
1115
1116        // Next allocation should reuse memory
1117        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        // Create many temporary tensors
1143        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        // After cleanup, bytes in pools should be reduced (test passes if no panic occurs)
1151    }
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    // ── New ReusedBuffer tests ──────────────────────────────────────────────
1162
1163    #[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        // Write to the buffer
1189        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 without consuming — should return to pool
1205            drop(buf);
1206        }
1207
1208        // Second acquire should be a pool hit (same size class)
1209        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    // ── Aligned-acquire tests ───────────────────────────────────────────────
1230
1231    #[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        // 32-byte alignment (AVX2 / scirs2_core::simd_aligned::SIMD_ALIGNMENT).
1237        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        // Pointer must still be aligned after pool reuse.
1267        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        // 32-byte aligned acquire/release for a size that maps to a particular size class.
1277        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        // Natural-alignment acquire of the same (T, size) — must NOT collide with the
1282        // 32-aligned bucket; it should produce a fresh allocation.
1283        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    // ── Memory-mapped allocation path ───────────────────────────────────────
1301    // These exercise the real disk-backed `scirs2_core::memory_efficient::MemoryMappedArray`
1302    // path and are gated on the `memory_efficient` feature. Run with:
1303    //   cargo test -p torsh-tensor --features memory_efficient
1304
1305    /// Round-trips KNOWN (non-default) data through the exact helper used by
1306    /// `create_memory_mapped_tensor`: write to a temp-dir backing file, read back via
1307    /// `as_slice()`, assert equality. Fails if the mmap wiring drops/garbles the data.
1308    #[cfg(feature = "memory_efficient")]
1309    #[test]
1310    fn test_map_through_mmap_file_roundtrips_known_data() {
1311        // Non-default values so a zero-init regression cannot accidentally pass.
1312        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        // Defensive cleanup in case the helper's best-effort removal failed.
1329        let _ = std::fs::remove_file(&backing_path);
1330    }
1331
1332    /// Directly drives `MemoryMappedArray::new(..)` + `as_slice()` with known `f64` data under
1333    /// the temp directory to pin the exact scirs2-core API contract the wiring relies on.
1334    #[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    /// Drives the production method `create_memory_mapped_tensor` end-to-end through the
1358    /// memory-mapped path and verifies the resulting tensor's shape and contents.
1359    #[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        // Data was staged through the memory map and read back via as_slice();
1373        // a freshly-allocated tensor holds default (zero) values.
1374        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}