use std::any::TypeId;
use std::cmp::Reverse;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::env;
use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::thread;
use std::time::{Duration, Instant};
use crate::arbiter::{
inherited_or_new_execution_owner, with_execution_owner, ResourceArbiter, ResourceOwner,
ResourcePermit,
};
use crate::buffer_pool::{BufferPool, BufferPoolStats, PoolScalar};
use crate::dot_runtime::{
CpuProviderBundle, CpuProviderBundleInstallError, CpuProviderDomainContract,
};
use crate::engine::{CpuEngine, EngineResources};
use crate::indexed_plan_cache::{
IndexedPlanCache, IndexedPlanCacheLimits, DEFAULT_INDEXED_PLAN_CACHE_LIMITS,
};
use crate::placement::{
resolve_placement, resolve_placement_with_affinity, CpuEngineConstructionError,
ResolvedCpuExecution,
};
use crate::provider::{CpuExecutionContext, CpuOperationEntry, ParallelMode};
use crate::{
discover_cpu_topology, CpuAdmissionMode, CpuDomainId, CpuDomainOwnership, CpuExecutorAffinity,
CpuExecutorShutdown, CpuId, CpuPlacement, CpuPlacementError, CpuPlacementGuarantee, CpuSet,
CpuTopology, CpuTopologyError, ExternalCpuDomain, NumaNodeId, ResolvedCpuPlacement,
};
use crate::{
CacheStats, Tensor, TensorRank, TensorRead, TensorScalar, TensorValue, TensorWrite,
TypedTensor, TypedTensorView, TypedTensorViewMut,
};
use tenferro_tensor::backend::{ElementwiseFusionPlan, GroupedGemmConfig};
use tenferro_tensor::SharedTensorAllocationDomain;
use tenferro_tensor::{
AllocationDomainId, BackendCachedDot, BackendRuntimeCache, BackendSession, BackendSessionHost,
ContractionScalar, DotGeneralAccumulation, ElementwiseReadOp, TensorAnalytic, TensorBackend,
TensorBuffer, TensorDeviceTransfer, TensorDot, TensorElementwise, TensorFusion, TensorIndexing,
TensorReduction, TensorStructural, TensorViewCanonicalization,
};
use tenferro_tensor::{
CompareDir, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
};
use super::exec_session::CpuExecSession;
use super::{
analytic, copy_tensor_read_into, elementwise, gemm, indexing, materialize_tensor_read,
reduction, structural, CpuContext,
};
pub(crate) fn tag_fresh_output(output: &mut Tensor, domain: CpuDomainId) {
macro_rules! tag {
($tensor:expr) => {{
$tensor.set_cpu_affinity(Some(domain));
}};
}
match output {
Tensor::F32(tensor) => tag!(tensor),
Tensor::F64(tensor) => tag!(tensor),
Tensor::I32(tensor) => tag!(tensor),
Tensor::I64(tensor) => tag!(tensor),
Tensor::Bool(tensor) => tag!(tensor),
Tensor::C32(tensor) => tag!(tensor),
Tensor::C64(tensor) => tag!(tensor),
}
}
pub(crate) fn elementwise_read_into_fallback_with_pool(
buffers: &mut BufferPool,
op: ElementwiseReadOp,
inputs: &[TensorRead<'_>],
out: TensorWrite<'_>,
) -> crate::Result<()> {
let result = match op {
ElementwiseReadOp::Add => {
elementwise::add_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
}
ElementwiseReadOp::Subtract => {
elementwise::sub_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
}
ElementwiseReadOp::Multiply => {
elementwise::mul_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
}
ElementwiseReadOp::Negate => elementwise::neg_read_with_pool(buffers, inputs[0].clone())?,
ElementwiseReadOp::Conj => elementwise::conj_read_with_pool(buffers, inputs[0].clone())?,
ElementwiseReadOp::Divide => {
elementwise::div_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
}
_ => {
return Err(crate::Error::unsupported(
"CpuBackend::elementwise_read_into",
format!("CPU backend does not implement {op:?}"),
))
}
};
copy_tensor_read_into(
"CpuBackend::elementwise_read_into",
TensorRead::from_tensor(&result),
out,
)
}
pub(crate) trait FreshCpuOutput {
fn tag_fresh(&mut self, domain: CpuDomainId);
}
impl FreshCpuOutput for Tensor {
fn tag_fresh(&mut self, domain: CpuDomainId) {
tag_fresh_output(self, domain);
}
}
impl<T, R: TensorRank> FreshCpuOutput for TypedTensor<T, R> {
fn tag_fresh(&mut self, domain: CpuDomainId) {
self.set_cpu_affinity(Some(domain));
}
}
impl<T: FreshCpuOutput> FreshCpuOutput for Option<T> {
fn tag_fresh(&mut self, domain: CpuDomainId) {
if let Some(output) = self {
output.tag_fresh(domain);
}
}
}
impl<T: FreshCpuOutput> FreshCpuOutput for Vec<T> {
fn tag_fresh(&mut self, domain: CpuDomainId) {
for output in self {
output.tag_fresh(domain);
}
}
}
#[derive(Debug, Default, Clone)]
struct CpuSessionProfileEntry {
calls: usize,
total_time: Duration,
}
fn cpu_session_profile_enabled() -> bool {
static ENABLED: OnceLock<bool> = OnceLock::new();
*ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_CPU_SESSION").is_ok())
}
fn cpu_session_profile_print_every() -> Option<usize> {
static PRINT_EVERY: OnceLock<Option<usize>> = OnceLock::new();
*PRINT_EVERY.get_or_init(|| {
env::var("TENFERRO_PROFILE_CPU_SESSION_PRINT_EVERY")
.ok()
.and_then(|value| value.parse::<usize>().ok())
.filter(|&value| value > 0)
})
}
fn cpu_session_profile_state() -> &'static Mutex<HashMap<&'static str, CpuSessionProfileEntry>> {
static STATE: OnceLock<Mutex<HashMap<&'static str, CpuSessionProfileEntry>>> = OnceLock::new();
STATE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn record_cpu_session_profile(section: &'static str, elapsed: Duration) {
if !cpu_session_profile_enabled() {
return;
}
let Ok(mut state) = cpu_session_profile_state().lock() else {
return;
};
let entry = state.entry(section).or_default();
entry.calls += 1;
entry.total_time += elapsed;
}
fn profile_cpu_session_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
if !cpu_session_profile_enabled() {
return f();
}
let started = Instant::now();
let result = f();
record_cpu_session_profile(section, started.elapsed());
result
}
fn maybe_print_cpu_session_profile() {
let Some(print_every) = cpu_session_profile_print_every() else {
return;
};
let should_print = {
let Ok(state) = cpu_session_profile_state().lock() else {
return;
};
state
.get("with_backend_session_cached.total")
.is_some_and(|entry| entry.calls % print_every == 0)
};
if !should_print {
return;
}
let mut entries = {
let Ok(mut state) = cpu_session_profile_state().lock() else {
return;
};
let entries = state
.iter()
.map(|(section, entry)| (*section, entry.clone()))
.collect::<Vec<_>>();
state.clear();
entries
};
entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
eprintln!("=== tenferro CPU session profile ===");
for (section, entry) in entries {
eprintln!(
"{section}: calls={} total={:.6}ms per_call={:.3}us",
entry.calls,
entry.total_time.as_secs_f64() * 1.0e3,
entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64,
);
}
}
struct BufferPoolLoan<'a> {
buffers: &'a mut BufferPool,
}
impl<'a> BufferPoolLoan<'a> {
fn new(buffers: &'a mut BufferPool) -> Self {
Self { buffers }
}
fn get_mut(&mut self) -> &mut BufferPool {
self.buffers
}
}
impl Drop for BufferPoolLoan<'_> {
fn drop(&mut self) {
if thread::panicking() {
self.buffers.replenish_in_flight_retained();
} else {
self.buffers.clear_in_flight_retained();
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum CpuBackendKind {
Faer,
Blas,
}
impl CpuBackendKind {
pub fn default_compiled() -> Self {
#[cfg(feature = "cpu-blas")]
{
Self::Blas
}
#[cfg(all(not(feature = "cpu-blas"), feature = "cpu-faer"))]
{
Self::Faer
}
}
#[allow(dead_code)]
pub(crate) fn name(self) -> &'static str {
match self {
Self::Faer => "faer",
Self::Blas => "blas",
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum CpuExecutionMode {
Managed,
ExternalManaged,
CallerManaged,
ProviderDefaultExclusive,
Compatibility,
}
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
pub enum ExternalCpuDomainRegistryError {
#[error("externally managed CPU registry must contain at least one domain")]
EmptyRegistry,
#[error("CPU domain ID {id:?} is registered more than once")]
DuplicateDomainId {
id: CpuDomainId,
},
#[error("CPU placement {placement:?} is registered more than once")]
DuplicatePlacementIdentity {
placement: CpuPlacement,
},
#[error("CPU domain {domain:?} declares process-disallowed CPU {cpu}")]
CpuOutsideAllowedSet {
domain: CpuDomainId,
cpu: CpuId,
},
#[error("default CPU domain {default_domain:?} is not registered")]
MissingDefaultDomain {
default_domain: CpuDomainId,
},
#[error(
"exact all-allowed CPU domain {domain:?} declares {declared:?}, but the process allows {allowed:?}"
)]
ExactAllAllowedMismatch {
domain: CpuDomainId,
declared: CpuSet,
allowed: CpuSet,
},
}
#[derive(Debug, thiserror::Error)]
pub enum CpuBackendError {
#[error(transparent)]
Tensor(#[from] crate::Error),
#[error("{op}: {source}")]
Placement {
op: &'static str,
#[source]
source: CpuPlacementError,
},
#[error(transparent)]
ExternalRegistry(#[from] ExternalCpuDomainRegistryError),
}
impl CpuBackendError {
fn placement(op: &'static str, source: CpuPlacementError) -> Self {
Self::Placement { op, source }
}
pub fn placement_error(&self) -> Option<&CpuPlacementError> {
match self {
Self::Tensor(_) => None,
Self::Placement { source, .. } => Some(source),
Self::ExternalRegistry(_) => None,
}
}
}
impl From<CpuBackendError> for crate::Error {
fn from(error: CpuBackendError) -> Self {
match error {
CpuBackendError::Tensor(error) => error,
CpuBackendError::ExternalRegistry(source) => Self::extension(
"CpuBackend::from_external_managed_domains",
"cpu",
crate::ErrorKind::Validation(crate::ValidationKind::InvalidArgument),
source,
),
CpuBackendError::Placement { op, source } => match source {
CpuPlacementError::TopologyDiscovery { .. }
| CpuPlacementError::ManagedAffinityUnavailable { .. }
| CpuPlacementError::NumaDiscoveryUnavailable { .. }
| CpuPlacementError::UnknownNumaNode { .. }
| CpuPlacementError::UnregisteredExternalPlacement { .. }
| CpuPlacementError::UnregisteredExternalDomain { .. } => {
Self::runtime_state_source(op, source)
}
CpuPlacementError::ExternalProviderAffinityUnmanaged { .. } => {
Self::extension(op, "cpu", crate::ErrorKind::Unsupported, source)
}
CpuPlacementError::EngineConstruction { .. } => Self::backend_source(op, source),
CpuPlacementError::InternalState { .. } => {
Self::extension(op, "cpu", crate::ErrorKind::Internal, source)
}
},
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CpuExecutionInfo {
backend_kind: CpuBackendKind,
execution_mode: CpuExecutionMode,
requested_placement: CpuPlacement,
resolved_placement: Option<ResolvedCpuPlacement>,
topology: CpuTopology,
domain_id: CpuDomainId,
domain_cpus: Option<CpuSet>,
worker_count: usize,
thread_budget: usize,
placement_guarantee: Option<CpuPlacementGuarantee>,
admission_mode: CpuAdmissionMode,
domain_ownership: CpuDomainOwnership,
executor_affinity: CpuExecutorAffinity,
executor_shutdown: CpuExecutorShutdown,
provider_diagnostic: &'static str,
}
impl CpuExecutionInfo {
pub fn backend_kind(&self) -> CpuBackendKind {
self.backend_kind
}
pub fn execution_mode(&self) -> CpuExecutionMode {
self.execution_mode
}
pub fn requested_placement(&self) -> CpuPlacement {
self.requested_placement
}
pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
self.resolved_placement.as_ref()
}
pub fn topology(&self) -> &CpuTopology {
&self.topology
}
pub fn domain_id(&self) -> CpuDomainId {
self.domain_id
}
pub fn domain_cpus(&self) -> Option<&CpuSet> {
self.domain_cpus.as_ref()
}
pub fn worker_count(&self) -> usize {
self.worker_count
}
pub fn thread_budget(&self) -> usize {
self.thread_budget
}
pub fn placement_guarantee(&self) -> Option<CpuPlacementGuarantee> {
self.placement_guarantee
}
pub fn admission_mode(&self) -> CpuAdmissionMode {
self.admission_mode
}
pub fn domain_ownership(&self) -> CpuDomainOwnership {
self.domain_ownership
}
pub fn executor_affinity(&self) -> CpuExecutorAffinity {
self.executor_affinity
}
pub fn executor_shutdown(&self) -> CpuExecutorShutdown {
self.executor_shutdown
}
pub fn provider_diagnostic(&self) -> &'static str {
self.provider_diagnostic
}
}
fn provider_diagnostic(
kind: CpuBackendKind,
ownership: CpuDomainOwnership,
admission_mode: CpuAdmissionMode,
) -> &'static str {
if ownership == CpuDomainOwnership::ExternalManaged {
if admission_mode == CpuAdmissionMode::CallerManaged {
debug_assert_eq!(kind, CpuBackendKind::Faer);
return "faer (caller-managed CPU executor and admission)";
}
return match kind {
CpuBackendKind::Faer => "faer (externally managed CPU executor)",
CpuBackendKind::Blas => "BLAS/LAPACK (externally managed CPU executor)",
};
}
match kind {
CpuBackendKind::Faer => "faer (tenferro-managed Rayon affinity)",
CpuBackendKind::Blas => {
#[cfg(feature = "blas-openblas")]
return "OpenBLAS (external worker affinity)";
#[cfg(feature = "blas-mkl")]
return "Intel MKL (external worker affinity)";
#[cfg(feature = "blas-accelerate")]
return "Apple Accelerate (external worker affinity)";
#[cfg(feature = "provider-inject")]
return "runtime-injected BLAS/LAPACK (external worker affinity)";
#[cfg(not(any(
feature = "blas-openblas",
feature = "blas-mkl",
feature = "blas-accelerate",
feature = "provider-inject"
)))]
return "linked BLAS/LAPACK provider (identity unknown; external worker affinity)";
}
}
}
fn ensure_cpu_backend_kind_available(kind: CpuBackendKind, op: &'static str) -> crate::Result<()> {
let _ = op;
match kind {
CpuBackendKind::Faer => {
#[cfg(feature = "cpu-faer")]
{
Ok(())
}
#[cfg(not(feature = "cpu-faer"))]
{
Err(crate::Error::invalid_argument(
op,
"configuration",
"CpuBackendKind::Faer requires the cpu-faer feature".to_string(),
))
}
}
CpuBackendKind::Blas => {
#[cfg(feature = "cpu-blas")]
{
Ok(())
}
#[cfg(not(feature = "cpu-blas"))]
{
Err(crate::Error::invalid_argument(
op,
"configuration",
"CpuBackendKind::Blas requires the cpu-blas feature".to_string(),
))
}
}
}
}
fn constructor_tensor_error(op: &'static str, error: crate::Error) -> CpuBackendError {
CpuBackendError::Tensor(match error {
crate::Error::Validation { source, .. } => crate::Error::validation(op, source),
error => error,
})
}
#[allow(dead_code)]
pub(super) fn unavailable_cpu_backend_kind(kind: CpuBackendKind, op: &'static str) -> crate::Error {
crate::Error::invalid_argument(
op,
"configuration",
format!("CPU backend kind {} is not compiled in", kind.name()),
)
}
struct ManagedEngineRegistry {
node_engines: Mutex<BTreeMap<NumaNodeId, Arc<CpuEngine>>>,
node_domain_ids: BTreeMap<NumaNodeId, CpuDomainId>,
all_allowed: OnceLock<Arc<CpuEngine>>,
all_allowed_build: Mutex<()>,
base_engine: Arc<CpuEngine>,
thread_budget: usize,
}
struct ExternalEngineRegistry {
by_id: BTreeMap<CpuDomainId, Arc<CpuEngine>>,
by_node: BTreeMap<NumaNodeId, Arc<CpuEngine>>,
all_allowed: Option<Arc<CpuEngine>>,
default_domain: CpuDomainId,
}
enum CpuEngineRegistry {
ManagedLazy(ManagedEngineRegistry),
ExternalPrebuilt(ExternalEngineRegistry),
}
struct CpuBackendState {
topology: CpuTopology,
engines: CpuEngineRegistry,
arbiter: ResourceArbiter,
kind: CpuBackendKind,
buffer_limit: AtomicUsize,
indexed_plan_cache_limits: Mutex<IndexedPlanCacheLimits>,
}
impl CpuBackendState {
fn managed_engine_for(
&self,
placement: &ResolvedCpuPlacement,
requested: CpuPlacement,
) -> Result<Arc<CpuEngine>, CpuPlacementError> {
let cache_configuration = self.indexed_plan_cache_limits.lock().map_err(|_| {
CpuPlacementError::InternalState {
requested,
backend: self.kind,
message: "CPU indexed-plan cache configuration lock is poisoned",
}
})?;
let cache_limits = *cache_configuration;
let CpuEngineRegistry::ManagedLazy(registry) = &self.engines else {
return Err(CpuPlacementError::InternalState {
requested,
backend: self.kind,
message: "managed placement requested from an external engine registry",
});
};
match placement {
ResolvedCpuPlacement::NumaNode { id, .. } => {
let mut engines = registry
.node_engines
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(engine) = engines.get(id) {
return Ok(Arc::clone(engine));
}
let Some(domain_id) = registry.node_domain_ids.get(id).copied() else {
return Err(CpuPlacementError::InternalState {
requested,
backend: self.kind,
message: "managed NUMA node has no coordinator-stable domain ID",
});
};
let engine = Arc::new(
CpuEngine::new_managed(
domain_id,
placement.clone(),
registry.thread_budget,
self.buffer_limit.load(Ordering::Relaxed),
)
.map_err(|error| {
CpuPlacementError::EngineConstruction {
requested,
backend: self.kind,
source: CpuEngineConstructionError::Context(error),
}
})?,
);
self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
engines.insert(*id, Arc::clone(&engine));
Ok(engine)
}
ResolvedCpuPlacement::AllAllowed { .. } => {
if let Some(engine) = registry.all_allowed.get() {
return Ok(Arc::clone(engine));
}
let _build = registry
.all_allowed_build
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(engine) = registry.all_allowed.get() {
return Ok(Arc::clone(engine));
}
let engine = Arc::new(
CpuEngine::new_managed(
CpuDomainId::new(0),
placement.clone(),
registry.thread_budget,
self.buffer_limit.load(Ordering::Relaxed),
)
.map_err(|error| {
CpuPlacementError::EngineConstruction {
requested,
backend: self.kind,
source: CpuEngineConstructionError::Context(error),
}
})?,
);
self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
let _ = registry.all_allowed.set(Arc::clone(&engine));
Ok(engine)
}
}
}
fn configure_new_indexed_plan_cache(
&self,
engine: &CpuEngine,
requested: CpuPlacement,
limits: IndexedPlanCacheLimits,
) -> Result<(), CpuPlacementError> {
let mut resources =
engine
.resources
.lock()
.map_err(|_| CpuPlacementError::InternalState {
requested,
backend: self.kind,
message: "new CPU engine indexed-plan cache lock is poisoned",
})?;
resources.indexed_plan_cache.set_limits(limits);
Ok(())
}
fn managed_base_engine(
&self,
requested: CpuPlacement,
) -> Result<Arc<CpuEngine>, CpuPlacementError> {
match &self.engines {
CpuEngineRegistry::ManagedLazy(registry) => Ok(Arc::clone(®istry.base_engine)),
CpuEngineRegistry::ExternalPrebuilt(_) => Err(CpuPlacementError::InternalState {
requested,
backend: self.kind,
message: "managed compatibility placement requested from an external registry",
}),
}
}
fn external_engine_for(
&self,
requested: CpuPlacement,
) -> Result<Arc<CpuEngine>, CpuPlacementError> {
let CpuEngineRegistry::ExternalPrebuilt(registry) = &self.engines else {
return Err(CpuPlacementError::InternalState {
requested,
backend: self.kind,
message: "external placement requested from a managed engine registry",
});
};
let engine = match requested {
CpuPlacement::Auto => registry.by_id.get(®istry.default_domain),
CpuPlacement::NumaNode(id) => registry.by_node.get(&id),
CpuPlacement::AllAllowed => registry.all_allowed.as_ref(),
};
engine
.cloned()
.ok_or(CpuPlacementError::UnregisteredExternalPlacement { requested })
}
fn external_engine_for_id(
&self,
domain: CpuDomainId,
) -> Result<Arc<CpuEngine>, CpuPlacementError> {
let CpuEngineRegistry::ExternalPrebuilt(registry) = &self.engines else {
return Err(CpuPlacementError::UnregisteredExternalDomain { domain });
};
registry
.by_id
.get(&domain)
.cloned()
.ok_or(CpuPlacementError::UnregisteredExternalDomain { domain })
}
fn is_external(&self) -> bool {
matches!(&self.engines, CpuEngineRegistry::ExternalPrebuilt(_))
}
fn initialized_engines(&self, op: &'static str) -> crate::Result<Vec<Arc<CpuEngine>>> {
let mut engines = match &self.engines {
CpuEngineRegistry::ManagedLazy(registry) => {
let mut engines = vec![Arc::clone(®istry.base_engine)];
if let Some(engine) = registry.all_allowed.get() {
engines.push(Arc::clone(engine));
}
engines.extend(
registry
.node_engines
.lock()
.map_err(|_| poisoned_cpu_lock(op, "CPU engine registry"))?
.values()
.cloned(),
);
engines
}
CpuEngineRegistry::ExternalPrebuilt(registry) => {
registry.by_id.values().cloned().collect()
}
};
if engines.len() > 1 {
engines.sort_unstable_by_key(|engine| Arc::as_ptr(engine) as usize);
engines.dedup_by(|left, right| Arc::ptr_eq(left, right));
}
Ok(engines)
}
}
fn poisoned_cpu_lock(op: &'static str, lock: &'static str) -> crate::Error {
crate::Error::runtime_state(op, format!("{lock} lock poisoned"))
}
fn lock_engine_resources<'a>(
engine: &'a CpuEngine,
op: &'static str,
) -> crate::Result<std::sync::MutexGuard<'a, EngineResources>> {
engine
.resources
.lock()
.map_err(|_| poisoned_cpu_lock(op, "CPU engine resources"))
}
fn saturating_add_tensor_cache_stats(total: &mut CacheStats, value: CacheStats) {
total.entries = total.entries.saturating_add(value.entries);
total.retained_bytes = total.retained_bytes.saturating_add(value.retained_bytes);
total.hits = total.hits.saturating_add(value.hits);
total.misses = total.misses.saturating_add(value.misses);
total.evictions = total.evictions.saturating_add(value.evictions);
total.clears = total.clears.saturating_add(value.clears);
}
#[doc(hidden)]
struct CpuBackendSessionMarker;
#[derive(Clone)]
pub struct CpuBackend {
runtime_identity: CpuRuntimeIdentity,
shared: Arc<CpuBackendState>,
requested: CpuPlacement,
resolved: ResolvedCpuExecution,
engine: Arc<CpuEngine>,
provider_bundle: CpuProviderBundle,
allocation_domain: Option<Arc<dyn SharedTensorAllocationDomain>>,
}
#[derive(Clone, Debug)]
pub struct CpuRuntimeIdentity {
marker: Arc<()>,
}
impl CpuRuntimeIdentity {
fn fresh() -> Self {
Self {
marker: Arc::new(()),
}
}
}
impl PartialEq for CpuRuntimeIdentity {
fn eq(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.marker, &other.marker)
}
}
impl Eq for CpuRuntimeIdentity {}
fn resolve_discovered_topology(
kind: CpuBackendKind,
topology: Result<CpuTopology, CpuTopologyError>,
) -> Result<CpuTopology, CpuPlacementError> {
topology.map_err(|source| CpuPlacementError::TopologyDiscovery {
requested: CpuPlacement::Auto,
backend: kind,
source,
})
}
fn external_engine_resolution(
engine: &CpuEngine,
requested: CpuPlacement,
kind: CpuBackendKind,
) -> Result<ResolvedCpuExecution, CpuPlacementError> {
match engine.domain().admission_mode() {
CpuAdmissionMode::CooperativeCpuSet => engine
.placement()
.cloned()
.map(ResolvedCpuExecution::ExternalManaged)
.ok_or(CpuPlacementError::InternalState {
requested,
backend: kind,
message: "cooperative external domain has no placement",
}),
CpuAdmissionMode::CallerManaged => Ok(ResolvedCpuExecution::ExternalCallerManaged),
}
}
fn external_domain_backend_kind(
op: &'static str,
domains: &[ExternalCpuDomain],
) -> Result<CpuBackendKind, CpuBackendError> {
let kind = if domains
.iter()
.any(|domain| domain.admission_mode() == CpuAdmissionMode::CallerManaged)
{
CpuBackendKind::Faer
} else {
CpuBackendKind::default_compiled()
};
ensure_cpu_backend_kind_available(kind, op)
.map_err(|error| constructor_tensor_error(op, error))?;
Ok(kind)
}
fn coordinator_node_domain_ids(topology: &CpuTopology) -> BTreeMap<NumaNodeId, CpuDomainId> {
topology
.nodes()
.iter()
.enumerate()
.filter_map(|(index, node)| {
u64::try_from(index)
.ok()
.and_then(|index| index.checked_add(1))
.map(|id| (node.id(), CpuDomainId::new(id)))
})
.collect()
}
impl fmt::Debug for CpuBackend {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CpuBackend")
.field("kind", &self.kind())
.field("provider_bundle", &self.provider_bundle)
.field("requested_placement", &self.requested)
.field("resolved_execution", &self.resolved)
.field("engine_placement", &self.engine.placement())
.field("num_threads", &self.num_threads())
.field("allocation_domain", &self.allocation_domain())
.field("buffer_pool_cache_stats", &self.buffer_pool_cache_stats())
.field("buffer_pool_limit_bytes", &self.buffer_pool_limit_bytes())
.finish_non_exhaustive()
}
}
impl CpuBackend {
fn from_thread_budget_and_kind(
thread_budget: usize,
kind: CpuBackendKind,
max_retained_capacity_bytes: usize,
) -> Result<Self, CpuPlacementError> {
let topology = resolve_discovered_topology(kind, discover_cpu_topology())?;
let resolved = resolve_placement(kind, CpuPlacement::Auto, &topology)?;
#[cfg(not(any(target_os = "linux", target_os = "android")))]
{
let context = CpuContext::with_threads(thread_budget).map_err(|error| {
CpuPlacementError::EngineConstruction {
requested: CpuPlacement::Auto,
backend: kind,
source: CpuEngineConstructionError::Tensor(error),
}
})?;
Ok(Self::compatibility_with_topology(
Arc::new(context),
max_retained_capacity_bytes,
kind,
topology,
resolved,
))
}
#[cfg(any(target_os = "linux", target_os = "android"))]
{
let engine_placement = ResolvedCpuPlacement::AllAllowed {
cpus: topology.allowed_cpus().clone(),
};
let engine = Arc::new(
CpuEngine::new_managed(
CpuDomainId::new(0),
engine_placement,
thread_budget,
max_retained_capacity_bytes,
)
.map_err(|error| CpuPlacementError::EngineConstruction {
requested: CpuPlacement::Auto,
backend: kind,
source: CpuEngineConstructionError::Context(error),
})?,
);
let all_allowed = OnceLock::new();
let _ = all_allowed.set(Arc::clone(&engine));
Ok(Self {
shared: Arc::new(CpuBackendState {
engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
node_engines: Mutex::new(BTreeMap::new()),
node_domain_ids: coordinator_node_domain_ids(&topology),
all_allowed,
all_allowed_build: Mutex::new(()),
base_engine: Arc::clone(&engine),
thread_budget,
}),
topology,
arbiter: ResourceArbiter::global(),
kind,
buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
}),
runtime_identity: CpuRuntimeIdentity::fresh(),
requested: CpuPlacement::Auto,
resolved,
engine,
provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
allocation_domain: None,
})
}
}
fn compatibility(
ctx: Arc<CpuContext>,
max_retained_capacity_bytes: usize,
kind: CpuBackendKind,
) -> Self {
let topology = discover_cpu_topology().unwrap_or_else(|_| {
let allowed = crate::process_cpu_affinity().unwrap_or_else(|| {
CpuSet::new((0..crate::available_parallelism()).map(CpuId::new))
.unwrap_or_else(|_| CpuSet::singleton(CpuId::new(0)))
});
CpuTopology::all_allowed(allowed)
});
let resolved = if kind == CpuBackendKind::Blas {
ResolvedCpuExecution::ProviderDefaultExclusive
} else {
ResolvedCpuExecution::Compatibility
};
Self::compatibility_with_topology(
ctx,
max_retained_capacity_bytes,
kind,
topology,
resolved,
)
}
fn compatibility_with_topology(
ctx: Arc<CpuContext>,
max_retained_capacity_bytes: usize,
kind: CpuBackendKind,
topology: CpuTopology,
resolved: ResolvedCpuExecution,
) -> Self {
let placement = ResolvedCpuPlacement::AllAllowed {
cpus: topology.allowed_cpus().clone(),
};
let base_engine = Arc::new(CpuEngine::from_context(
CpuDomainId::new(0),
placement,
ctx,
max_retained_capacity_bytes,
));
Self {
shared: Arc::new(CpuBackendState {
engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
node_engines: Mutex::new(BTreeMap::new()),
node_domain_ids: coordinator_node_domain_ids(&topology),
all_allowed: OnceLock::new(),
all_allowed_build: Mutex::new(()),
base_engine: Arc::clone(&base_engine),
thread_budget: base_engine.domain().thread_budget().get(),
}),
topology,
arbiter: ResourceArbiter::global(),
kind,
buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
}),
runtime_identity: CpuRuntimeIdentity::fresh(),
requested: CpuPlacement::Auto,
resolved,
engine: base_engine,
provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
allocation_domain: None,
}
}
pub fn new() -> Self {
let context = Arc::new(CpuContext::from_env());
Self::from_thread_budget_and_kind(
context.num_threads(),
CpuBackendKind::default_compiled(),
crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
)
.unwrap_or_else(|error| {
eprintln!(
"tenferro_cpu: using the unpinned compatibility context after placement error: {error}"
);
Self::from_context(context)
})
}
pub fn from_external_managed_domains(
default_domain: CpuDomainId,
domains: impl IntoIterator<Item = ExternalCpuDomain>,
) -> Result<Self, CpuBackendError> {
let op = "CpuBackend::from_external_managed_domains";
let domains: Vec<_> = domains.into_iter().collect();
let kind = external_domain_backend_kind(op, &domains)?;
let topology = resolve_discovered_topology(kind, discover_cpu_topology())
.map_err(|source| CpuBackendError::placement(op, source))?;
Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
default_domain,
domains,
topology,
ResourceArbiter::global(),
kind,
CpuProviderBundle::standard(kind, false),
)
}
pub fn from_external_managed_domains_with_provider_bundle(
default_domain: CpuDomainId,
domains: impl IntoIterator<Item = ExternalCpuDomain>,
provider_bundle: CpuProviderBundle,
) -> Result<Self, CpuBackendError> {
let op = "CpuBackend::from_external_managed_domains_with_provider_bundle";
let domains: Vec<_> = domains.into_iter().collect();
let kind = external_domain_backend_kind(op, &domains)?;
let topology = resolve_discovered_topology(kind, discover_cpu_topology())
.map_err(|source| CpuBackendError::placement(op, source))?;
Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
default_domain,
domains,
topology,
ResourceArbiter::global(),
kind,
provider_bundle,
)
}
fn from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
default_domain: CpuDomainId,
domains: impl IntoIterator<Item = ExternalCpuDomain>,
topology: CpuTopology,
arbiter: ResourceArbiter,
kind: CpuBackendKind,
provider_bundle: CpuProviderBundle,
) -> Result<Self, CpuBackendError> {
let domains: Vec<_> = domains.into_iter().collect();
if domains.is_empty() {
return Err(ExternalCpuDomainRegistryError::EmptyRegistry.into());
}
let mut domain_ids = BTreeSet::new();
let mut node_ids = BTreeSet::new();
let mut has_all_allowed = false;
for domain in &domains {
if !domain_ids.insert(domain.id()) {
return Err(
ExternalCpuDomainRegistryError::DuplicateDomainId { id: domain.id() }.into(),
);
}
if let Some(placement) = domain.placement() {
match placement {
ResolvedCpuPlacement::NumaNode { id, .. } => {
if !node_ids.insert(*id) {
return Err(
ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
placement: CpuPlacement::NumaNode(*id),
}
.into(),
);
}
}
ResolvedCpuPlacement::AllAllowed { cpus } => {
if has_all_allowed {
return Err(
ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
placement: CpuPlacement::AllAllowed,
}
.into(),
);
}
has_all_allowed = true;
if domain.placement_guarantee()
== Some(CpuPlacementGuarantee::ExactDeclared)
&& cpus != topology.allowed_cpus()
{
return Err(ExternalCpuDomainRegistryError::ExactAllAllowedMismatch {
domain: domain.id(),
declared: cpus.clone(),
allowed: topology.allowed_cpus().clone(),
}
.into());
}
}
}
if let Some(cpu) = placement
.cpus()
.as_slice()
.iter()
.copied()
.find(|cpu| !topology.allowed_cpus().contains(*cpu))
{
return Err(ExternalCpuDomainRegistryError::CpuOutsideAllowedSet {
domain: domain.id(),
cpu,
}
.into());
}
}
}
let buffer_limit = crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES;
let mut by_id = BTreeMap::new();
let mut by_node = BTreeMap::new();
let mut all_allowed = None;
for domain in domains {
let id = domain.id();
let placement = domain.placement().cloned();
let engine = Arc::new(CpuEngine::from_external(domain, buffer_limit));
match placement {
Some(ResolvedCpuPlacement::NumaNode { id, .. }) => {
by_node.insert(id, Arc::clone(&engine));
}
Some(ResolvedCpuPlacement::AllAllowed { .. }) => {
all_allowed = Some(Arc::clone(&engine));
}
None => {}
}
by_id.insert(id, engine);
}
let Some(engine) = by_id.get(&default_domain).cloned() else {
return Err(
ExternalCpuDomainRegistryError::MissingDefaultDomain { default_domain }.into(),
);
};
let resolved = match engine.domain().admission_mode() {
CpuAdmissionMode::CooperativeCpuSet => ResolvedCpuExecution::ExternalManaged(
engine.placement().cloned().ok_or_else(|| {
CpuBackendError::placement(
"CpuBackend external domain resolution",
CpuPlacementError::InternalState {
requested: CpuPlacement::Auto,
backend: kind,
message: "cooperative external domain has no placement",
},
)
})?,
),
CpuAdmissionMode::CallerManaged => ResolvedCpuExecution::ExternalCallerManaged,
};
let backend = Self {
runtime_identity: CpuRuntimeIdentity::fresh(),
shared: Arc::new(CpuBackendState {
topology,
engines: CpuEngineRegistry::ExternalPrebuilt(ExternalEngineRegistry {
by_id,
by_node,
all_allowed,
default_domain,
}),
arbiter,
kind,
buffer_limit: AtomicUsize::new(buffer_limit),
indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
}),
requested: CpuPlacement::Auto,
resolved,
engine,
provider_bundle,
allocation_domain: None,
};
backend
.validate_provider_bundle_for_domains(&backend.provider_bundle)
.map_err(|source| {
CpuBackendError::Tensor(crate::Error::backend_source(
"CpuBackend ExternalManaged provider validation",
source,
))
})?;
Ok(backend)
}
pub fn with_kind(kind: CpuBackendKind) -> Result<Self, CpuBackendError> {
let op = "CpuBackend::with_kind";
ensure_cpu_backend_kind_available(kind, op)
.map_err(|error| constructor_tensor_error(op, error))?;
let context = CpuContext::from_env();
Self::from_thread_budget_and_kind(
context.num_threads(),
kind,
crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
)
.map_err(|error| CpuBackendError::placement(op, error))
}
pub fn try_new() -> Result<Self, CpuBackendError> {
let op = "CpuBackend::try_new";
let context =
CpuContext::try_from_env().map_err(|error| constructor_tensor_error(op, error))?;
Self::from_thread_budget_and_kind(
context.num_threads(),
CpuBackendKind::default_compiled(),
crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
)
.map_err(|error| CpuBackendError::placement(op, error))
}
pub fn from_context(ctx: Arc<CpuContext>) -> Self {
Self::compatibility(
ctx,
crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
CpuBackendKind::default_compiled(),
)
}
pub fn from_context_with_buffer_pool_limit(
ctx: Arc<CpuContext>,
max_retained_capacity_bytes: usize,
) -> Self {
Self::from_context_with_buffer_pool_limit_and_kind(
ctx,
max_retained_capacity_bytes,
CpuBackendKind::default_compiled(),
)
}
fn from_context_with_buffer_pool_limit_and_kind(
ctx: Arc<CpuContext>,
max_retained_capacity_bytes: usize,
kind: CpuBackendKind,
) -> Self {
Self::compatibility(ctx, max_retained_capacity_bytes, kind)
}
pub fn with_threads(num_threads: usize) -> Result<Self, CpuBackendError> {
let op = "CpuBackend::with_threads";
let context = CpuContext::with_threads(num_threads)
.map_err(|error| constructor_tensor_error(op, error))?;
Self::from_thread_budget_and_kind(
context.num_threads(),
CpuBackendKind::default_compiled(),
crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
)
.map_err(|error| CpuBackendError::placement(op, error))
}
pub fn with_threads_and_kind(
num_threads: usize,
kind: CpuBackendKind,
) -> Result<Self, CpuBackendError> {
let op = "CpuBackend::with_threads_and_kind";
ensure_cpu_backend_kind_available(kind, op)
.map_err(|error| constructor_tensor_error(op, error))?;
let context = CpuContext::with_threads(num_threads)
.map_err(|error| constructor_tensor_error(op, error))?;
Self::from_thread_budget_and_kind(
context.num_threads(),
kind,
crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
)
.map_err(|error| CpuBackendError::placement(op, error))
}
pub fn for_placement(&self, requested: CpuPlacement) -> Result<Self, CpuPlacementError> {
self.for_placement_with_affinity(
requested,
cfg!(any(target_os = "linux", target_os = "android")),
)
}
pub fn for_domain(&self, domain: CpuDomainId) -> Result<Self, CpuPlacementError> {
let engine = self.shared.external_engine_for_id(domain)?;
let resolved = external_engine_resolution(&engine, CpuPlacement::Auto, self.kind())?;
Ok(Self {
runtime_identity: CpuRuntimeIdentity::fresh(),
shared: Arc::clone(&self.shared),
requested: CpuPlacement::Auto,
resolved,
engine,
provider_bundle: self.provider_bundle.clone(),
allocation_domain: self.allocation_domain.clone(),
})
}
fn for_placement_with_affinity(
&self,
requested: CpuPlacement,
managed_affinity_available: bool,
) -> Result<Self, CpuPlacementError> {
if self.shared.is_external() {
let engine = self.shared.external_engine_for(requested)?;
let resolved = external_engine_resolution(&engine, requested, self.kind())?;
return Ok(Self {
runtime_identity: CpuRuntimeIdentity::fresh(),
shared: Arc::clone(&self.shared),
requested,
resolved,
engine,
provider_bundle: self.provider_bundle.clone(),
allocation_domain: self.allocation_domain.clone(),
});
}
let resolved = resolve_placement_with_affinity(
self.kind(),
requested,
&self.shared.topology,
managed_affinity_available,
)?;
if requested == CpuPlacement::Auto && !managed_affinity_available {
return Ok(Self {
runtime_identity: CpuRuntimeIdentity::fresh(),
shared: Arc::clone(&self.shared),
requested,
resolved,
engine: self.shared.managed_base_engine(requested)?,
provider_bundle: self.provider_bundle.clone(),
allocation_domain: self.allocation_domain.clone(),
});
}
let engine_placement = match &resolved {
ResolvedCpuExecution::Managed(placement) => placement.clone(),
ResolvedCpuExecution::ExternalManaged(_)
| ResolvedCpuExecution::ExternalCallerManaged => {
return Err(CpuPlacementError::InternalState {
requested,
backend: self.kind(),
message: "managed resolver returned an external execution mode",
});
}
ResolvedCpuExecution::ProviderDefaultExclusive => ResolvedCpuPlacement::AllAllowed {
cpus: self.shared.topology.allowed_cpus().clone(),
},
ResolvedCpuExecution::Compatibility => {
return Err(CpuPlacementError::InternalState {
requested,
backend: self.kind(),
message: "placement resolution returned an internal compatibility mode",
});
}
};
let engine = self
.shared
.managed_engine_for(&engine_placement, requested)?;
Ok(Self {
runtime_identity: CpuRuntimeIdentity::fresh(),
shared: Arc::clone(&self.shared),
requested,
resolved,
engine,
provider_bundle: self.provider_bundle.clone(),
allocation_domain: self.allocation_domain.clone(),
})
}
pub fn placement(&self) -> CpuPlacement {
self.requested
}
pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
match &self.resolved {
ResolvedCpuExecution::Managed(placement)
| ResolvedCpuExecution::ExternalManaged(placement) => Some(placement),
ResolvedCpuExecution::Compatibility
| ResolvedCpuExecution::ExternalCallerManaged
| ResolvedCpuExecution::ProviderDefaultExclusive => None,
}
}
pub fn topology(&self) -> &CpuTopology {
&self.shared.topology
}
pub fn supports_placement(&self, placement: CpuPlacement) -> bool {
if self.shared.is_external() {
self.shared.external_engine_for(placement).is_ok()
} else {
resolve_placement(self.kind(), placement, &self.shared.topology).is_ok()
}
}
pub fn execution_info(&self) -> CpuExecutionInfo {
let domain = self.engine.domain();
let capabilities = domain.executor_capabilities();
let (executor_affinity, executor_shutdown) =
match (domain.ownership(), domain.admission_mode()) {
(CpuDomainOwnership::ExternalManaged, CpuAdmissionMode::CooperativeCpuSet) => (
CpuExecutorAffinity::CallerDeclaredUnverified,
CpuExecutorShutdown::CallerOwned,
),
(CpuDomainOwnership::ExternalManaged, CpuAdmissionMode::CallerManaged) => {
(capabilities.affinity, CpuExecutorShutdown::CallerOwned)
}
(CpuDomainOwnership::Managed, _) => (capabilities.affinity, capabilities.shutdown),
};
CpuExecutionInfo {
backend_kind: self.kind(),
execution_mode: match &self.resolved {
ResolvedCpuExecution::Managed(_) => CpuExecutionMode::Managed,
ResolvedCpuExecution::ExternalManaged(_) => CpuExecutionMode::ExternalManaged,
ResolvedCpuExecution::ExternalCallerManaged => CpuExecutionMode::CallerManaged,
ResolvedCpuExecution::ProviderDefaultExclusive => {
CpuExecutionMode::ProviderDefaultExclusive
}
ResolvedCpuExecution::Compatibility => CpuExecutionMode::Compatibility,
},
requested_placement: self.requested,
resolved_placement: self.resolved_placement().cloned(),
topology: self.shared.topology.clone(),
domain_id: domain.id(),
domain_cpus: domain.cpus().cloned(),
worker_count: capabilities.worker_count.get(),
thread_budget: domain.thread_budget().get(),
placement_guarantee: domain.placement_guarantee(),
admission_mode: domain.admission_mode(),
domain_ownership: domain.ownership(),
executor_affinity,
executor_shutdown,
provider_diagnostic: provider_diagnostic(
self.kind(),
domain.ownership(),
domain.admission_mode(),
),
}
}
#[cfg(all(
test,
feature = "cpu-faer",
any(target_os = "linux", target_os = "android")
))]
fn coordinator_id_for_test(&self) -> usize {
Arc::as_ptr(&self.shared) as usize
}
#[cfg(test)]
pub(crate) fn context_id_for_test(&self) -> usize {
Arc::as_ptr(self.engine.domain().executor()) as *const () as usize
}
pub fn kind(&self) -> CpuBackendKind {
self.shared.kind
}
pub fn provider_bundle(&self) -> &CpuProviderBundle {
&self.provider_bundle
}
pub fn runtime_identity(&self) -> CpuRuntimeIdentity {
self.runtime_identity.clone()
}
pub fn with_provider_bundle(
mut self,
bundle: CpuProviderBundle,
) -> Result<Self, CpuProviderBundleInstallError> {
self.validate_provider_bundle_for_domains(&bundle)?;
self.provider_bundle = bundle;
self.runtime_identity = CpuRuntimeIdentity::fresh();
Ok(self)
}
fn validate_provider_bundle_for_domains(
&self,
bundle: &CpuProviderBundle,
) -> Result<(), CpuProviderBundleInstallError> {
let allowed = self.shared.topology.allowed_cpus();
let validate_engine = |engine: &CpuEngine| {
let domain = engine.domain();
let contract = match (domain.placement_guarantee(), domain.cpus()) {
(Some(placement_guarantee), Some(domain_cpus)) => {
CpuProviderDomainContract::CooperativeCpuSet {
placement_guarantee,
domain_cpus,
process_allowed_cpus: allowed,
}
}
(None, None) => CpuProviderDomainContract::CallerManaged,
_ => unreachable!("CPU domain placement and guarantee must match"),
};
bundle.validate_for_domain(domain.id(), domain.thread_budget(), contract)
};
match &self.shared.engines {
CpuEngineRegistry::ExternalPrebuilt(registry) => {
for engine in registry.by_id.values() {
validate_engine(engine)?;
}
}
CpuEngineRegistry::ManagedLazy(registry) => {
validate_engine(®istry.base_engine)?;
#[cfg(any(target_os = "linux", target_os = "android"))]
for node in self.shared.topology.nodes() {
let Some(domain_id) = registry.node_domain_ids.get(&node.id()).copied() else {
continue;
};
let budget =
std::num::NonZeroUsize::new(registry.thread_budget.min(node.cpus().len()))
.expect("usable topology nodes have non-empty CPU sets");
bundle.validate_for_domain(
domain_id,
budget,
CpuProviderDomainContract::CooperativeCpuSet {
placement_guarantee: CpuPlacementGuarantee::ExactDeclared,
domain_cpus: node.cpus(),
process_allowed_cpus: allowed,
},
)?;
}
}
}
Ok(())
}
pub fn num_threads(&self) -> usize {
self.engine.domain().thread_budget().get()
}
pub fn buffer_pool_len(&self) -> crate::Result<usize> {
self.shared
.initialized_engines("CpuBackend::buffer_pool_len")?
.iter()
.try_fold(0, |total, engine| {
Ok(total
+ lock_engine_resources(engine, "CpuBackend::buffer_pool_len")?
.buffers
.len())
})
}
pub fn buffer_pool_stats(&self) -> crate::Result<BufferPoolStats> {
self.shared
.initialized_engines("CpuBackend::buffer_pool_stats")?
.iter()
.try_fold(BufferPoolStats::default(), |mut total, engine| {
let stats = lock_engine_resources(engine, "CpuBackend::buffer_pool_stats")?
.buffers
.stats();
total.buffers += stats.buffers;
total.capacity_bytes += stats.capacity_bytes;
Ok(total)
})
}
pub fn buffer_pool_cache_stats(&self) -> crate::Result<CacheStats> {
let stats = self.buffer_pool_stats()?;
Ok(CacheStats {
entries: stats.buffers,
retained_bytes: stats.capacity_bytes,
hits: 0,
misses: 0,
evictions: 0,
clears: 0,
})
}
pub fn indexed_plan_cache_limits(&self) -> crate::Result<IndexedPlanCacheLimits> {
self.shared
.indexed_plan_cache_limits
.lock()
.map(|limits| *limits)
.map_err(|_| {
poisoned_cpu_lock(
"CpuBackend::indexed_plan_cache_limits",
"CPU indexed-plan cache configuration",
)
})
}
pub fn set_indexed_plan_cache_limits(
&mut self,
limits: IndexedPlanCacheLimits,
) -> crate::Result<()> {
let mut configured_limits = self.shared.indexed_plan_cache_limits.lock().map_err(|_| {
poisoned_cpu_lock(
"CpuBackend::set_indexed_plan_cache_limits",
"CPU indexed-plan cache configuration",
)
})?;
let engines = self
.shared
.initialized_engines("CpuBackend::set_indexed_plan_cache_limits")?;
let mut resources = engines
.iter()
.map(|engine| {
lock_engine_resources(engine, "CpuBackend::set_indexed_plan_cache_limits")
})
.collect::<crate::Result<Vec<_>>>()?;
*configured_limits = limits;
for resource in &mut resources {
resource.indexed_plan_cache.set_limits(limits);
}
Ok(())
}
pub fn indexed_plan_cache_stats(&self) -> crate::Result<CacheStats> {
self.shared
.initialized_engines("CpuBackend::indexed_plan_cache_stats")?
.iter()
.try_fold(CacheStats::default(), |mut total, engine| {
let stats = lock_engine_resources(engine, "CpuBackend::indexed_plan_cache_stats")?
.indexed_plan_cache
.stats();
saturating_add_tensor_cache_stats(&mut total, stats);
Ok(total)
})
}
pub fn clear_indexed_plan_cache(&mut self) -> crate::Result<()> {
let engines = self
.shared
.initialized_engines("CpuBackend::clear_indexed_plan_cache")?;
let mut resources = engines
.iter()
.map(|engine| lock_engine_resources(engine, "CpuBackend::clear_indexed_plan_cache"))
.collect::<crate::Result<Vec<_>>>()?;
for resource in &mut resources {
resource.indexed_plan_cache.clear();
}
Ok(())
}
pub fn buffer_pool_limit_bytes(&self) -> usize {
self.shared.buffer_limit.load(Ordering::Relaxed)
}
pub fn set_buffer_pool_limit_bytes(
&mut self,
max_retained_capacity_bytes: usize,
) -> crate::Result<()> {
let engines = self
.shared
.initialized_engines("CpuBackend::set_buffer_pool_limit_bytes")?;
let mut resources = engines
.iter()
.map(|engine| lock_engine_resources(engine, "CpuBackend::set_buffer_pool_limit_bytes"))
.collect::<crate::Result<Vec<_>>>()?;
self.shared
.buffer_limit
.store(max_retained_capacity_bytes, Ordering::Relaxed);
for resource in &mut resources {
resource
.buffers
.set_max_retained_capacity_bytes(max_retained_capacity_bytes);
}
Ok(())
}
pub fn reset_buffer_pool(&mut self) -> crate::Result<()> {
let engines = self
.shared
.initialized_engines("CpuBackend::reset_buffer_pool")?;
let mut resources = engines
.iter()
.map(|engine| lock_engine_resources(engine, "CpuBackend::reset_buffer_pool"))
.collect::<crate::Result<Vec<_>>>()?;
for resource in &mut resources {
resource.buffers.clear();
}
Ok(())
}
pub(crate) fn runtime_cache_stats(
&self,
) -> crate::Result<tenferro_runtime::runtime::CacheStats> {
let resources = lock_engine_resources(&self.engine, "CpuBackend::runtime_cache_stats")?;
let buffers = resources.buffers.cache_stats();
let gemm = tenferro_tensor::RuntimeCacheControl::stats(&resources.gemm_analysis_cache);
let indexed = resources.indexed_plan_cache.stats();
Ok(tenferro_runtime::runtime::CacheStats {
entries: buffers
.entries
.saturating_add(gemm.entries)
.saturating_add(indexed.entries),
retained_bytes: buffers
.retained_bytes
.saturating_add(gemm.retained_bytes)
.saturating_add(indexed.retained_bytes),
hits: indexed.hits,
misses: indexed.misses,
evictions: indexed.evictions,
clears: indexed.clears,
})
}
pub(crate) fn clear_runtime_caches(&self) -> crate::Result<()> {
let mut resources =
lock_engine_resources(&self.engine, "CpuBackend::clear_runtime_caches")?;
resources.buffers.clear();
tenferro_tensor::RuntimeCacheControl::clear(&mut resources.gemm_analysis_cache);
resources.indexed_plan_cache.clear();
Ok(())
}
pub fn install<R: Send>(&self, op: impl FnOnce() -> R + Send) -> R {
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
match entry.enter(ParallelMode::Sequential, |_| op()) {
Ok(result) => result,
Err(error) => panic!("CpuBackend::install executor failed: {error}"),
}
}
fn try_install<R: Send>(
&self,
op: impl FnOnce() -> crate::Result<R> + Send,
) -> crate::Result<R> {
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
let mode = entry.preferred_engine_mode();
entry
.enter(mode, |context| context.with_native_parallelism(op))
.map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
}
fn try_install_with_context<R: Send>(
&self,
op: impl FnOnce(&CpuExecutionContext<'_>) -> crate::Result<R> + Send,
) -> crate::Result<R> {
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
let mode = entry.preferred_engine_mode();
entry
.enter(mode, |context| {
context.with_native_parallelism(|| op(context))
})
.map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
}
fn try_install_fresh<R: FreshCpuOutput + Send>(
&self,
op: impl FnOnce() -> crate::Result<R> + Send,
) -> crate::Result<R> {
let domain = self.engine.domain().id();
let mut output = self.try_install(op)?;
output.tag_fresh(domain);
Ok(output)
}
fn try_install_fresh_with_context<R: FreshCpuOutput + Send>(
&self,
op: impl FnOnce(&CpuExecutionContext<'_>) -> crate::Result<R> + Send,
) -> crate::Result<R> {
let domain = self.engine.domain().id();
let mut output = self.try_install_with_context(op)?;
output.tag_fresh(domain);
Ok(output)
}
fn install_with_pool_unmarked<R: Send>(
&mut self,
op: impl FnOnce(&mut BufferPool) -> crate::Result<R> + Send,
) -> crate::Result<R> {
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
let mode = entry.preferred_engine_mode();
entry
.enter(mode, |context| {
context.with_native_parallelism(|| {
self.with_execution_resources(&permit, |resources| {
let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
op(buffers.get_mut())
})
})
})
.map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
}
fn install_with_pool_context_unmarked<R: Send>(
&mut self,
op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
) -> crate::Result<R> {
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
let mode = entry.preferred_engine_mode();
entry
.enter(mode, |context| {
context.with_native_parallelism(|| {
self.with_execution_resources(&permit, |resources| {
let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
op(context, buffers.get_mut())
})
})
})
.map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
}
fn install_with_indexed_pool_context_unmarked<R: Send>(
&mut self,
op: impl FnOnce(
&CpuExecutionContext<'_>,
&mut BufferPool,
&mut IndexedPlanCache,
) -> crate::Result<R>
+ Send,
) -> crate::Result<R> {
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
let mode = entry.preferred_engine_mode();
entry
.enter(mode, |context| {
context.with_native_parallelism(|| {
self.with_execution_resources(&permit, |resources| {
let EngineResources {
buffers,
indexed_plan_cache,
..
} = resources;
let mut buffers = BufferPoolLoan::new(buffers);
op(context, buffers.get_mut(), indexed_plan_cache)
})
})
})
.map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
}
fn install_with_pool<R: FreshCpuOutput + Send>(
&mut self,
op: impl FnOnce(&mut BufferPool) -> crate::Result<R> + Send,
) -> crate::Result<R> {
let domain = self.engine.domain().id();
let mut output = self.install_with_pool_unmarked(op)?;
output.tag_fresh(domain);
Ok(output)
}
fn install_with_pool_context<R: FreshCpuOutput + Send>(
&mut self,
op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
) -> crate::Result<R> {
let domain = self.engine.domain().id();
let mut output = self.install_with_pool_context_unmarked(op)?;
output.tag_fresh(domain);
Ok(output)
}
fn install_with_indexed_pool_context<R: FreshCpuOutput + Send>(
&mut self,
op: impl FnOnce(
&CpuExecutionContext<'_>,
&mut BufferPool,
&mut IndexedPlanCache,
) -> crate::Result<R>
+ Send,
) -> crate::Result<R> {
let domain = self.engine.domain().id();
let mut output = self.install_with_indexed_pool_context_unmarked(op)?;
output.tag_fresh(domain);
Ok(output)
}
#[doc(hidden)]
pub fn with_linalg_pool<R: Send>(
&mut self,
op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
) -> crate::Result<R> {
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
let mode = entry.preferred_linalg_mode(self.kind());
entry
.enter(mode, |context| {
context.with_native_parallelism(|| {
self.with_execution_resources(&permit, |resources| {
let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
op(context, buffers.get_mut())
})
})
})
.map_err(|error| crate::Error::backend_source("CPU linalg execution", error))?
}
fn with_execution_resources<R>(
&self,
permit: &ResourcePermit,
op: impl FnOnce(&mut EngineResources) -> R,
) -> R {
if permit.is_reentrant() {
let mut resources =
EngineResources::new(self.shared.buffer_limit.load(Ordering::Relaxed));
return op(&mut resources);
}
let mut resources = self
.engine
.resources
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
op(&mut resources)
}
fn acquire_execution_permit(&self, owner: ResourceOwner) -> ResourcePermit {
match &self.resolved {
ResolvedCpuExecution::Managed(placement)
| ResolvedCpuExecution::ExternalManaged(placement) => self
.shared
.arbiter
.acquire_recovering(placement.cpus().clone(), owner),
ResolvedCpuExecution::ExternalCallerManaged => {
let active = self
.engine
.domain()
.caller_managed_active()
.unwrap_or_else(|| {
unreachable!("caller-managed execution needs a local admission guard")
});
ResourcePermit::caller_managed(active, owner)
}
ResolvedCpuExecution::Compatibility => self
.shared
.arbiter
.acquire_recovering(self.shared.topology.allowed_cpus().clone(), owner),
ResolvedCpuExecution::ProviderDefaultExclusive => self
.shared
.arbiter
.acquire_provider_exclusive_recovering(owner),
}
}
#[cfg(test)]
fn try_acquire_execution_permit_for_test(
&self,
) -> Result<Option<ResourcePermit>, crate::arbiter::ResourceArbiterError> {
match &self.resolved {
ResolvedCpuExecution::Managed(placement)
| ResolvedCpuExecution::ExternalManaged(placement) => {
self.shared.arbiter.try_acquire(placement.cpus().clone())
}
ResolvedCpuExecution::ExternalCallerManaged => Ok(None),
ResolvedCpuExecution::Compatibility => self
.shared
.arbiter
.try_acquire(self.shared.topology.allowed_cpus().clone()),
ResolvedCpuExecution::ProviderDefaultExclusive => {
self.shared.arbiter.try_acquire_provider_exclusive()
}
}
}
}
impl BackendSession for CpuBackend {
fn vdot_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.run_backend_session_cached(None, move |session| session.vdot_read(lhs, rhs))
}
fn norm_squared_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.run_backend_session_cached(None, move |session| session.norm_squared_read(input))
}
fn axpby_read_into_accum(
&mut self,
alpha: ContractionScalar,
x: TensorRead<'_>,
beta: ContractionScalar,
y: TensorWrite<'_>,
) -> crate::Result<()> {
self.run_backend_session_cached(None, move |session| {
session.axpby_read_into_accum(alpha, x, beta, y)
})
}
fn session_type_id(&self) -> TypeId {
TypeId::of::<CpuBackendSessionMarker>()
}
unsafe fn session_data_mut(&mut self) -> *mut () {
self as *mut Self as *mut ()
}
}
impl BackendRuntimeCache for CpuBackend {
type RuntimeCache = gemm::GemmAnalysisCache;
}
impl TensorElementwise for CpuBackend {
fn elementwise_read_into(
&mut self,
op: ElementwiseReadOp,
inputs: &[TensorRead<'_>],
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.install_with_pool_context_unmarked(|context, buffers| {
let exec_context = context.strided_exec_context();
tenferro_tensor::backend::elementwise_read_into_with_context(
op,
inputs,
out,
&exec_context,
|inputs, out| elementwise_read_into_fallback_with_pool(buffers, op, inputs, out),
)
})
}
fn add(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::add_with_pool(buffers, lhs, rhs))
}
fn add_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::add_read_with_pool(buffers, lhs, rhs))
}
fn sub(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::sub_with_pool(buffers, lhs, rhs))
}
fn sub_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::sub_read_with_pool(buffers, lhs, rhs))
}
fn mul(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::mul_with_pool(buffers, lhs, rhs))
}
fn mul_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::mul_read_with_pool(buffers, lhs, rhs))
}
fn neg(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::neg_with_pool(buffers, input))
}
fn neg_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::neg_read_with_pool(buffers, input))
}
fn conj(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::conj_with_pool(buffers, input))
}
fn conj_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::conj_read_with_pool(buffers, input))
}
fn div(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::div_with_pool(buffers, lhs, rhs))
}
fn div_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::div_read_with_pool(buffers, lhs, rhs))
}
fn rem(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::rem_with_pool(buffers, lhs, rhs))
}
fn rem_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::rem_read_with_pool(buffers, lhs, rhs))
}
fn abs(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::abs_with_pool(buffers, input))
}
fn abs_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::abs_read_with_pool(buffers, input))
}
fn sign(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::sign_with_pool(buffers, input))
}
fn sign_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::sign_read_with_pool(buffers, input))
}
fn maximum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::maximum_with_pool(buffers, lhs, rhs))
}
fn maximum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::maximum_read_with_pool(buffers, lhs, rhs))
}
fn minimum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::minimum_with_pool(buffers, lhs, rhs))
}
fn minimum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::minimum_read_with_pool(buffers, lhs, rhs))
}
fn compare(&mut self, lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::compare_with_pool(buffers, lhs, rhs, dir))
}
fn compare_read(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
dir: &CompareDir,
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
elementwise::compare_read_with_pool(buffers, lhs, rhs, dir)
})
}
fn select(
&mut self,
pred: &Tensor,
on_true: &Tensor,
on_false: &Tensor,
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
elementwise::select_with_pool(buffers, pred, on_true, on_false)
})
}
fn select_read(
&mut self,
pred: TensorRead<'_>,
on_true: TensorRead<'_>,
on_false: TensorRead<'_>,
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
elementwise::select_read_with_pool(buffers, pred, on_true, on_false)
})
}
fn clamp(&mut self, input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| elementwise::clamp_with_pool(buffers, input, lower, upper))
}
fn clamp_read(
&mut self,
input: TensorRead<'_>,
lower: TensorRead<'_>,
upper: TensorRead<'_>,
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
elementwise::clamp_read_with_pool(buffers, input, lower, upper)
})
}
}
impl TensorAnalytic for CpuBackend {
fn exp(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::exp_with_pool(buffers, input))
}
fn exp_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::exp_read_with_pool(buffers, input))
}
fn log(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::log_with_pool(buffers, input))
}
fn log_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::log_read_with_pool(buffers, input))
}
fn sin(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::sin_with_pool(buffers, input))
}
fn sin_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::sin_read_with_pool(buffers, input))
}
fn cos(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::cos_with_pool(buffers, input))
}
fn cos_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::cos_read_with_pool(buffers, input))
}
fn tanh(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::tanh_with_pool(buffers, input))
}
fn tanh_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::tanh_read_with_pool(buffers, input))
}
fn sqrt(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::sqrt_with_pool(buffers, input))
}
fn sqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::sqrt_read_with_pool(buffers, input))
}
fn rsqrt(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::rsqrt_with_pool(buffers, input))
}
fn rsqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::rsqrt_read_with_pool(buffers, input))
}
fn pow(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::pow_with_pool(buffers, lhs, rhs))
}
fn pow_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::pow_read_with_pool(buffers, lhs, rhs))
}
fn expm1(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::expm1_with_pool(buffers, input))
}
fn expm1_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::expm1_read_with_pool(buffers, input))
}
fn log1p(&mut self, input: &Tensor) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::log1p_with_pool(buffers, input))
}
fn log1p_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| analytic::log1p_read_with_pool(buffers, input))
}
}
impl TensorStructural for CpuBackend {
fn to_contiguous_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
materialize_tensor_read(buffers, "CpuBackend::to_contiguous_read", input)
})
}
fn copy_read_into(&mut self, src: TensorRead<'_>, dst: TensorWrite<'_>) -> crate::Result<()> {
self.try_install(|| copy_tensor_read_into("CpuBackend::copy_read_into", src, dst))
}
fn transpose(&mut self, input: &Tensor, perm: &[usize]) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| structural::transpose_with_pool(buffers, input, perm))
}
fn transpose_read(&mut self, input: TensorRead<'_>, perm: &[usize]) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| structural::transpose_read_with_pool(buffers, input, perm))
}
fn reshape(&mut self, input: &Tensor, shape: &[usize]) -> crate::Result<Tensor> {
structural::reshape(input, shape)
}
fn reshape_read(&mut self, input: TensorRead<'_>, shape: &[usize]) -> crate::Result<Tensor> {
match &input {
TensorRead::Tensor(tensor) => structural::reshape(tensor, shape),
TensorRead::View(_) => self.install_with_pool(|buffers| {
structural::reshape_read_with_pool(buffers, input, shape)
}),
}
}
fn broadcast_in_dim(
&mut self,
input: &Tensor,
shape: &[usize],
dims: &[usize],
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
structural::broadcast_in_dim_with_pool(buffers, input, shape, dims)
})
}
fn broadcast_in_dim_read(
&mut self,
input: TensorRead<'_>,
shape: &[usize],
dims: &[usize],
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
structural::broadcast_in_dim_read_with_pool(buffers, input, shape, dims)
})
}
fn cast(&mut self, input: &Tensor, to: crate::DType) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| structural::cast_with_pool(buffers, input, to))
}
fn extract_diagonal(
&mut self,
input: &Tensor,
axis_a: usize,
axis_b: usize,
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
structural::extract_diagonal_with_pool(buffers, input, axis_a, axis_b)
})
}
fn embed_diagonal(
&mut self,
input: &Tensor,
axis_a: usize,
axis_b: usize,
) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| {
structural::embed_diagonal_with_pool(buffers, input, axis_a, axis_b)
})
}
fn tril(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| structural::tril_with_pool(buffers, input, k))
}
fn triu(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| structural::triu_with_pool(buffers, input, k))
}
}
impl TensorReduction for CpuBackend {
fn reduce_sum(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
self.try_install_fresh_with_context(|context| {
let exec_context = context.strided_exec_context();
reduction::reduce_sum(input, axes, &exec_context)
})
}
fn reduce_sum_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
reduction::reduce_sum_read(buffers, input, axes, &exec_context)
})
}
fn reduce_sum_squares_read(
&mut self,
input: TensorRead<'_>,
axes: &[usize],
) -> crate::Result<Tensor> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
reduction::reduce_sum_squares_read(buffers, input, axes, &exec_context)
})
}
fn reduce_prod(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
self.try_install_fresh_with_context(|context| {
let exec_context = context.strided_exec_context();
reduction::reduce_prod(input, axes, &exec_context)
})
}
fn reduce_prod_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
reduction::reduce_prod_read(buffers, input, axes, &exec_context)
})
}
fn reduce_max(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
self.try_install_fresh(|| reduction::reduce_max(input, axes))
}
fn reduce_max_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| reduction::reduce_max_read(buffers, input, axes))
}
fn reduce_min(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
self.try_install_fresh(|| reduction::reduce_min(input, axes))
}
fn reduce_min_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
self.install_with_pool(|buffers| reduction::reduce_min_read(buffers, input, axes))
}
}
impl TensorDot for CpuBackend {
fn dot_general(
&mut self,
lhs: &Tensor,
rhs: &Tensor,
config: &DotGeneralConfig,
) -> crate::Result<Tensor> {
self.run_backend_session_cached(None, move |session| session.dot_general(lhs, rhs, config))
}
fn dot_general_read(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
) -> crate::Result<Tensor> {
self.run_backend_session_cached(None, move |session| {
session.dot_general_read(lhs, rhs, config)
})
}
fn dot_general_read_into(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.run_backend_session_cached(None, move |session| {
session.dot_general_read_into(lhs, rhs, config, out)
})
}
fn dot_general_read_into_accum(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
accumulation: DotGeneralAccumulation,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.run_backend_session_cached(None, move |session| {
session.dot_general_read_into_accum(lhs, rhs, config, accumulation, out)
})
}
fn dot_general_with_conj(
&mut self,
lhs: &Tensor,
rhs: &Tensor,
config: &DotGeneralConfig,
lhs_conj: bool,
rhs_conj: bool,
) -> crate::Result<Tensor> {
self.run_backend_session_cached(None, move |session| {
session.dot_general_with_conj(lhs, rhs, config, lhs_conj, rhs_conj)
})
}
}
impl BackendCachedDot for CpuBackend {
fn dot_general_cached(
&mut self,
cache: &mut Self::RuntimeCache,
cache_slot: Option<usize>,
lhs: &Tensor,
rhs: &Tensor,
config: &DotGeneralConfig,
) -> crate::Result<Tensor> {
self.run_backend_session_cached(Some(cache), move |session| {
session.dot_general_cached(cache_slot, lhs, rhs, config)
})
}
fn dot_general_with_conj_cached(
&mut self,
cache: &mut Self::RuntimeCache,
cache_slot: Option<usize>,
lhs: &Tensor,
rhs: &Tensor,
config: &DotGeneralConfig,
lhs_conj: bool,
rhs_conj: bool,
) -> crate::Result<Tensor> {
self.run_backend_session_cached(Some(cache), move |session| {
session.dot_general_with_conj_cached(cache_slot, lhs, rhs, config, lhs_conj, rhs_conj)
})
}
fn dot_general_read_into_accum_cached(
&mut self,
cache: &mut Self::RuntimeCache,
cache_slot: Option<usize>,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
accumulation: DotGeneralAccumulation,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.run_backend_session_cached(Some(cache), move |session| {
session.dot_general_read_into_accum_cached(
cache_slot,
lhs,
rhs,
config,
accumulation,
out,
)
})
}
fn grouped_gemm_cached(
&mut self,
cache: &mut Self::RuntimeCache,
cache_slot: Option<usize>,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &GroupedGemmConfig<'_>,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.run_backend_session_cached(Some(cache), move |session| {
session.grouped_gemm_cached(cache_slot, lhs, rhs, config, out)
})
}
}
impl TensorIndexing for CpuBackend {
fn gather(
&mut self,
operand: &Tensor,
start_indices: &Tensor,
config: &GatherConfig,
) -> crate::Result<Tensor> {
self.install_with_indexed_pool_context(|context, buffers, cache| {
let exec_context = context.strided_exec_context();
indexing::gather_with_pool(
buffers,
cache,
&exec_context,
operand,
start_indices,
config,
)
})
}
fn scatter(
&mut self,
operand: &Tensor,
scatter_indices: &Tensor,
updates: &Tensor,
config: &ScatterConfig,
) -> crate::Result<Tensor> {
self.install_with_indexed_pool_context(|context, buffers, cache| {
let exec_context = context.strided_exec_context();
indexing::scatter_with_pool(
buffers,
cache,
&exec_context,
operand,
scatter_indices,
updates,
config,
)
})
}
fn slice(&mut self, input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
indexing::try_slice_with_pool(buffers, &exec_context, input, config)
})
}
fn dynamic_slice(
&mut self,
input: &Tensor,
starts: &Tensor,
slice_sizes: &[usize],
) -> crate::Result<Tensor> {
self.install_with_indexed_pool_context(|context, buffers, cache| {
let exec_context = context.strided_exec_context();
indexing::dynamic_slice_with_pool(
buffers,
cache,
&exec_context,
input,
starts,
slice_sizes,
)
})
}
fn dynamic_update_slice(
&mut self,
operand: &Tensor,
update: &Tensor,
starts: &Tensor,
) -> crate::Result<Tensor> {
self.install_with_indexed_pool_context(|context, buffers, cache| {
let exec_context = context.strided_exec_context();
indexing::dynamic_update_slice_with_pool(
buffers,
cache,
&exec_context,
operand,
update,
starts,
)
})
}
fn pad(&mut self, input: &Tensor, config: &PadConfig) -> crate::Result<Tensor> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
indexing::try_pad_with_pool(buffers, &exec_context, input, config)
})
}
fn concatenate(&mut self, inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
indexing::try_concatenate_with_pool(buffers, &exec_context, inputs, axis)
})
}
fn reverse(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
indexing::reverse_with_pool(buffers, &exec_context, input, axes)
})
}
}
impl CpuBackend {
pub fn with_allocation_domain(mut self, domain: Arc<dyn SharedTensorAllocationDomain>) -> Self {
self.allocation_domain = Some(domain);
self.runtime_identity = CpuRuntimeIdentity::fresh();
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()
}
fn run_backend_session_cached<R: Send>(
&mut self,
cache: Option<&mut gemm::GemmAnalysisCache>,
f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
) -> R {
let providers = self.provider_bundle.clone();
let owner = inherited_or_new_execution_owner();
let permit = self.acquire_execution_permit(owner);
let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
let enter_managed_session = entry.supports_infallible_session_entry()
&& !matches!(
&self.resolved,
ResolvedCpuExecution::ProviderDefaultExclusive
);
let run = |entered| {
self.with_execution_resources(&permit, |resources| {
let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
let cache = cache.unwrap_or(&mut resources.gemm_analysis_cache);
let session_started = Instant::now();
let mut session = CpuExecSession {
entry,
entered,
buffers: buffers.get_mut(),
gemm_analysis_cache: cache,
indexed_plan_cache: &mut resources.indexed_plan_cache,
providers: &providers,
backend_kind: self.kind(),
allocation_domain: self.allocation_domain.as_ref(),
};
record_cpu_session_profile(
"with_backend_session_cached.session_construct",
session_started.elapsed(),
);
let exec_started = Instant::now();
let result = f(&mut session);
record_cpu_session_profile(
"with_backend_session_cached.exec_body",
exec_started.elapsed(),
);
result
})
};
if enter_managed_session {
entry.enter_managed_session(|context| run(Some(context)))
} else {
with_execution_owner(owner, || run(None))
}
}
}
impl BackendSessionHost for CpuBackend {
fn with_backend_session<R: Send>(
&mut self,
f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
) -> R {
self.run_backend_session_cached(None, f)
}
fn with_backend_session_cached<R: Send>(
&mut self,
cache: &mut Self::RuntimeCache,
f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
) -> R {
if !cpu_session_profile_enabled() {
return self.run_backend_session_cached(Some(cache), f);
}
let total_started = Instant::now();
let result =
profile_cpu_session_section("with_backend_session_cached.exec_session", || {
self.run_backend_session_cached(Some(cache), f)
});
record_cpu_session_profile("with_backend_session_cached.total", total_started.elapsed());
maybe_print_cpu_session_profile();
result
}
}
impl TensorBuffer for CpuBackend {
fn reclaim_buffer(&mut self, tensor: Tensor) {
let owner = inherited_or_new_execution_owner();
with_execution_owner(owner, || {
let permit = self.acquire_execution_permit(owner);
self.with_execution_resources(&permit, |resources| {
let buffers = &mut resources.buffers;
match tensor {
Tensor::F32(t) => reclaim_typed(buffers, t),
Tensor::F64(t) => reclaim_typed(buffers, t),
Tensor::I32(t) => reclaim_typed(buffers, t),
Tensor::I64(t) => reclaim_typed(buffers, t),
Tensor::Bool(t) => reclaim_typed(buffers, t),
Tensor::C32(t) => reclaim_typed(buffers, t),
Tensor::C64(t) => reclaim_typed(buffers, t),
}
})
})
}
}
impl<T, R> TensorViewCanonicalization<T, R> for CpuBackend
where
T: TensorScalar + PoolScalar,
R: TensorRank,
R::Shape: Send + Sync,
R::Strides: Send + Sync,
{
fn to_contiguous(
&mut self,
view: &TypedTensorView<'_, T, R>,
) -> crate::Result<TypedTensor<T, R>> {
self.install_with_pool(|buffers| {
structural::typed_materialize_view_with_pool(buffers, view, "CpuBackend::to_contiguous")
})
}
fn copy_into(
&mut self,
src: &TypedTensorView<'_, T, R>,
dst: &mut TypedTensorViewMut<'_, T, R>,
) -> crate::Result<()> {
self.try_install(|| structural::typed_copy_view_into(src, dst, "CpuBackend::copy_into"))
}
}
impl TensorFusion for CpuBackend {
fn execute_elementwise_fusion(
&mut self,
inputs: &[&Tensor],
plan: &ElementwiseFusionPlan,
) -> crate::Result<Option<Vec<Tensor>>> {
self.install_with_pool_context(|context, buffers| {
let exec_context = context.strided_exec_context();
elementwise::elementwise_fusion_with_pool(buffers, &exec_context, inputs, plan)
})
}
fn execute_broadcast_multiply(
&mut self,
lhs: TensorRead<'_>,
lhs_shape: &[usize],
lhs_dims: &[usize],
rhs: TensorRead<'_>,
rhs_shape: &[usize],
rhs_dims: &[usize],
) -> crate::Result<Option<Tensor>> {
self.install_with_pool(|buffers| {
elementwise::broadcast_multiply_read_with_pool(
buffers, lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims,
)
})
}
fn execute_broadcast_multiply_value(
&mut self,
lhs: TensorRead<'_>,
lhs_shape: &[usize],
lhs_dims: &[usize],
rhs: TensorRead<'_>,
rhs_shape: &[usize],
rhs_dims: &[usize],
) -> crate::Result<Option<TensorValue>> {
let domain = self.engine.domain().id();
self.install_with_pool_unmarked(|buffers| {
elementwise::broadcast_multiply_value_with_pool_and_tag(
buffers,
lhs,
lhs_shape,
lhs_dims,
rhs,
rhs_shape,
rhs_dims,
|tensor| tag_fresh_output(tensor, domain),
)
})
}
}
impl TensorDeviceTransfer for CpuBackend {
fn download_to_host(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
if tensor.backend_family().is_some() {
return Err(crate::Error::runtime_state(
"CpuBackend::download_to_host",
"CPU backend received a backend buffer; download the tensor to host with its owning backend before CPU execution",
));
}
tensor.tensor_view().duplicate()
}
fn upload_host_tensor(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
if tensor.backend_family().is_some() {
return Err(crate::Error::runtime_state(
"CpuBackend::upload_host_tensor",
"CPU backend upload_host_tensor expects a host tensor; download backend buffers to host before CPU execution",
));
}
tensor.tensor_view().duplicate()
}
}
impl TensorBackend for CpuBackend {}
pub(crate) fn reclaim_typed<T: PoolScalar>(pool: &mut BufferPool, typed: TypedTensor<T>) {
if typed.backend_buffer().is_some() {
return;
}
if let Ok(data) = typed.into_host_vec() {
T::pool_release(pool, data);
}
}
impl Default for CpuBackend {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests;