1use num_complex::{Complex32, Complex64};
2use std::cmp::Reverse;
3use std::collections::{BTreeMap, BTreeSet, HashMap};
4use std::env;
5use std::fmt;
6use std::sync::atomic::{AtomicUsize, Ordering};
7use std::sync::{Arc, Mutex, OnceLock};
8use std::thread;
9use std::time::{Duration, Instant};
10use strided_kernel::ExecContext;
11use tenferro_tensor::DType;
12
13use crate::arbiter::{with_execution_owner, ResourceArbiter, ResourceOwner, ResourcePermit};
14use crate::buffer_pool::{BufferPool, BufferPoolStats, PoolScalar};
15use crate::dot_runtime::{
16 CpuProviderBundle, CpuProviderBundleInstallError, CpuProviderDomainContract,
17};
18use crate::engine::{CpuEngine, EngineResources};
19use crate::indexed_plan_cache::{IndexedPlanCacheLimits, DEFAULT_INDEXED_PLAN_CACHE_LIMITS};
20use crate::placement::{
21 resolve_placement, resolve_placement_with_affinity, CpuEngineConstructionError,
22 ResolvedCpuExecution,
23};
24use crate::provider::{CpuOperationEntry, ParallelMode};
25use crate::{
26 discover_cpu_topology, CpuAdmissionMode, CpuDomainId, CpuDomainOwnership, CpuExecutorAffinity,
27 CpuExecutorShutdown, CpuId, CpuPlacement, CpuPlacementError, CpuSet, CpuTopology,
28 CpuTopologyError, ExternalCpuDomain, NumaNodeId, ResolvedCpuPlacement,
29};
30use crate::{CacheStats, Tensor, TensorRank, TensorRead, TensorScalar, TensorWrite, TypedTensor};
31use tenferro_tensor::{
32 AllocationDomainId, BackendRuntimeCache, BackendSession, BackendSessionHost, ElementwiseReadOp,
33 TensorBackend, TensorDeviceTransfer,
34};
35use tenferro_tensor::{SessionEntryError, SharedTensorAllocationDomain};
36
37use super::exec_session::CpuExecSession;
38use super::{copy_tensor_read_into, elementwise, gemm, CpuContext};
39
40pub(crate) fn tag_fresh_output(output: &mut Tensor, domain: CpuDomainId) {
41 match output.dtype() {
42 DType::F32 => tag_fresh_typed::<f32>(output, domain),
43 DType::F64 => tag_fresh_typed::<f64>(output, domain),
44 DType::I32 => tag_fresh_typed::<i32>(output, domain),
45 DType::I64 => tag_fresh_typed::<i64>(output, domain),
46 DType::Bool => tag_fresh_typed::<bool>(output, domain),
47 DType::C32 => tag_fresh_typed::<Complex32>(output, domain),
48 DType::C64 => tag_fresh_typed::<Complex64>(output, domain),
49 DType::External(_) => {}
52 }
53}
54
55fn tag_fresh_typed<T: TensorScalar>(output: &mut Tensor, domain: CpuDomainId) {
57 if let Some(tensor) = output.as_typed_mut::<T>() {
58 tensor.set_cpu_affinity(Some(domain));
59 }
60}
61
62pub(crate) fn elementwise_read_into_fallback_with_pool(
63 buffers: &mut BufferPool,
64 ctx: &ExecContext,
65 op: ElementwiseReadOp,
66 inputs: &[TensorRead<'_>],
67 out: TensorWrite<'_>,
68) -> crate::Result<()> {
69 let result = match op {
70 ElementwiseReadOp::Add => {
71 elementwise::add_read_with_pool(buffers, ctx, inputs[0].clone(), inputs[1].clone())?
72 }
73 ElementwiseReadOp::Subtract => {
74 elementwise::sub_read_with_pool(buffers, ctx, inputs[0].clone(), inputs[1].clone())?
75 }
76 ElementwiseReadOp::Multiply => {
77 elementwise::mul_read_with_pool(buffers, ctx, inputs[0].clone(), inputs[1].clone())?
78 }
79 ElementwiseReadOp::Negate => {
80 elementwise::neg_read_with_pool(buffers, ctx, inputs[0].clone())?
81 }
82 ElementwiseReadOp::Conj => {
83 elementwise::conj_read_with_pool(buffers, ctx, inputs[0].clone())?
84 }
85 ElementwiseReadOp::Divide => {
86 elementwise::div_read_with_pool(buffers, ctx, inputs[0].clone(), inputs[1].clone())?
87 }
88 _ => {
89 return Err(crate::Error::unsupported(
90 "CpuBackend::elementwise_read_into",
91 format!("CPU backend does not implement {op:?}"),
92 ))
93 }
94 };
95 let copied = copy_tensor_read_into(
96 "CpuBackend::elementwise_read_into",
97 TensorRead::from_tensor(&result),
98 out,
99 );
100 reclaim_tensor(buffers, result);
102 copied
103}
104
105pub(crate) trait FreshCpuOutput {
106 fn tag_fresh(&mut self, domain: CpuDomainId);
107}
108
109impl FreshCpuOutput for Tensor {
110 fn tag_fresh(&mut self, domain: CpuDomainId) {
111 tag_fresh_output(self, domain);
112 }
113}
114
115impl<T, R: TensorRank> FreshCpuOutput for TypedTensor<T, R> {
116 fn tag_fresh(&mut self, domain: CpuDomainId) {
117 self.set_cpu_affinity(Some(domain));
118 }
119}
120
121impl<T: FreshCpuOutput> FreshCpuOutput for Option<T> {
122 fn tag_fresh(&mut self, domain: CpuDomainId) {
123 if let Some(output) = self {
124 output.tag_fresh(domain);
125 }
126 }
127}
128
129impl<T: FreshCpuOutput> FreshCpuOutput for Vec<T> {
130 fn tag_fresh(&mut self, domain: CpuDomainId) {
131 for output in self {
132 output.tag_fresh(domain);
133 }
134 }
135}
136
137#[derive(Debug, Default, Clone)]
138struct CpuSessionProfileEntry {
139 calls: usize,
140 total_time: Duration,
141}
142
143fn cpu_session_profile_enabled() -> bool {
144 static ENABLED: OnceLock<bool> = OnceLock::new();
145 *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_CPU_SESSION").is_ok())
146}
147
148fn cpu_session_profile_print_every() -> Option<usize> {
149 static PRINT_EVERY: OnceLock<Option<usize>> = OnceLock::new();
150 *PRINT_EVERY.get_or_init(|| {
151 env::var("TENFERRO_PROFILE_CPU_SESSION_PRINT_EVERY")
152 .ok()
153 .and_then(|value| value.parse::<usize>().ok())
154 .filter(|&value| value > 0)
155 })
156}
157
158fn cpu_session_profile_state() -> &'static Mutex<HashMap<&'static str, CpuSessionProfileEntry>> {
159 static STATE: OnceLock<Mutex<HashMap<&'static str, CpuSessionProfileEntry>>> = OnceLock::new();
160 STATE.get_or_init(|| Mutex::new(HashMap::new()))
161}
162
163fn record_cpu_session_profile(section: &'static str, elapsed: Duration) {
164 if !cpu_session_profile_enabled() {
165 return;
166 }
167 let Ok(mut state) = cpu_session_profile_state().lock() else {
168 return;
169 };
170 let entry = state.entry(section).or_default();
171 entry.calls += 1;
172 entry.total_time += elapsed;
173}
174
175fn profile_cpu_session_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
176 if !cpu_session_profile_enabled() {
177 return f();
178 }
179 let started = Instant::now();
180 let result = f();
181 record_cpu_session_profile(section, started.elapsed());
182 result
183}
184
185fn maybe_print_cpu_session_profile() {
186 let Some(print_every) = cpu_session_profile_print_every() else {
187 return;
188 };
189 let should_print = {
190 let Ok(state) = cpu_session_profile_state().lock() else {
191 return;
192 };
193 state
194 .get("with_backend_session_cached.total")
195 .is_some_and(|entry| entry.calls % print_every == 0)
196 };
197 if !should_print {
198 return;
199 }
200 let mut entries = {
201 let Ok(mut state) = cpu_session_profile_state().lock() else {
202 return;
203 };
204 let entries = state
205 .iter()
206 .map(|(section, entry)| (*section, entry.clone()))
207 .collect::<Vec<_>>();
208 state.clear();
209 entries
210 };
211 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
212 eprintln!("=== tenferro CPU session profile ===");
213 for (section, entry) in entries {
214 eprintln!(
215 "{section}: calls={} total={:.6}ms per_call={:.3}us",
216 entry.calls,
217 entry.total_time.as_secs_f64() * 1.0e3,
218 entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64,
219 );
220 }
221}
222
223pub(crate) const CPU_BACKEND: &str = "CpuBackend";
225
226struct BufferPoolLoan<'a> {
227 buffers: &'a mut BufferPool,
228}
229
230impl<'a> BufferPoolLoan<'a> {
231 fn new(buffers: &'a mut BufferPool) -> Self {
232 Self { buffers }
233 }
234
235 fn get_mut(&mut self) -> &mut BufferPool {
236 self.buffers
237 }
238}
239
240impl Drop for BufferPoolLoan<'_> {
241 fn drop(&mut self) {
242 if thread::panicking() {
243 self.buffers.replenish_in_flight_retained();
244 } else {
245 self.buffers.clear_in_flight_retained();
246 }
247 }
248}
249
250#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
265pub enum CpuBackendKind {
266 Faer,
268 Blas,
270}
271
272impl CpuBackendKind {
273 pub fn default_compiled() -> Self {
287 #[cfg(feature = "cpu-blas")]
288 {
289 Self::Blas
290 }
291 #[cfg(all(not(feature = "cpu-blas"), feature = "cpu-faer"))]
292 {
293 Self::Faer
294 }
295 }
296
297 #[allow(dead_code)]
300 pub(crate) fn name(self) -> &'static str {
301 match self {
302 Self::Faer => "faer",
303 Self::Blas => "blas",
304 }
305 }
306}
307
308#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
326pub enum CpuExecutionMode {
327 Managed,
329 ExternalManaged,
331 CallerManaged,
333 ProviderDefaultExclusive,
335 Compatibility,
337}
338
339#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
350pub enum ExternalCpuDomainRegistryError {
351 #[error("externally managed CPU registry must contain at least one domain")]
353 EmptyRegistry,
354 #[error("CPU domain ID {id:?} is registered more than once")]
356 DuplicateDomainId {
357 id: CpuDomainId,
359 },
360 #[error("CPU placement {placement:?} is registered more than once")]
362 DuplicatePlacementIdentity {
363 placement: CpuPlacement,
365 },
366 #[error("CPU domain {domain:?} declares process-disallowed CPU {cpu}")]
368 CpuOutsideAllowedSet {
369 domain: CpuDomainId,
371 cpu: CpuId,
373 },
374 #[error("default CPU domain {default_domain:?} is not registered")]
376 MissingDefaultDomain {
377 default_domain: CpuDomainId,
379 },
380 #[error(
382 "exact all-allowed CPU domain {domain:?} declares {declared:?}, but the process allows {allowed:?}"
383 )]
384 ExactAllAllowedMismatch {
385 domain: CpuDomainId,
387 declared: CpuSet,
389 allowed: CpuSet,
391 },
392}
393
394#[derive(Debug, thiserror::Error)]
409pub enum CpuBackendError {
410 #[error(transparent)]
412 Tensor(#[from] crate::Error),
413 #[error("{op}: {source}")]
415 Placement {
416 op: &'static str,
418 #[source]
420 source: CpuPlacementError,
421 },
422 #[error(transparent)]
424 ExternalRegistry(#[from] ExternalCpuDomainRegistryError),
425}
426
427impl CpuBackendError {
428 fn placement(op: &'static str, source: CpuPlacementError) -> Self {
429 Self::Placement { op, source }
430 }
431
432 pub fn placement_error(&self) -> Option<&CpuPlacementError> {
445 match self {
446 Self::Tensor(_) => None,
447 Self::Placement { source, .. } => Some(source),
448 Self::ExternalRegistry(_) => None,
449 }
450 }
451}
452
453impl From<CpuBackendError> for crate::Error {
454 fn from(error: CpuBackendError) -> Self {
455 match error {
456 CpuBackendError::Tensor(error) => error,
457 CpuBackendError::ExternalRegistry(source) => Self::extension(
458 "CpuBackend::from_external_managed_domains",
459 "cpu",
460 crate::ErrorKind::Validation(crate::ValidationKind::InvalidArgument),
461 source,
462 ),
463 CpuBackendError::Placement { op, source } => match source {
464 CpuPlacementError::TopologyDiscovery { .. }
465 | CpuPlacementError::ManagedAffinityUnavailable { .. }
466 | CpuPlacementError::NumaDiscoveryUnavailable { .. }
467 | CpuPlacementError::UnknownNumaNode { .. }
468 | CpuPlacementError::UnregisteredExternalPlacement { .. }
469 | CpuPlacementError::UnregisteredExternalDomain { .. } => {
470 Self::runtime_state_source(op, source)
471 }
472 CpuPlacementError::ExternalProviderAffinityUnmanaged { .. } => {
473 Self::extension(op, "cpu", crate::ErrorKind::Unsupported, source)
474 }
475 CpuPlacementError::EngineConstruction { .. } => Self::backend_source(op, source),
476 CpuPlacementError::InternalState { .. } => {
477 Self::extension(op, "cpu", crate::ErrorKind::Internal, source)
478 }
479 },
480 }
481 }
482}
483
484#[derive(Clone, Debug, PartialEq, Eq)]
499pub struct CpuExecutionInfo {
500 backend_kind: CpuBackendKind,
501 execution_mode: CpuExecutionMode,
502 requested_placement: CpuPlacement,
503 resolved_placement: Option<ResolvedCpuPlacement>,
504 topology: CpuTopology,
505 domain_id: CpuDomainId,
506 domain_cpus: Option<CpuSet>,
507 worker_count: usize,
508 thread_budget: usize,
509 admission_mode: CpuAdmissionMode,
510 domain_ownership: CpuDomainOwnership,
511 executor_affinity: CpuExecutorAffinity,
512 executor_shutdown: CpuExecutorShutdown,
513 provider_diagnostic: &'static str,
514}
515
516impl CpuExecutionInfo {
517 pub fn backend_kind(&self) -> CpuBackendKind {
526 self.backend_kind
527 }
528
529 pub fn execution_mode(&self) -> CpuExecutionMode {
540 self.execution_mode
541 }
542
543 pub fn requested_placement(&self) -> CpuPlacement {
552 self.requested_placement
553 }
554
555 pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
564 self.resolved_placement.as_ref()
565 }
566
567 pub fn topology(&self) -> &CpuTopology {
576 &self.topology
577 }
578
579 pub fn domain_id(&self) -> CpuDomainId {
588 self.domain_id
589 }
590
591 pub fn domain_cpus(&self) -> Option<&CpuSet> {
602 self.domain_cpus.as_ref()
603 }
604
605 pub fn worker_count(&self) -> usize {
614 self.worker_count
615 }
616
617 pub fn thread_budget(&self) -> usize {
630 self.thread_budget
631 }
632
633 pub fn admission_mode(&self) -> CpuAdmissionMode {
644 self.admission_mode
645 }
646
647 pub fn domain_ownership(&self) -> CpuDomainOwnership {
658 self.domain_ownership
659 }
660
661 pub fn executor_affinity(&self) -> CpuExecutorAffinity {
672 self.executor_affinity
673 }
674
675 pub fn executor_shutdown(&self) -> CpuExecutorShutdown {
686 self.executor_shutdown
687 }
688
689 pub fn provider_diagnostic(&self) -> &'static str {
702 self.provider_diagnostic
703 }
704}
705
706fn provider_diagnostic(
707 kind: CpuBackendKind,
708 ownership: CpuDomainOwnership,
709 admission_mode: CpuAdmissionMode,
710) -> &'static str {
711 if ownership == CpuDomainOwnership::ExternalManaged {
712 if admission_mode == CpuAdmissionMode::CallerManaged {
713 debug_assert_eq!(kind, CpuBackendKind::Faer);
716 return "faer (caller-managed CPU executor and admission)";
717 }
718 return match kind {
719 CpuBackendKind::Faer => "faer (externally managed CPU executor)",
720 CpuBackendKind::Blas => "BLAS/LAPACK (externally managed CPU executor)",
721 };
722 }
723 match kind {
724 CpuBackendKind::Faer => "faer (tenferro-managed Rayon affinity)",
725 CpuBackendKind::Blas => {
726 #[cfg(feature = "blas-openblas")]
727 return "OpenBLAS (external worker affinity)";
728 #[cfg(feature = "blas-mkl")]
729 return "Intel MKL (external worker affinity)";
730 #[cfg(feature = "blas-accelerate")]
731 return "Apple Accelerate (external worker affinity)";
732 #[cfg(feature = "provider-inject")]
733 return "runtime-injected BLAS/LAPACK (external worker affinity)";
734 #[cfg(not(any(
735 feature = "blas-openblas",
736 feature = "blas-mkl",
737 feature = "blas-accelerate",
738 feature = "provider-inject"
739 )))]
740 return "linked BLAS/LAPACK provider (identity unknown; external worker affinity)";
741 }
742 }
743}
744
745fn ensure_cpu_backend_kind_available(kind: CpuBackendKind, op: &'static str) -> crate::Result<()> {
746 let _ = op;
747 match kind {
748 CpuBackendKind::Faer => {
749 #[cfg(feature = "cpu-faer")]
750 {
751 Ok(())
752 }
753 #[cfg(not(feature = "cpu-faer"))]
754 {
755 Err(crate::Error::invalid_argument(
756 op,
757 "configuration",
758 "CpuBackendKind::Faer requires the cpu-faer feature".to_string(),
759 ))
760 }
761 }
762 CpuBackendKind::Blas => {
763 #[cfg(feature = "cpu-blas")]
764 {
765 Ok(())
766 }
767 #[cfg(not(feature = "cpu-blas"))]
768 {
769 Err(crate::Error::invalid_argument(
770 op,
771 "configuration",
772 "CpuBackendKind::Blas requires the cpu-blas feature".to_string(),
773 ))
774 }
775 }
776 }
777}
778
779fn constructor_tensor_error(op: &'static str, error: crate::Error) -> CpuBackendError {
780 CpuBackendError::Tensor(match error {
781 crate::Error::Validation { source, .. } => crate::Error::validation(op, source),
782 error => error,
783 })
784}
785
786#[allow(dead_code)]
789pub(super) fn unavailable_cpu_backend_kind(kind: CpuBackendKind, op: &'static str) -> crate::Error {
790 crate::Error::invalid_argument(
791 op,
792 "configuration",
793 format!("CPU backend kind {} is not compiled in", kind.name()),
794 )
795}
796
797struct ManagedEngineRegistry {
798 node_engines: Mutex<BTreeMap<NumaNodeId, Arc<CpuEngine>>>,
799 node_domain_ids: BTreeMap<NumaNodeId, CpuDomainId>,
800 all_allowed: OnceLock<Arc<CpuEngine>>,
801 all_allowed_build: Mutex<()>,
802 base_engine: Arc<CpuEngine>,
803 thread_budget: usize,
804}
805
806struct ExternalEngineRegistry {
807 by_id: BTreeMap<CpuDomainId, Arc<CpuEngine>>,
808 by_node: BTreeMap<NumaNodeId, Arc<CpuEngine>>,
809 all_allowed: Option<Arc<CpuEngine>>,
810 default_domain: CpuDomainId,
811}
812
813enum CpuEngineRegistry {
814 ManagedLazy(ManagedEngineRegistry),
815 ExternalPrebuilt(ExternalEngineRegistry),
816}
817
818struct CpuBackendState {
819 topology: CpuTopology,
820 engines: CpuEngineRegistry,
821 arbiter: ResourceArbiter,
822 kind: CpuBackendKind,
823 buffer_limit: AtomicUsize,
824 indexed_plan_cache_limits: Mutex<IndexedPlanCacheLimits>,
825}
826
827impl CpuBackendState {
828 fn managed_engine_for(
829 &self,
830 placement: &ResolvedCpuPlacement,
831 requested: CpuPlacement,
832 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
833 let cache_configuration = self.indexed_plan_cache_limits.lock().map_err(|_| {
837 CpuPlacementError::InternalState {
838 requested,
839 backend: self.kind,
840 message: "CPU indexed-plan cache configuration lock is poisoned",
841 }
842 })?;
843 let cache_limits = *cache_configuration;
844 let CpuEngineRegistry::ManagedLazy(registry) = &self.engines else {
845 return Err(CpuPlacementError::InternalState {
846 requested,
847 backend: self.kind,
848 message: "managed placement requested from an external engine registry",
849 });
850 };
851 match placement {
852 ResolvedCpuPlacement::NumaNode { id, .. } => {
853 let mut engines = registry
854 .node_engines
855 .lock()
856 .unwrap_or_else(std::sync::PoisonError::into_inner);
857 if let Some(engine) = engines.get(id) {
858 return Ok(Arc::clone(engine));
859 }
860 let Some(domain_id) = registry.node_domain_ids.get(id).copied() else {
861 return Err(CpuPlacementError::InternalState {
862 requested,
863 backend: self.kind,
864 message: "managed NUMA node has no coordinator-stable domain ID",
865 });
866 };
867 let engine = Arc::new(
868 CpuEngine::new_managed(
869 domain_id,
870 placement.clone(),
871 registry.thread_budget,
872 self.buffer_limit.load(Ordering::Relaxed),
873 )
874 .map_err(|error| {
875 CpuPlacementError::EngineConstruction {
876 requested,
877 backend: self.kind,
878 source: CpuEngineConstructionError::Context(error),
879 }
880 })?,
881 );
882 self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
883 engines.insert(*id, Arc::clone(&engine));
884 Ok(engine)
885 }
886 ResolvedCpuPlacement::AllAllowed { .. } => {
887 if let Some(engine) = registry.all_allowed.get() {
888 return Ok(Arc::clone(engine));
889 }
890 let _build = registry
891 .all_allowed_build
892 .lock()
893 .unwrap_or_else(std::sync::PoisonError::into_inner);
894 if let Some(engine) = registry.all_allowed.get() {
895 return Ok(Arc::clone(engine));
896 }
897 let engine = Arc::new(
898 CpuEngine::new_managed(
899 CpuDomainId::new(0),
900 placement.clone(),
901 registry.thread_budget,
902 self.buffer_limit.load(Ordering::Relaxed),
903 )
904 .map_err(|error| {
905 CpuPlacementError::EngineConstruction {
906 requested,
907 backend: self.kind,
908 source: CpuEngineConstructionError::Context(error),
909 }
910 })?,
911 );
912 self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
913 let _ = registry.all_allowed.set(Arc::clone(&engine));
914 Ok(engine)
915 }
916 }
917 }
918
919 fn configure_new_indexed_plan_cache(
920 &self,
921 engine: &CpuEngine,
922 requested: CpuPlacement,
923 limits: IndexedPlanCacheLimits,
924 ) -> Result<(), CpuPlacementError> {
925 let mut resources =
926 engine
927 .resources
928 .lock()
929 .map_err(|_| CpuPlacementError::InternalState {
930 requested,
931 backend: self.kind,
932 message: "new CPU engine indexed-plan cache lock is poisoned",
933 })?;
934 resources.indexed_plan_cache.set_limits(limits);
935 Ok(())
936 }
937
938 fn managed_base_engine(
939 &self,
940 requested: CpuPlacement,
941 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
942 match &self.engines {
943 CpuEngineRegistry::ManagedLazy(registry) => Ok(Arc::clone(®istry.base_engine)),
944 CpuEngineRegistry::ExternalPrebuilt(_) => Err(CpuPlacementError::InternalState {
945 requested,
946 backend: self.kind,
947 message: "managed compatibility placement requested from an external registry",
948 }),
949 }
950 }
951
952 fn external_engine_for(
953 &self,
954 requested: CpuPlacement,
955 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
956 let CpuEngineRegistry::ExternalPrebuilt(registry) = &self.engines else {
957 return Err(CpuPlacementError::InternalState {
958 requested,
959 backend: self.kind,
960 message: "external placement requested from a managed engine registry",
961 });
962 };
963 let engine = match requested {
964 CpuPlacement::Auto => registry.by_id.get(®istry.default_domain),
965 CpuPlacement::NumaNode(id) => registry.by_node.get(&id),
966 CpuPlacement::AllAllowed => registry.all_allowed.as_ref(),
967 };
968 engine
969 .cloned()
970 .ok_or(CpuPlacementError::UnregisteredExternalPlacement { requested })
971 }
972
973 fn external_engine_for_id(
974 &self,
975 domain: CpuDomainId,
976 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
977 let CpuEngineRegistry::ExternalPrebuilt(registry) = &self.engines else {
978 return Err(CpuPlacementError::UnregisteredExternalDomain { domain });
979 };
980 registry
981 .by_id
982 .get(&domain)
983 .cloned()
984 .ok_or(CpuPlacementError::UnregisteredExternalDomain { domain })
985 }
986
987 fn is_external(&self) -> bool {
988 matches!(&self.engines, CpuEngineRegistry::ExternalPrebuilt(_))
989 }
990
991 fn initialized_engines(&self, op: &'static str) -> crate::Result<Vec<Arc<CpuEngine>>> {
992 let mut engines = match &self.engines {
993 CpuEngineRegistry::ManagedLazy(registry) => {
994 let mut engines = vec![Arc::clone(®istry.base_engine)];
995 if let Some(engine) = registry.all_allowed.get() {
996 engines.push(Arc::clone(engine));
997 }
998 engines.extend(
999 registry
1000 .node_engines
1001 .lock()
1002 .map_err(|_| poisoned_cpu_lock(op, "CPU engine registry"))?
1003 .values()
1004 .cloned(),
1005 );
1006 engines
1007 }
1008 CpuEngineRegistry::ExternalPrebuilt(registry) => {
1009 registry.by_id.values().cloned().collect()
1010 }
1011 };
1012 if engines.len() > 1 {
1013 engines.sort_unstable_by_key(|engine| Arc::as_ptr(engine) as usize);
1014 engines.dedup_by(|left, right| Arc::ptr_eq(left, right));
1015 }
1016 Ok(engines)
1017 }
1018}
1019
1020fn poisoned_cpu_lock(op: &'static str, lock: &'static str) -> crate::Error {
1021 crate::Error::runtime_state(op, format!("{lock} lock poisoned"))
1022}
1023
1024fn lock_engine_resources<'a>(
1025 engine: &'a CpuEngine,
1026 op: &'static str,
1027) -> crate::Result<std::sync::MutexGuard<'a, EngineResources>> {
1028 engine
1029 .resources
1030 .lock()
1031 .map_err(|_| poisoned_cpu_lock(op, "CPU engine resources"))
1032}
1033
1034fn saturating_add_tensor_cache_stats(total: &mut CacheStats, value: CacheStats) {
1035 total.entries = total.entries.saturating_add(value.entries);
1036 total.retained_bytes = total.retained_bytes.saturating_add(value.retained_bytes);
1037 total.hits = total.hits.saturating_add(value.hits);
1038 total.misses = total.misses.saturating_add(value.misses);
1039 total.evictions = total.evictions.saturating_add(value.evictions);
1040 total.clears = total.clears.saturating_add(value.clears);
1041}
1042
1043#[derive(Clone)]
1057pub struct CpuBackend {
1058 runtime_identity: CpuRuntimeIdentity,
1059 shared: Arc<CpuBackendState>,
1060 requested: CpuPlacement,
1061 resolved: ResolvedCpuExecution,
1062 engine: Arc<CpuEngine>,
1063 provider_bundle: CpuProviderBundle,
1064 allocation_domain: Option<Arc<dyn SharedTensorAllocationDomain>>,
1065 batch_policy: crate::CpuBatchPolicy,
1066}
1067
1068#[derive(Clone, Debug)]
1084pub struct CpuRuntimeIdentity {
1085 marker: Arc<()>,
1086}
1087
1088impl CpuRuntimeIdentity {
1089 fn fresh() -> Self {
1090 Self {
1091 marker: Arc::new(()),
1092 }
1093 }
1094}
1095
1096impl PartialEq for CpuRuntimeIdentity {
1097 fn eq(&self, other: &Self) -> bool {
1098 Arc::ptr_eq(&self.marker, &other.marker)
1099 }
1100}
1101
1102impl Eq for CpuRuntimeIdentity {}
1103
1104fn resolve_discovered_topology(
1105 kind: CpuBackendKind,
1106 topology: Result<CpuTopology, CpuTopologyError>,
1107) -> Result<CpuTopology, CpuPlacementError> {
1108 topology.map_err(|source| CpuPlacementError::TopologyDiscovery {
1109 requested: CpuPlacement::Auto,
1110 backend: kind,
1111 source,
1112 })
1113}
1114
1115fn external_engine_resolution(
1116 engine: &CpuEngine,
1117 requested: CpuPlacement,
1118 kind: CpuBackendKind,
1119) -> Result<ResolvedCpuExecution, CpuPlacementError> {
1120 match engine.domain().admission_mode() {
1121 CpuAdmissionMode::CooperativeCpuSet => engine
1122 .placement()
1123 .cloned()
1124 .map(ResolvedCpuExecution::ExternalManaged)
1125 .ok_or(CpuPlacementError::InternalState {
1126 requested,
1127 backend: kind,
1128 message: "cooperative external domain has no placement",
1129 }),
1130 CpuAdmissionMode::CallerManaged => Ok(ResolvedCpuExecution::ExternalCallerManaged),
1131 }
1132}
1133
1134#[allow(clippy::result_large_err)]
1137fn external_domain_backend_kind(
1138 op: &'static str,
1139 domains: &[ExternalCpuDomain],
1140) -> Result<CpuBackendKind, CpuBackendError> {
1141 let kind = if domains
1142 .iter()
1143 .any(|domain| domain.admission_mode() == CpuAdmissionMode::CallerManaged)
1144 {
1145 CpuBackendKind::Faer
1146 } else {
1147 CpuBackendKind::default_compiled()
1148 };
1149 ensure_cpu_backend_kind_available(kind, op)
1150 .map_err(|error| constructor_tensor_error(op, error))?;
1151 Ok(kind)
1152}
1153
1154fn coordinator_node_domain_ids(topology: &CpuTopology) -> BTreeMap<NumaNodeId, CpuDomainId> {
1155 topology
1156 .nodes()
1157 .iter()
1158 .enumerate()
1159 .filter_map(|(index, node)| {
1160 u64::try_from(index)
1161 .ok()
1162 .and_then(|index| index.checked_add(1))
1163 .map(|id| (node.id(), CpuDomainId::new(id)))
1164 })
1165 .collect()
1166}
1167
1168impl fmt::Debug for CpuBackend {
1169 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1170 f.debug_struct("CpuBackend")
1171 .field("kind", &self.kind())
1172 .field("provider_bundle", &self.provider_bundle)
1173 .field("requested_placement", &self.requested)
1174 .field("resolved_execution", &self.resolved)
1175 .field("engine_placement", &self.engine.placement())
1176 .field("num_threads", &self.num_threads())
1177 .field("allocation_domain", &self.allocation_domain())
1178 .field("buffer_pool_cache_stats", &self.buffer_pool_cache_stats())
1179 .field("buffer_pool_limit_bytes", &self.buffer_pool_limit_bytes())
1180 .finish_non_exhaustive()
1181 }
1182}
1183
1184impl CpuBackend {
1185 fn from_thread_budget_and_kind(
1186 thread_budget: usize,
1187 kind: CpuBackendKind,
1188 max_retained_capacity_bytes: usize,
1189 ) -> Result<Self, CpuPlacementError> {
1190 let topology = resolve_discovered_topology(kind, discover_cpu_topology())?;
1191 let resolved = resolve_placement(kind, CpuPlacement::Auto, &topology)?;
1192 #[cfg(not(any(target_os = "linux", target_os = "android")))]
1193 {
1194 let context = CpuContext::with_threads(thread_budget).map_err(|error| {
1195 CpuPlacementError::EngineConstruction {
1196 requested: CpuPlacement::Auto,
1197 backend: kind,
1198 source: CpuEngineConstructionError::Tensor(error),
1199 }
1200 })?;
1201 Ok(Self::compatibility_with_topology(
1202 Arc::new(context),
1203 max_retained_capacity_bytes,
1204 kind,
1205 topology,
1206 resolved,
1207 ))
1208 }
1209 #[cfg(any(target_os = "linux", target_os = "android"))]
1210 {
1211 let engine_placement = ResolvedCpuPlacement::AllAllowed {
1212 cpus: topology.allowed_cpus().clone(),
1213 };
1214 let engine = Arc::new(
1215 CpuEngine::new_managed(
1216 CpuDomainId::new(0),
1217 engine_placement,
1218 thread_budget,
1219 max_retained_capacity_bytes,
1220 )
1221 .map_err(|error| CpuPlacementError::EngineConstruction {
1222 requested: CpuPlacement::Auto,
1223 backend: kind,
1224 source: CpuEngineConstructionError::Context(error),
1225 })?,
1226 );
1227 let all_allowed = OnceLock::new();
1228 let _ = all_allowed.set(Arc::clone(&engine));
1229 Ok(Self {
1230 shared: Arc::new(CpuBackendState {
1231 engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
1232 node_engines: Mutex::new(BTreeMap::new()),
1233 node_domain_ids: coordinator_node_domain_ids(&topology),
1234 all_allowed,
1235 all_allowed_build: Mutex::new(()),
1236 base_engine: Arc::clone(&engine),
1237 thread_budget,
1238 }),
1239 topology,
1240 arbiter: ResourceArbiter::global(),
1241 kind,
1242 buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
1243 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1244 }),
1245 runtime_identity: CpuRuntimeIdentity::fresh(),
1246 requested: CpuPlacement::Auto,
1247 resolved,
1248 engine,
1249 provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
1250 allocation_domain: None,
1251 batch_policy: crate::CpuBatchPolicy::default(),
1252 })
1253 }
1254 }
1255
1256 fn compatibility(
1257 ctx: Arc<CpuContext>,
1258 max_retained_capacity_bytes: usize,
1259 kind: CpuBackendKind,
1260 ) -> Self {
1261 let topology = discover_cpu_topology().unwrap_or_else(|_| {
1262 let allowed = crate::process_cpu_affinity().unwrap_or_else(|| {
1263 CpuSet::new((0..crate::available_parallelism()).map(CpuId::new))
1264 .unwrap_or_else(|_| CpuSet::singleton(CpuId::new(0)))
1265 });
1266 CpuTopology::all_allowed(allowed)
1267 });
1268 let resolved = if kind == CpuBackendKind::Blas {
1269 ResolvedCpuExecution::ProviderDefaultExclusive
1270 } else {
1271 ResolvedCpuExecution::Compatibility
1272 };
1273 Self::compatibility_with_topology(
1274 ctx,
1275 max_retained_capacity_bytes,
1276 kind,
1277 topology,
1278 resolved,
1279 )
1280 }
1281
1282 fn compatibility_with_topology(
1283 ctx: Arc<CpuContext>,
1284 max_retained_capacity_bytes: usize,
1285 kind: CpuBackendKind,
1286 topology: CpuTopology,
1287 resolved: ResolvedCpuExecution,
1288 ) -> Self {
1289 let placement = ResolvedCpuPlacement::AllAllowed {
1290 cpus: topology.allowed_cpus().clone(),
1291 };
1292 let base_engine = Arc::new(CpuEngine::from_context(
1293 CpuDomainId::new(0),
1294 placement,
1295 ctx,
1296 max_retained_capacity_bytes,
1297 ));
1298 Self {
1299 shared: Arc::new(CpuBackendState {
1300 engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
1301 node_engines: Mutex::new(BTreeMap::new()),
1302 node_domain_ids: coordinator_node_domain_ids(&topology),
1303 all_allowed: OnceLock::new(),
1304 all_allowed_build: Mutex::new(()),
1305 base_engine: Arc::clone(&base_engine),
1306 thread_budget: base_engine.domain().thread_budget().get(),
1307 }),
1308 topology,
1309 arbiter: ResourceArbiter::global(),
1310 kind,
1311 buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
1312 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1313 }),
1314 runtime_identity: CpuRuntimeIdentity::fresh(),
1315 requested: CpuPlacement::Auto,
1316 resolved,
1317 engine: base_engine,
1318 provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
1319 allocation_domain: None,
1320 batch_policy: crate::CpuBatchPolicy::default(),
1321 }
1322 }
1323
1324 pub fn new() -> Self {
1334 let context = Arc::new(CpuContext::from_env());
1335 Self::from_thread_budget_and_kind(
1336 context.num_threads(),
1337 CpuBackendKind::default_compiled(),
1338 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1339 )
1340 .unwrap_or_else(|error| {
1341 eprintln!(
1342 "tenferro_cpu: using the unpinned compatibility context after placement error: {error}"
1343 );
1344 Self::from_context(context)
1345 })
1346 }
1347
1348 #[allow(clippy::result_large_err)]
1351 pub fn from_external_managed_domains(
1408 default_domain: CpuDomainId,
1409 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1410 ) -> Result<Self, CpuBackendError> {
1411 let op = "CpuBackend::from_external_managed_domains";
1412 let domains: Vec<_> = domains.into_iter().collect();
1413 let kind = external_domain_backend_kind(op, &domains)?;
1414 let topology = resolve_discovered_topology(kind, discover_cpu_topology())
1415 .map_err(|source| CpuBackendError::placement(op, source))?;
1416 Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1417 default_domain,
1418 domains,
1419 topology,
1420 ResourceArbiter::global(),
1421 kind,
1422 CpuProviderBundle::standard(kind, false),
1423 )
1424 }
1425
1426 #[allow(clippy::result_large_err)]
1429 pub fn from_external_managed_domains_with_provider_bundle(
1483 default_domain: CpuDomainId,
1484 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1485 provider_bundle: CpuProviderBundle,
1486 ) -> Result<Self, CpuBackendError> {
1487 let op = "CpuBackend::from_external_managed_domains_with_provider_bundle";
1488 let domains: Vec<_> = domains.into_iter().collect();
1489 let kind = external_domain_backend_kind(op, &domains)?;
1490 let topology = resolve_discovered_topology(kind, discover_cpu_topology())
1491 .map_err(|source| CpuBackendError::placement(op, source))?;
1492 Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1493 default_domain,
1494 domains,
1495 topology,
1496 ResourceArbiter::global(),
1497 kind,
1498 provider_bundle,
1499 )
1500 }
1501
1502 #[allow(clippy::result_large_err)]
1505 fn from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1506 default_domain: CpuDomainId,
1507 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1508 topology: CpuTopology,
1509 arbiter: ResourceArbiter,
1510 kind: CpuBackendKind,
1511 provider_bundle: CpuProviderBundle,
1512 ) -> Result<Self, CpuBackendError> {
1513 let domains: Vec<_> = domains.into_iter().collect();
1514 if domains.is_empty() {
1515 return Err(ExternalCpuDomainRegistryError::EmptyRegistry.into());
1516 }
1517
1518 let mut domain_ids = BTreeSet::new();
1519 let mut node_ids = BTreeSet::new();
1520 let mut has_all_allowed = false;
1521 for domain in &domains {
1522 if !domain_ids.insert(domain.id()) {
1523 return Err(
1524 ExternalCpuDomainRegistryError::DuplicateDomainId { id: domain.id() }.into(),
1525 );
1526 }
1527 if let Some(placement) = domain.placement() {
1528 match placement {
1529 ResolvedCpuPlacement::NumaNode { id, .. } => {
1530 if !node_ids.insert(*id) {
1531 return Err(
1532 ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
1533 placement: CpuPlacement::NumaNode(*id),
1534 }
1535 .into(),
1536 );
1537 }
1538 }
1539 ResolvedCpuPlacement::AllAllowed { cpus } => {
1540 if has_all_allowed {
1541 return Err(
1542 ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
1543 placement: CpuPlacement::AllAllowed,
1544 }
1545 .into(),
1546 );
1547 }
1548 has_all_allowed = true;
1549 if cpus != topology.allowed_cpus() {
1553 return Err(ExternalCpuDomainRegistryError::ExactAllAllowedMismatch {
1554 domain: domain.id(),
1555 declared: cpus.clone(),
1556 allowed: topology.allowed_cpus().clone(),
1557 }
1558 .into());
1559 }
1560 }
1561 }
1562 if let Some(cpu) = placement
1563 .cpus()
1564 .as_slice()
1565 .iter()
1566 .copied()
1567 .find(|cpu| !topology.allowed_cpus().contains(*cpu))
1568 {
1569 return Err(ExternalCpuDomainRegistryError::CpuOutsideAllowedSet {
1570 domain: domain.id(),
1571 cpu,
1572 }
1573 .into());
1574 }
1575 }
1576 }
1577 let buffer_limit = crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES;
1578 let mut by_id = BTreeMap::new();
1579 let mut by_node = BTreeMap::new();
1580 let mut all_allowed = None;
1581 for domain in domains {
1582 let id = domain.id();
1583 let placement = domain.placement().cloned();
1584 let engine = Arc::new(CpuEngine::from_external(domain, buffer_limit));
1585 match placement {
1586 Some(ResolvedCpuPlacement::NumaNode { id, .. }) => {
1587 by_node.insert(id, Arc::clone(&engine));
1588 }
1589 Some(ResolvedCpuPlacement::AllAllowed { .. }) => {
1590 all_allowed = Some(Arc::clone(&engine));
1591 }
1592 None => {}
1593 }
1594 by_id.insert(id, engine);
1595 }
1596 let Some(engine) = by_id.get(&default_domain).cloned() else {
1597 return Err(
1598 ExternalCpuDomainRegistryError::MissingDefaultDomain { default_domain }.into(),
1599 );
1600 };
1601 let resolved = match engine.domain().admission_mode() {
1602 CpuAdmissionMode::CooperativeCpuSet => ResolvedCpuExecution::ExternalManaged(
1603 engine.placement().cloned().ok_or_else(|| {
1604 CpuBackendError::placement(
1605 "CpuBackend external domain resolution",
1606 CpuPlacementError::InternalState {
1607 requested: CpuPlacement::Auto,
1608 backend: kind,
1609 message: "cooperative external domain has no placement",
1610 },
1611 )
1612 })?,
1613 ),
1614 CpuAdmissionMode::CallerManaged => ResolvedCpuExecution::ExternalCallerManaged,
1615 };
1616 let backend = Self {
1617 runtime_identity: CpuRuntimeIdentity::fresh(),
1618 shared: Arc::new(CpuBackendState {
1619 topology,
1620 engines: CpuEngineRegistry::ExternalPrebuilt(ExternalEngineRegistry {
1621 by_id,
1622 by_node,
1623 all_allowed,
1624 default_domain,
1625 }),
1626 arbiter,
1627 kind,
1628 buffer_limit: AtomicUsize::new(buffer_limit),
1629 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1630 }),
1631 requested: CpuPlacement::Auto,
1632 resolved,
1633 engine,
1634 provider_bundle,
1635 allocation_domain: None,
1636 batch_policy: crate::CpuBatchPolicy::default(),
1637 };
1638 backend
1639 .validate_provider_bundle_for_domains(&backend.provider_bundle)
1640 .map_err(|source| {
1641 CpuBackendError::Tensor(crate::Error::backend_source(
1642 "CpuBackend ExternalManaged provider validation",
1643 source,
1644 ))
1645 })?;
1646 Ok(backend)
1647 }
1648
1649 #[allow(clippy::result_large_err)]
1652 pub fn with_kind(kind: CpuBackendKind) -> Result<Self, CpuBackendError> {
1669 let op = "CpuBackend::with_kind";
1670 ensure_cpu_backend_kind_available(kind, op)
1671 .map_err(|error| constructor_tensor_error(op, error))?;
1672 let context = CpuContext::from_env();
1673 Self::from_thread_budget_and_kind(
1674 context.num_threads(),
1675 kind,
1676 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1677 )
1678 .map_err(|error| CpuBackendError::placement(op, error))
1679 }
1680
1681 #[allow(clippy::result_large_err)]
1684 pub fn try_new() -> Result<Self, CpuBackendError> {
1703 let op = "CpuBackend::try_new";
1704 let context =
1705 CpuContext::try_from_env().map_err(|error| constructor_tensor_error(op, error))?;
1706 Self::from_thread_budget_and_kind(
1707 context.num_threads(),
1708 CpuBackendKind::default_compiled(),
1709 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1710 )
1711 .map_err(|error| CpuBackendError::placement(op, error))
1712 }
1713
1714 pub fn from_context(ctx: Arc<CpuContext>) -> Self {
1727 Self::compatibility(
1728 ctx,
1729 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1730 CpuBackendKind::default_compiled(),
1731 )
1732 }
1733
1734 pub fn from_context_with_buffer_pool_limit(
1750 ctx: Arc<CpuContext>,
1751 max_retained_capacity_bytes: usize,
1752 ) -> Self {
1753 Self::from_context_with_buffer_pool_limit_and_kind(
1754 ctx,
1755 max_retained_capacity_bytes,
1756 CpuBackendKind::default_compiled(),
1757 )
1758 }
1759
1760 fn from_context_with_buffer_pool_limit_and_kind(
1761 ctx: Arc<CpuContext>,
1762 max_retained_capacity_bytes: usize,
1763 kind: CpuBackendKind,
1764 ) -> Self {
1765 Self::compatibility(ctx, max_retained_capacity_bytes, kind)
1766 }
1767
1768 #[allow(clippy::result_large_err)]
1771 pub fn with_threads(num_threads: usize) -> Result<Self, CpuBackendError> {
1788 let op = "CpuBackend::with_threads";
1789 let context = CpuContext::with_threads(num_threads)
1790 .map_err(|error| constructor_tensor_error(op, error))?;
1791 Self::from_thread_budget_and_kind(
1792 context.num_threads(),
1793 CpuBackendKind::default_compiled(),
1794 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1795 )
1796 .map_err(|error| CpuBackendError::placement(op, error))
1797 }
1798
1799 #[allow(clippy::result_large_err)]
1802 pub fn with_threads_and_kind(
1823 num_threads: usize,
1824 kind: CpuBackendKind,
1825 ) -> Result<Self, CpuBackendError> {
1826 let op = "CpuBackend::with_threads_and_kind";
1827 ensure_cpu_backend_kind_available(kind, op)
1828 .map_err(|error| constructor_tensor_error(op, error))?;
1829 let context = CpuContext::with_threads(num_threads)
1830 .map_err(|error| constructor_tensor_error(op, error))?;
1831 Self::from_thread_budget_and_kind(
1832 context.num_threads(),
1833 kind,
1834 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1835 )
1836 .map_err(|error| CpuBackendError::placement(op, error))
1837 }
1838
1839 pub fn for_placement(&self, requested: CpuPlacement) -> Result<Self, CpuPlacementError> {
1863 self.for_placement_with_affinity(
1864 requested,
1865 cfg!(any(target_os = "linux", target_os = "android")),
1866 )
1867 }
1868
1869 pub fn for_domain(&self, domain: CpuDomainId) -> Result<Self, CpuPlacementError> {
1896 let engine = self.shared.external_engine_for_id(domain)?;
1897 let resolved = external_engine_resolution(&engine, CpuPlacement::Auto, self.kind())?;
1898 Ok(Self {
1899 runtime_identity: CpuRuntimeIdentity::fresh(),
1900 shared: Arc::clone(&self.shared),
1901 requested: CpuPlacement::Auto,
1902 resolved,
1903 engine,
1904 provider_bundle: self.provider_bundle.clone(),
1905 allocation_domain: self.allocation_domain.clone(),
1906 batch_policy: self.batch_policy,
1907 })
1908 }
1909
1910 fn for_placement_with_affinity(
1911 &self,
1912 requested: CpuPlacement,
1913 managed_affinity_available: bool,
1914 ) -> Result<Self, CpuPlacementError> {
1915 if self.shared.is_external() {
1916 let engine = self.shared.external_engine_for(requested)?;
1917 let resolved = external_engine_resolution(&engine, requested, self.kind())?;
1918 return Ok(Self {
1919 runtime_identity: CpuRuntimeIdentity::fresh(),
1920 shared: Arc::clone(&self.shared),
1921 requested,
1922 resolved,
1923 engine,
1924 provider_bundle: self.provider_bundle.clone(),
1925 allocation_domain: self.allocation_domain.clone(),
1926 batch_policy: self.batch_policy,
1927 });
1928 }
1929 let resolved = resolve_placement_with_affinity(
1930 self.kind(),
1931 requested,
1932 &self.shared.topology,
1933 managed_affinity_available,
1934 )?;
1935 if requested == CpuPlacement::Auto && !managed_affinity_available {
1936 return Ok(Self {
1937 runtime_identity: CpuRuntimeIdentity::fresh(),
1938 shared: Arc::clone(&self.shared),
1939 requested,
1940 resolved,
1941 engine: self.shared.managed_base_engine(requested)?,
1942 provider_bundle: self.provider_bundle.clone(),
1943 allocation_domain: self.allocation_domain.clone(),
1944 batch_policy: self.batch_policy,
1945 });
1946 }
1947 let engine_placement = match &resolved {
1948 ResolvedCpuExecution::Managed(placement) => placement.clone(),
1949 ResolvedCpuExecution::ExternalManaged(_)
1950 | ResolvedCpuExecution::ExternalCallerManaged => {
1951 return Err(CpuPlacementError::InternalState {
1952 requested,
1953 backend: self.kind(),
1954 message: "managed resolver returned an external execution mode",
1955 });
1956 }
1957 ResolvedCpuExecution::ProviderDefaultExclusive => ResolvedCpuPlacement::AllAllowed {
1958 cpus: self.shared.topology.allowed_cpus().clone(),
1959 },
1960 ResolvedCpuExecution::Compatibility => {
1961 return Err(CpuPlacementError::InternalState {
1962 requested,
1963 backend: self.kind(),
1964 message: "placement resolution returned an internal compatibility mode",
1965 });
1966 }
1967 };
1968 let engine = self
1969 .shared
1970 .managed_engine_for(&engine_placement, requested)?;
1971 Ok(Self {
1972 runtime_identity: CpuRuntimeIdentity::fresh(),
1973 shared: Arc::clone(&self.shared),
1974 requested,
1975 resolved,
1976 engine,
1977 provider_bundle: self.provider_bundle.clone(),
1978 allocation_domain: self.allocation_domain.clone(),
1979 batch_policy: self.batch_policy,
1980 })
1981 }
1982
1983 pub fn placement(&self) -> CpuPlacement {
1993 self.requested
1994 }
1995
1996 pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
2015 match &self.resolved {
2016 ResolvedCpuExecution::Managed(placement)
2017 | ResolvedCpuExecution::ExternalManaged(placement) => Some(placement),
2018 ResolvedCpuExecution::Compatibility
2019 | ResolvedCpuExecution::ExternalCallerManaged
2020 | ResolvedCpuExecution::ProviderDefaultExclusive => None,
2021 }
2022 }
2023
2024 pub fn topology(&self) -> &CpuTopology {
2034 &self.shared.topology
2035 }
2036
2037 pub fn supports_placement(&self, placement: CpuPlacement) -> bool {
2047 if self.shared.is_external() {
2048 self.shared.external_engine_for(placement).is_ok()
2049 } else {
2050 resolve_placement(self.kind(), placement, &self.shared.topology).is_ok()
2051 }
2052 }
2053
2054 pub fn execution_info(&self) -> CpuExecutionInfo {
2063 let domain = self.engine.domain();
2064 let capabilities = domain.executor_capabilities();
2065 let (executor_affinity, executor_shutdown) =
2066 match (domain.ownership(), domain.admission_mode()) {
2067 (CpuDomainOwnership::ExternalManaged, CpuAdmissionMode::CooperativeCpuSet) => (
2068 CpuExecutorAffinity::CallerDeclaredUnverified,
2069 CpuExecutorShutdown::CallerOwned,
2070 ),
2071 (CpuDomainOwnership::ExternalManaged, CpuAdmissionMode::CallerManaged) => {
2072 (capabilities.affinity, CpuExecutorShutdown::CallerOwned)
2073 }
2074 (CpuDomainOwnership::Managed, _) => (capabilities.affinity, capabilities.shutdown),
2075 };
2076 CpuExecutionInfo {
2077 backend_kind: self.kind(),
2078 execution_mode: match &self.resolved {
2079 ResolvedCpuExecution::Managed(_) => CpuExecutionMode::Managed,
2080 ResolvedCpuExecution::ExternalManaged(_) => CpuExecutionMode::ExternalManaged,
2081 ResolvedCpuExecution::ExternalCallerManaged => CpuExecutionMode::CallerManaged,
2082 ResolvedCpuExecution::ProviderDefaultExclusive => {
2083 CpuExecutionMode::ProviderDefaultExclusive
2084 }
2085 ResolvedCpuExecution::Compatibility => CpuExecutionMode::Compatibility,
2086 },
2087 requested_placement: self.requested,
2088 resolved_placement: self.resolved_placement().cloned(),
2089 topology: self.shared.topology.clone(),
2090 domain_id: domain.id(),
2091 domain_cpus: domain.cpus().cloned(),
2092 worker_count: capabilities.worker_count.get(),
2093 thread_budget: domain.thread_budget().get(),
2094 admission_mode: domain.admission_mode(),
2095 domain_ownership: domain.ownership(),
2096 executor_affinity,
2097 executor_shutdown,
2098 provider_diagnostic: provider_diagnostic(
2099 self.kind(),
2100 domain.ownership(),
2101 domain.admission_mode(),
2102 ),
2103 }
2104 }
2105
2106 #[cfg(all(
2107 test,
2108 feature = "cpu-faer",
2109 any(target_os = "linux", target_os = "android")
2110 ))]
2111 fn coordinator_id_for_test(&self) -> usize {
2112 Arc::as_ptr(&self.shared) as usize
2113 }
2114
2115 #[cfg(test)]
2116 pub(crate) fn context_id_for_test(&self) -> usize {
2117 Arc::as_ptr(self.engine.domain().executor()) as *const () as usize
2118 }
2119
2120 pub fn kind(&self) -> CpuBackendKind {
2131 self.shared.kind
2132 }
2133
2134 pub fn provider_bundle(&self) -> &CpuProviderBundle {
2136 &self.provider_bundle
2137 }
2138
2139 pub fn runtime_identity(&self) -> CpuRuntimeIdentity {
2146 self.runtime_identity.clone()
2147 }
2148
2149 pub fn with_provider_bundle(
2169 mut self,
2170 bundle: CpuProviderBundle,
2171 ) -> Result<Self, CpuProviderBundleInstallError> {
2172 self.validate_provider_bundle_for_domains(&bundle)?;
2173 self.provider_bundle = bundle;
2174 self.runtime_identity = CpuRuntimeIdentity::fresh();
2175 Ok(self)
2176 }
2177
2178 fn validate_provider_bundle_for_domains(
2179 &self,
2180 bundle: &CpuProviderBundle,
2181 ) -> Result<(), CpuProviderBundleInstallError> {
2182 let validate_engine = |engine: &CpuEngine| {
2183 let domain = engine.domain();
2184 let contract = if domain.cpus().is_some() {
2185 CpuProviderDomainContract::CooperativeCpuSet
2186 } else {
2187 CpuProviderDomainContract::CallerManaged
2188 };
2189 bundle.validate_for_domain(domain.id(), domain.thread_budget(), contract)
2190 };
2191
2192 match &self.shared.engines {
2193 CpuEngineRegistry::ExternalPrebuilt(registry) => {
2194 for engine in registry.by_id.values() {
2195 validate_engine(engine)?;
2196 }
2197 }
2198 CpuEngineRegistry::ManagedLazy(registry) => {
2199 validate_engine(®istry.base_engine)?;
2200
2201 #[cfg(any(target_os = "linux", target_os = "android"))]
2206 for node in self.shared.topology.nodes() {
2207 let Some(domain_id) = registry.node_domain_ids.get(&node.id()).copied() else {
2208 continue;
2209 };
2210 let budget =
2211 std::num::NonZeroUsize::new(registry.thread_budget.min(node.cpus().len()))
2212 .expect("usable topology nodes have non-empty CPU sets");
2213 bundle.validate_for_domain(
2214 domain_id,
2215 budget,
2216 CpuProviderDomainContract::CooperativeCpuSet,
2217 )?;
2218 }
2219 }
2220 }
2221 Ok(())
2222 }
2223
2224 pub fn num_threads(&self) -> usize {
2235 self.engine.domain().thread_budget().get()
2236 }
2237
2238 pub fn buffer_pool_len(&self) -> crate::Result<usize> {
2255 self.shared
2256 .initialized_engines("CpuBackend::buffer_pool_len")?
2257 .iter()
2258 .try_fold(0, |total, engine| {
2259 Ok(total
2260 + lock_engine_resources(engine, "CpuBackend::buffer_pool_len")?
2261 .buffers
2262 .len())
2263 })
2264 }
2265
2266 pub fn buffer_pool_stats(&self) -> crate::Result<BufferPoolStats> {
2285 self.shared
2286 .initialized_engines("CpuBackend::buffer_pool_stats")?
2287 .iter()
2288 .try_fold(BufferPoolStats::default(), |mut total, engine| {
2289 let stats = lock_engine_resources(engine, "CpuBackend::buffer_pool_stats")?
2290 .buffers
2291 .stats();
2292 total.buffers += stats.buffers;
2293 total.capacity_bytes += stats.capacity_bytes;
2294 Ok(total)
2295 })
2296 }
2297
2298 pub fn buffer_pool_cache_stats(&self) -> crate::Result<CacheStats> {
2317 let stats = self.buffer_pool_stats()?;
2318 Ok(CacheStats {
2319 entries: stats.buffers,
2320 retained_bytes: stats.capacity_bytes,
2321 hits: 0,
2322 misses: 0,
2323 evictions: 0,
2324 clears: 0,
2325 })
2326 }
2327
2328 pub fn indexed_plan_cache_limits(&self) -> crate::Result<IndexedPlanCacheLimits> {
2345 self.shared
2346 .indexed_plan_cache_limits
2347 .lock()
2348 .map(|limits| *limits)
2349 .map_err(|_| {
2350 poisoned_cpu_lock(
2351 "CpuBackend::indexed_plan_cache_limits",
2352 "CPU indexed-plan cache configuration",
2353 )
2354 })
2355 }
2356
2357 pub fn set_indexed_plan_cache_limits(
2378 &mut self,
2379 limits: IndexedPlanCacheLimits,
2380 ) -> crate::Result<()> {
2381 let mut configured_limits = self.shared.indexed_plan_cache_limits.lock().map_err(|_| {
2385 poisoned_cpu_lock(
2386 "CpuBackend::set_indexed_plan_cache_limits",
2387 "CPU indexed-plan cache configuration",
2388 )
2389 })?;
2390 let engines = self
2391 .shared
2392 .initialized_engines("CpuBackend::set_indexed_plan_cache_limits")?;
2393 let mut resources = engines
2394 .iter()
2395 .map(|engine| {
2396 lock_engine_resources(engine, "CpuBackend::set_indexed_plan_cache_limits")
2397 })
2398 .collect::<crate::Result<Vec<_>>>()?;
2399 *configured_limits = limits;
2400 for resource in &mut resources {
2401 resource.indexed_plan_cache.set_limits(limits);
2402 }
2403 Ok(())
2404 }
2405
2406 pub fn indexed_plan_cache_stats(&self) -> crate::Result<CacheStats> {
2423 self.shared
2424 .initialized_engines("CpuBackend::indexed_plan_cache_stats")?
2425 .iter()
2426 .try_fold(CacheStats::default(), |mut total, engine| {
2427 let stats = lock_engine_resources(engine, "CpuBackend::indexed_plan_cache_stats")?
2428 .indexed_plan_cache
2429 .stats();
2430 saturating_add_tensor_cache_stats(&mut total, stats);
2431 Ok(total)
2432 })
2433 }
2434
2435 pub fn clear_indexed_plan_cache(&mut self) -> crate::Result<()> {
2453 let engines = self
2454 .shared
2455 .initialized_engines("CpuBackend::clear_indexed_plan_cache")?;
2456 let mut resources = engines
2457 .iter()
2458 .map(|engine| lock_engine_resources(engine, "CpuBackend::clear_indexed_plan_cache"))
2459 .collect::<crate::Result<Vec<_>>>()?;
2460 for resource in &mut resources {
2461 resource.indexed_plan_cache.clear();
2462 }
2463 Ok(())
2464 }
2465
2466 pub fn buffer_pool_limit_bytes(&self) -> usize {
2481 self.shared.buffer_limit.load(Ordering::Relaxed)
2482 }
2483
2484 pub fn set_buffer_pool_limit_bytes(
2507 &mut self,
2508 max_retained_capacity_bytes: usize,
2509 ) -> crate::Result<()> {
2510 let engines = self
2511 .shared
2512 .initialized_engines("CpuBackend::set_buffer_pool_limit_bytes")?;
2513 let mut resources = engines
2514 .iter()
2515 .map(|engine| lock_engine_resources(engine, "CpuBackend::set_buffer_pool_limit_bytes"))
2516 .collect::<crate::Result<Vec<_>>>()?;
2517 self.shared
2518 .buffer_limit
2519 .store(max_retained_capacity_bytes, Ordering::Relaxed);
2520 for resource in &mut resources {
2521 resource
2522 .buffers
2523 .set_max_retained_capacity_bytes(max_retained_capacity_bytes);
2524 }
2525 Ok(())
2526 }
2527
2528 pub fn reset_buffer_pool(&mut self) -> crate::Result<()> {
2551 let engines = self
2552 .shared
2553 .initialized_engines("CpuBackend::reset_buffer_pool")?;
2554 let mut resources = engines
2555 .iter()
2556 .map(|engine| lock_engine_resources(engine, "CpuBackend::reset_buffer_pool"))
2557 .collect::<crate::Result<Vec<_>>>()?;
2558 for resource in &mut resources {
2559 resource.buffers.clear();
2560 }
2561 Ok(())
2562 }
2563
2564 pub(crate) fn runtime_cache_stats(
2565 &self,
2566 ) -> crate::Result<tenferro_runtime::runtime::CacheStats> {
2567 let resources = lock_engine_resources(&self.engine, "CpuBackend::runtime_cache_stats")?;
2568 let buffers = resources.buffers.cache_stats();
2569 let gemm = tenferro_tensor::RuntimeCacheControl::stats(&resources.gemm_analysis_cache);
2570 let indexed = resources.indexed_plan_cache.stats();
2571 Ok(tenferro_runtime::runtime::CacheStats {
2572 entries: buffers
2573 .entries
2574 .saturating_add(gemm.entries)
2575 .saturating_add(indexed.entries),
2576 retained_bytes: buffers
2577 .retained_bytes
2578 .saturating_add(gemm.retained_bytes)
2579 .saturating_add(indexed.retained_bytes),
2580 hits: indexed.hits,
2581 misses: indexed.misses,
2582 evictions: indexed.evictions,
2583 clears: indexed.clears,
2584 })
2585 }
2586
2587 pub(crate) fn clear_runtime_caches(&self) -> crate::Result<()> {
2588 let mut resources =
2589 lock_engine_resources(&self.engine, "CpuBackend::clear_runtime_caches")?;
2590 resources.buffers.clear();
2591 tenferro_tensor::RuntimeCacheControl::clear(&mut resources.gemm_analysis_cache);
2592 resources.indexed_plan_cache.clear();
2593 Ok(())
2594 }
2595
2596 pub fn install<R: Send>(&self, op: impl FnOnce() -> R + Send) -> crate::Result<R> {
2620 let admission = self.execution_admission()?;
2621 let permit = admission.permit();
2622 let entry = CpuOperationEntry::new(self.engine.domain(), permit)
2623 .with_batch_policy(self.batch_policy);
2624 entry
2625 .enter(ParallelMode::Sequential, |_| op())
2626 .map_err(|error| crate::Error::backend_source("CpuBackend::install", error))
2627 }
2628
2629 fn with_execution_resources<R>(
2630 &self,
2631 permit: &ResourcePermit,
2632 op: impl FnOnce(&mut EngineResources) -> R,
2633 ) -> R {
2634 if permit.is_reentrant() {
2635 let mut resources =
2636 EngineResources::new(self.shared.buffer_limit.load(Ordering::Relaxed));
2637 return op(&mut resources);
2638 }
2639 let mut resources = self
2645 .engine
2646 .resources
2647 .lock()
2648 .unwrap_or_else(std::sync::PoisonError::into_inner);
2649 op(&mut resources)
2650 }
2651
2652 fn acquire_execution_permit(
2653 &self,
2654 owner: ResourceOwner,
2655 ) -> Result<ResourcePermit, SessionEntryError> {
2656 let arbiter_poisoned = |_| SessionEntryError::ResourcePoisoned {
2657 backend: CPU_BACKEND,
2658 resource: "the CPU resource arbiter",
2659 };
2660 match &self.resolved {
2661 ResolvedCpuExecution::Managed(placement)
2662 | ResolvedCpuExecution::ExternalManaged(placement) => self
2663 .shared
2664 .arbiter
2665 .acquire_waiting(placement.cpus().clone(), owner)
2666 .map_err(arbiter_poisoned),
2667 ResolvedCpuExecution::ExternalCallerManaged => {
2668 let active = self
2671 .engine
2672 .domain()
2673 .caller_managed_active()
2674 .unwrap_or_else(|| {
2675 unreachable!("caller-managed execution needs a local admission guard")
2676 });
2677 ResourcePermit::caller_managed(active, owner).ok_or_else(|| {
2678 SessionEntryError::Contended {
2679 backend: CPU_BACKEND,
2680 message: "the caller-managed CPU domain is already executing; \
2681 serialize entries to a caller-managed domain"
2682 .to_owned(),
2683 }
2684 })
2685 }
2686 ResolvedCpuExecution::Compatibility => self
2687 .shared
2688 .arbiter
2689 .acquire_waiting(self.shared.topology.allowed_cpus().clone(), owner)
2690 .map_err(arbiter_poisoned),
2691 ResolvedCpuExecution::ProviderDefaultExclusive => self
2692 .shared
2693 .arbiter
2694 .acquire_provider_exclusive_waiting(owner)
2695 .map_err(arbiter_poisoned),
2696 }
2697 }
2698
2699 #[cfg(test)]
2700 fn try_acquire_execution_permit_for_test(
2701 &self,
2702 ) -> Result<Option<ResourcePermit>, crate::arbiter::ResourceArbiterError> {
2703 match &self.resolved {
2704 ResolvedCpuExecution::Managed(placement)
2705 | ResolvedCpuExecution::ExternalManaged(placement) => {
2706 self.shared.arbiter.try_acquire(placement.cpus().clone())
2707 }
2708 ResolvedCpuExecution::ExternalCallerManaged => Ok(None),
2709 ResolvedCpuExecution::Compatibility => self
2710 .shared
2711 .arbiter
2712 .try_acquire(self.shared.topology.allowed_cpus().clone()),
2713 ResolvedCpuExecution::ProviderDefaultExclusive => {
2714 self.shared.arbiter.try_acquire_provider_exclusive()
2715 }
2716 }
2717 }
2718}
2719
2720impl BackendRuntimeCache for CpuBackend {
2721 type RuntimeCache = gemm::GemmAnalysisCache;
2722}
2723
2724impl CpuBackend {
2725 #[must_use]
2745 pub fn with_batch_policy(mut self, policy: crate::CpuBatchPolicy) -> Self {
2746 self.batch_policy = policy;
2747 self
2748 }
2749
2750 #[must_use]
2760 pub fn batch_policy(&self) -> crate::CpuBatchPolicy {
2761 self.batch_policy
2762 }
2763
2764 pub fn with_allocation_domain(mut self, domain: Arc<dyn SharedTensorAllocationDomain>) -> Self {
2789 self.allocation_domain = Some(domain);
2790 self.runtime_identity = CpuRuntimeIdentity::fresh();
2791 self
2792 }
2793
2794 pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
2804 self.allocation_domain.as_ref().map(|domain| domain.id())
2805 }
2806
2807 pub fn shared_allocation_domain(&self) -> Option<&Arc<dyn SharedTensorAllocationDomain>> {
2817 self.allocation_domain.as_ref()
2818 }
2819
2820 fn run_backend_session_cached<R: Send>(
2821 &mut self,
2822 cache: Option<&mut gemm::GemmAnalysisCache>,
2823 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
2824 ) -> Result<R, SessionEntryError> {
2825 let providers = self.provider_bundle.clone();
2826 let admission = self.execution_admission()?;
2827 let permit = admission.permit();
2828 let owner = permit.owner();
2829 let entry = CpuOperationEntry::new(self.engine.domain(), permit)
2830 .with_batch_policy(self.batch_policy);
2831 let enter_managed_session = entry.enters_executor_per_session();
2834 let run = |entered| {
2835 self.with_execution_resources(permit, |resources| {
2836 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
2837 let cache = cache.unwrap_or(&mut resources.gemm_analysis_cache);
2838 let session_started = Instant::now();
2839 let mut session = CpuExecSession {
2840 entry,
2841 entered,
2842 buffers: buffers.get_mut(),
2843 gemm_analysis_cache: cache,
2844 indexed_plan_cache: &mut resources.indexed_plan_cache,
2845 providers: &providers,
2846 backend_kind: self.kind(),
2847 allocation_domain: self.allocation_domain.as_ref(),
2848 };
2849 record_cpu_session_profile(
2850 "with_backend_session_cached.session_construct",
2851 session_started.elapsed(),
2852 );
2853 let exec_started = Instant::now();
2854 let result = f(&mut session);
2855 record_cpu_session_profile(
2856 "with_backend_session_cached.exec_body",
2857 exec_started.elapsed(),
2858 );
2859 result
2860 })
2861 };
2862 if enter_managed_session {
2863 entry.enter_managed_session(|context| run(Some(context)))
2864 } else {
2865 Ok(with_execution_owner(owner, || run(None)))
2867 }
2868 }
2869}
2870
2871impl BackendSessionHost for CpuBackend {
2872 fn with_backend_session<R: Send>(
2873 &mut self,
2874 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
2875 ) -> Result<R, SessionEntryError> {
2876 self.run_backend_session_cached(None, f)
2877 }
2878
2879 fn with_backend_session_cached<R: Send>(
2880 &mut self,
2881 cache: &mut Self::RuntimeCache,
2882 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
2883 ) -> Result<R, SessionEntryError> {
2884 if !cpu_session_profile_enabled() {
2885 return self.run_backend_session_cached(Some(cache), f);
2886 }
2887 let total_started = Instant::now();
2888 let result =
2889 profile_cpu_session_section("with_backend_session_cached.exec_session", || {
2890 self.run_backend_session_cached(Some(cache), f)
2891 });
2892 record_cpu_session_profile("with_backend_session_cached.total", total_started.elapsed());
2893 maybe_print_cpu_session_profile();
2894 result
2895 }
2896}
2897
2898pub(crate) fn reclaim_tensor(buffers: &mut BufferPool, tensor: Tensor) {
2903 match tensor.dtype() {
2904 DType::F32 => reclaim_tensor_typed::<f32>(buffers, tensor),
2905 DType::F64 => reclaim_tensor_typed::<f64>(buffers, tensor),
2906 DType::I32 => reclaim_tensor_typed::<i32>(buffers, tensor),
2907 DType::I64 => reclaim_tensor_typed::<i64>(buffers, tensor),
2908 DType::Bool => reclaim_tensor_typed::<bool>(buffers, tensor),
2909 DType::C32 => reclaim_tensor_typed::<Complex32>(buffers, tensor),
2910 DType::C64 => reclaim_tensor_typed::<Complex64>(buffers, tensor),
2911 DType::External(_) => {}
2912 }
2913}
2914
2915fn reclaim_tensor_typed<T: tenferro_cpu_basic::PoolScalar>(
2917 buffers: &mut BufferPool,
2918 tensor: Tensor,
2919) {
2920 if let Ok(typed) = tensor.into_typed::<T>() {
2921 reclaim_typed(buffers, typed);
2922 }
2923}
2924
2925impl TensorDeviceTransfer for CpuBackend {
2926 fn download_to_host(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
2927 if tensor.backend_family().is_some() {
2928 return Err(crate::Error::runtime_state(
2929 "CpuBackend::download_to_host",
2930 "CPU backend received a backend buffer; download the tensor to host with its owning backend before CPU execution",
2931 ));
2932 }
2933 tensor.tensor_view().duplicate()
2934 }
2935
2936 fn upload_host_tensor(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
2937 if tensor.backend_family().is_some() {
2938 return Err(crate::Error::runtime_state(
2939 "CpuBackend::upload_host_tensor",
2940 "CPU backend upload_host_tensor expects a host tensor; download backend buffers to host before CPU execution",
2941 ));
2942 }
2943 tensor.tensor_view().duplicate()
2944 }
2945}
2946
2947impl TensorBackend for CpuBackend {}
2948
2949pub(crate) fn reclaim_typed<T: PoolScalar>(pool: &mut BufferPool, typed: TypedTensor<T>) {
2950 if typed.backend_buffer().is_some() {
2951 return;
2952 }
2953 if let Ok(data) = typed.into_host_vec() {
2954 T::pool_release(pool, data);
2955 }
2956}
2957
2958impl Default for CpuBackend {
2959 fn default() -> Self {
2960 Self::new()
2961 }
2962}
2963
2964pub(crate) mod execution_scope;
2965
2966#[cfg(test)]
2967mod tests;