Skip to main content

tenferro_cpu/
engine.rs

1//! The CPU execution engine: a caller-held value that owns one
2//! [`threadteam::Pool`] (or none) and the workspace every operation on it
3//! writes.
4//!
5//! tenferro never builds threads. The caller builds the pool with
6//! [`Pool::builder`](crate::Pool::builder) and hands it to [`CpuEngine::new`]; a
7//! [`CpuEngine::serial`] engine has no pool at all. Every operation takes the
8//! engine through `&mut`, so one computation uses it at a time without run-time
9//! locks. An operation starts on the calling thread with one placement token
10//! minted from the engine's pool; the kernels split that token and enter the pool
11//! only after a provider has chosen more than one task for it.
12
13use std::num::NonZeroUsize;
14use std::sync::Arc;
15
16use tenferro_tensor::{AllocationDomainId, SharedTensorAllocationDomain};
17
18use crate::buffer_pool::{BufferPool, BufferPoolStats};
19use crate::contraction::ExecutionResources;
20use crate::{CacheStats, CpuPlanCache};
21
22/// Typed refusals reported by a [`CpuEngine`] before an operation runs.
23///
24/// They reach callers as the source of a [`crate::Error::RuntimeState`].
25///
26/// # Examples
27///
28/// ```
29/// use tenferro_cpu::CpuEngineError;
30///
31/// assert!(CpuEngineError::Broken.to_string().contains("panic"));
32/// ```
33#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
34#[non_exhaustive]
35pub enum CpuEngineError {
36    /// An earlier operation panicked while it was using the engine.
37    ///
38    /// The engine's workspace may be half updated, so later operations are
39    /// refused; replace the engine with a new one.
40    #[error("the CPU engine is broken by an earlier panic; build a new engine")]
41    Broken,
42}
43
44/// A caller-held CPU execution engine.
45///
46/// The engine owns one [`threadteam::Pool`] (or none for [`CpuEngine::serial`]) and the
47/// workspace operations write: the buffer pool that recycles outputs, the
48/// contraction scratch, and a handle to a [`CpuPlanCache`]. It implements
49/// [`tenferro_tensor::Engine`], so every operation takes it as `&mut`.
50///
51/// # Examples
52///
53/// ```
54/// use tenferro_cpu::CpuEngine;
55/// use tenferro_tensor::{Tensor, TensorElementwise, TensorRead};
56///
57/// let mut engine = CpuEngine::serial();
58/// let a = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
59/// let b = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0])?;
60/// let c = engine.add_read(TensorRead::from_tensor(&a), TensorRead::from_tensor(&b))?;
61/// assert_eq!(c.as_slice::<f64>()?, &[4.0, 6.0]);
62/// # Ok::<(), tenferro_tensor::Error>(())
63/// ```
64pub struct CpuEngine {
65    pub(crate) resources: ExecutionResources,
66    pub(crate) buffers: BufferPool,
67    pub(crate) plans: CpuPlanCache,
68    pub(crate) batch_policy: crate::CpuBatchPolicy,
69    pub(crate) allocation_domain: Option<Arc<dyn SharedTensorAllocationDomain>>,
70    runtime_clears: u64,
71    pub(crate) broken: bool,
72    /// Runtime engine id this engine runs prepared programs as.
73    pub(crate) runtime_engine_id: tenferro_runtime::EngineId,
74    /// Extension caches the runtime lends to extension executors during a run.
75    pub(crate) extension_caches: tenferro_runtime::ExtensionCacheStore,
76}
77
78impl std::fmt::Debug for CpuEngine {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        f.debug_struct("CpuEngine")
81            .field("num_threads", &self.num_threads())
82            .field("has_team", &self.has_team())
83            .field("batch_policy", &self.batch_policy)
84            .field("allocation_domain", &self.allocation_domain())
85            .field("broken", &self.broken)
86            .finish_non_exhaustive()
87    }
88}
89
90impl CpuEngine {
91    fn with_resources(resources: ExecutionResources) -> Self {
92        Self {
93            resources,
94            buffers: BufferPool::new(),
95            plans: CpuPlanCache::new(),
96            batch_policy: crate::CpuBatchPolicy::default(),
97            allocation_domain: None,
98            runtime_clears: 0,
99            broken: false,
100            runtime_engine_id: crate::runtime_adapter::default_runtime_engine_id(),
101            extension_caches: tenferro_runtime::ExtensionCacheStore::default(),
102        }
103    }
104
105    /// An engine that computes on `pool`, which it takes over.
106    ///
107    /// Every operation may use all workers of the pool; a caller that wants
108    /// fewer threads builds a smaller pool. [`Pool::builder`](crate::Pool::builder)
109    /// reserves the worker stack the providers need by default.
110    ///
111    /// # Examples
112    ///
113    /// ```
114    /// let pool = tenferro_cpu::Pool::builder().threads(4).build()?;
115    /// let engine = tenferro_cpu::CpuEngine::new(pool);
116    /// assert_eq!(engine.num_threads(), 4);
117    /// # Ok::<(), Box<dyn std::error::Error>>(())
118    /// ```
119    pub fn new(pool: crate::Pool) -> Self {
120        Self::with_resources(ExecutionResources::new(Some(pool)))
121    }
122
123    /// An engine with no pool: every operation runs on the calling thread.
124    ///
125    /// A serial engine never starts or enters a pool, so it is safe to use from
126    /// inside any other pool.
127    ///
128    /// # Examples
129    ///
130    /// ```
131    /// let engine = tenferro_cpu::CpuEngine::serial();
132    /// assert_eq!(engine.num_threads(), 1);
133    /// ```
134    pub fn serial() -> Self {
135        Self::with_resources(ExecutionResources::new(None))
136    }
137
138    /// Use `plans` as this engine's plan cache, shared with every other engine
139    /// that holds a clone of it.
140    ///
141    /// # Examples
142    ///
143    /// ```
144    /// use tenferro_cpu::{CpuEngine, CpuPlanCache};
145    ///
146    /// let plans = CpuPlanCache::new();
147    /// let engine = CpuEngine::serial().with_plan_cache(plans.clone());
148    /// assert!(engine.plan_cache().shares_with(&plans));
149    /// ```
150    #[must_use]
151    pub fn with_plan_cache(mut self, plans: CpuPlanCache) -> Self {
152        self.plans = plans;
153        self
154    }
155
156    /// The plan cache this engine reads plans from.
157    ///
158    /// # Examples
159    ///
160    /// ```
161    /// let engine = tenferro_cpu::CpuEngine::serial();
162    /// assert_eq!(engine.plan_cache().stats().entries, 0);
163    /// ```
164    pub fn plan_cache(&self) -> &CpuPlanCache {
165        &self.plans
166    }
167
168    /// Run prepared programs under a runtime engine id other than the default
169    /// [`crate::runtime_engine_id`]; it must match the id this engine kind was
170    /// registered under.
171    ///
172    /// # Examples
173    ///
174    /// ```
175    /// use tenferro_runtime::EngineId;
176    /// let id = EngineId::new("example.cpu.v1")?;
177    /// let engine = tenferro_cpu::CpuEngine::serial().with_runtime_engine_id(id.clone());
178    /// assert_eq!(tenferro_runtime::ExecutionEngine::target(&engine).engine_id(), &id);
179    /// # Ok::<(), Box<dyn std::error::Error>>(())
180    /// ```
181    #[must_use]
182    pub fn with_runtime_engine_id(mut self, id: tenferro_runtime::EngineId) -> Self {
183        self.runtime_engine_id = id;
184        self
185    }
186
187    /// Bind outputs and managed inputs to a shared allocation domain.
188    ///
189    /// # Examples
190    ///
191    /// ```
192    /// let engine = tenferro_cpu::CpuEngine::serial();
193    /// assert!(engine.allocation_domain().is_none());
194    /// ```
195    #[must_use]
196    pub fn with_allocation_domain(mut self, domain: Arc<dyn SharedTensorAllocationDomain>) -> Self {
197        self.allocation_domain = Some(domain);
198        self
199    }
200
201    /// The shared allocation domain this engine is bound to, if any.
202    ///
203    /// # Examples
204    ///
205    /// ```
206    /// assert_eq!(tenferro_cpu::CpuEngine::serial().allocation_domain(), None);
207    /// ```
208    pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
209        self.allocation_domain.as_ref().map(|domain| domain.id())
210    }
211
212    /// The shared allocator this engine is bound to, if any.
213    ///
214    /// # Examples
215    ///
216    /// ```
217    /// assert!(tenferro_cpu::CpuEngine::serial().shared_allocation_domain().is_none());
218    /// ```
219    pub fn shared_allocation_domain(&self) -> Option<&Arc<dyn SharedTensorAllocationDomain>> {
220        self.allocation_domain.as_ref()
221    }
222
223    /// This engine with a different default batch policy.
224    ///
225    /// # Examples
226    ///
227    /// ```
228    /// use tenferro_cpu::{CpuBatchPolicy, CpuBatchStrategy, CpuEngine};
229    ///
230    /// let engine = CpuEngine::serial()
231    ///     .with_batch_policy(CpuBatchPolicy::new(CpuBatchStrategy::Sequential));
232    /// assert_eq!(engine.batch_policy().strategy, CpuBatchStrategy::Sequential);
233    /// ```
234    #[must_use]
235    pub fn with_batch_policy(mut self, policy: crate::CpuBatchPolicy) -> Self {
236        self.batch_policy = policy;
237        self
238    }
239
240    /// Replace the default batch policy, returning the previous one.
241    ///
242    /// For a one-operation override use [`crate::with_batch_policy`].
243    ///
244    /// # Examples
245    ///
246    /// ```
247    /// use tenferro_cpu::{CpuBatchPolicy, CpuBatchStrategy, CpuEngine};
248    ///
249    /// let mut engine = CpuEngine::serial();
250    /// let previous = engine.set_batch_policy(CpuBatchPolicy::new(CpuBatchStrategy::Outer));
251    /// assert_eq!(previous.strategy, CpuBatchStrategy::Auto);
252    /// ```
253    pub fn set_batch_policy(&mut self, policy: crate::CpuBatchPolicy) -> crate::CpuBatchPolicy {
254        std::mem::replace(&mut self.batch_policy, policy)
255    }
256
257    /// The default batch policy of this engine.
258    ///
259    /// # Examples
260    ///
261    /// ```
262    /// use tenferro_cpu::{CpuBatchStrategy, CpuEngine};
263    /// assert_eq!(CpuEngine::serial().batch_policy().strategy, CpuBatchStrategy::Auto);
264    /// ```
265    #[must_use]
266    pub fn batch_policy(&self) -> crate::CpuBatchPolicy {
267        self.batch_policy
268    }
269
270    /// The number of workers operations may use: the pool size, or one.
271    ///
272    /// # Examples
273    ///
274    /// ```
275    /// assert_eq!(tenferro_cpu::CpuEngine::serial().num_threads(), 1);
276    /// ```
277    pub fn num_threads(&self) -> usize {
278        self.thread_budget().get()
279    }
280
281    /// Whether the engine runs its parallel work on a reserved
282    /// `threadteam::Team`.
283    ///
284    /// `true` for an engine built from a pool of two or more workers whose
285    /// members were reserved. `false` for [`CpuEngine::serial`], for a
286    /// one-worker pool, and for a pool whose members could not be reserved
287    /// within about ten seconds of construction (the engine then runs serially
288    /// and [`CpuEngine::num_threads`] reports one). Benchmarks should assert
289    /// this after construction.
290    ///
291    /// # Examples
292    ///
293    /// ```
294    /// assert!(!tenferro_cpu::CpuEngine::serial().has_team());
295    /// let pool = tenferro_cpu::Pool::builder().threads(2).build()?;
296    /// assert!(tenferro_cpu::CpuEngine::new(pool).has_team());
297    /// # Ok::<(), Box<dyn std::error::Error>>(())
298    /// ```
299    pub fn has_team(&self) -> bool {
300        self.resources.has_team()
301    }
302
303    pub(crate) fn thread_budget(&self) -> NonZeroUsize {
304        self.resources.thread_budget()
305    }
306
307    pub(crate) fn dispatch_costs(&self) -> &threadteam::DispatchCosts {
308        self.resources.dispatch_costs()
309    }
310
311    /// Whether an earlier operation broke this engine (see
312    /// [`CpuEngineError::Broken`]).
313    ///
314    /// # Examples
315    ///
316    /// ```
317    /// assert!(!tenferro_cpu::CpuEngine::serial().is_broken());
318    /// ```
319    pub fn is_broken(&self) -> bool {
320        self.broken
321    }
322
323    /// The provider execution bundle for this engine: a placement token over
324    /// every worker of its pool, its workspace and its batch policy.
325    ///
326    /// A downstream kernel that calls a provider (cpueinsum, tcontract,
327    /// tlinalg) directly passes this bundle by `&mut`, so it uses the same
328    /// threads, budget and batch policy as tenferro. The bundle is not `Clone`:
329    /// it is the engine's whole thread budget for that one call, and the provider
330    /// splits it. Getting it applies the same checks as an operation.
331    ///
332    /// # Errors
333    ///
334    /// Returns [`crate::Error::RuntimeState`] with a [`CpuEngineError`] source
335    /// when the engine is broken.
336    ///
337    /// # Examples
338    ///
339    /// ```
340    /// let mut engine = tenferro_cpu::CpuEngine::serial();
341    /// let exec = engine.exec()?;
342    /// assert_eq!(exec.budget(), 1);
343    /// # Ok::<(), tenferro_tensor::Error>(())
344    /// ```
345    pub fn exec(&mut self) -> crate::Result<cpueinsum::Exec<'_>> {
346        self.check_entry("CpuEngine::exec")?;
347        Ok(self.resources.exec(self.batch_policy))
348    }
349
350    /// Run `op` on the engine's provider bundle, keeping the engine's guarantees.
351    ///
352    /// `op` receives the same bundle an operation mints (a placement token over
353    /// every worker of the pool, the workspace and the batch policy; see
354    /// [`CpuEngine::exec`]) and runs on the calling thread. Unlike
355    /// [`CpuEngine::exec`], the call is an engine operation: a panic inside `op`
356    /// marks the engine broken, and contraction scratch retained above
357    /// [`CpuEngine::buffer_pool_limit_bytes`] is released when `op` returns.
358    ///
359    /// The bundle is borrowed for the call only, and the engine stays mutably
360    /// borrowed while `op` runs, so neither can be stored or re-entered from
361    /// inside `op`; the compiler rejects both:
362    ///
363    /// ```compile_fail,E0501
364    /// let mut engine = tenferro_cpu::CpuEngine::serial();
365    /// let engine_ref = &mut engine;
366    /// engine_ref.with_exec(|_| engine_ref.with_exec(|exec| exec.budget()))?;
367    /// # Ok::<(), tenferro_tensor::Error>(())
368    /// ```
369    ///
370    /// ```compile_fail,E0521
371    /// let mut engine = tenferro_cpu::CpuEngine::serial();
372    /// let mut kept = None;
373    /// engine.with_exec(|exec| kept = Some(exec))?;
374    /// # Ok::<(), tenferro_tensor::Error>(())
375    /// ```
376    ///
377    /// # Errors
378    ///
379    /// Returns [`crate::Error::RuntimeState`] with a [`CpuEngineError::Broken`]
380    /// source, without calling `op`, when an earlier operation panicked.
381    ///
382    /// # Panics
383    ///
384    /// Resumes a panic raised by `op` after marking the engine broken.
385    ///
386    /// # Examples
387    ///
388    /// ```
389    /// let mut engine = tenferro_cpu::CpuEngine::serial();
390    /// let budget = engine.with_exec(|exec| exec.budget())?;
391    /// assert_eq!(budget, 1);
392    /// # Ok::<(), tenferro_tensor::Error>(())
393    /// ```
394    ///
395    /// A panic inside `op` breaks the engine, and a broken engine refuses later
396    /// calls with an error:
397    ///
398    /// ```
399    /// use tenferro_cpu::{CpuEngine, CpuEngineError};
400    ///
401    /// let mut engine = CpuEngine::serial();
402    /// let unwind = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
403    ///     engine.with_exec(|_| panic!("a failing block"))
404    /// }));
405    /// assert!(unwind.is_err());
406    /// assert!(engine.is_broken());
407    ///
408    /// let refused = engine.with_exec(|exec| exec.budget()).unwrap_err();
409    /// let source = std::error::Error::source(&refused).and_then(|s| s.downcast_ref());
410    /// assert_eq!(source, Some(&CpuEngineError::Broken));
411    /// ```
412    pub fn with_exec<R>(
413        &mut self,
414        op: impl FnOnce(&mut cpueinsum::Exec<'_>) -> R,
415    ) -> crate::Result<R> {
416        self.enter("CpuEngine::with_exec", true, |exec, _, _| Ok(op(exec)))
417    }
418
419    /// Borrow the provider bundle, the buffer pool and the N-ary contraction
420    /// scratch for one downstream call.
421    ///
422    /// # Errors
423    ///
424    /// Returns the entry refusals of [`CpuEngine::exec`], or `op`'s own error.
425    ///
426    /// # Examples
427    ///
428    /// ```
429    /// let mut engine = tenferro_cpu::CpuEngine::serial();
430    /// let len = engine.with_contraction_exec(|_, buffers, workspaces| {
431    ///     workspaces.with_scratch::<f64, _>(4, buffers.max_retained_capacity_bytes(), |scratch| {
432    ///         Ok(scratch.len())
433    ///     })
434    /// })?;
435    /// assert_eq!(len, 4);
436    /// # Ok::<(), tenferro_tensor::Error>(())
437    /// ```
438    #[doc(hidden)]
439    pub fn with_contraction_exec<R>(
440        &mut self,
441        op: impl FnOnce(
442            &mut cpueinsum::Exec<'_>,
443            &mut BufferPool,
444            &crate::ContractionWorkspaces,
445        ) -> crate::Result<R>,
446    ) -> crate::Result<R> {
447        self.enter("CpuEngine::with_contraction_exec", true, op)
448    }
449
450    /// Borrow the provider bundle and buffer pool for a linalg provider call.
451    ///
452    /// The provider chooses its own execution through the bundle; this
453    /// callback itself runs on the calling thread.
454    ///
455    /// # Errors
456    ///
457    /// Returns the entry refusals of [`CpuEngine::exec`], or `op`'s own error.
458    ///
459    /// # Examples
460    ///
461    /// ```
462    /// use tenferro_cpu::{CpuBatchStrategy, CpuEngine};
463    ///
464    /// let mut engine = CpuEngine::serial();
465    /// let strategy = engine.with_linalg_pool(|exec, _| Ok(exec.batch_policy().strategy))?;
466    /// assert_eq!(strategy, CpuBatchStrategy::Auto);
467    /// # Ok::<(), tenferro_tensor::Error>(())
468    /// ```
469    #[doc(hidden)]
470    pub fn with_linalg_pool<R>(
471        &mut self,
472        op: impl FnOnce(&mut cpueinsum::Exec<'_>, &mut BufferPool) -> crate::Result<R>,
473    ) -> crate::Result<R> {
474        self.enter("CpuEngine::with_linalg_pool", false, |exec, buffers, _| {
475            op(exec, buffers)
476        })
477    }
478
479    /// The one guarded entry behind the closure-scoped entries: refuse a broken
480    /// engine, lend the bundle, the buffer pool and the N-ary scratch, mark the
481    /// engine broken if `op` unwinds, and with `bound_retention` release
482    /// workspace retained above the buffer limit afterwards.
483    fn enter<R>(
484        &mut self,
485        name: &'static str,
486        bound_retention: bool,
487        op: impl FnOnce(
488            &mut cpueinsum::Exec<'_>,
489            &mut BufferPool,
490            &crate::ContractionWorkspaces,
491        ) -> crate::Result<R>,
492    ) -> crate::Result<R> {
493        self.check_entry(name)?;
494        let _retention = bound_retention.then(|| {
495            self.resources
496                .retention_guard(self.buffers.max_retained_capacity_bytes())
497        });
498        let mut exec = self.resources.exec(self.batch_policy);
499        let nary = &self.resources.nary;
500        guard_unwind(&mut self.buffers, &mut self.broken, |buffers| {
501            op(&mut exec, buffers, nary)
502        })
503    }
504
505    /// Refuse an operation that cannot run on this engine now.
506    pub(crate) fn check_entry(&self, op: &'static str) -> crate::Result<()> {
507        if self.broken {
508            return Err(crate::Error::runtime_state_source(
509                op,
510                CpuEngineError::Broken,
511            ));
512        }
513        Ok(())
514    }
515
516    /// Ledger roots of every placement token the engine minted since the last
517    /// call, so a test can read `threadteam::ledger` for the operation it just
518    /// ran (`max_live`, `installs`).
519    ///
520    /// Available with the `ledger` cargo feature, which is test support.
521    #[cfg(feature = "ledger")]
522    #[doc(hidden)]
523    pub fn take_ledger_roots(&self) -> Vec<u64> {
524        self.resources.take_ledger_roots()
525    }
526
527    /// Pool entries (installs from a plain thread or a foreign pool) made by
528    /// the tokens minted since the last `take_ledger_*` call, which it consumes.
529    ///
530    /// Zero means every operation since then stayed on the calling thread.
531    /// Available with the `ledger` cargo feature, which is test support.
532    #[cfg(feature = "ledger")]
533    #[doc(hidden)]
534    pub fn take_ledger_installs(&self) -> usize {
535        self.take_ledger_roots()
536            .into_iter()
537            .map(threadteam::ledger::installs)
538            .sum()
539    }
540
541    /// Largest number of tasks live at once under any token minted since the
542    /// last `take_ledger_*` call, which it consumes. It never exceeds
543    /// [`CpuEngine::num_threads`].
544    ///
545    /// Available with the `ledger` cargo feature, which is test support.
546    #[cfg(feature = "ledger")]
547    #[doc(hidden)]
548    pub fn take_ledger_max_live(&self) -> usize {
549        self.take_ledger_roots()
550            .into_iter()
551            .map(threadteam::ledger::max_live)
552            .max()
553            .unwrap_or(0)
554    }
555
556    /// Number of retained buffers in the engine's buffer pool.
557    ///
558    /// # Examples
559    ///
560    /// ```
561    /// assert_eq!(tenferro_cpu::CpuEngine::serial().buffer_pool_len(), 0);
562    /// ```
563    pub fn buffer_pool_len(&self) -> usize {
564        self.buffers.len()
565    }
566
567    /// Retained-buffer statistics of the engine's buffer pool.
568    ///
569    /// # Examples
570    ///
571    /// ```
572    /// assert_eq!(tenferro_cpu::CpuEngine::serial().buffer_pool_stats().buffers, 0);
573    /// ```
574    pub fn buffer_pool_stats(&self) -> BufferPoolStats {
575        self.buffers.stats()
576    }
577
578    /// The buffer pool as cache statistics (entries and retained bytes).
579    ///
580    /// # Examples
581    ///
582    /// ```
583    /// assert_eq!(tenferro_cpu::CpuEngine::serial().buffer_pool_cache_stats().entries, 0);
584    /// ```
585    pub fn buffer_pool_cache_stats(&self) -> CacheStats {
586        let stats = self.buffers.stats();
587        CacheStats {
588            entries: stats.buffers,
589            retained_bytes: stats.capacity_bytes,
590            ..CacheStats::default()
591        }
592    }
593
594    /// The retention limit of recycled buffers, in retained capacity bytes.
595    ///
596    /// # Examples
597    ///
598    /// ```
599    /// let mut engine = tenferro_cpu::CpuEngine::serial();
600    /// engine.set_buffer_pool_limit_bytes(0);
601    /// assert_eq!(engine.buffer_pool_limit_bytes(), 0);
602    /// ```
603    pub fn buffer_pool_limit_bytes(&self) -> usize {
604        self.buffers.max_retained_capacity_bytes()
605    }
606
607    /// Bound the bytes of recycled buffers and contraction scratch retained
608    /// between operations; zero disables retention.
609    ///
610    /// Destinations larger than the retention limit cannot be recycled, so
611    /// repeated calls allocate again and may pay first-touch costs. The default
612    /// is 100 MiB; raise it when a run keeps producing outputs above that size.
613    ///
614    /// # Examples
615    ///
616    /// ```
617    /// let mut engine = tenferro_cpu::CpuEngine::serial();
618    /// engine.set_buffer_pool_limit_bytes(1 << 20);
619    /// assert_eq!(engine.buffer_pool_limit_bytes(), 1 << 20);
620    /// ```
621    pub fn set_buffer_pool_limit_bytes(&mut self, max_retained_capacity_bytes: usize) {
622        self.buffers
623            .set_max_retained_capacity_bytes(max_retained_capacity_bytes);
624        self.resources.nary.trim_to(max_retained_capacity_bytes);
625        if self.resources.retained_bytes() > max_retained_capacity_bytes {
626            self.resources.trim();
627        }
628    }
629
630    /// Drop every retained buffer and contraction scratch.
631    ///
632    /// # Examples
633    ///
634    /// ```
635    /// let mut engine = tenferro_cpu::CpuEngine::serial();
636    /// engine.reset_buffer_pool();
637    /// assert_eq!(engine.buffer_pool_len(), 0);
638    /// ```
639    pub fn reset_buffer_pool(&mut self) {
640        self.buffers.clear();
641        self.resources.nary.clear_all();
642        self.resources.trim();
643    }
644
645    /// Aggregate statistics of everything this engine retains between
646    /// operations: buffers, contraction scratch and its plan cache.
647    ///
648    /// # Examples
649    ///
650    /// ```
651    /// assert_eq!(tenferro_cpu::CpuEngine::serial().runtime_cache_stats().entries, 0);
652    /// ```
653    #[doc(hidden)]
654    pub fn runtime_cache_stats(&self) -> CacheStats {
655        let (workspace_entries, nary_bytes) = self.resources.nary.total_stats();
656        let workspace_bytes = nary_bytes.saturating_add(self.resources.retained_bytes());
657        let buffers = self.buffers.cache_stats();
658        let plans = self.plans.stats();
659        CacheStats {
660            entries: buffers
661                .entries
662                .saturating_add(plans.entries)
663                .saturating_add(workspace_entries),
664            retained_bytes: buffers
665                .retained_bytes
666                .saturating_add(plans.retained_bytes)
667                .saturating_add(workspace_bytes),
668            hits: plans.hits,
669            misses: plans.misses,
670            evictions: plans.evictions,
671            clears: self.runtime_clears,
672        }
673    }
674
675    /// Drop everything this engine retains between operations, including its
676    /// plan cache (and therefore the plans of every engine sharing it).
677    ///
678    /// # Examples
679    ///
680    /// ```
681    /// let mut engine = tenferro_cpu::CpuEngine::serial();
682    /// engine.clear_runtime_caches();
683    /// assert_eq!(engine.runtime_cache_stats().clears, 1);
684    /// ```
685    #[doc(hidden)]
686    pub fn clear_runtime_caches(&mut self) {
687        self.reset_buffer_pool();
688        self.plans.clear();
689        self.runtime_clears = self.runtime_clears.saturating_add(1);
690    }
691}
692
693/// Run one operation's use of the buffer pool, recovering the pool's
694/// accounting and marking the engine broken if it unwinds.
695///
696/// Buffers checked out during the operation are counted as in flight. A normal
697/// exit forgets that accounting (the buffers left as outputs); an unwind
698/// replenishes the pool with replacements for the buffers it lost and marks the
699/// engine broken, since other workspace may be half updated.
700pub(crate) fn guard_unwind<R>(
701    buffers: &mut BufferPool,
702    broken: &mut bool,
703    op: impl FnOnce(&mut BufferPool) -> R,
704) -> R {
705    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| op(&mut *buffers))) {
706        Ok(value) => value,
707        Err(payload) => {
708            // The engine is never used again after a panic, so the buffers the
709            // operation had checked out need no accounting repair.
710            *broken = true;
711            std::panic::resume_unwind(payload)
712        }
713    }
714}
715
716impl CpuEngine {
717    /// Split borrows for the operation helpers: the provider bundle minted for
718    /// this operation (it borrows only the pool resources), the buffer pool, the
719    /// plan cache and the broken flag.
720    pub(crate) fn parts(
721        &mut self,
722    ) -> (
723        cpueinsum::Exec<'_>,
724        &mut BufferPool,
725        &CpuPlanCache,
726        &mut bool,
727    ) {
728        (
729            self.resources.exec(self.batch_policy),
730            &mut self.buffers,
731            &self.plans,
732            &mut self.broken,
733        )
734    }
735}
736
737// Uses the crate-level test helpers, which `provider-inject` builds leave out (see `lib.rs`).
738#[cfg(all(test, not(feature = "provider-inject")))]
739mod tests;