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;