1use std::any::TypeId;
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};
10
11use crate::arbiter::{with_execution_owner, ResourceArbiter, ResourceOwner, ResourcePermit};
12use crate::buffer_pool::{BufferPool, BufferPoolStats, PoolScalar};
13use crate::dot_runtime::{
14 CpuProviderBundle, CpuProviderBundleInstallError, CpuProviderDomainContract,
15};
16use crate::engine::{CpuEngine, EngineResources};
17use crate::indexed_plan_cache::{
18 IndexedPlanCache, IndexedPlanCacheLimits, DEFAULT_INDEXED_PLAN_CACHE_LIMITS,
19};
20use crate::placement::{
21 resolve_placement, resolve_placement_with_affinity, CpuEngineConstructionError,
22 ResolvedCpuExecution,
23};
24use crate::provider::{CpuExecutionContext, CpuOperationEntry, ParallelMode};
25use crate::{
26 discover_cpu_topology, CpuAdmissionMode, CpuDomainId, CpuDomainOwnership, CpuExecutorAffinity,
27 CpuExecutorShutdown, CpuId, CpuPlacement, CpuPlacementError, CpuPlacementGuarantee, CpuSet,
28 CpuTopology, CpuTopologyError, ExternalCpuDomain, NumaNodeId, ResolvedCpuPlacement,
29};
30use crate::{
31 CacheStats, Tensor, TensorRank, TensorRead, TensorScalar, TensorValue, TensorWrite,
32 TypedTensor, TypedTensorView, TypedTensorViewMut,
33};
34use tenferro_tensor::backend::{ElementwiseFusionPlan, GroupedGemmConfig};
35use tenferro_tensor::SharedTensorAllocationDomain;
36use tenferro_tensor::{
37 AllocationDomainId, BackendCachedDot, BackendRuntimeCache, BackendSession, BackendSessionHost,
38 ContractionScalar, DotGeneralAccumulation, ElementwiseReadOp, TensorAnalytic, TensorBackend,
39 TensorBuffer, TensorDeviceTransfer, TensorDot, TensorElementwise, TensorFusion, TensorIndexing,
40 TensorReduction, TensorStructural, TensorViewCanonicalization,
41};
42use tenferro_tensor::{
43 CompareDir, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
44};
45
46use super::exec_session::CpuExecSession;
47use super::{
48 analytic, copy_tensor_read_into, elementwise, gemm, indexing, materialize_tensor_read,
49 reduction, structural, CpuContext,
50};
51
52pub(crate) fn tag_fresh_output(output: &mut Tensor, domain: CpuDomainId) {
53 macro_rules! tag {
54 ($tensor:expr) => {{
55 $tensor.set_cpu_affinity(Some(domain));
56 }};
57 }
58 match output {
59 Tensor::F32(tensor) => tag!(tensor),
60 Tensor::F64(tensor) => tag!(tensor),
61 Tensor::I32(tensor) => tag!(tensor),
62 Tensor::I64(tensor) => tag!(tensor),
63 Tensor::Bool(tensor) => tag!(tensor),
64 Tensor::C32(tensor) => tag!(tensor),
65 Tensor::C64(tensor) => tag!(tensor),
66 }
67}
68
69pub(crate) fn elementwise_read_into_fallback_with_pool(
70 buffers: &mut BufferPool,
71 op: ElementwiseReadOp,
72 inputs: &[TensorRead<'_>],
73 out: TensorWrite<'_>,
74) -> crate::Result<()> {
75 let result = match op {
76 ElementwiseReadOp::Add => {
77 elementwise::add_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
78 }
79 ElementwiseReadOp::Subtract => {
80 elementwise::sub_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
81 }
82 ElementwiseReadOp::Multiply => {
83 elementwise::mul_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
84 }
85 ElementwiseReadOp::Negate => elementwise::neg_read_with_pool(buffers, inputs[0].clone())?,
86 ElementwiseReadOp::Conj => elementwise::conj_read_with_pool(buffers, inputs[0].clone())?,
87 ElementwiseReadOp::Divide => {
88 elementwise::div_read_with_pool(buffers, inputs[0].clone(), inputs[1].clone())?
89 }
90 _ => {
91 return Err(crate::Error::unsupported(
92 "CpuBackend::elementwise_read_into",
93 format!("CPU backend does not implement {op:?}"),
94 ))
95 }
96 };
97 copy_tensor_read_into(
98 "CpuBackend::elementwise_read_into",
99 TensorRead::from_tensor(&result),
100 out,
101 )
102}
103
104pub(crate) trait FreshCpuOutput {
105 fn tag_fresh(&mut self, domain: CpuDomainId);
106}
107
108impl FreshCpuOutput for Tensor {
109 fn tag_fresh(&mut self, domain: CpuDomainId) {
110 tag_fresh_output(self, domain);
111 }
112}
113
114impl<T, R: TensorRank> FreshCpuOutput for TypedTensor<T, R> {
115 fn tag_fresh(&mut self, domain: CpuDomainId) {
116 self.set_cpu_affinity(Some(domain));
117 }
118}
119
120impl<T: FreshCpuOutput> FreshCpuOutput for Option<T> {
121 fn tag_fresh(&mut self, domain: CpuDomainId) {
122 if let Some(output) = self {
123 output.tag_fresh(domain);
124 }
125 }
126}
127
128impl<T: FreshCpuOutput> FreshCpuOutput for Vec<T> {
129 fn tag_fresh(&mut self, domain: CpuDomainId) {
130 for output in self {
131 output.tag_fresh(domain);
132 }
133 }
134}
135
136#[derive(Debug, Default, Clone)]
137struct CpuSessionProfileEntry {
138 calls: usize,
139 total_time: Duration,
140}
141
142fn cpu_session_profile_enabled() -> bool {
143 static ENABLED: OnceLock<bool> = OnceLock::new();
144 *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_CPU_SESSION").is_ok())
145}
146
147fn cpu_session_profile_print_every() -> Option<usize> {
148 static PRINT_EVERY: OnceLock<Option<usize>> = OnceLock::new();
149 *PRINT_EVERY.get_or_init(|| {
150 env::var("TENFERRO_PROFILE_CPU_SESSION_PRINT_EVERY")
151 .ok()
152 .and_then(|value| value.parse::<usize>().ok())
153 .filter(|&value| value > 0)
154 })
155}
156
157fn cpu_session_profile_state() -> &'static Mutex<HashMap<&'static str, CpuSessionProfileEntry>> {
158 static STATE: OnceLock<Mutex<HashMap<&'static str, CpuSessionProfileEntry>>> = OnceLock::new();
159 STATE.get_or_init(|| Mutex::new(HashMap::new()))
160}
161
162fn record_cpu_session_profile(section: &'static str, elapsed: Duration) {
163 if !cpu_session_profile_enabled() {
164 return;
165 }
166 let Ok(mut state) = cpu_session_profile_state().lock() else {
167 return;
168 };
169 let entry = state.entry(section).or_default();
170 entry.calls += 1;
171 entry.total_time += elapsed;
172}
173
174fn profile_cpu_session_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
175 if !cpu_session_profile_enabled() {
176 return f();
177 }
178 let started = Instant::now();
179 let result = f();
180 record_cpu_session_profile(section, started.elapsed());
181 result
182}
183
184fn maybe_print_cpu_session_profile() {
185 let Some(print_every) = cpu_session_profile_print_every() else {
186 return;
187 };
188 let should_print = {
189 let Ok(state) = cpu_session_profile_state().lock() else {
190 return;
191 };
192 state
193 .get("with_backend_session_cached.total")
194 .is_some_and(|entry| entry.calls % print_every == 0)
195 };
196 if !should_print {
197 return;
198 }
199 let mut entries = {
200 let Ok(mut state) = cpu_session_profile_state().lock() else {
201 return;
202 };
203 let entries = state
204 .iter()
205 .map(|(section, entry)| (*section, entry.clone()))
206 .collect::<Vec<_>>();
207 state.clear();
208 entries
209 };
210 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
211 eprintln!("=== tenferro CPU session profile ===");
212 for (section, entry) in entries {
213 eprintln!(
214 "{section}: calls={} total={:.6}ms per_call={:.3}us",
215 entry.calls,
216 entry.total_time.as_secs_f64() * 1.0e3,
217 entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64,
218 );
219 }
220}
221
222struct BufferPoolLoan<'a> {
223 buffers: &'a mut BufferPool,
224}
225
226impl<'a> BufferPoolLoan<'a> {
227 fn new(buffers: &'a mut BufferPool) -> Self {
228 Self { buffers }
229 }
230
231 fn get_mut(&mut self) -> &mut BufferPool {
232 self.buffers
233 }
234}
235
236impl Drop for BufferPoolLoan<'_> {
237 fn drop(&mut self) {
238 if thread::panicking() {
239 self.buffers.replenish_in_flight_retained();
240 } else {
241 self.buffers.clear_in_flight_retained();
242 }
243 }
244}
245
246#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
261pub enum CpuBackendKind {
262 Faer,
264 Blas,
266}
267
268impl CpuBackendKind {
269 pub fn default_compiled() -> Self {
283 #[cfg(feature = "cpu-blas")]
284 {
285 Self::Blas
286 }
287 #[cfg(all(not(feature = "cpu-blas"), feature = "cpu-faer"))]
288 {
289 Self::Faer
290 }
291 }
292
293 #[allow(dead_code)]
296 pub(crate) fn name(self) -> &'static str {
297 match self {
298 Self::Faer => "faer",
299 Self::Blas => "blas",
300 }
301 }
302}
303
304#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
322pub enum CpuExecutionMode {
323 Managed,
325 ExternalManaged,
327 CallerManaged,
329 ProviderDefaultExclusive,
331 Compatibility,
333}
334
335#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
346pub enum ExternalCpuDomainRegistryError {
347 #[error("externally managed CPU registry must contain at least one domain")]
349 EmptyRegistry,
350 #[error("CPU domain ID {id:?} is registered more than once")]
352 DuplicateDomainId {
353 id: CpuDomainId,
355 },
356 #[error("CPU placement {placement:?} is registered more than once")]
358 DuplicatePlacementIdentity {
359 placement: CpuPlacement,
361 },
362 #[error("CPU domain {domain:?} declares process-disallowed CPU {cpu}")]
364 CpuOutsideAllowedSet {
365 domain: CpuDomainId,
367 cpu: CpuId,
369 },
370 #[error("default CPU domain {default_domain:?} is not registered")]
372 MissingDefaultDomain {
373 default_domain: CpuDomainId,
375 },
376 #[error(
378 "exact all-allowed CPU domain {domain:?} declares {declared:?}, but the process allows {allowed:?}"
379 )]
380 ExactAllAllowedMismatch {
381 domain: CpuDomainId,
383 declared: CpuSet,
385 allowed: CpuSet,
387 },
388}
389
390#[derive(Debug, thiserror::Error)]
405pub enum CpuBackendError {
406 #[error(transparent)]
408 Tensor(#[from] crate::Error),
409 #[error("{op}: {source}")]
411 Placement {
412 op: &'static str,
414 #[source]
416 source: CpuPlacementError,
417 },
418 #[error(transparent)]
420 ExternalRegistry(#[from] ExternalCpuDomainRegistryError),
421}
422
423impl CpuBackendError {
424 fn placement(op: &'static str, source: CpuPlacementError) -> Self {
425 Self::Placement { op, source }
426 }
427
428 pub fn placement_error(&self) -> Option<&CpuPlacementError> {
441 match self {
442 Self::Tensor(_) => None,
443 Self::Placement { source, .. } => Some(source),
444 Self::ExternalRegistry(_) => None,
445 }
446 }
447}
448
449impl From<CpuBackendError> for crate::Error {
450 fn from(error: CpuBackendError) -> Self {
451 match error {
452 CpuBackendError::Tensor(error) => error,
453 CpuBackendError::ExternalRegistry(source) => Self::extension(
454 "CpuBackend::from_external_managed_domains",
455 "cpu",
456 crate::ErrorKind::Validation(crate::ValidationKind::InvalidArgument),
457 source,
458 ),
459 CpuBackendError::Placement { op, source } => match source {
460 CpuPlacementError::TopologyDiscovery { .. }
461 | CpuPlacementError::ManagedAffinityUnavailable { .. }
462 | CpuPlacementError::NumaDiscoveryUnavailable { .. }
463 | CpuPlacementError::UnknownNumaNode { .. }
464 | CpuPlacementError::UnregisteredExternalPlacement { .. }
465 | CpuPlacementError::UnregisteredExternalDomain { .. } => {
466 Self::runtime_state_source(op, source)
467 }
468 CpuPlacementError::ExternalProviderAffinityUnmanaged { .. } => {
469 Self::extension(op, "cpu", crate::ErrorKind::Unsupported, source)
470 }
471 CpuPlacementError::EngineConstruction { .. } => Self::backend_source(op, source),
472 CpuPlacementError::InternalState { .. } => {
473 Self::extension(op, "cpu", crate::ErrorKind::Internal, source)
474 }
475 },
476 }
477 }
478}
479
480#[derive(Clone, Debug, PartialEq, Eq)]
495pub struct CpuExecutionInfo {
496 backend_kind: CpuBackendKind,
497 execution_mode: CpuExecutionMode,
498 requested_placement: CpuPlacement,
499 resolved_placement: Option<ResolvedCpuPlacement>,
500 topology: CpuTopology,
501 domain_id: CpuDomainId,
502 domain_cpus: Option<CpuSet>,
503 worker_count: usize,
504 thread_budget: usize,
505 placement_guarantee: Option<CpuPlacementGuarantee>,
506 admission_mode: CpuAdmissionMode,
507 domain_ownership: CpuDomainOwnership,
508 executor_affinity: CpuExecutorAffinity,
509 executor_shutdown: CpuExecutorShutdown,
510 provider_diagnostic: &'static str,
511}
512
513impl CpuExecutionInfo {
514 pub fn backend_kind(&self) -> CpuBackendKind {
523 self.backend_kind
524 }
525
526 pub fn execution_mode(&self) -> CpuExecutionMode {
537 self.execution_mode
538 }
539
540 pub fn requested_placement(&self) -> CpuPlacement {
549 self.requested_placement
550 }
551
552 pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
561 self.resolved_placement.as_ref()
562 }
563
564 pub fn topology(&self) -> &CpuTopology {
573 &self.topology
574 }
575
576 pub fn domain_id(&self) -> CpuDomainId {
585 self.domain_id
586 }
587
588 pub fn domain_cpus(&self) -> Option<&CpuSet> {
599 self.domain_cpus.as_ref()
600 }
601
602 pub fn worker_count(&self) -> usize {
611 self.worker_count
612 }
613
614 pub fn thread_budget(&self) -> usize {
627 self.thread_budget
628 }
629
630 pub fn placement_guarantee(&self) -> Option<CpuPlacementGuarantee> {
641 self.placement_guarantee
642 }
643
644 pub fn admission_mode(&self) -> CpuAdmissionMode {
655 self.admission_mode
656 }
657
658 pub fn domain_ownership(&self) -> CpuDomainOwnership {
669 self.domain_ownership
670 }
671
672 pub fn executor_affinity(&self) -> CpuExecutorAffinity {
683 self.executor_affinity
684 }
685
686 pub fn executor_shutdown(&self) -> CpuExecutorShutdown {
697 self.executor_shutdown
698 }
699
700 pub fn provider_diagnostic(&self) -> &'static str {
713 self.provider_diagnostic
714 }
715}
716
717fn provider_diagnostic(
718 kind: CpuBackendKind,
719 ownership: CpuDomainOwnership,
720 admission_mode: CpuAdmissionMode,
721) -> &'static str {
722 if ownership == CpuDomainOwnership::ExternalManaged {
723 if admission_mode == CpuAdmissionMode::CallerManaged {
724 debug_assert_eq!(kind, CpuBackendKind::Faer);
727 return "faer (caller-managed CPU executor and admission)";
728 }
729 return match kind {
730 CpuBackendKind::Faer => "faer (externally managed CPU executor)",
731 CpuBackendKind::Blas => "BLAS/LAPACK (externally managed CPU executor)",
732 };
733 }
734 match kind {
735 CpuBackendKind::Faer => "faer (tenferro-managed Rayon affinity)",
736 CpuBackendKind::Blas => {
737 #[cfg(feature = "blas-openblas")]
738 return "OpenBLAS (external worker affinity)";
739 #[cfg(feature = "blas-mkl")]
740 return "Intel MKL (external worker affinity)";
741 #[cfg(feature = "blas-accelerate")]
742 return "Apple Accelerate (external worker affinity)";
743 #[cfg(feature = "provider-inject")]
744 return "runtime-injected BLAS/LAPACK (external worker affinity)";
745 #[cfg(not(any(
746 feature = "blas-openblas",
747 feature = "blas-mkl",
748 feature = "blas-accelerate",
749 feature = "provider-inject"
750 )))]
751 return "linked BLAS/LAPACK provider (identity unknown; external worker affinity)";
752 }
753 }
754}
755
756fn ensure_cpu_backend_kind_available(kind: CpuBackendKind, op: &'static str) -> crate::Result<()> {
757 let _ = op;
758 match kind {
759 CpuBackendKind::Faer => {
760 #[cfg(feature = "cpu-faer")]
761 {
762 Ok(())
763 }
764 #[cfg(not(feature = "cpu-faer"))]
765 {
766 Err(crate::Error::invalid_argument(
767 op,
768 "configuration",
769 "CpuBackendKind::Faer requires the cpu-faer feature".to_string(),
770 ))
771 }
772 }
773 CpuBackendKind::Blas => {
774 #[cfg(feature = "cpu-blas")]
775 {
776 Ok(())
777 }
778 #[cfg(not(feature = "cpu-blas"))]
779 {
780 Err(crate::Error::invalid_argument(
781 op,
782 "configuration",
783 "CpuBackendKind::Blas requires the cpu-blas feature".to_string(),
784 ))
785 }
786 }
787 }
788}
789
790fn constructor_tensor_error(op: &'static str, error: crate::Error) -> CpuBackendError {
791 CpuBackendError::Tensor(match error {
792 crate::Error::Validation { source, .. } => crate::Error::validation(op, source),
793 error => error,
794 })
795}
796
797#[allow(dead_code)]
800pub(super) fn unavailable_cpu_backend_kind(kind: CpuBackendKind, op: &'static str) -> crate::Error {
801 crate::Error::invalid_argument(
802 op,
803 "configuration",
804 format!("CPU backend kind {} is not compiled in", kind.name()),
805 )
806}
807
808struct ManagedEngineRegistry {
809 node_engines: Mutex<BTreeMap<NumaNodeId, Arc<CpuEngine>>>,
810 node_domain_ids: BTreeMap<NumaNodeId, CpuDomainId>,
811 all_allowed: OnceLock<Arc<CpuEngine>>,
812 all_allowed_build: Mutex<()>,
813 base_engine: Arc<CpuEngine>,
814 thread_budget: usize,
815}
816
817struct ExternalEngineRegistry {
818 by_id: BTreeMap<CpuDomainId, Arc<CpuEngine>>,
819 by_node: BTreeMap<NumaNodeId, Arc<CpuEngine>>,
820 all_allowed: Option<Arc<CpuEngine>>,
821 default_domain: CpuDomainId,
822}
823
824enum CpuEngineRegistry {
825 ManagedLazy(ManagedEngineRegistry),
826 ExternalPrebuilt(ExternalEngineRegistry),
827}
828
829struct CpuBackendState {
830 topology: CpuTopology,
831 engines: CpuEngineRegistry,
832 arbiter: ResourceArbiter,
833 kind: CpuBackendKind,
834 buffer_limit: AtomicUsize,
835 indexed_plan_cache_limits: Mutex<IndexedPlanCacheLimits>,
836}
837
838impl CpuBackendState {
839 fn managed_engine_for(
840 &self,
841 placement: &ResolvedCpuPlacement,
842 requested: CpuPlacement,
843 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
844 let cache_configuration = self.indexed_plan_cache_limits.lock().map_err(|_| {
848 CpuPlacementError::InternalState {
849 requested,
850 backend: self.kind,
851 message: "CPU indexed-plan cache configuration lock is poisoned",
852 }
853 })?;
854 let cache_limits = *cache_configuration;
855 let CpuEngineRegistry::ManagedLazy(registry) = &self.engines else {
856 return Err(CpuPlacementError::InternalState {
857 requested,
858 backend: self.kind,
859 message: "managed placement requested from an external engine registry",
860 });
861 };
862 match placement {
863 ResolvedCpuPlacement::NumaNode { id, .. } => {
864 let mut engines = registry
865 .node_engines
866 .lock()
867 .unwrap_or_else(std::sync::PoisonError::into_inner);
868 if let Some(engine) = engines.get(id) {
869 return Ok(Arc::clone(engine));
870 }
871 let Some(domain_id) = registry.node_domain_ids.get(id).copied() else {
872 return Err(CpuPlacementError::InternalState {
873 requested,
874 backend: self.kind,
875 message: "managed NUMA node has no coordinator-stable domain ID",
876 });
877 };
878 let engine = Arc::new(
879 CpuEngine::new_managed(
880 domain_id,
881 placement.clone(),
882 registry.thread_budget,
883 self.buffer_limit.load(Ordering::Relaxed),
884 )
885 .map_err(|error| {
886 CpuPlacementError::EngineConstruction {
887 requested,
888 backend: self.kind,
889 source: CpuEngineConstructionError::Context(error),
890 }
891 })?,
892 );
893 self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
894 engines.insert(*id, Arc::clone(&engine));
895 Ok(engine)
896 }
897 ResolvedCpuPlacement::AllAllowed { .. } => {
898 if let Some(engine) = registry.all_allowed.get() {
899 return Ok(Arc::clone(engine));
900 }
901 let _build = registry
902 .all_allowed_build
903 .lock()
904 .unwrap_or_else(std::sync::PoisonError::into_inner);
905 if let Some(engine) = registry.all_allowed.get() {
906 return Ok(Arc::clone(engine));
907 }
908 let engine = Arc::new(
909 CpuEngine::new_managed(
910 CpuDomainId::new(0),
911 placement.clone(),
912 registry.thread_budget,
913 self.buffer_limit.load(Ordering::Relaxed),
914 )
915 .map_err(|error| {
916 CpuPlacementError::EngineConstruction {
917 requested,
918 backend: self.kind,
919 source: CpuEngineConstructionError::Context(error),
920 }
921 })?,
922 );
923 self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
924 let _ = registry.all_allowed.set(Arc::clone(&engine));
925 Ok(engine)
926 }
927 }
928 }
929
930 fn configure_new_indexed_plan_cache(
931 &self,
932 engine: &CpuEngine,
933 requested: CpuPlacement,
934 limits: IndexedPlanCacheLimits,
935 ) -> Result<(), CpuPlacementError> {
936 let mut resources =
937 engine
938 .resources
939 .lock()
940 .map_err(|_| CpuPlacementError::InternalState {
941 requested,
942 backend: self.kind,
943 message: "new CPU engine indexed-plan cache lock is poisoned",
944 })?;
945 resources.indexed_plan_cache.set_limits(limits);
946 Ok(())
947 }
948
949 fn managed_base_engine(
950 &self,
951 requested: CpuPlacement,
952 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
953 match &self.engines {
954 CpuEngineRegistry::ManagedLazy(registry) => Ok(Arc::clone(®istry.base_engine)),
955 CpuEngineRegistry::ExternalPrebuilt(_) => Err(CpuPlacementError::InternalState {
956 requested,
957 backend: self.kind,
958 message: "managed compatibility placement requested from an external registry",
959 }),
960 }
961 }
962
963 fn external_engine_for(
964 &self,
965 requested: CpuPlacement,
966 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
967 let CpuEngineRegistry::ExternalPrebuilt(registry) = &self.engines else {
968 return Err(CpuPlacementError::InternalState {
969 requested,
970 backend: self.kind,
971 message: "external placement requested from a managed engine registry",
972 });
973 };
974 let engine = match requested {
975 CpuPlacement::Auto => registry.by_id.get(®istry.default_domain),
976 CpuPlacement::NumaNode(id) => registry.by_node.get(&id),
977 CpuPlacement::AllAllowed => registry.all_allowed.as_ref(),
978 };
979 engine
980 .cloned()
981 .ok_or(CpuPlacementError::UnregisteredExternalPlacement { requested })
982 }
983
984 fn external_engine_for_id(
985 &self,
986 domain: CpuDomainId,
987 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
988 let CpuEngineRegistry::ExternalPrebuilt(registry) = &self.engines else {
989 return Err(CpuPlacementError::UnregisteredExternalDomain { domain });
990 };
991 registry
992 .by_id
993 .get(&domain)
994 .cloned()
995 .ok_or(CpuPlacementError::UnregisteredExternalDomain { domain })
996 }
997
998 fn is_external(&self) -> bool {
999 matches!(&self.engines, CpuEngineRegistry::ExternalPrebuilt(_))
1000 }
1001
1002 fn initialized_engines(&self, op: &'static str) -> crate::Result<Vec<Arc<CpuEngine>>> {
1003 let mut engines = match &self.engines {
1004 CpuEngineRegistry::ManagedLazy(registry) => {
1005 let mut engines = vec![Arc::clone(®istry.base_engine)];
1006 if let Some(engine) = registry.all_allowed.get() {
1007 engines.push(Arc::clone(engine));
1008 }
1009 engines.extend(
1010 registry
1011 .node_engines
1012 .lock()
1013 .map_err(|_| poisoned_cpu_lock(op, "CPU engine registry"))?
1014 .values()
1015 .cloned(),
1016 );
1017 engines
1018 }
1019 CpuEngineRegistry::ExternalPrebuilt(registry) => {
1020 registry.by_id.values().cloned().collect()
1021 }
1022 };
1023 if engines.len() > 1 {
1024 engines.sort_unstable_by_key(|engine| Arc::as_ptr(engine) as usize);
1025 engines.dedup_by(|left, right| Arc::ptr_eq(left, right));
1026 }
1027 Ok(engines)
1028 }
1029}
1030
1031fn poisoned_cpu_lock(op: &'static str, lock: &'static str) -> crate::Error {
1032 crate::Error::runtime_state(op, format!("{lock} lock poisoned"))
1033}
1034
1035fn lock_engine_resources<'a>(
1036 engine: &'a CpuEngine,
1037 op: &'static str,
1038) -> crate::Result<std::sync::MutexGuard<'a, EngineResources>> {
1039 engine
1040 .resources
1041 .lock()
1042 .map_err(|_| poisoned_cpu_lock(op, "CPU engine resources"))
1043}
1044
1045fn saturating_add_tensor_cache_stats(total: &mut CacheStats, value: CacheStats) {
1046 total.entries = total.entries.saturating_add(value.entries);
1047 total.retained_bytes = total.retained_bytes.saturating_add(value.retained_bytes);
1048 total.hits = total.hits.saturating_add(value.hits);
1049 total.misses = total.misses.saturating_add(value.misses);
1050 total.evictions = total.evictions.saturating_add(value.evictions);
1051 total.clears = total.clears.saturating_add(value.clears);
1052}
1053
1054#[doc(hidden)]
1069struct CpuBackendSessionMarker;
1070
1071#[derive(Clone)]
1072pub struct CpuBackend {
1073 runtime_identity: CpuRuntimeIdentity,
1074 shared: Arc<CpuBackendState>,
1075 requested: CpuPlacement,
1076 resolved: ResolvedCpuExecution,
1077 engine: Arc<CpuEngine>,
1078 provider_bundle: CpuProviderBundle,
1079 allocation_domain: Option<Arc<dyn SharedTensorAllocationDomain>>,
1080}
1081
1082#[derive(Clone, Debug)]
1098pub struct CpuRuntimeIdentity {
1099 marker: Arc<()>,
1100}
1101
1102impl CpuRuntimeIdentity {
1103 fn fresh() -> Self {
1104 Self {
1105 marker: Arc::new(()),
1106 }
1107 }
1108}
1109
1110impl PartialEq for CpuRuntimeIdentity {
1111 fn eq(&self, other: &Self) -> bool {
1112 Arc::ptr_eq(&self.marker, &other.marker)
1113 }
1114}
1115
1116impl Eq for CpuRuntimeIdentity {}
1117
1118fn resolve_discovered_topology(
1119 kind: CpuBackendKind,
1120 topology: Result<CpuTopology, CpuTopologyError>,
1121) -> Result<CpuTopology, CpuPlacementError> {
1122 topology.map_err(|source| CpuPlacementError::TopologyDiscovery {
1123 requested: CpuPlacement::Auto,
1124 backend: kind,
1125 source,
1126 })
1127}
1128
1129fn external_engine_resolution(
1130 engine: &CpuEngine,
1131 requested: CpuPlacement,
1132 kind: CpuBackendKind,
1133) -> Result<ResolvedCpuExecution, CpuPlacementError> {
1134 match engine.domain().admission_mode() {
1135 CpuAdmissionMode::CooperativeCpuSet => engine
1136 .placement()
1137 .cloned()
1138 .map(ResolvedCpuExecution::ExternalManaged)
1139 .ok_or(CpuPlacementError::InternalState {
1140 requested,
1141 backend: kind,
1142 message: "cooperative external domain has no placement",
1143 }),
1144 CpuAdmissionMode::CallerManaged => Ok(ResolvedCpuExecution::ExternalCallerManaged),
1145 }
1146}
1147
1148fn external_domain_backend_kind(
1149 op: &'static str,
1150 domains: &[ExternalCpuDomain],
1151) -> Result<CpuBackendKind, CpuBackendError> {
1152 let kind = if domains
1153 .iter()
1154 .any(|domain| domain.admission_mode() == CpuAdmissionMode::CallerManaged)
1155 {
1156 CpuBackendKind::Faer
1157 } else {
1158 CpuBackendKind::default_compiled()
1159 };
1160 ensure_cpu_backend_kind_available(kind, op)
1161 .map_err(|error| constructor_tensor_error(op, error))?;
1162 Ok(kind)
1163}
1164
1165fn coordinator_node_domain_ids(topology: &CpuTopology) -> BTreeMap<NumaNodeId, CpuDomainId> {
1166 topology
1167 .nodes()
1168 .iter()
1169 .enumerate()
1170 .filter_map(|(index, node)| {
1171 u64::try_from(index)
1172 .ok()
1173 .and_then(|index| index.checked_add(1))
1174 .map(|id| (node.id(), CpuDomainId::new(id)))
1175 })
1176 .collect()
1177}
1178
1179impl fmt::Debug for CpuBackend {
1180 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1181 f.debug_struct("CpuBackend")
1182 .field("kind", &self.kind())
1183 .field("provider_bundle", &self.provider_bundle)
1184 .field("requested_placement", &self.requested)
1185 .field("resolved_execution", &self.resolved)
1186 .field("engine_placement", &self.engine.placement())
1187 .field("num_threads", &self.num_threads())
1188 .field("allocation_domain", &self.allocation_domain())
1189 .field("buffer_pool_cache_stats", &self.buffer_pool_cache_stats())
1190 .field("buffer_pool_limit_bytes", &self.buffer_pool_limit_bytes())
1191 .finish_non_exhaustive()
1192 }
1193}
1194
1195impl CpuBackend {
1196 fn from_thread_budget_and_kind(
1197 thread_budget: usize,
1198 kind: CpuBackendKind,
1199 max_retained_capacity_bytes: usize,
1200 ) -> Result<Self, CpuPlacementError> {
1201 let topology = resolve_discovered_topology(kind, discover_cpu_topology())?;
1202 let resolved = resolve_placement(kind, CpuPlacement::Auto, &topology)?;
1203 #[cfg(not(any(target_os = "linux", target_os = "android")))]
1204 {
1205 let context = CpuContext::with_threads(thread_budget).map_err(|error| {
1206 CpuPlacementError::EngineConstruction {
1207 requested: CpuPlacement::Auto,
1208 backend: kind,
1209 source: CpuEngineConstructionError::Tensor(error),
1210 }
1211 })?;
1212 Ok(Self::compatibility_with_topology(
1213 Arc::new(context),
1214 max_retained_capacity_bytes,
1215 kind,
1216 topology,
1217 resolved,
1218 ))
1219 }
1220 #[cfg(any(target_os = "linux", target_os = "android"))]
1221 {
1222 let engine_placement = ResolvedCpuPlacement::AllAllowed {
1223 cpus: topology.allowed_cpus().clone(),
1224 };
1225 let engine = Arc::new(
1226 CpuEngine::new_managed(
1227 CpuDomainId::new(0),
1228 engine_placement,
1229 thread_budget,
1230 max_retained_capacity_bytes,
1231 )
1232 .map_err(|error| CpuPlacementError::EngineConstruction {
1233 requested: CpuPlacement::Auto,
1234 backend: kind,
1235 source: CpuEngineConstructionError::Context(error),
1236 })?,
1237 );
1238 let all_allowed = OnceLock::new();
1239 let _ = all_allowed.set(Arc::clone(&engine));
1240 Ok(Self {
1241 shared: Arc::new(CpuBackendState {
1242 engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
1243 node_engines: Mutex::new(BTreeMap::new()),
1244 node_domain_ids: coordinator_node_domain_ids(&topology),
1245 all_allowed,
1246 all_allowed_build: Mutex::new(()),
1247 base_engine: Arc::clone(&engine),
1248 thread_budget,
1249 }),
1250 topology,
1251 arbiter: ResourceArbiter::global(),
1252 kind,
1253 buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
1254 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1255 }),
1256 runtime_identity: CpuRuntimeIdentity::fresh(),
1257 requested: CpuPlacement::Auto,
1258 resolved,
1259 engine,
1260 provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
1261 allocation_domain: None,
1262 })
1263 }
1264 }
1265
1266 fn compatibility(
1267 ctx: Arc<CpuContext>,
1268 max_retained_capacity_bytes: usize,
1269 kind: CpuBackendKind,
1270 ) -> Self {
1271 let topology = discover_cpu_topology().unwrap_or_else(|_| {
1272 let allowed = crate::process_cpu_affinity().unwrap_or_else(|| {
1273 CpuSet::new((0..crate::available_parallelism()).map(CpuId::new))
1274 .unwrap_or_else(|_| CpuSet::singleton(CpuId::new(0)))
1275 });
1276 CpuTopology::all_allowed(allowed)
1277 });
1278 let resolved = if kind == CpuBackendKind::Blas {
1279 ResolvedCpuExecution::ProviderDefaultExclusive
1280 } else {
1281 ResolvedCpuExecution::Compatibility
1282 };
1283 Self::compatibility_with_topology(
1284 ctx,
1285 max_retained_capacity_bytes,
1286 kind,
1287 topology,
1288 resolved,
1289 )
1290 }
1291
1292 fn compatibility_with_topology(
1293 ctx: Arc<CpuContext>,
1294 max_retained_capacity_bytes: usize,
1295 kind: CpuBackendKind,
1296 topology: CpuTopology,
1297 resolved: ResolvedCpuExecution,
1298 ) -> Self {
1299 let placement = ResolvedCpuPlacement::AllAllowed {
1300 cpus: topology.allowed_cpus().clone(),
1301 };
1302 let base_engine = Arc::new(CpuEngine::from_context(
1303 CpuDomainId::new(0),
1304 placement,
1305 ctx,
1306 max_retained_capacity_bytes,
1307 ));
1308 Self {
1309 shared: Arc::new(CpuBackendState {
1310 engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
1311 node_engines: Mutex::new(BTreeMap::new()),
1312 node_domain_ids: coordinator_node_domain_ids(&topology),
1313 all_allowed: OnceLock::new(),
1314 all_allowed_build: Mutex::new(()),
1315 base_engine: Arc::clone(&base_engine),
1316 thread_budget: base_engine.domain().thread_budget().get(),
1317 }),
1318 topology,
1319 arbiter: ResourceArbiter::global(),
1320 kind,
1321 buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
1322 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1323 }),
1324 runtime_identity: CpuRuntimeIdentity::fresh(),
1325 requested: CpuPlacement::Auto,
1326 resolved,
1327 engine: base_engine,
1328 provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
1329 allocation_domain: None,
1330 }
1331 }
1332
1333 pub fn new() -> Self {
1343 let context = Arc::new(CpuContext::from_env());
1344 Self::from_thread_budget_and_kind(
1345 context.num_threads(),
1346 CpuBackendKind::default_compiled(),
1347 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1348 )
1349 .unwrap_or_else(|error| {
1350 eprintln!(
1351 "tenferro_cpu: using the unpinned compatibility context after placement error: {error}"
1352 );
1353 Self::from_context(context)
1354 })
1355 }
1356
1357 pub fn from_external_managed_domains(
1415 default_domain: CpuDomainId,
1416 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1417 ) -> Result<Self, CpuBackendError> {
1418 let op = "CpuBackend::from_external_managed_domains";
1419 let domains: Vec<_> = domains.into_iter().collect();
1420 let kind = external_domain_backend_kind(op, &domains)?;
1421 let topology = resolve_discovered_topology(kind, discover_cpu_topology())
1422 .map_err(|source| CpuBackendError::placement(op, source))?;
1423 Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1424 default_domain,
1425 domains,
1426 topology,
1427 ResourceArbiter::global(),
1428 kind,
1429 CpuProviderBundle::standard(kind, false),
1430 )
1431 }
1432
1433 pub fn from_external_managed_domains_with_provider_bundle(
1488 default_domain: CpuDomainId,
1489 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1490 provider_bundle: CpuProviderBundle,
1491 ) -> Result<Self, CpuBackendError> {
1492 let op = "CpuBackend::from_external_managed_domains_with_provider_bundle";
1493 let domains: Vec<_> = domains.into_iter().collect();
1494 let kind = external_domain_backend_kind(op, &domains)?;
1495 let topology = resolve_discovered_topology(kind, discover_cpu_topology())
1496 .map_err(|source| CpuBackendError::placement(op, source))?;
1497 Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1498 default_domain,
1499 domains,
1500 topology,
1501 ResourceArbiter::global(),
1502 kind,
1503 provider_bundle,
1504 )
1505 }
1506
1507 fn from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1508 default_domain: CpuDomainId,
1509 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1510 topology: CpuTopology,
1511 arbiter: ResourceArbiter,
1512 kind: CpuBackendKind,
1513 provider_bundle: CpuProviderBundle,
1514 ) -> Result<Self, CpuBackendError> {
1515 let domains: Vec<_> = domains.into_iter().collect();
1516 if domains.is_empty() {
1517 return Err(ExternalCpuDomainRegistryError::EmptyRegistry.into());
1518 }
1519
1520 let mut domain_ids = BTreeSet::new();
1521 let mut node_ids = BTreeSet::new();
1522 let mut has_all_allowed = false;
1523 for domain in &domains {
1524 if !domain_ids.insert(domain.id()) {
1525 return Err(
1526 ExternalCpuDomainRegistryError::DuplicateDomainId { id: domain.id() }.into(),
1527 );
1528 }
1529 if let Some(placement) = domain.placement() {
1530 match placement {
1531 ResolvedCpuPlacement::NumaNode { id, .. } => {
1532 if !node_ids.insert(*id) {
1533 return Err(
1534 ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
1535 placement: CpuPlacement::NumaNode(*id),
1536 }
1537 .into(),
1538 );
1539 }
1540 }
1541 ResolvedCpuPlacement::AllAllowed { cpus } => {
1542 if has_all_allowed {
1543 return Err(
1544 ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
1545 placement: CpuPlacement::AllAllowed,
1546 }
1547 .into(),
1548 );
1549 }
1550 has_all_allowed = true;
1551 if domain.placement_guarantee()
1552 == Some(CpuPlacementGuarantee::ExactDeclared)
1553 && cpus != topology.allowed_cpus()
1554 {
1555 return Err(ExternalCpuDomainRegistryError::ExactAllAllowedMismatch {
1556 domain: domain.id(),
1557 declared: cpus.clone(),
1558 allowed: topology.allowed_cpus().clone(),
1559 }
1560 .into());
1561 }
1562 }
1563 }
1564 if let Some(cpu) = placement
1565 .cpus()
1566 .as_slice()
1567 .iter()
1568 .copied()
1569 .find(|cpu| !topology.allowed_cpus().contains(*cpu))
1570 {
1571 return Err(ExternalCpuDomainRegistryError::CpuOutsideAllowedSet {
1572 domain: domain.id(),
1573 cpu,
1574 }
1575 .into());
1576 }
1577 }
1578 }
1579 let buffer_limit = crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES;
1580 let mut by_id = BTreeMap::new();
1581 let mut by_node = BTreeMap::new();
1582 let mut all_allowed = None;
1583 for domain in domains {
1584 let id = domain.id();
1585 let placement = domain.placement().cloned();
1586 let engine = Arc::new(CpuEngine::from_external(domain, buffer_limit));
1587 match placement {
1588 Some(ResolvedCpuPlacement::NumaNode { id, .. }) => {
1589 by_node.insert(id, Arc::clone(&engine));
1590 }
1591 Some(ResolvedCpuPlacement::AllAllowed { .. }) => {
1592 all_allowed = Some(Arc::clone(&engine));
1593 }
1594 None => {}
1595 }
1596 by_id.insert(id, engine);
1597 }
1598 let Some(engine) = by_id.get(&default_domain).cloned() else {
1599 return Err(
1600 ExternalCpuDomainRegistryError::MissingDefaultDomain { default_domain }.into(),
1601 );
1602 };
1603 let resolved = match engine.domain().admission_mode() {
1604 CpuAdmissionMode::CooperativeCpuSet => ResolvedCpuExecution::ExternalManaged(
1605 engine.placement().cloned().ok_or_else(|| {
1606 CpuBackendError::placement(
1607 "CpuBackend external domain resolution",
1608 CpuPlacementError::InternalState {
1609 requested: CpuPlacement::Auto,
1610 backend: kind,
1611 message: "cooperative external domain has no placement",
1612 },
1613 )
1614 })?,
1615 ),
1616 CpuAdmissionMode::CallerManaged => ResolvedCpuExecution::ExternalCallerManaged,
1617 };
1618 let backend = Self {
1619 runtime_identity: CpuRuntimeIdentity::fresh(),
1620 shared: Arc::new(CpuBackendState {
1621 topology,
1622 engines: CpuEngineRegistry::ExternalPrebuilt(ExternalEngineRegistry {
1623 by_id,
1624 by_node,
1625 all_allowed,
1626 default_domain,
1627 }),
1628 arbiter,
1629 kind,
1630 buffer_limit: AtomicUsize::new(buffer_limit),
1631 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1632 }),
1633 requested: CpuPlacement::Auto,
1634 resolved,
1635 engine,
1636 provider_bundle,
1637 allocation_domain: None,
1638 };
1639 backend
1640 .validate_provider_bundle_for_domains(&backend.provider_bundle)
1641 .map_err(|source| {
1642 CpuBackendError::Tensor(crate::Error::backend_source(
1643 "CpuBackend ExternalManaged provider validation",
1644 source,
1645 ))
1646 })?;
1647 Ok(backend)
1648 }
1649
1650 pub fn with_kind(kind: CpuBackendKind) -> Result<Self, CpuBackendError> {
1667 let op = "CpuBackend::with_kind";
1668 ensure_cpu_backend_kind_available(kind, op)
1669 .map_err(|error| constructor_tensor_error(op, error))?;
1670 let context = CpuContext::from_env();
1671 Self::from_thread_budget_and_kind(
1672 context.num_threads(),
1673 kind,
1674 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1675 )
1676 .map_err(|error| CpuBackendError::placement(op, error))
1677 }
1678
1679 pub fn try_new() -> Result<Self, CpuBackendError> {
1698 let op = "CpuBackend::try_new";
1699 let context =
1700 CpuContext::try_from_env().map_err(|error| constructor_tensor_error(op, error))?;
1701 Self::from_thread_budget_and_kind(
1702 context.num_threads(),
1703 CpuBackendKind::default_compiled(),
1704 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1705 )
1706 .map_err(|error| CpuBackendError::placement(op, error))
1707 }
1708
1709 pub fn from_context(ctx: Arc<CpuContext>) -> Self {
1722 Self::compatibility(
1723 ctx,
1724 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1725 CpuBackendKind::default_compiled(),
1726 )
1727 }
1728
1729 pub fn from_context_with_buffer_pool_limit(
1745 ctx: Arc<CpuContext>,
1746 max_retained_capacity_bytes: usize,
1747 ) -> Self {
1748 Self::from_context_with_buffer_pool_limit_and_kind(
1749 ctx,
1750 max_retained_capacity_bytes,
1751 CpuBackendKind::default_compiled(),
1752 )
1753 }
1754
1755 fn from_context_with_buffer_pool_limit_and_kind(
1756 ctx: Arc<CpuContext>,
1757 max_retained_capacity_bytes: usize,
1758 kind: CpuBackendKind,
1759 ) -> Self {
1760 Self::compatibility(ctx, max_retained_capacity_bytes, kind)
1761 }
1762
1763 pub fn with_threads(num_threads: usize) -> Result<Self, CpuBackendError> {
1780 let op = "CpuBackend::with_threads";
1781 let context = CpuContext::with_threads(num_threads)
1782 .map_err(|error| constructor_tensor_error(op, error))?;
1783 Self::from_thread_budget_and_kind(
1784 context.num_threads(),
1785 CpuBackendKind::default_compiled(),
1786 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1787 )
1788 .map_err(|error| CpuBackendError::placement(op, error))
1789 }
1790
1791 pub fn with_threads_and_kind(
1812 num_threads: usize,
1813 kind: CpuBackendKind,
1814 ) -> Result<Self, CpuBackendError> {
1815 let op = "CpuBackend::with_threads_and_kind";
1816 ensure_cpu_backend_kind_available(kind, op)
1817 .map_err(|error| constructor_tensor_error(op, error))?;
1818 let context = CpuContext::with_threads(num_threads)
1819 .map_err(|error| constructor_tensor_error(op, error))?;
1820 Self::from_thread_budget_and_kind(
1821 context.num_threads(),
1822 kind,
1823 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1824 )
1825 .map_err(|error| CpuBackendError::placement(op, error))
1826 }
1827
1828 pub fn for_placement(&self, requested: CpuPlacement) -> Result<Self, CpuPlacementError> {
1852 self.for_placement_with_affinity(
1853 requested,
1854 cfg!(any(target_os = "linux", target_os = "android")),
1855 )
1856 }
1857
1858 pub fn for_domain(&self, domain: CpuDomainId) -> Result<Self, CpuPlacementError> {
1885 let engine = self.shared.external_engine_for_id(domain)?;
1886 let resolved = external_engine_resolution(&engine, CpuPlacement::Auto, self.kind())?;
1887 Ok(Self {
1888 runtime_identity: CpuRuntimeIdentity::fresh(),
1889 shared: Arc::clone(&self.shared),
1890 requested: CpuPlacement::Auto,
1891 resolved,
1892 engine,
1893 provider_bundle: self.provider_bundle.clone(),
1894 allocation_domain: self.allocation_domain.clone(),
1895 })
1896 }
1897
1898 fn for_placement_with_affinity(
1899 &self,
1900 requested: CpuPlacement,
1901 managed_affinity_available: bool,
1902 ) -> Result<Self, CpuPlacementError> {
1903 if self.shared.is_external() {
1904 let engine = self.shared.external_engine_for(requested)?;
1905 let resolved = external_engine_resolution(&engine, requested, self.kind())?;
1906 return Ok(Self {
1907 runtime_identity: CpuRuntimeIdentity::fresh(),
1908 shared: Arc::clone(&self.shared),
1909 requested,
1910 resolved,
1911 engine,
1912 provider_bundle: self.provider_bundle.clone(),
1913 allocation_domain: self.allocation_domain.clone(),
1914 });
1915 }
1916 let resolved = resolve_placement_with_affinity(
1917 self.kind(),
1918 requested,
1919 &self.shared.topology,
1920 managed_affinity_available,
1921 )?;
1922 if requested == CpuPlacement::Auto && !managed_affinity_available {
1923 return Ok(Self {
1924 runtime_identity: CpuRuntimeIdentity::fresh(),
1925 shared: Arc::clone(&self.shared),
1926 requested,
1927 resolved,
1928 engine: self.shared.managed_base_engine(requested)?,
1929 provider_bundle: self.provider_bundle.clone(),
1930 allocation_domain: self.allocation_domain.clone(),
1931 });
1932 }
1933 let engine_placement = match &resolved {
1934 ResolvedCpuExecution::Managed(placement) => placement.clone(),
1935 ResolvedCpuExecution::ExternalManaged(_)
1936 | ResolvedCpuExecution::ExternalCallerManaged => {
1937 return Err(CpuPlacementError::InternalState {
1938 requested,
1939 backend: self.kind(),
1940 message: "managed resolver returned an external execution mode",
1941 });
1942 }
1943 ResolvedCpuExecution::ProviderDefaultExclusive => ResolvedCpuPlacement::AllAllowed {
1944 cpus: self.shared.topology.allowed_cpus().clone(),
1945 },
1946 ResolvedCpuExecution::Compatibility => {
1947 return Err(CpuPlacementError::InternalState {
1948 requested,
1949 backend: self.kind(),
1950 message: "placement resolution returned an internal compatibility mode",
1951 });
1952 }
1953 };
1954 let engine = self
1955 .shared
1956 .managed_engine_for(&engine_placement, requested)?;
1957 Ok(Self {
1958 runtime_identity: CpuRuntimeIdentity::fresh(),
1959 shared: Arc::clone(&self.shared),
1960 requested,
1961 resolved,
1962 engine,
1963 provider_bundle: self.provider_bundle.clone(),
1964 allocation_domain: self.allocation_domain.clone(),
1965 })
1966 }
1967
1968 pub fn placement(&self) -> CpuPlacement {
1978 self.requested
1979 }
1980
1981 pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
2000 match &self.resolved {
2001 ResolvedCpuExecution::Managed(placement)
2002 | ResolvedCpuExecution::ExternalManaged(placement) => Some(placement),
2003 ResolvedCpuExecution::Compatibility
2004 | ResolvedCpuExecution::ExternalCallerManaged
2005 | ResolvedCpuExecution::ProviderDefaultExclusive => None,
2006 }
2007 }
2008
2009 pub fn topology(&self) -> &CpuTopology {
2019 &self.shared.topology
2020 }
2021
2022 pub fn supports_placement(&self, placement: CpuPlacement) -> bool {
2032 if self.shared.is_external() {
2033 self.shared.external_engine_for(placement).is_ok()
2034 } else {
2035 resolve_placement(self.kind(), placement, &self.shared.topology).is_ok()
2036 }
2037 }
2038
2039 pub fn execution_info(&self) -> CpuExecutionInfo {
2048 let domain = self.engine.domain();
2049 let capabilities = domain.executor_capabilities();
2050 let (executor_affinity, executor_shutdown) =
2051 match (domain.ownership(), domain.admission_mode()) {
2052 (CpuDomainOwnership::ExternalManaged, CpuAdmissionMode::CooperativeCpuSet) => (
2053 CpuExecutorAffinity::CallerDeclaredUnverified,
2054 CpuExecutorShutdown::CallerOwned,
2055 ),
2056 (CpuDomainOwnership::ExternalManaged, CpuAdmissionMode::CallerManaged) => {
2057 (capabilities.affinity, CpuExecutorShutdown::CallerOwned)
2058 }
2059 (CpuDomainOwnership::Managed, _) => (capabilities.affinity, capabilities.shutdown),
2060 };
2061 CpuExecutionInfo {
2062 backend_kind: self.kind(),
2063 execution_mode: match &self.resolved {
2064 ResolvedCpuExecution::Managed(_) => CpuExecutionMode::Managed,
2065 ResolvedCpuExecution::ExternalManaged(_) => CpuExecutionMode::ExternalManaged,
2066 ResolvedCpuExecution::ExternalCallerManaged => CpuExecutionMode::CallerManaged,
2067 ResolvedCpuExecution::ProviderDefaultExclusive => {
2068 CpuExecutionMode::ProviderDefaultExclusive
2069 }
2070 ResolvedCpuExecution::Compatibility => CpuExecutionMode::Compatibility,
2071 },
2072 requested_placement: self.requested,
2073 resolved_placement: self.resolved_placement().cloned(),
2074 topology: self.shared.topology.clone(),
2075 domain_id: domain.id(),
2076 domain_cpus: domain.cpus().cloned(),
2077 worker_count: capabilities.worker_count.get(),
2078 thread_budget: domain.thread_budget().get(),
2079 placement_guarantee: domain.placement_guarantee(),
2080 admission_mode: domain.admission_mode(),
2081 domain_ownership: domain.ownership(),
2082 executor_affinity,
2083 executor_shutdown,
2084 provider_diagnostic: provider_diagnostic(
2085 self.kind(),
2086 domain.ownership(),
2087 domain.admission_mode(),
2088 ),
2089 }
2090 }
2091
2092 #[cfg(all(
2093 test,
2094 feature = "cpu-faer",
2095 any(target_os = "linux", target_os = "android")
2096 ))]
2097 fn coordinator_id_for_test(&self) -> usize {
2098 Arc::as_ptr(&self.shared) as usize
2099 }
2100
2101 #[cfg(test)]
2102 pub(crate) fn context_id_for_test(&self) -> usize {
2103 Arc::as_ptr(self.engine.domain().executor()) as *const () as usize
2104 }
2105
2106 pub fn kind(&self) -> CpuBackendKind {
2117 self.shared.kind
2118 }
2119
2120 pub fn provider_bundle(&self) -> &CpuProviderBundle {
2122 &self.provider_bundle
2123 }
2124
2125 pub fn runtime_identity(&self) -> CpuRuntimeIdentity {
2132 self.runtime_identity.clone()
2133 }
2134
2135 pub fn with_provider_bundle(
2155 mut self,
2156 bundle: CpuProviderBundle,
2157 ) -> Result<Self, CpuProviderBundleInstallError> {
2158 self.validate_provider_bundle_for_domains(&bundle)?;
2159 self.provider_bundle = bundle;
2160 self.runtime_identity = CpuRuntimeIdentity::fresh();
2161 Ok(self)
2162 }
2163
2164 fn validate_provider_bundle_for_domains(
2165 &self,
2166 bundle: &CpuProviderBundle,
2167 ) -> Result<(), CpuProviderBundleInstallError> {
2168 let allowed = self.shared.topology.allowed_cpus();
2169 let validate_engine = |engine: &CpuEngine| {
2170 let domain = engine.domain();
2171 let contract = match (domain.placement_guarantee(), domain.cpus()) {
2172 (Some(placement_guarantee), Some(domain_cpus)) => {
2173 CpuProviderDomainContract::CooperativeCpuSet {
2174 placement_guarantee,
2175 domain_cpus,
2176 process_allowed_cpus: allowed,
2177 }
2178 }
2179 (None, None) => CpuProviderDomainContract::CallerManaged,
2180 _ => unreachable!("CPU domain placement and guarantee must match"),
2183 };
2184 bundle.validate_for_domain(domain.id(), domain.thread_budget(), contract)
2185 };
2186
2187 match &self.shared.engines {
2188 CpuEngineRegistry::ExternalPrebuilt(registry) => {
2189 for engine in registry.by_id.values() {
2190 validate_engine(engine)?;
2191 }
2192 }
2193 CpuEngineRegistry::ManagedLazy(registry) => {
2194 validate_engine(®istry.base_engine)?;
2195
2196 #[cfg(any(target_os = "linux", target_os = "android"))]
2201 for node in self.shared.topology.nodes() {
2202 let Some(domain_id) = registry.node_domain_ids.get(&node.id()).copied() else {
2203 continue;
2204 };
2205 let budget =
2206 std::num::NonZeroUsize::new(registry.thread_budget.min(node.cpus().len()))
2207 .expect("usable topology nodes have non-empty CPU sets");
2208 bundle.validate_for_domain(
2209 domain_id,
2210 budget,
2211 CpuProviderDomainContract::CooperativeCpuSet {
2212 placement_guarantee: CpuPlacementGuarantee::ExactDeclared,
2213 domain_cpus: node.cpus(),
2214 process_allowed_cpus: allowed,
2215 },
2216 )?;
2217 }
2218 }
2219 }
2220 Ok(())
2221 }
2222
2223 pub fn num_threads(&self) -> usize {
2234 self.engine.domain().thread_budget().get()
2235 }
2236
2237 pub fn buffer_pool_len(&self) -> crate::Result<usize> {
2254 self.shared
2255 .initialized_engines("CpuBackend::buffer_pool_len")?
2256 .iter()
2257 .try_fold(0, |total, engine| {
2258 Ok(total
2259 + lock_engine_resources(engine, "CpuBackend::buffer_pool_len")?
2260 .buffers
2261 .len())
2262 })
2263 }
2264
2265 pub fn buffer_pool_stats(&self) -> crate::Result<BufferPoolStats> {
2284 self.shared
2285 .initialized_engines("CpuBackend::buffer_pool_stats")?
2286 .iter()
2287 .try_fold(BufferPoolStats::default(), |mut total, engine| {
2288 let stats = lock_engine_resources(engine, "CpuBackend::buffer_pool_stats")?
2289 .buffers
2290 .stats();
2291 total.buffers += stats.buffers;
2292 total.capacity_bytes += stats.capacity_bytes;
2293 Ok(total)
2294 })
2295 }
2296
2297 pub fn buffer_pool_cache_stats(&self) -> crate::Result<CacheStats> {
2316 let stats = self.buffer_pool_stats()?;
2317 Ok(CacheStats {
2318 entries: stats.buffers,
2319 retained_bytes: stats.capacity_bytes,
2320 hits: 0,
2321 misses: 0,
2322 evictions: 0,
2323 clears: 0,
2324 })
2325 }
2326
2327 pub fn indexed_plan_cache_limits(&self) -> crate::Result<IndexedPlanCacheLimits> {
2344 self.shared
2345 .indexed_plan_cache_limits
2346 .lock()
2347 .map(|limits| *limits)
2348 .map_err(|_| {
2349 poisoned_cpu_lock(
2350 "CpuBackend::indexed_plan_cache_limits",
2351 "CPU indexed-plan cache configuration",
2352 )
2353 })
2354 }
2355
2356 pub fn set_indexed_plan_cache_limits(
2377 &mut self,
2378 limits: IndexedPlanCacheLimits,
2379 ) -> crate::Result<()> {
2380 let mut configured_limits = self.shared.indexed_plan_cache_limits.lock().map_err(|_| {
2384 poisoned_cpu_lock(
2385 "CpuBackend::set_indexed_plan_cache_limits",
2386 "CPU indexed-plan cache configuration",
2387 )
2388 })?;
2389 let engines = self
2390 .shared
2391 .initialized_engines("CpuBackend::set_indexed_plan_cache_limits")?;
2392 let mut resources = engines
2393 .iter()
2394 .map(|engine| {
2395 lock_engine_resources(engine, "CpuBackend::set_indexed_plan_cache_limits")
2396 })
2397 .collect::<crate::Result<Vec<_>>>()?;
2398 *configured_limits = limits;
2399 for resource in &mut resources {
2400 resource.indexed_plan_cache.set_limits(limits);
2401 }
2402 Ok(())
2403 }
2404
2405 pub fn indexed_plan_cache_stats(&self) -> crate::Result<CacheStats> {
2422 self.shared
2423 .initialized_engines("CpuBackend::indexed_plan_cache_stats")?
2424 .iter()
2425 .try_fold(CacheStats::default(), |mut total, engine| {
2426 let stats = lock_engine_resources(engine, "CpuBackend::indexed_plan_cache_stats")?
2427 .indexed_plan_cache
2428 .stats();
2429 saturating_add_tensor_cache_stats(&mut total, stats);
2430 Ok(total)
2431 })
2432 }
2433
2434 pub fn clear_indexed_plan_cache(&mut self) -> crate::Result<()> {
2452 let engines = self
2453 .shared
2454 .initialized_engines("CpuBackend::clear_indexed_plan_cache")?;
2455 let mut resources = engines
2456 .iter()
2457 .map(|engine| lock_engine_resources(engine, "CpuBackend::clear_indexed_plan_cache"))
2458 .collect::<crate::Result<Vec<_>>>()?;
2459 for resource in &mut resources {
2460 resource.indexed_plan_cache.clear();
2461 }
2462 Ok(())
2463 }
2464
2465 pub fn buffer_pool_limit_bytes(&self) -> usize {
2480 self.shared.buffer_limit.load(Ordering::Relaxed)
2481 }
2482
2483 pub fn set_buffer_pool_limit_bytes(
2506 &mut self,
2507 max_retained_capacity_bytes: usize,
2508 ) -> crate::Result<()> {
2509 let engines = self
2510 .shared
2511 .initialized_engines("CpuBackend::set_buffer_pool_limit_bytes")?;
2512 let mut resources = engines
2513 .iter()
2514 .map(|engine| lock_engine_resources(engine, "CpuBackend::set_buffer_pool_limit_bytes"))
2515 .collect::<crate::Result<Vec<_>>>()?;
2516 self.shared
2517 .buffer_limit
2518 .store(max_retained_capacity_bytes, Ordering::Relaxed);
2519 for resource in &mut resources {
2520 resource
2521 .buffers
2522 .set_max_retained_capacity_bytes(max_retained_capacity_bytes);
2523 }
2524 Ok(())
2525 }
2526
2527 pub fn reset_buffer_pool(&mut self) -> crate::Result<()> {
2550 let engines = self
2551 .shared
2552 .initialized_engines("CpuBackend::reset_buffer_pool")?;
2553 let mut resources = engines
2554 .iter()
2555 .map(|engine| lock_engine_resources(engine, "CpuBackend::reset_buffer_pool"))
2556 .collect::<crate::Result<Vec<_>>>()?;
2557 for resource in &mut resources {
2558 resource.buffers.clear();
2559 }
2560 Ok(())
2561 }
2562
2563 pub(crate) fn runtime_cache_stats(
2564 &self,
2565 ) -> crate::Result<tenferro_runtime::runtime::CacheStats> {
2566 let resources = lock_engine_resources(&self.engine, "CpuBackend::runtime_cache_stats")?;
2567 let buffers = resources.buffers.cache_stats();
2568 let gemm = tenferro_tensor::RuntimeCacheControl::stats(&resources.gemm_analysis_cache);
2569 let indexed = resources.indexed_plan_cache.stats();
2570 Ok(tenferro_runtime::runtime::CacheStats {
2571 entries: buffers
2572 .entries
2573 .saturating_add(gemm.entries)
2574 .saturating_add(indexed.entries),
2575 retained_bytes: buffers
2576 .retained_bytes
2577 .saturating_add(gemm.retained_bytes)
2578 .saturating_add(indexed.retained_bytes),
2579 hits: indexed.hits,
2580 misses: indexed.misses,
2581 evictions: indexed.evictions,
2582 clears: indexed.clears,
2583 })
2584 }
2585
2586 pub(crate) fn clear_runtime_caches(&self) -> crate::Result<()> {
2587 let mut resources =
2588 lock_engine_resources(&self.engine, "CpuBackend::clear_runtime_caches")?;
2589 resources.buffers.clear();
2590 tenferro_tensor::RuntimeCacheControl::clear(&mut resources.gemm_analysis_cache);
2591 resources.indexed_plan_cache.clear();
2592 Ok(())
2593 }
2594
2595 pub fn install<R: Send>(&self, op: impl FnOnce() -> R + Send) -> R {
2616 let admission = self.infallible_execution_admission();
2617 let permit = admission.permit();
2618 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
2619 match entry.enter(ParallelMode::Sequential, |_| op()) {
2620 Ok(result) => result,
2621 Err(error) => panic!("CpuBackend::install executor failed: {error}"),
2622 }
2623 }
2624
2625 fn try_install<R: Send>(
2626 &self,
2627 op: impl FnOnce() -> crate::Result<R> + Send,
2628 ) -> crate::Result<R> {
2629 let admission = self.execution_admission()?;
2630 let permit = admission.permit();
2631 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
2632 let mode = entry.preferred_engine_mode();
2633 entry
2634 .enter(mode, |context| context.with_native_parallelism(op))
2635 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2636 }
2637
2638 fn try_install_with_context<R: Send>(
2639 &self,
2640 op: impl FnOnce(&CpuExecutionContext<'_>) -> crate::Result<R> + Send,
2641 ) -> crate::Result<R> {
2642 let admission = self.execution_admission()?;
2643 let permit = admission.permit();
2644 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
2645 let mode = entry.preferred_engine_mode();
2646 entry
2647 .enter(mode, |context| {
2648 context.with_native_parallelism(|| op(context))
2649 })
2650 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2651 }
2652
2653 fn try_install_fresh<R: FreshCpuOutput + Send>(
2654 &self,
2655 op: impl FnOnce() -> crate::Result<R> + Send,
2656 ) -> crate::Result<R> {
2657 let domain = self.engine.domain().id();
2658 let mut output = self.try_install(op)?;
2659 output.tag_fresh(domain);
2660 Ok(output)
2661 }
2662
2663 fn try_install_fresh_with_context<R: FreshCpuOutput + Send>(
2664 &self,
2665 op: impl FnOnce(&CpuExecutionContext<'_>) -> crate::Result<R> + Send,
2666 ) -> crate::Result<R> {
2667 let domain = self.engine.domain().id();
2668 let mut output = self.try_install_with_context(op)?;
2669 output.tag_fresh(domain);
2670 Ok(output)
2671 }
2672
2673 fn install_with_pool_unmarked<R: Send>(
2674 &mut self,
2675 op: impl FnOnce(&mut BufferPool) -> crate::Result<R> + Send,
2676 ) -> crate::Result<R> {
2677 let admission = self.execution_admission()?;
2678 let permit = admission.permit();
2679 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
2680 let mode = entry.preferred_engine_mode();
2681 entry
2682 .enter(mode, |context| {
2683 context.with_native_parallelism(|| {
2684 self.with_execution_resources(permit, |resources| {
2685 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
2686 op(buffers.get_mut())
2687 })
2688 })
2689 })
2690 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2691 }
2692
2693 fn install_with_pool_context_unmarked<R: Send>(
2694 &mut self,
2695 op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
2696 ) -> crate::Result<R> {
2697 let admission = self.execution_admission()?;
2698 let permit = admission.permit();
2699 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
2700 let mode = entry.preferred_engine_mode();
2701 entry
2702 .enter(mode, |context| {
2703 context.with_native_parallelism(|| {
2704 self.with_execution_resources(permit, |resources| {
2705 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
2706 op(context, buffers.get_mut())
2707 })
2708 })
2709 })
2710 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2711 }
2712
2713 fn install_with_indexed_pool_context_unmarked<R: Send>(
2714 &mut self,
2715 op: impl FnOnce(
2716 &CpuExecutionContext<'_>,
2717 &mut BufferPool,
2718 &mut IndexedPlanCache,
2719 ) -> crate::Result<R>
2720 + Send,
2721 ) -> crate::Result<R> {
2722 let admission = self.execution_admission()?;
2723 let permit = admission.permit();
2724 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
2725 let mode = entry.preferred_engine_mode();
2726 entry
2727 .enter(mode, |context| {
2728 context.with_native_parallelism(|| {
2729 self.with_execution_resources(permit, |resources| {
2730 let EngineResources {
2731 buffers,
2732 indexed_plan_cache,
2733 ..
2734 } = resources;
2735 let mut buffers = BufferPoolLoan::new(buffers);
2736 op(context, buffers.get_mut(), indexed_plan_cache)
2737 })
2738 })
2739 })
2740 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2741 }
2742
2743 fn install_with_pool<R: FreshCpuOutput + Send>(
2744 &mut self,
2745 op: impl FnOnce(&mut BufferPool) -> crate::Result<R> + Send,
2746 ) -> crate::Result<R> {
2747 let domain = self.engine.domain().id();
2748 let mut output = self.install_with_pool_unmarked(op)?;
2749 output.tag_fresh(domain);
2750 Ok(output)
2751 }
2752
2753 fn install_with_pool_context<R: FreshCpuOutput + Send>(
2754 &mut self,
2755 op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
2756 ) -> crate::Result<R> {
2757 let domain = self.engine.domain().id();
2758 let mut output = self.install_with_pool_context_unmarked(op)?;
2759 output.tag_fresh(domain);
2760 Ok(output)
2761 }
2762
2763 fn install_with_indexed_pool_context<R: FreshCpuOutput + Send>(
2764 &mut self,
2765 op: impl FnOnce(
2766 &CpuExecutionContext<'_>,
2767 &mut BufferPool,
2768 &mut IndexedPlanCache,
2769 ) -> crate::Result<R>
2770 + Send,
2771 ) -> crate::Result<R> {
2772 let domain = self.engine.domain().id();
2773 let mut output = self.install_with_indexed_pool_context_unmarked(op)?;
2774 output.tag_fresh(domain);
2775 Ok(output)
2776 }
2777
2778 #[doc(hidden)]
2803 pub fn with_linalg_pool<R: Send>(
2804 &mut self,
2805 op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
2806 ) -> crate::Result<R> {
2807 let admission = self.execution_admission()?;
2808 let permit = admission.permit();
2809 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
2810 let mode = entry.preferred_linalg_mode(self.kind());
2811 entry
2812 .enter(mode, |context| {
2813 context.with_native_parallelism(|| {
2814 self.with_execution_resources(permit, |resources| {
2815 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
2816 op(context, buffers.get_mut())
2817 })
2818 })
2819 })
2820 .map_err(|error| crate::Error::backend_source("CPU linalg execution", error))?
2821 }
2822
2823 fn with_execution_resources<R>(
2824 &self,
2825 permit: &ResourcePermit,
2826 op: impl FnOnce(&mut EngineResources) -> R,
2827 ) -> R {
2828 if permit.is_reentrant() {
2829 let mut resources =
2830 EngineResources::new(self.shared.buffer_limit.load(Ordering::Relaxed));
2831 return op(&mut resources);
2832 }
2833 let mut resources = self
2834 .engine
2835 .resources
2836 .lock()
2837 .unwrap_or_else(std::sync::PoisonError::into_inner);
2838 op(&mut resources)
2839 }
2840
2841 fn acquire_execution_permit(&self, owner: ResourceOwner) -> ResourcePermit {
2842 match &self.resolved {
2843 ResolvedCpuExecution::Managed(placement)
2844 | ResolvedCpuExecution::ExternalManaged(placement) => self
2845 .shared
2846 .arbiter
2847 .acquire_recovering(placement.cpus().clone(), owner),
2848 ResolvedCpuExecution::ExternalCallerManaged => {
2849 let active = self
2852 .engine
2853 .domain()
2854 .caller_managed_active()
2855 .unwrap_or_else(|| {
2856 unreachable!("caller-managed execution needs a local admission guard")
2857 });
2858 ResourcePermit::caller_managed(active, owner)
2859 }
2860 ResolvedCpuExecution::Compatibility => self
2861 .shared
2862 .arbiter
2863 .acquire_recovering(self.shared.topology.allowed_cpus().clone(), owner),
2864 ResolvedCpuExecution::ProviderDefaultExclusive => self
2865 .shared
2866 .arbiter
2867 .acquire_provider_exclusive_recovering(owner),
2868 }
2869 }
2870
2871 #[cfg(test)]
2872 fn try_acquire_execution_permit_for_test(
2873 &self,
2874 ) -> Result<Option<ResourcePermit>, crate::arbiter::ResourceArbiterError> {
2875 match &self.resolved {
2876 ResolvedCpuExecution::Managed(placement)
2877 | ResolvedCpuExecution::ExternalManaged(placement) => {
2878 self.shared.arbiter.try_acquire(placement.cpus().clone())
2879 }
2880 ResolvedCpuExecution::ExternalCallerManaged => Ok(None),
2881 ResolvedCpuExecution::Compatibility => self
2882 .shared
2883 .arbiter
2884 .try_acquire(self.shared.topology.allowed_cpus().clone()),
2885 ResolvedCpuExecution::ProviderDefaultExclusive => {
2886 self.shared.arbiter.try_acquire_provider_exclusive()
2887 }
2888 }
2889 }
2890}
2891
2892impl BackendSession for CpuBackend {
2893 fn vdot_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2894 self.run_backend_session_cached(None, move |session| session.vdot_read(lhs, rhs))
2895 }
2896
2897 fn norm_squared_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2898 self.run_backend_session_cached(None, move |session| session.norm_squared_read(input))
2899 }
2900
2901 fn axpby_read_into_accum(
2902 &mut self,
2903 alpha: ContractionScalar,
2904 x: TensorRead<'_>,
2905 beta: ContractionScalar,
2906 y: TensorWrite<'_>,
2907 ) -> crate::Result<()> {
2908 self.run_backend_session_cached(None, move |session| {
2909 session.axpby_read_into_accum(alpha, x, beta, y)
2910 })
2911 }
2912
2913 fn session_type_id(&self) -> TypeId {
2914 TypeId::of::<CpuBackendSessionMarker>()
2915 }
2916
2917 unsafe fn session_data_mut(&mut self) -> *mut () {
2918 self as *mut Self as *mut ()
2919 }
2920}
2921
2922impl BackendRuntimeCache for CpuBackend {
2923 type RuntimeCache = gemm::GemmAnalysisCache;
2924}
2925
2926impl TensorElementwise for CpuBackend {
2927 fn elementwise_read_into(
2928 &mut self,
2929 op: ElementwiseReadOp,
2930 inputs: &[TensorRead<'_>],
2931 out: TensorWrite<'_>,
2932 ) -> crate::Result<()> {
2933 self.install_with_pool_context_unmarked(|context, buffers| {
2934 let exec_context = context.strided_exec_context();
2935 tenferro_internal_cpu_kernels::elementwise_read_into_with_context(
2936 op,
2937 inputs,
2938 out,
2939 &exec_context,
2940 |inputs, out| elementwise_read_into_fallback_with_pool(buffers, op, inputs, out),
2941 )
2942 })
2943 }
2944
2945 fn add(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2946 self.install_with_pool(|buffers| elementwise::add_with_pool(buffers, lhs, rhs))
2947 }
2948
2949 fn add_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2950 self.install_with_pool(|buffers| elementwise::add_read_with_pool(buffers, lhs, rhs))
2951 }
2952
2953 fn sub(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2954 self.install_with_pool(|buffers| elementwise::sub_with_pool(buffers, lhs, rhs))
2955 }
2956
2957 fn sub_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2958 self.install_with_pool(|buffers| elementwise::sub_read_with_pool(buffers, lhs, rhs))
2959 }
2960
2961 fn mul(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2962 self.install_with_pool(|buffers| elementwise::mul_with_pool(buffers, lhs, rhs))
2963 }
2964
2965 fn mul_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2966 self.install_with_pool(|buffers| elementwise::mul_read_with_pool(buffers, lhs, rhs))
2967 }
2968
2969 fn neg(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2970 self.install_with_pool(|buffers| elementwise::neg_with_pool(buffers, input))
2971 }
2972
2973 fn neg_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2974 self.install_with_pool(|buffers| elementwise::neg_read_with_pool(buffers, input))
2975 }
2976
2977 fn conj(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2978 self.install_with_pool(|buffers| elementwise::conj_with_pool(buffers, input))
2979 }
2980
2981 fn conj_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2982 self.install_with_pool(|buffers| elementwise::conj_read_with_pool(buffers, input))
2983 }
2984
2985 fn div(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2986 self.install_with_pool(|buffers| elementwise::div_with_pool(buffers, lhs, rhs))
2987 }
2988
2989 fn div_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2990 self.install_with_pool(|buffers| elementwise::div_read_with_pool(buffers, lhs, rhs))
2991 }
2992
2993 fn rem(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2994 self.install_with_pool(|buffers| elementwise::rem_with_pool(buffers, lhs, rhs))
2995 }
2996
2997 fn rem_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2998 self.install_with_pool(|buffers| elementwise::rem_read_with_pool(buffers, lhs, rhs))
2999 }
3000
3001 fn abs(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3002 self.install_with_pool(|buffers| elementwise::abs_with_pool(buffers, input))
3003 }
3004
3005 fn abs_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3006 self.install_with_pool(|buffers| elementwise::abs_read_with_pool(buffers, input))
3007 }
3008
3009 fn sign(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3010 self.install_with_pool(|buffers| elementwise::sign_with_pool(buffers, input))
3011 }
3012
3013 fn sign_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3014 self.install_with_pool(|buffers| elementwise::sign_read_with_pool(buffers, input))
3015 }
3016
3017 fn maximum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
3018 self.install_with_pool(|buffers| elementwise::maximum_with_pool(buffers, lhs, rhs))
3019 }
3020
3021 fn maximum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
3022 self.install_with_pool(|buffers| elementwise::maximum_read_with_pool(buffers, lhs, rhs))
3023 }
3024
3025 fn minimum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
3026 self.install_with_pool(|buffers| elementwise::minimum_with_pool(buffers, lhs, rhs))
3027 }
3028
3029 fn minimum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
3030 self.install_with_pool(|buffers| elementwise::minimum_read_with_pool(buffers, lhs, rhs))
3031 }
3032
3033 fn compare(&mut self, lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor> {
3034 self.install_with_pool(|buffers| elementwise::compare_with_pool(buffers, lhs, rhs, dir))
3035 }
3036
3037 fn compare_read(
3038 &mut self,
3039 lhs: TensorRead<'_>,
3040 rhs: TensorRead<'_>,
3041 dir: &CompareDir,
3042 ) -> crate::Result<Tensor> {
3043 self.install_with_pool(|buffers| {
3044 elementwise::compare_read_with_pool(buffers, lhs, rhs, dir)
3045 })
3046 }
3047
3048 fn select(
3049 &mut self,
3050 pred: &Tensor,
3051 on_true: &Tensor,
3052 on_false: &Tensor,
3053 ) -> crate::Result<Tensor> {
3054 self.install_with_pool(|buffers| {
3055 elementwise::select_with_pool(buffers, pred, on_true, on_false)
3056 })
3057 }
3058
3059 fn select_read(
3060 &mut self,
3061 pred: TensorRead<'_>,
3062 on_true: TensorRead<'_>,
3063 on_false: TensorRead<'_>,
3064 ) -> crate::Result<Tensor> {
3065 self.install_with_pool(|buffers| {
3066 elementwise::select_read_with_pool(buffers, pred, on_true, on_false)
3067 })
3068 }
3069
3070 fn clamp(&mut self, input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor> {
3071 self.install_with_pool(|buffers| elementwise::clamp_with_pool(buffers, input, lower, upper))
3072 }
3073
3074 fn clamp_read(
3075 &mut self,
3076 input: TensorRead<'_>,
3077 lower: TensorRead<'_>,
3078 upper: TensorRead<'_>,
3079 ) -> crate::Result<Tensor> {
3080 self.install_with_pool(|buffers| {
3081 elementwise::clamp_read_with_pool(buffers, input, lower, upper)
3082 })
3083 }
3084}
3085
3086impl TensorAnalytic for CpuBackend {
3087 fn exp(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3088 self.install_with_pool(|buffers| analytic::exp_with_pool(buffers, input))
3089 }
3090
3091 fn exp_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3092 self.install_with_pool(|buffers| analytic::exp_read_with_pool(buffers, input))
3093 }
3094
3095 fn log(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3096 self.install_with_pool(|buffers| analytic::log_with_pool(buffers, input))
3097 }
3098
3099 fn log_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3100 self.install_with_pool(|buffers| analytic::log_read_with_pool(buffers, input))
3101 }
3102
3103 fn sin(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3104 self.install_with_pool(|buffers| analytic::sin_with_pool(buffers, input))
3105 }
3106
3107 fn sin_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3108 self.install_with_pool(|buffers| analytic::sin_read_with_pool(buffers, input))
3109 }
3110
3111 fn cos(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3112 self.install_with_pool(|buffers| analytic::cos_with_pool(buffers, input))
3113 }
3114
3115 fn cos_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3116 self.install_with_pool(|buffers| analytic::cos_read_with_pool(buffers, input))
3117 }
3118
3119 fn tanh(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3120 self.install_with_pool(|buffers| analytic::tanh_with_pool(buffers, input))
3121 }
3122
3123 fn tanh_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3124 self.install_with_pool(|buffers| analytic::tanh_read_with_pool(buffers, input))
3125 }
3126
3127 fn sqrt(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3128 self.install_with_pool(|buffers| analytic::sqrt_with_pool(buffers, input))
3129 }
3130
3131 fn sqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3132 self.install_with_pool(|buffers| analytic::sqrt_read_with_pool(buffers, input))
3133 }
3134
3135 fn rsqrt(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3136 self.install_with_pool(|buffers| analytic::rsqrt_with_pool(buffers, input))
3137 }
3138
3139 fn rsqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3140 self.install_with_pool(|buffers| analytic::rsqrt_read_with_pool(buffers, input))
3141 }
3142
3143 fn pow(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
3144 self.install_with_pool(|buffers| analytic::pow_with_pool(buffers, lhs, rhs))
3145 }
3146
3147 fn pow_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
3148 self.install_with_pool(|buffers| analytic::pow_read_with_pool(buffers, lhs, rhs))
3149 }
3150
3151 fn expm1(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3152 self.install_with_pool(|buffers| analytic::expm1_with_pool(buffers, input))
3153 }
3154
3155 fn expm1_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3156 self.install_with_pool(|buffers| analytic::expm1_read_with_pool(buffers, input))
3157 }
3158
3159 fn log1p(&mut self, input: &Tensor) -> crate::Result<Tensor> {
3160 self.install_with_pool(|buffers| analytic::log1p_with_pool(buffers, input))
3161 }
3162
3163 fn log1p_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3164 self.install_with_pool(|buffers| analytic::log1p_read_with_pool(buffers, input))
3165 }
3166}
3167
3168impl TensorStructural for CpuBackend {
3169 fn to_contiguous_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
3170 self.install_with_pool(|buffers| {
3171 materialize_tensor_read(buffers, "CpuBackend::to_contiguous_read", input)
3172 })
3173 }
3174
3175 fn copy_read_into(&mut self, src: TensorRead<'_>, dst: TensorWrite<'_>) -> crate::Result<()> {
3176 self.try_install(|| copy_tensor_read_into("CpuBackend::copy_read_into", src, dst))
3177 }
3178
3179 fn transpose(&mut self, input: &Tensor, perm: &[usize]) -> crate::Result<Tensor> {
3180 self.install_with_pool(|buffers| structural::transpose_with_pool(buffers, input, perm))
3181 }
3182
3183 fn transpose_read(&mut self, input: TensorRead<'_>, perm: &[usize]) -> crate::Result<Tensor> {
3184 self.install_with_pool(|buffers| structural::transpose_read_with_pool(buffers, input, perm))
3185 }
3186
3187 fn reshape(&mut self, input: &Tensor, shape: &[usize]) -> crate::Result<Tensor> {
3188 structural::reshape(input, shape)
3192 }
3193
3194 fn reshape_read(&mut self, input: TensorRead<'_>, shape: &[usize]) -> crate::Result<Tensor> {
3195 match &input {
3196 TensorRead::Tensor(tensor) => structural::reshape(tensor, shape),
3200 TensorRead::View(_) => self.install_with_pool(|buffers| {
3201 structural::reshape_read_with_pool(buffers, input, shape)
3202 }),
3203 }
3204 }
3205
3206 fn broadcast_in_dim(
3207 &mut self,
3208 input: &Tensor,
3209 shape: &[usize],
3210 dims: &[usize],
3211 ) -> crate::Result<Tensor> {
3212 self.install_with_pool(|buffers| {
3213 structural::broadcast_in_dim_with_pool(buffers, input, shape, dims)
3214 })
3215 }
3216
3217 fn broadcast_in_dim_read(
3218 &mut self,
3219 input: TensorRead<'_>,
3220 shape: &[usize],
3221 dims: &[usize],
3222 ) -> crate::Result<Tensor> {
3223 self.install_with_pool(|buffers| {
3224 structural::broadcast_in_dim_read_with_pool(buffers, input, shape, dims)
3225 })
3226 }
3227
3228 fn cast(&mut self, input: &Tensor, to: crate::DType) -> crate::Result<Tensor> {
3229 self.install_with_pool(|buffers| structural::cast_with_pool(buffers, input, to))
3230 }
3231
3232 fn extract_diagonal(
3233 &mut self,
3234 input: &Tensor,
3235 axis_a: usize,
3236 axis_b: usize,
3237 ) -> crate::Result<Tensor> {
3238 self.install_with_pool(|buffers| {
3239 structural::extract_diagonal_with_pool(buffers, input, axis_a, axis_b)
3240 })
3241 }
3242
3243 fn embed_diagonal(
3244 &mut self,
3245 input: &Tensor,
3246 axis_a: usize,
3247 axis_b: usize,
3248 ) -> crate::Result<Tensor> {
3249 self.install_with_pool(|buffers| {
3250 structural::embed_diagonal_with_pool(buffers, input, axis_a, axis_b)
3251 })
3252 }
3253
3254 fn tril(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor> {
3255 self.install_with_pool(|buffers| structural::tril_with_pool(buffers, input, k))
3256 }
3257
3258 fn triu(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor> {
3259 self.install_with_pool(|buffers| structural::triu_with_pool(buffers, input, k))
3260 }
3261}
3262
3263impl TensorReduction for CpuBackend {
3264 fn reduce_sum(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3265 self.try_install_fresh_with_context(|context| {
3266 let exec_context = context.strided_exec_context();
3267 reduction::reduce_sum(input, axes, &exec_context)
3268 })
3269 }
3270
3271 fn reduce_sum_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3272 self.install_with_pool_context(|context, buffers| {
3273 let exec_context = context.strided_exec_context();
3274 reduction::reduce_sum_read(buffers, input, axes, &exec_context)
3275 })
3276 }
3277
3278 fn reduce_sum_squares_read(
3279 &mut self,
3280 input: TensorRead<'_>,
3281 axes: &[usize],
3282 ) -> crate::Result<Tensor> {
3283 self.install_with_pool_context(|context, buffers| {
3284 let exec_context = context.strided_exec_context();
3285 reduction::reduce_sum_squares_read(buffers, input, axes, &exec_context)
3286 })
3287 }
3288
3289 fn reduce_prod(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3290 self.try_install_fresh_with_context(|context| {
3291 let exec_context = context.strided_exec_context();
3292 reduction::reduce_prod(input, axes, &exec_context)
3293 })
3294 }
3295
3296 fn reduce_prod_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3297 self.install_with_pool_context(|context, buffers| {
3298 let exec_context = context.strided_exec_context();
3299 reduction::reduce_prod_read(buffers, input, axes, &exec_context)
3300 })
3301 }
3302
3303 fn reduce_max(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3304 self.try_install_fresh(|| reduction::reduce_max(input, axes))
3305 }
3306
3307 fn reduce_max_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3308 self.install_with_pool(|buffers| reduction::reduce_max_read(buffers, input, axes))
3309 }
3310
3311 fn reduce_min(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3312 self.try_install_fresh(|| reduction::reduce_min(input, axes))
3313 }
3314
3315 fn reduce_min_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3316 self.install_with_pool(|buffers| reduction::reduce_min_read(buffers, input, axes))
3317 }
3318}
3319
3320impl TensorDot for CpuBackend {
3321 fn dot_general(
3322 &mut self,
3323 lhs: &Tensor,
3324 rhs: &Tensor,
3325 config: &DotGeneralConfig,
3326 ) -> crate::Result<Tensor> {
3327 self.run_backend_session_cached(None, move |session| session.dot_general(lhs, rhs, config))
3328 }
3329
3330 fn dot_general_read(
3331 &mut self,
3332 lhs: TensorRead<'_>,
3333 rhs: TensorRead<'_>,
3334 config: &DotGeneralConfig,
3335 ) -> crate::Result<Tensor> {
3336 self.run_backend_session_cached(None, move |session| {
3337 session.dot_general_read(lhs, rhs, config)
3338 })
3339 }
3340
3341 fn dot_general_read_into(
3342 &mut self,
3343 lhs: TensorRead<'_>,
3344 rhs: TensorRead<'_>,
3345 config: &DotGeneralConfig,
3346 out: TensorWrite<'_>,
3347 ) -> crate::Result<()> {
3348 self.run_backend_session_cached(None, move |session| {
3349 session.dot_general_read_into(lhs, rhs, config, out)
3350 })
3351 }
3352
3353 fn dot_general_read_into_accum(
3354 &mut self,
3355 lhs: TensorRead<'_>,
3356 rhs: TensorRead<'_>,
3357 config: &DotGeneralConfig,
3358 accumulation: DotGeneralAccumulation,
3359 out: TensorWrite<'_>,
3360 ) -> crate::Result<()> {
3361 self.run_backend_session_cached(None, move |session| {
3362 session.dot_general_read_into_accum(lhs, rhs, config, accumulation, out)
3363 })
3364 }
3365
3366 fn dot_general_with_conj(
3367 &mut self,
3368 lhs: &Tensor,
3369 rhs: &Tensor,
3370 config: &DotGeneralConfig,
3371 lhs_conj: bool,
3372 rhs_conj: bool,
3373 ) -> crate::Result<Tensor> {
3374 self.run_backend_session_cached(None, move |session| {
3375 session.dot_general_with_conj(lhs, rhs, config, lhs_conj, rhs_conj)
3376 })
3377 }
3378}
3379
3380impl BackendCachedDot for CpuBackend {
3381 fn dot_general_cached(
3382 &mut self,
3383 cache: &mut Self::RuntimeCache,
3384 cache_slot: Option<usize>,
3385 lhs: &Tensor,
3386 rhs: &Tensor,
3387 config: &DotGeneralConfig,
3388 ) -> crate::Result<Tensor> {
3389 self.run_backend_session_cached(Some(cache), move |session| {
3390 session.dot_general_cached(cache_slot, lhs, rhs, config)
3391 })
3392 }
3393
3394 fn dot_general_with_conj_cached(
3395 &mut self,
3396 cache: &mut Self::RuntimeCache,
3397 cache_slot: Option<usize>,
3398 lhs: &Tensor,
3399 rhs: &Tensor,
3400 config: &DotGeneralConfig,
3401 lhs_conj: bool,
3402 rhs_conj: bool,
3403 ) -> crate::Result<Tensor> {
3404 self.run_backend_session_cached(Some(cache), move |session| {
3405 session.dot_general_with_conj_cached(cache_slot, lhs, rhs, config, lhs_conj, rhs_conj)
3406 })
3407 }
3408
3409 fn dot_general_read_into_accum_cached(
3410 &mut self,
3411 cache: &mut Self::RuntimeCache,
3412 cache_slot: Option<usize>,
3413 lhs: TensorRead<'_>,
3414 rhs: TensorRead<'_>,
3415 config: &DotGeneralConfig,
3416 accumulation: DotGeneralAccumulation,
3417 out: TensorWrite<'_>,
3418 ) -> crate::Result<()> {
3419 self.run_backend_session_cached(Some(cache), move |session| {
3420 session.dot_general_read_into_accum_cached(
3421 cache_slot,
3422 lhs,
3423 rhs,
3424 config,
3425 accumulation,
3426 out,
3427 )
3428 })
3429 }
3430
3431 fn grouped_gemm_cached(
3432 &mut self,
3433 cache: &mut Self::RuntimeCache,
3434 cache_slot: Option<usize>,
3435 lhs: TensorRead<'_>,
3436 rhs: TensorRead<'_>,
3437 config: &GroupedGemmConfig<'_>,
3438 out: TensorWrite<'_>,
3439 ) -> crate::Result<()> {
3440 self.run_backend_session_cached(Some(cache), move |session| {
3441 session.grouped_gemm_cached(cache_slot, lhs, rhs, config, out)
3442 })
3443 }
3444}
3445
3446impl TensorIndexing for CpuBackend {
3447 fn gather(
3448 &mut self,
3449 operand: &Tensor,
3450 start_indices: &Tensor,
3451 config: &GatherConfig,
3452 ) -> crate::Result<Tensor> {
3453 self.install_with_indexed_pool_context(|context, buffers, cache| {
3454 let exec_context = context.strided_exec_context();
3455 indexing::gather_with_pool(
3456 buffers,
3457 cache,
3458 &exec_context,
3459 operand,
3460 start_indices,
3461 config,
3462 )
3463 })
3464 }
3465
3466 fn scatter(
3467 &mut self,
3468 operand: &Tensor,
3469 scatter_indices: &Tensor,
3470 updates: &Tensor,
3471 config: &ScatterConfig,
3472 ) -> crate::Result<Tensor> {
3473 self.install_with_indexed_pool_context(|context, buffers, cache| {
3474 let exec_context = context.strided_exec_context();
3475 indexing::scatter_with_pool(
3476 buffers,
3477 cache,
3478 &exec_context,
3479 operand,
3480 scatter_indices,
3481 updates,
3482 config,
3483 )
3484 })
3485 }
3486
3487 fn slice(&mut self, input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor> {
3488 self.install_with_pool_context(|context, buffers| {
3489 let exec_context = context.strided_exec_context();
3490 indexing::try_slice_with_pool(buffers, &exec_context, input, config)
3491 })
3492 }
3493
3494 fn dynamic_slice(
3495 &mut self,
3496 input: &Tensor,
3497 starts: &Tensor,
3498 slice_sizes: &[usize],
3499 ) -> crate::Result<Tensor> {
3500 self.install_with_indexed_pool_context(|context, buffers, cache| {
3501 let exec_context = context.strided_exec_context();
3502 indexing::dynamic_slice_with_pool(
3503 buffers,
3504 cache,
3505 &exec_context,
3506 input,
3507 starts,
3508 slice_sizes,
3509 )
3510 })
3511 }
3512
3513 fn dynamic_update_slice(
3514 &mut self,
3515 operand: &Tensor,
3516 update: &Tensor,
3517 starts: &Tensor,
3518 ) -> crate::Result<Tensor> {
3519 self.install_with_indexed_pool_context(|context, buffers, cache| {
3520 let exec_context = context.strided_exec_context();
3521 indexing::dynamic_update_slice_with_pool(
3522 buffers,
3523 cache,
3524 &exec_context,
3525 operand,
3526 update,
3527 starts,
3528 )
3529 })
3530 }
3531
3532 fn pad(&mut self, input: &Tensor, config: &PadConfig) -> crate::Result<Tensor> {
3533 self.install_with_pool_context(|context, buffers| {
3534 let exec_context = context.strided_exec_context();
3535 indexing::try_pad_with_pool(buffers, &exec_context, input, config)
3536 })
3537 }
3538
3539 fn concatenate(&mut self, inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor> {
3540 self.install_with_pool_context(|context, buffers| {
3541 let exec_context = context.strided_exec_context();
3542 indexing::try_concatenate_with_pool(buffers, &exec_context, inputs, axis)
3543 })
3544 }
3545
3546 fn reverse(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3547 self.install_with_pool_context(|context, buffers| {
3548 let exec_context = context.strided_exec_context();
3549 indexing::reverse_with_pool(buffers, &exec_context, input, axes)
3550 })
3551 }
3552}
3553
3554impl CpuBackend {
3555 pub fn with_allocation_domain(mut self, domain: Arc<dyn SharedTensorAllocationDomain>) -> Self {
3580 self.allocation_domain = Some(domain);
3581 self.runtime_identity = CpuRuntimeIdentity::fresh();
3582 self
3583 }
3584
3585 pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
3595 self.allocation_domain.as_ref().map(|domain| domain.id())
3596 }
3597
3598 pub fn shared_allocation_domain(&self) -> Option<&Arc<dyn SharedTensorAllocationDomain>> {
3608 self.allocation_domain.as_ref()
3609 }
3610
3611 fn run_backend_session_cached<R: Send>(
3612 &mut self,
3613 cache: Option<&mut gemm::GemmAnalysisCache>,
3614 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
3615 ) -> R {
3616 let providers = self.provider_bundle.clone();
3617 let admission = self.infallible_execution_admission();
3618 let permit = admission.permit();
3619 let owner = permit.owner();
3620 let entry = CpuOperationEntry::new(self.engine.domain(), permit);
3621 let enter_managed_session = entry.supports_infallible_session_entry();
3624 let run = |entered| {
3625 self.with_execution_resources(permit, |resources| {
3626 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
3627 let cache = cache.unwrap_or(&mut resources.gemm_analysis_cache);
3628 let session_started = Instant::now();
3629 let mut session = CpuExecSession {
3630 entry,
3631 entered,
3632 buffers: buffers.get_mut(),
3633 gemm_analysis_cache: cache,
3634 indexed_plan_cache: &mut resources.indexed_plan_cache,
3635 providers: &providers,
3636 backend_kind: self.kind(),
3637 allocation_domain: self.allocation_domain.as_ref(),
3638 };
3639 record_cpu_session_profile(
3640 "with_backend_session_cached.session_construct",
3641 session_started.elapsed(),
3642 );
3643 let exec_started = Instant::now();
3644 let result = f(&mut session);
3645 record_cpu_session_profile(
3646 "with_backend_session_cached.exec_body",
3647 exec_started.elapsed(),
3648 );
3649 result
3650 })
3651 };
3652 if enter_managed_session {
3653 entry.enter_managed_session(|context| run(Some(context)))
3654 } else {
3655 with_execution_owner(owner, || run(None))
3656 }
3657 }
3658}
3659
3660impl BackendSessionHost for CpuBackend {
3661 fn with_backend_session<R: Send>(
3662 &mut self,
3663 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
3664 ) -> R {
3665 self.run_backend_session_cached(None, f)
3666 }
3667
3668 fn with_backend_session_cached<R: Send>(
3669 &mut self,
3670 cache: &mut Self::RuntimeCache,
3671 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
3672 ) -> R {
3673 if !cpu_session_profile_enabled() {
3674 return self.run_backend_session_cached(Some(cache), f);
3675 }
3676 let total_started = Instant::now();
3677 let result =
3678 profile_cpu_session_section("with_backend_session_cached.exec_session", || {
3679 self.run_backend_session_cached(Some(cache), f)
3680 });
3681 record_cpu_session_profile("with_backend_session_cached.total", total_started.elapsed());
3682 maybe_print_cpu_session_profile();
3683 result
3684 }
3685}
3686
3687impl TensorBuffer for CpuBackend {
3688 fn reclaim_buffer(&mut self, tensor: Tensor) {
3689 let admission = self.infallible_execution_admission();
3690 let permit = admission.permit();
3691 with_execution_owner(permit.owner(), || {
3692 self.with_execution_resources(permit, |resources| {
3693 let buffers = &mut resources.buffers;
3694 match tensor {
3695 Tensor::F32(t) => reclaim_typed(buffers, t),
3696 Tensor::F64(t) => reclaim_typed(buffers, t),
3697 Tensor::I32(t) => reclaim_typed(buffers, t),
3698 Tensor::I64(t) => reclaim_typed(buffers, t),
3699 Tensor::Bool(t) => reclaim_typed(buffers, t),
3700 Tensor::C32(t) => reclaim_typed(buffers, t),
3701 Tensor::C64(t) => reclaim_typed(buffers, t),
3702 }
3703 })
3704 })
3705 }
3706}
3707
3708impl<T, R> TensorViewCanonicalization<T, R> for CpuBackend
3709where
3710 T: TensorScalar + PoolScalar,
3711 R: TensorRank,
3712 R::Shape: Send + Sync,
3713 R::Strides: Send + Sync,
3714{
3715 fn to_contiguous(
3716 &mut self,
3717 view: &TypedTensorView<'_, T, R>,
3718 ) -> crate::Result<TypedTensor<T, R>> {
3719 self.install_with_pool(|buffers| {
3720 structural::typed_materialize_view_with_pool(buffers, view, "CpuBackend::to_contiguous")
3721 })
3722 }
3723
3724 fn copy_into(
3725 &mut self,
3726 src: &TypedTensorView<'_, T, R>,
3727 dst: &mut TypedTensorViewMut<'_, T, R>,
3728 ) -> crate::Result<()> {
3729 self.try_install(|| structural::typed_copy_view_into(src, dst, "CpuBackend::copy_into"))
3730 }
3731}
3732
3733impl TensorFusion for CpuBackend {
3734 fn execute_elementwise_fusion(
3735 &mut self,
3736 inputs: &[&Tensor],
3737 plan: &ElementwiseFusionPlan,
3738 ) -> crate::Result<Option<Vec<Tensor>>> {
3739 self.install_with_pool_context(|context, buffers| {
3740 let exec_context = context.strided_exec_context();
3741 tenferro_cpu_fused::elementwise_fusion_with_pool(buffers, &exec_context, inputs, plan)
3742 })
3743 }
3744
3745 fn execute_broadcast_multiply(
3746 &mut self,
3747 lhs: TensorRead<'_>,
3748 lhs_shape: &[usize],
3749 lhs_dims: &[usize],
3750 rhs: TensorRead<'_>,
3751 rhs_shape: &[usize],
3752 rhs_dims: &[usize],
3753 ) -> crate::Result<Option<Tensor>> {
3754 self.install_with_pool(|buffers| {
3755 elementwise::broadcast_multiply_read_with_pool(
3756 buffers, lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims,
3757 )
3758 })
3759 }
3760
3761 fn execute_broadcast_multiply_value(
3762 &mut self,
3763 lhs: TensorRead<'_>,
3764 lhs_shape: &[usize],
3765 lhs_dims: &[usize],
3766 rhs: TensorRead<'_>,
3767 rhs_shape: &[usize],
3768 rhs_dims: &[usize],
3769 ) -> crate::Result<Option<TensorValue>> {
3770 let domain = self.engine.domain().id();
3771 self.install_with_pool_unmarked(|buffers| {
3772 elementwise::broadcast_multiply_value_with_pool_and_tag(
3773 buffers,
3774 lhs,
3775 lhs_shape,
3776 lhs_dims,
3777 rhs,
3778 rhs_shape,
3779 rhs_dims,
3780 |tensor| tag_fresh_output(tensor, domain),
3781 )
3782 })
3783 }
3784}
3785
3786impl TensorDeviceTransfer for CpuBackend {
3787 fn download_to_host(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
3788 if tensor.backend_family().is_some() {
3789 return Err(crate::Error::runtime_state(
3790 "CpuBackend::download_to_host",
3791 "CPU backend received a backend buffer; download the tensor to host with its owning backend before CPU execution",
3792 ));
3793 }
3794 tensor.tensor_view().duplicate()
3795 }
3796
3797 fn upload_host_tensor(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor> {
3798 if tensor.backend_family().is_some() {
3799 return Err(crate::Error::runtime_state(
3800 "CpuBackend::upload_host_tensor",
3801 "CPU backend upload_host_tensor expects a host tensor; download backend buffers to host before CPU execution",
3802 ));
3803 }
3804 tensor.tensor_view().duplicate()
3805 }
3806}
3807
3808impl TensorBackend for CpuBackend {}
3809
3810pub(crate) fn reclaim_typed<T: PoolScalar>(pool: &mut BufferPool, typed: TypedTensor<T>) {
3811 if typed.backend_buffer().is_some() {
3812 return;
3813 }
3814 if let Ok(data) = typed.into_host_vec() {
3815 T::pool_release(pool, data);
3816 }
3817}
3818
3819impl Default for CpuBackend {
3820 fn default() -> Self {
3821 Self::new()
3822 }
3823}
3824
3825pub(crate) mod execution_scope;
3826
3827#[cfg(test)]
3828mod tests;