tenferro_cpu_basic/
buffer_pool.rs

1//! Typed host buffer pooling for reusable tensor allocations.
2//!
3//! # Examples
4//!
5//! ```rust
6//! use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
7//!
8//! let mut pool = BufferPool::new();
9//! let mut buf = pool.acquire_zeroed::<f64>(4);
10//! buf.fill(1.0);
11//! <f64 as PoolScalar>::pool_release(&mut pool, buf);
12//! assert_eq!(pool.len(), 1);
13//! ```
14
15use std::collections::BTreeMap;
16use std::env;
17use std::ffi::OsString;
18use std::fmt;
19use std::mem::{size_of, ManuallyDrop, MaybeUninit};
20use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
21use tenferro_tensor::HostBufferRecycler;
22
23use num_complex::{Complex32, Complex64};
24
25use crate::CacheStats;
26
27/// Non-Copy proof of a tracked pool checkout.
28#[derive(Debug)]
29pub(crate) enum UninitCheckoutToken {
30    /// The allocation was freshly allocated.
31    Fresh { actual_capacity: usize },
32    /// Retained storage was removed, with its actual capacity.
33    Reused { actual_capacity: usize },
34}
35
36/// Environment variable overriding the CPU buffer-pool retention cap in bytes.
37///
38/// The value is parsed as an unsigned integer. Invalid values fall back to
39/// [`DEFAULT_MAX_RETAINED_CAPACITY_BYTES`].
40pub const BUFFER_POOL_MAX_RETAINED_BYTES_ENV: &str = "TENFERRO_BUFFER_POOL_MAX_RETAINED_BYTES";
41
42/// Default retained CPU buffer capacity per backend.
43///
44/// The cap keeps long-running workloads from accumulating obsolete buffer
45/// sizes as tensor shapes grow while still preserving reuse for hot working
46/// sets.
47pub const DEFAULT_MAX_RETAINED_CAPACITY_BYTES: usize = 100 * 1024 * 1024;
48
49static DEFAULT_MAX_RETAINED_CAPACITY_FROM_ENV: OnceLock<usize> = OnceLock::new();
50
51/// Snapshot of typed host buffers retained by a [`BufferPool`].
52///
53/// `buffers` counts retained `Vec` allocations, while `capacity_bytes` counts
54/// their total element capacity in bytes. Allocators may keep freed memory in
55/// process-local arenas after a pool is cleared, so this reports memory that is
56/// still live in the pool rather than operating-system RSS.
57#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
58pub struct BufferPoolStats {
59    /// Number of retained vector allocations.
60    pub buffers: usize,
61    /// Total retained vector capacity in bytes.
62    pub capacity_bytes: usize,
63}
64
65/// Typed buffer pool keyed by element capacity and separated by scalar type.
66///
67/// Each supported dtype has an independent best-fit pool. Acquired buffers are
68/// returned without zero-initialization so kernels can avoid redundant writes
69/// when they fully overwrite the output. Use [`PoolScalar::pool_acquire_zeroed`]
70/// when the caller may read the buffer before writing every element.
71///
72/// # Examples
73///
74/// ```rust
75/// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
76///
77/// let mut pool = BufferPool::new();
78/// let buf = pool.acquire_zeroed::<f32>(8);
79/// <f32 as PoolScalar>::pool_release(&mut pool, buf);
80/// assert_eq!(pool.len(), 1);
81/// ```
82pub struct BufferPool {
83    state: Arc<SharedPool>,
84}
85
86#[derive(Debug)]
87struct SharedPool(Mutex<PoolState>);
88
89#[derive(Debug)]
90struct PoolState {
91    f64_pool: BTreeMap<usize, Vec<Vec<f64>>>,
92    f32_pool: BTreeMap<usize, Vec<Vec<f32>>>,
93    i32_pool: BTreeMap<usize, Vec<Vec<i32>>>,
94    i64_pool: BTreeMap<usize, Vec<Vec<i64>>>,
95    bool_pool: BTreeMap<usize, Vec<Vec<bool>>>,
96    c64_pool: BTreeMap<usize, Vec<Vec<Complex64>>>,
97    c32_pool: BTreeMap<usize, Vec<Vec<Complex32>>>,
98    f64_in_flight: BTreeMap<usize, usize>,
99    f32_in_flight: BTreeMap<usize, usize>,
100    i32_in_flight: BTreeMap<usize, usize>,
101    i64_in_flight: BTreeMap<usize, usize>,
102    bool_in_flight: BTreeMap<usize, usize>,
103    c64_in_flight: BTreeMap<usize, usize>,
104    c32_in_flight: BTreeMap<usize, usize>,
105    retained_capacity_bytes: usize,
106    max_retained_capacity_bytes: usize,
107}
108
109impl fmt::Debug for BufferPool {
110    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111        f.debug_struct("BufferPool")
112            .field("stats", &self.stats())
113            .field(
114                "max_retained_capacity_bytes",
115                &self.max_retained_capacity_bytes(),
116            )
117            .finish_non_exhaustive()
118    }
119}
120
121/// Scalar types supported by [`BufferPool`].
122///
123/// The trait is sealed to the scalar dtypes that tenferro currently pools for
124/// CPU execution.
125///
126/// # Examples
127///
128/// ```rust
129/// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
130///
131/// let mut pool = BufferPool::new();
132/// let mut buf = pool.acquire_zeroed::<f64>(2);
133/// buf.copy_from_slice(&[3.0, 4.0]);
134/// <f64 as PoolScalar>::pool_release(&mut pool, buf);
135/// ```
136pub trait PoolScalar:
137    Copy + Sized + Send + Sync + tenferro_tensor::TensorScalar + private::Sealed
138{
139    /// Zero value used to initialize acquired buffers.
140    fn pool_zero() -> Self;
141
142    /// Acquire a buffer with length `len` and every element set to zero.
143    ///
144    /// This is the safe path for callers that may read the buffer before every
145    /// element is overwritten. Full-overwrite kernels should use
146    /// [`crate::PooledUninitOutput`] or an operation-specific uninitialized
147    /// destination guard.
148    ///
149    /// # Examples
150    ///
151    /// ```rust
152    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
153    ///
154    /// let mut pool = BufferPool::new();
155    /// let buf = <f64 as PoolScalar>::pool_acquire_zeroed(&mut pool, 2);
156    /// assert_eq!(buf, vec![0.0, 0.0]);
157    /// ```
158    fn pool_acquire_zeroed(pool: &mut BufferPool, len: usize) -> Vec<Self>;
159
160    /// Return a buffer to the typed pool for later reuse.
161    ///
162    /// Zero-capacity buffers are ignored.
163    ///
164    /// # Examples
165    ///
166    /// ```rust
167    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
168    ///
169    /// let mut pool = BufferPool::new();
170    /// let buf = vec![1.0_f32; 4];
171    /// <f32 as PoolScalar>::pool_release(&mut pool, buf);
172    /// assert_eq!(pool.len(), 1);
173    /// ```
174    fn pool_release(pool: &mut BufferPool, buf: Vec<Self>);
175}
176
177pub(crate) mod private {
178    use std::mem::MaybeUninit;
179
180    // INVARIANT: this sealed crate-private trait is the only implementation
181    // boundary for tracked guard tokens; it is not part of the public API.
182    #[allow(private_interfaces)]
183    pub trait Sealed {
184        /// Checks out uninitialized storage with an exact cleanup token.
185        ///
186        /// # Errors
187        /// Returns `Error::Validation` if retained-capacity byte accounting
188        /// overflows, or `Error::BackendSource` if fresh allocation fails.
189        fn pool_acquire_uninit_tracked(
190            pool: &mut super::BufferPool,
191            len: usize,
192        ) -> crate::Result<(Vec<MaybeUninit<Self>>, super::UninitCheckoutToken)>
193        where
194            Self: Sized;
195        fn pool_finish_recycled(
196            pool: &mut super::BufferPool,
197            checkout: super::UninitCheckoutToken,
198        ) -> std::sync::Weak<dyn tenferro_tensor::HostBufferRecycler<Self>>
199        where
200            Self: Sized;
201
202        fn pool_discard_uninit(
203            pool: &mut super::BufferPool,
204            data: Vec<MaybeUninit<Self>>,
205            checkout: super::UninitCheckoutToken,
206        ) where
207            Self: Sized;
208    }
209}
210
211fn take_best_fit<T>(pool: &mut BTreeMap<usize, Vec<Vec<T>>>, len: usize) -> Option<Vec<T>> {
212    let key = *pool.range(len..).next()?.0;
213    let buf = {
214        let vecs = pool.get_mut(&key)?;
215        vecs.pop()
216    };
217    if pool.get(&key).is_some_and(Vec::is_empty) {
218        pool.remove(&key);
219    }
220    buf
221}
222
223fn pool_len<T>(pool: &BTreeMap<usize, Vec<Vec<T>>>) -> usize {
224    pool.values().map(Vec::len).sum()
225}
226
227fn evict_one_from_pool<T>(pool: &mut BTreeMap<usize, Vec<Vec<T>>>) -> Option<usize> {
228    let key = *pool.keys().next()?;
229    let vecs = pool.get_mut(&key)?;
230    let _ = vecs.pop()?;
231    if vecs.is_empty() {
232        pool.remove(&key);
233    }
234    Some(key.saturating_mul(size_of::<T>()))
235}
236
237#[derive(Clone, Copy)]
238enum TypedPoolKind {
239    F64,
240    F32,
241    I32,
242    I64,
243    Bool,
244    C64,
245    C32,
246}
247
248fn smallest_pool_candidate<T>(
249    pool: &BTreeMap<usize, Vec<Vec<T>>>,
250    kind: TypedPoolKind,
251) -> Option<(usize, TypedPoolKind)> {
252    pool.keys()
253        .next()
254        .map(|&capacity| (capacity.saturating_mul(size_of::<T>()), kind))
255}
256
257fn increment_in_flight(in_flight: &mut BTreeMap<usize, usize>, cap: usize) {
258    if cap > 0 {
259        *in_flight.entry(cap).or_default() += 1;
260    }
261}
262
263fn decrement_in_flight(in_flight: &mut BTreeMap<usize, usize>, cap: usize) {
264    if cap == 0 {
265        return;
266    }
267    let Some(count) = in_flight.get_mut(&cap) else {
268        return;
269    };
270    *count -= 1;
271    if *count == 0 {
272        in_flight.remove(&cap);
273    }
274}
275
276fn replenish_in_flight_for<T>(
277    pool: &mut BTreeMap<usize, Vec<Vec<T>>>,
278    in_flight: &mut BTreeMap<usize, usize>,
279    retained_capacity_bytes: &mut usize,
280) {
281    for (&cap, &count) in in_flight.iter() {
282        for _ in 0..count {
283            let mut replacement = Vec::new();
284            if replacement.try_reserve_exact(cap).is_err() {
285                continue;
286            }
287            let actual_cap = replacement.capacity();
288            *retained_capacity_bytes =
289                retained_capacity_bytes.saturating_add(actual_cap.saturating_mul(size_of::<T>()));
290            pool.entry(actual_cap).or_default().push(replacement);
291        }
292    }
293    in_flight.clear();
294}
295
296macro_rules! impl_pool_scalar {
297    ($ty:ty, $field:ident, $in_flight:ident, $zero:expr) => {
298        // INVARIANT: this implementation is reachable only through the sealed
299        // crate-private helper and cannot be named by sibling crates.
300        #[allow(private_interfaces)]
301        impl private::Sealed for $ty {
302            fn pool_acquire_uninit_tracked(
303                pool: &mut BufferPool,
304                len: usize,
305            ) -> crate::Result<(Vec<MaybeUninit<Self>>, UninitCheckoutToken)> {
306                let mut state = lock_pool(&pool.state);
307                let pool = &mut *state;
308                match take_best_fit(&mut pool.$field, len) {
309                    Some(buf) => {
310                        let cap = buf.capacity();
311                        let bytes = cap.checked_mul(size_of::<Self>()).ok_or_else(|| {
312                            crate::Error::invalid_argument(
313                                "pooled_uninit_output",
314                                "length",
315                                "pool capacity byte length overflow",
316                            )
317                        })?;
318                        pool.retained_capacity_bytes -= bytes;
319                        increment_in_flight(&mut pool.$in_flight, cap);
320                        let mut buf = ManuallyDrop::new(buf);
321                        // SAFETY: MaybeUninit<Self> has the same layout as Self and len <= cap.
322                        let buf = unsafe { Vec::from_raw_parts(buf.as_mut_ptr().cast(), len, cap) };
323                        Ok((
324                            buf,
325                            UninitCheckoutToken::Reused {
326                                actual_capacity: cap,
327                            },
328                        ))
329                    }
330                    None => {
331                        let mut buf = Vec::new();
332                        buf.try_reserve_exact(len).map_err(|err| {
333                            crate::Error::backend_source("pooled_uninit_output", err)
334                        })?;
335                        // SAFETY: every bit pattern is valid for MaybeUninit.
336                        unsafe { buf.set_len(len) };
337                        let actual_capacity = buf.capacity();
338                        Ok((buf, UninitCheckoutToken::Fresh { actual_capacity }))
339                    }
340                }
341            }
342
343            fn pool_finish_recycled(
344                pool: &mut BufferPool,
345                checkout: UninitCheckoutToken,
346            ) -> Weak<dyn HostBufferRecycler<Self>> {
347                if let UninitCheckoutToken::Reused { actual_capacity } = checkout {
348                    decrement_in_flight(&mut lock_pool(&pool.state).$in_flight, actual_capacity);
349                }
350                let owner: Arc<dyn HostBufferRecycler<Self>> = pool.state.clone();
351                Arc::downgrade(&owner)
352            }
353
354            fn pool_discard_uninit(
355                pool: &mut BufferPool,
356                data: Vec<MaybeUninit<Self>>,
357                checkout: UninitCheckoutToken,
358            ) {
359                drop(data);
360                let mut state = lock_pool(&pool.state);
361                let pool = &mut *state;
362                if let UninitCheckoutToken::Reused { actual_capacity } = checkout {
363                    decrement_in_flight(&mut pool.$in_flight, actual_capacity);
364                }
365            }
366        }
367
368        impl PoolScalar for $ty {
369            fn pool_zero() -> Self {
370                $zero
371            }
372
373            fn pool_acquire_zeroed(pool: &mut BufferPool, len: usize) -> Vec<Self> {
374                let mut state = lock_pool(&pool.state);
375                let pool = &mut *state;
376                match take_best_fit(&mut pool.$field, len) {
377                    Some(mut buf) => {
378                        pool.retained_capacity_bytes = pool
379                            .retained_capacity_bytes
380                            .saturating_sub(buf.capacity().saturating_mul(size_of::<Self>()));
381                        increment_in_flight(&mut pool.$in_flight, buf.capacity());
382                        buf.resize(len, Self::pool_zero());
383                        buf.fill(Self::pool_zero());
384                        buf
385                    }
386                    None => vec![Self::pool_zero(); len],
387                }
388            }
389
390            fn pool_release(pool: &mut BufferPool, buf: Vec<Self>) {
391                decrement_in_flight(&mut lock_pool(&pool.state).$in_flight, buf.capacity());
392                pool.state.recycle(buf);
393            }
394        }
395
396        impl HostBufferRecycler<$ty> for SharedPool {
397            fn recycle(&self, buf: Vec<$ty>) {
398                let mut state = lock_pool(self);
399                let pool = &mut *state;
400                let cap = buf.capacity();
401                if cap > 0 {
402                    pool.retained_capacity_bytes = pool
403                        .retained_capacity_bytes
404                        .saturating_add(cap.saturating_mul(size_of::<$ty>()));
405                    pool.$field.entry(cap).or_default().push(buf);
406                    pool.enforce_retention_limit();
407                }
408            }
409        }
410    };
411}
412
413impl_pool_scalar!(f64, f64_pool, f64_in_flight, 0.0);
414impl_pool_scalar!(f32, f32_pool, f32_in_flight, 0.0);
415impl_pool_scalar!(i32, i32_pool, i32_in_flight, 0);
416impl_pool_scalar!(i64, i64_pool, i64_in_flight, 0);
417impl_pool_scalar!(bool, bool_pool, bool_in_flight, false);
418impl_pool_scalar!(Complex64, c64_pool, c64_in_flight, Complex64::new(0.0, 0.0));
419impl_pool_scalar!(Complex32, c32_pool, c32_in_flight, Complex32::new(0.0, 0.0));
420
421impl BufferPool {
422    #[cfg(test)]
423    pub(crate) fn in_flight_is_empty(&self) -> bool {
424        let state = lock_pool(&self.state);
425        state.f64_in_flight.is_empty()
426            && state.f32_in_flight.is_empty()
427            && state.i32_in_flight.is_empty()
428            && state.i64_in_flight.is_empty()
429            && state.bool_in_flight.is_empty()
430            && state.c64_in_flight.is_empty()
431            && state.c32_in_flight.is_empty()
432    }
433    /// Create an empty typed buffer pool.
434    ///
435    /// # Examples
436    ///
437    /// ```rust
438    /// use tenferro_cpu_basic::buffer_pool::BufferPool;
439    ///
440    /// let pool = BufferPool::new();
441    /// assert!(pool.is_empty());
442    /// ```
443    pub fn new() -> Self {
444        Self::with_max_retained_capacity_bytes(default_max_retained_capacity_bytes())
445    }
446
447    /// Create an empty typed buffer pool with a specific retention cap.
448    ///
449    /// A cap of zero disables retention. Use [`BufferPool::unbounded`] only for
450    /// diagnostics or workloads that are externally memory-limited.
451    ///
452    /// # Examples
453    ///
454    /// ```rust
455    /// use tenferro_cpu_basic::buffer_pool::BufferPool;
456    ///
457    /// let pool = BufferPool::with_max_retained_capacity_bytes(1024);
458    /// assert_eq!(pool.max_retained_capacity_bytes(), 1024);
459    /// ```
460    pub fn with_max_retained_capacity_bytes(max_retained_capacity_bytes: usize) -> Self {
461        Self {
462            state: Arc::new(SharedPool(Mutex::new(PoolState {
463                f64_pool: BTreeMap::new(),
464                f32_pool: BTreeMap::new(),
465                i32_pool: BTreeMap::new(),
466                i64_pool: BTreeMap::new(),
467                bool_pool: BTreeMap::new(),
468                c64_pool: BTreeMap::new(),
469                c32_pool: BTreeMap::new(),
470                f64_in_flight: BTreeMap::new(),
471                f32_in_flight: BTreeMap::new(),
472                i32_in_flight: BTreeMap::new(),
473                i64_in_flight: BTreeMap::new(),
474                bool_in_flight: BTreeMap::new(),
475                c64_in_flight: BTreeMap::new(),
476                c32_in_flight: BTreeMap::new(),
477                retained_capacity_bytes: 0,
478                max_retained_capacity_bytes,
479            }))),
480        }
481    }
482
483    /// Create an empty typed buffer pool without a retention cap.
484    ///
485    /// This preserves the historical behavior and is mainly useful for
486    /// diagnostics or controlled benchmarks.
487    ///
488    /// # Examples
489    ///
490    /// ```rust
491    /// use tenferro_cpu_basic::buffer_pool::BufferPool;
492    ///
493    /// let pool = BufferPool::unbounded();
494    /// assert_eq!(pool.max_retained_capacity_bytes(), usize::MAX);
495    /// ```
496    pub fn unbounded() -> Self {
497        Self::with_max_retained_capacity_bytes(usize::MAX)
498    }
499
500    /// Maximum retained typed host-buffer capacity in bytes.
501    ///
502    /// # Examples
503    ///
504    /// ```rust
505    /// use tenferro_cpu_basic::buffer_pool::BufferPool;
506    ///
507    /// let pool = BufferPool::with_max_retained_capacity_bytes(4096);
508    /// assert_eq!(pool.max_retained_capacity_bytes(), 4096);
509    /// ```
510    pub fn max_retained_capacity_bytes(&self) -> usize {
511        lock_pool(&self.state).max_retained_capacity_bytes
512    }
513
514    /// Update the maximum retained typed host-buffer capacity in bytes.
515    ///
516    /// Shrinking below the currently retained capacity immediately evicts
517    /// retained buffers until the new cap is satisfied. A cap of zero disables
518    /// retention.
519    ///
520    /// # Examples
521    ///
522    /// ```
523    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
524    ///
525    /// let mut pool = BufferPool::with_max_retained_capacity_bytes(1024);
526    /// <f64 as PoolScalar>::pool_release(&mut pool, Vec::with_capacity(128));
527    /// pool.set_max_retained_capacity_bytes(0);
528    /// assert_eq!(pool.max_retained_capacity_bytes(), 0);
529    /// assert!(pool.is_empty());
530    /// ```
531    pub fn set_max_retained_capacity_bytes(&mut self, max_retained_capacity_bytes: usize) {
532        let mut state = lock_pool(&self.state);
533        state.max_retained_capacity_bytes = max_retained_capacity_bytes;
534        state.enforce_retention_limit();
535    }
536
537    /// Number of retained buffers across all typed pools.
538    ///
539    /// # Examples
540    ///
541    /// ```rust
542    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
543    ///
544    /// let mut pool = BufferPool::new();
545    /// <f64 as PoolScalar>::pool_release(&mut pool, vec![0.0; 2]);
546    /// assert_eq!(pool.len(), 1);
547    /// ```
548    pub fn len(&self) -> usize {
549        self.stats().buffers
550    }
551
552    /// Total retained typed host-buffer capacity in bytes.
553    ///
554    /// This counts capacity that is still live in the pool. The operating
555    /// system RSS may remain high after clearing the pool because the process
556    /// allocator can keep freed pages for future allocations.
557    ///
558    /// # Examples
559    ///
560    /// ```rust
561    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
562    ///
563    /// let mut pool = BufferPool::new();
564    /// <f64 as PoolScalar>::pool_release(&mut pool, Vec::with_capacity(2));
565    /// assert_eq!(pool.retained_capacity_bytes(), 16);
566    /// ```
567    pub fn retained_capacity_bytes(&self) -> usize {
568        self.stats().capacity_bytes
569    }
570
571    /// Snapshot retained-buffer count and capacity.
572    ///
573    /// # Examples
574    ///
575    /// ```rust
576    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
577    ///
578    /// let mut pool = BufferPool::new();
579    /// <f32 as PoolScalar>::pool_release(&mut pool, Vec::with_capacity(4));
580    /// let stats = pool.stats();
581    /// assert_eq!(stats.buffers, 1);
582    /// assert_eq!(stats.capacity_bytes, 16);
583    /// ```
584    pub fn stats(&self) -> BufferPoolStats {
585        let state = lock_pool(&self.state);
586        BufferPoolStats {
587            buffers: pool_len(&state.f64_pool)
588                + pool_len(&state.f32_pool)
589                + pool_len(&state.i32_pool)
590                + pool_len(&state.i64_pool)
591                + pool_len(&state.bool_pool)
592                + pool_len(&state.c64_pool)
593                + pool_len(&state.c32_pool),
594            capacity_bytes: state.retained_capacity_bytes,
595        }
596    }
597
598    /// Return cache-style stats for the buffers retained by this pool.
599    ///
600    /// `entries` is the number of retained buffers, and `retained_bytes` is the
601    /// total retained vector capacity in bytes.
602    ///
603    /// # Examples
604    ///
605    /// ```
606    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
607    ///
608    /// let mut pool = BufferPool::new();
609    /// <f32 as PoolScalar>::pool_release(&mut pool, Vec::with_capacity(4));
610    /// let stats = pool.cache_stats();
611    /// assert_eq!(stats.entries, 1);
612    /// assert_eq!(stats.retained_bytes, 16);
613    /// ```
614    pub fn cache_stats(&self) -> CacheStats {
615        let stats = self.stats();
616        CacheStats {
617            entries: stats.buffers,
618            retained_bytes: stats.capacity_bytes,
619            hits: 0,
620            misses: 0,
621            evictions: 0,
622            clears: 0,
623        }
624    }
625
626    /// Acquire a typed vector with length 0 and at least `cap` capacity.
627    ///
628    /// Returned buffers come from the typed pool when possible and are ready
629    /// for push-based population.
630    ///
631    /// # Examples
632    ///
633    /// ```rust
634    /// use tenferro_cpu_basic::buffer_pool::BufferPool;
635    ///
636    /// let mut pool = BufferPool::new();
637    /// let mut buf = pool.acquire_with_capacity::<f64>(4);
638    /// buf.extend_from_slice(&[1.0, 2.0]);
639    /// assert_eq!(buf.len(), 2);
640    /// assert!(buf.capacity() >= 4);
641    /// ```
642    pub fn acquire_with_capacity<T: PoolScalar>(&mut self, cap: usize) -> Vec<T> {
643        if cap == 0 {
644            return Vec::new();
645        }
646
647        let (data, _checkout) = <T as private::Sealed>::pool_acquire_uninit_tracked(self, cap)
648            .expect("validated typed pool capacity must be acquirable");
649        let mut data = ManuallyDrop::new(data);
650        let ptr = data.as_mut_ptr().cast::<T>();
651        let capacity = data.capacity();
652        // SAFETY: `MaybeUninit<T>` and `T` have identical layouts, and the
653        // returned vector has length zero, so no element is read before push.
654        unsafe { Vec::from_raw_parts(ptr, 0, capacity) }
655    }
656
657    /// Acquire a typed vector with length `len` initialized to zero.
658    ///
659    /// Use this only when the caller may read elements before overwriting the
660    /// entire buffer. Full-overwrite kernels should use
661    /// [`crate::PooledUninitOutput`] to avoid the initialization cost.
662    ///
663    /// # Examples
664    ///
665    /// ```rust
666    /// use tenferro_cpu_basic::buffer_pool::BufferPool;
667    ///
668    /// let mut pool = BufferPool::new();
669    /// let buf = pool.acquire_zeroed::<f32>(3);
670    /// assert_eq!(buf, vec![0.0, 0.0, 0.0]);
671    /// ```
672    pub fn acquire_zeroed<T: PoolScalar>(&mut self, len: usize) -> Vec<T> {
673        T::pool_acquire_zeroed(self, len)
674    }
675
676    /// Whether all typed pools are empty.
677    ///
678    /// # Examples
679    ///
680    /// ```rust
681    /// use tenferro_cpu_basic::buffer_pool::BufferPool;
682    ///
683    /// let pool = BufferPool::new();
684    /// assert!(pool.is_empty());
685    /// ```
686    pub fn is_empty(&self) -> bool {
687        self.len() == 0
688    }
689
690    /// Drop all retained buffers from the pool.
691    ///
692    /// This releases the vectors owned by the pool. The process allocator may
693    /// still keep freed pages mapped for reuse, so operating-system RSS is not
694    /// guaranteed to fall immediately.
695    ///
696    /// # Examples
697    ///
698    /// ```rust
699    /// use tenferro_cpu_basic::buffer_pool::{BufferPool, PoolScalar};
700    ///
701    /// let mut pool = BufferPool::new();
702    /// <f64 as PoolScalar>::pool_release(&mut pool, Vec::with_capacity(8));
703    /// pool.clear();
704    /// assert!(pool.is_empty());
705    /// ```
706    pub fn clear(&mut self) {
707        let mut state = lock_pool(&self.state);
708        state.clear();
709    }
710
711    #[doc(hidden)]
712    pub fn clear_in_flight_retained(&mut self) {
713        lock_pool(&self.state).clear_in_flight_retained();
714    }
715
716    #[doc(hidden)]
717    pub fn replenish_in_flight_retained(&mut self) {
718        lock_pool(&self.state).replenish_in_flight_retained();
719    }
720}
721
722// INVARIANT: no user callback or tensor destructor runs under this lock; it
723// protects only scalar Vec bins and their accounting, never execution admission.
724fn lock_pool(state: &SharedPool) -> MutexGuard<'_, PoolState> {
725    state
726        .0
727        .lock()
728        .unwrap_or_else(std::sync::PoisonError::into_inner)
729}
730
731impl PoolState {
732    fn clear(&mut self) {
733        self.f64_pool.clear();
734        self.f32_pool.clear();
735        self.i32_pool.clear();
736        self.i64_pool.clear();
737        self.bool_pool.clear();
738        self.c64_pool.clear();
739        self.c32_pool.clear();
740        self.clear_in_flight_retained();
741        self.retained_capacity_bytes = 0;
742    }
743
744    #[doc(hidden)]
745    pub fn clear_in_flight_retained(&mut self) {
746        self.f64_in_flight.clear();
747        self.f32_in_flight.clear();
748        self.i32_in_flight.clear();
749        self.i64_in_flight.clear();
750        self.bool_in_flight.clear();
751        self.c64_in_flight.clear();
752        self.c32_in_flight.clear();
753    }
754
755    #[doc(hidden)]
756    pub fn replenish_in_flight_retained(&mut self) {
757        replenish_in_flight_for(
758            &mut self.f64_pool,
759            &mut self.f64_in_flight,
760            &mut self.retained_capacity_bytes,
761        );
762        replenish_in_flight_for(
763            &mut self.f32_pool,
764            &mut self.f32_in_flight,
765            &mut self.retained_capacity_bytes,
766        );
767        replenish_in_flight_for(
768            &mut self.i32_pool,
769            &mut self.i32_in_flight,
770            &mut self.retained_capacity_bytes,
771        );
772        replenish_in_flight_for(
773            &mut self.i64_pool,
774            &mut self.i64_in_flight,
775            &mut self.retained_capacity_bytes,
776        );
777        replenish_in_flight_for(
778            &mut self.bool_pool,
779            &mut self.bool_in_flight,
780            &mut self.retained_capacity_bytes,
781        );
782        replenish_in_flight_for(
783            &mut self.c64_pool,
784            &mut self.c64_in_flight,
785            &mut self.retained_capacity_bytes,
786        );
787        replenish_in_flight_for(
788            &mut self.c32_pool,
789            &mut self.c32_in_flight,
790            &mut self.retained_capacity_bytes,
791        );
792        self.enforce_retention_limit();
793    }
794
795    fn enforce_retention_limit(&mut self) {
796        while self.retained_capacity_bytes > self.max_retained_capacity_bytes {
797            let Some(evicted_bytes) = self.evict_smallest_retained_buffer() else {
798                self.retained_capacity_bytes = 0;
799                return;
800            };
801            if evicted_bytes == 0 {
802                if self.retained_capacity_bytes == 0 {
803                    self.retained_capacity_bytes = 0;
804                    return;
805                }
806                continue;
807            }
808            self.retained_capacity_bytes =
809                self.retained_capacity_bytes.saturating_sub(evicted_bytes);
810        }
811    }
812
813    fn evict_smallest_retained_buffer(&mut self) -> Option<usize> {
814        let candidates = [
815            smallest_pool_candidate(&self.f64_pool, TypedPoolKind::F64),
816            smallest_pool_candidate(&self.f32_pool, TypedPoolKind::F32),
817            smallest_pool_candidate(&self.i32_pool, TypedPoolKind::I32),
818            smallest_pool_candidate(&self.i64_pool, TypedPoolKind::I64),
819            smallest_pool_candidate(&self.bool_pool, TypedPoolKind::Bool),
820            smallest_pool_candidate(&self.c64_pool, TypedPoolKind::C64),
821            smallest_pool_candidate(&self.c32_pool, TypedPoolKind::C32),
822        ];
823        let (_, kind) = candidates
824            .into_iter()
825            .flatten()
826            .min_by_key(|(bytes, _)| *bytes)?;
827        match kind {
828            TypedPoolKind::F64 => evict_one_from_pool(&mut self.f64_pool),
829            TypedPoolKind::F32 => evict_one_from_pool(&mut self.f32_pool),
830            TypedPoolKind::I32 => evict_one_from_pool(&mut self.i32_pool),
831            TypedPoolKind::I64 => evict_one_from_pool(&mut self.i64_pool),
832            TypedPoolKind::Bool => evict_one_from_pool(&mut self.bool_pool),
833            TypedPoolKind::C64 => evict_one_from_pool(&mut self.c64_pool),
834            TypedPoolKind::C32 => evict_one_from_pool(&mut self.c32_pool),
835        }
836    }
837}
838
839fn default_max_retained_capacity_bytes() -> usize {
840    *DEFAULT_MAX_RETAINED_CAPACITY_FROM_ENV.get_or_init(|| {
841        parse_default_max_retained_capacity_bytes(env::var_os(BUFFER_POOL_MAX_RETAINED_BYTES_ENV))
842    })
843}
844
845fn parse_default_max_retained_capacity_bytes(value: Option<OsString>) -> usize {
846    value
847        .and_then(|value| value.into_string().ok())
848        .and_then(|value| value.parse().ok())
849        .unwrap_or(DEFAULT_MAX_RETAINED_CAPACITY_BYTES)
850}
851
852impl Default for BufferPool {
853    fn default() -> Self {
854        Self::new()
855    }
856}
857
858#[cfg(test)]
859mod tests;
tenferro_cpu_basic/buffer_pool.rs

tenferro_cpu_basic/
buffer_pool.rs
