use std::num::NonZeroUsize;
use std::sync::Arc;
use tenferro_tensor::{AllocationDomainId, SharedTensorAllocationDomain};
use crate::buffer_pool::{BufferPool, BufferPoolStats};
use crate::contraction::ExecutionResources;
use crate::{CacheStats, CpuPlanCache};
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[non_exhaustive]
pub enum CpuEngineError {
#[error("the CPU engine is broken by an earlier panic; build a new engine")]
Broken,
}
pub struct CpuEngine {
pub(crate) resources: ExecutionResources,
pub(crate) buffers: BufferPool,
pub(crate) plans: CpuPlanCache,
pub(crate) batch_policy: crate::CpuBatchPolicy,
pub(crate) allocation_domain: Option<Arc<dyn SharedTensorAllocationDomain>>,
runtime_clears: u64,
pub(crate) broken: bool,
pub(crate) runtime_engine_id: tenferro_runtime::EngineId,
pub(crate) extension_caches: tenferro_runtime::ExtensionCacheStore,
}
impl std::fmt::Debug for CpuEngine {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CpuEngine")
.field("num_threads", &self.num_threads())
.field("has_team", &self.has_team())
.field("batch_policy", &self.batch_policy)
.field("allocation_domain", &self.allocation_domain())
.field("broken", &self.broken)
.finish_non_exhaustive()
}
}
impl CpuEngine {
fn with_resources(resources: ExecutionResources) -> Self {
Self {
resources,
buffers: BufferPool::new(),
plans: CpuPlanCache::new(),
batch_policy: crate::CpuBatchPolicy::default(),
allocation_domain: None,
runtime_clears: 0,
broken: false,
runtime_engine_id: crate::runtime_adapter::default_runtime_engine_id(),
extension_caches: tenferro_runtime::ExtensionCacheStore::default(),
}
}
pub fn new(pool: crate::Pool) -> Self {
Self::with_resources(ExecutionResources::new(Some(pool)))
}
pub fn serial() -> Self {
Self::with_resources(ExecutionResources::new(None))
}
#[must_use]
pub fn with_plan_cache(mut self, plans: CpuPlanCache) -> Self {
self.plans = plans;
self
}
pub fn plan_cache(&self) -> &CpuPlanCache {
&self.plans
}
#[must_use]
pub fn with_runtime_engine_id(mut self, id: tenferro_runtime::EngineId) -> Self {
self.runtime_engine_id = id;
self
}
#[must_use]
pub fn with_allocation_domain(mut self, domain: Arc<dyn SharedTensorAllocationDomain>) -> Self {
self.allocation_domain = Some(domain);
self
}
pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
self.allocation_domain.as_ref().map(|domain| domain.id())
}
pub fn shared_allocation_domain(&self) -> Option<&Arc<dyn SharedTensorAllocationDomain>> {
self.allocation_domain.as_ref()
}
#[must_use]
pub fn with_batch_policy(mut self, policy: crate::CpuBatchPolicy) -> Self {
self.batch_policy = policy;
self
}
pub fn set_batch_policy(&mut self, policy: crate::CpuBatchPolicy) -> crate::CpuBatchPolicy {
std::mem::replace(&mut self.batch_policy, policy)
}
#[must_use]
pub fn batch_policy(&self) -> crate::CpuBatchPolicy {
self.batch_policy
}
pub fn num_threads(&self) -> usize {
self.thread_budget().get()
}
pub fn has_team(&self) -> bool {
self.resources.has_team()
}
pub(crate) fn thread_budget(&self) -> NonZeroUsize {
self.resources.thread_budget()
}
pub(crate) fn dispatch_costs(&self) -> &threadteam::DispatchCosts {
self.resources.dispatch_costs()
}
pub fn is_broken(&self) -> bool {
self.broken
}
pub fn exec(&mut self) -> crate::Result<cpueinsum::Exec<'_>> {
self.check_entry("CpuEngine::exec")?;
Ok(self.resources.exec(self.batch_policy))
}
pub fn with_exec<R>(
&mut self,
op: impl FnOnce(&mut cpueinsum::Exec<'_>) -> R,
) -> crate::Result<R> {
self.enter("CpuEngine::with_exec", true, |exec, _, _| Ok(op(exec)))
}
#[doc(hidden)]
pub fn with_contraction_exec<R>(
&mut self,
op: impl FnOnce(
&mut cpueinsum::Exec<'_>,
&mut BufferPool,
&crate::ContractionWorkspaces,
) -> crate::Result<R>,
) -> crate::Result<R> {
self.enter("CpuEngine::with_contraction_exec", true, op)
}
#[doc(hidden)]
pub fn with_linalg_pool<R>(
&mut self,
op: impl FnOnce(&mut cpueinsum::Exec<'_>, &mut BufferPool) -> crate::Result<R>,
) -> crate::Result<R> {
self.enter("CpuEngine::with_linalg_pool", false, |exec, buffers, _| {
op(exec, buffers)
})
}
fn enter<R>(
&mut self,
name: &'static str,
bound_retention: bool,
op: impl FnOnce(
&mut cpueinsum::Exec<'_>,
&mut BufferPool,
&crate::ContractionWorkspaces,
) -> crate::Result<R>,
) -> crate::Result<R> {
self.check_entry(name)?;
let _retention = bound_retention.then(|| {
self.resources
.retention_guard(self.buffers.max_retained_capacity_bytes())
});
let mut exec = self.resources.exec(self.batch_policy);
let nary = &self.resources.nary;
guard_unwind(&mut self.buffers, &mut self.broken, |buffers| {
op(&mut exec, buffers, nary)
})
}
pub(crate) fn check_entry(&self, op: &'static str) -> crate::Result<()> {
if self.broken {
return Err(crate::Error::runtime_state_source(
op,
CpuEngineError::Broken,
));
}
Ok(())
}
#[cfg(feature = "ledger")]
#[doc(hidden)]
pub fn take_ledger_roots(&self) -> Vec<u64> {
self.resources.take_ledger_roots()
}
#[cfg(feature = "ledger")]
#[doc(hidden)]
pub fn take_ledger_installs(&self) -> usize {
self.take_ledger_roots()
.into_iter()
.map(threadteam::ledger::installs)
.sum()
}
#[cfg(feature = "ledger")]
#[doc(hidden)]
pub fn take_ledger_max_live(&self) -> usize {
self.take_ledger_roots()
.into_iter()
.map(threadteam::ledger::max_live)
.max()
.unwrap_or(0)
}
pub fn buffer_pool_len(&self) -> usize {
self.buffers.len()
}
pub fn buffer_pool_stats(&self) -> BufferPoolStats {
self.buffers.stats()
}
pub fn buffer_pool_cache_stats(&self) -> CacheStats {
let stats = self.buffers.stats();
CacheStats {
entries: stats.buffers,
retained_bytes: stats.capacity_bytes,
..CacheStats::default()
}
}
pub fn buffer_pool_limit_bytes(&self) -> usize {
self.buffers.max_retained_capacity_bytes()
}
pub fn set_buffer_pool_limit_bytes(&mut self, max_retained_capacity_bytes: usize) {
self.buffers
.set_max_retained_capacity_bytes(max_retained_capacity_bytes);
self.resources.nary.trim_to(max_retained_capacity_bytes);
if self.resources.retained_bytes() > max_retained_capacity_bytes {
self.resources.trim();
}
}
pub fn reset_buffer_pool(&mut self) {
self.buffers.clear();
self.resources.nary.clear_all();
self.resources.trim();
}
#[doc(hidden)]
pub fn runtime_cache_stats(&self) -> CacheStats {
let (workspace_entries, nary_bytes) = self.resources.nary.total_stats();
let workspace_bytes = nary_bytes.saturating_add(self.resources.retained_bytes());
let buffers = self.buffers.cache_stats();
let plans = self.plans.stats();
CacheStats {
entries: buffers
.entries
.saturating_add(plans.entries)
.saturating_add(workspace_entries),
retained_bytes: buffers
.retained_bytes
.saturating_add(plans.retained_bytes)
.saturating_add(workspace_bytes),
hits: plans.hits,
misses: plans.misses,
evictions: plans.evictions,
clears: self.runtime_clears,
}
}
#[doc(hidden)]
pub fn clear_runtime_caches(&mut self) {
self.reset_buffer_pool();
self.plans.clear();
self.runtime_clears = self.runtime_clears.saturating_add(1);
}
}
pub(crate) fn guard_unwind<R>(
buffers: &mut BufferPool,
broken: &mut bool,
op: impl FnOnce(&mut BufferPool) -> R,
) -> R {
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| op(&mut *buffers))) {
Ok(value) => value,
Err(payload) => {
*broken = true;
std::panic::resume_unwind(payload)
}
}
}
impl CpuEngine {
pub(crate) fn parts(
&mut self,
) -> (
cpueinsum::Exec<'_>,
&mut BufferPool,
&CpuPlanCache,
&mut bool,
) {
(
self.resources.exec(self.batch_policy),
&mut self.buffers,
&self.plans,
&mut self.broken,
)
}
}
#[cfg(all(test, not(feature = "provider-inject")))]
mod tests;