1use std::cmp::Reverse;
2use std::collections::{BTreeMap, BTreeSet, HashMap};
3use std::env;
4use std::fmt;
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::{Arc, Mutex, OnceLock};
7use std::thread;
8use std::time::{Duration, Instant};
9
10use crate::arbiter::{
11 inherited_or_new_execution_owner, with_execution_owner, ResourceArbiter, ResourceOwner,
12 ResourcePermit,
13};
14use crate::buffer_pool::{BufferPool, BufferPoolStats, PoolScalar};
15use crate::dot_runtime::{CpuProviderBundle, CpuProviderBundleInstallError};
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, CpuDomainId, CpuDomainOwnership, CpuExecutorAffinity,
27 CpuExecutorShutdown, CpuId, CpuPlacement, CpuPlacementError, CpuPlacementGuarantee, CpuSet,
28 CpuTopology, CpuTopologyError, ExternalCpuDomain, NumaNodeId, ResolvedCpuPlacement,
29};
30use crate::{
31 Buffer, 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 DotGeneralAccumulation, ElementwiseReadOp, TensorAnalytic, TensorBackend, TensorBuffer,
39 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)]
321pub enum CpuExecutionMode {
322 Managed,
324 ExternalManaged,
326 ProviderDefaultExclusive,
328 Compatibility,
330}
331
332#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
343pub enum ExternalCpuDomainRegistryError {
344 #[error("externally managed CPU registry must contain at least one domain")]
346 EmptyRegistry,
347 #[error("CPU domain ID {id:?} is registered more than once")]
349 DuplicateDomainId {
350 id: CpuDomainId,
352 },
353 #[error("CPU placement {placement:?} is registered more than once")]
355 DuplicatePlacementIdentity {
356 placement: CpuPlacement,
358 },
359 #[error("CPU domain {domain:?} declares process-disallowed CPU {cpu}")]
361 CpuOutsideAllowedSet {
362 domain: CpuDomainId,
364 cpu: CpuId,
366 },
367 #[error("default CPU domain {default_domain:?} is not registered")]
369 MissingDefaultDomain {
370 default_domain: CpuDomainId,
372 },
373 #[error(
375 "exact all-allowed CPU domain {domain:?} declares {declared:?}, but the process allows {allowed:?}"
376 )]
377 ExactAllAllowedMismatch {
378 domain: CpuDomainId,
380 declared: CpuSet,
382 allowed: CpuSet,
384 },
385}
386
387#[derive(Debug, thiserror::Error)]
402pub enum CpuBackendError {
403 #[error(transparent)]
405 Tensor(#[from] crate::Error),
406 #[error("{op}: {source}")]
408 Placement {
409 op: &'static str,
411 #[source]
413 source: CpuPlacementError,
414 },
415 #[error(transparent)]
417 ExternalRegistry(#[from] ExternalCpuDomainRegistryError),
418}
419
420impl CpuBackendError {
421 fn placement(op: &'static str, source: CpuPlacementError) -> Self {
422 Self::Placement { op, source }
423 }
424
425 pub fn placement_error(&self) -> Option<&CpuPlacementError> {
438 match self {
439 Self::Tensor(_) => None,
440 Self::Placement { source, .. } => Some(source),
441 Self::ExternalRegistry(_) => None,
442 }
443 }
444}
445
446impl From<CpuBackendError> for crate::Error {
447 fn from(error: CpuBackendError) -> Self {
448 match error {
449 CpuBackendError::Tensor(error) => error,
450 CpuBackendError::ExternalRegistry(source) => Self::extension(
451 "CpuBackend::from_external_managed_domains",
452 "cpu",
453 crate::ErrorKind::Validation(crate::ValidationKind::InvalidArgument),
454 source,
455 ),
456 CpuBackendError::Placement { op, source } => match source {
457 CpuPlacementError::TopologyDiscovery { .. }
458 | CpuPlacementError::ManagedAffinityUnavailable { .. }
459 | CpuPlacementError::NumaDiscoveryUnavailable { .. }
460 | CpuPlacementError::UnknownNumaNode { .. }
461 | CpuPlacementError::UnregisteredExternalPlacement { .. } => {
462 Self::runtime_state_source(op, source)
463 }
464 CpuPlacementError::ExternalProviderAffinityUnmanaged { .. } => {
465 Self::extension(op, "cpu", crate::ErrorKind::Unsupported, source)
466 }
467 CpuPlacementError::EngineConstruction { .. } => Self::backend_source(op, source),
468 CpuPlacementError::InternalState { .. } => {
469 Self::extension(op, "cpu", crate::ErrorKind::Internal, source)
470 }
471 },
472 }
473 }
474}
475
476#[derive(Clone, Debug, PartialEq, Eq)]
491pub struct CpuExecutionInfo {
492 backend_kind: CpuBackendKind,
493 execution_mode: CpuExecutionMode,
494 requested_placement: CpuPlacement,
495 resolved_placement: Option<ResolvedCpuPlacement>,
496 topology: CpuTopology,
497 domain_id: CpuDomainId,
498 domain_cpus: CpuSet,
499 worker_count: usize,
500 thread_budget: usize,
501 placement_guarantee: CpuPlacementGuarantee,
502 domain_ownership: CpuDomainOwnership,
503 executor_affinity: CpuExecutorAffinity,
504 executor_shutdown: CpuExecutorShutdown,
505 provider_diagnostic: &'static str,
506}
507
508impl CpuExecutionInfo {
509 pub fn backend_kind(&self) -> CpuBackendKind {
518 self.backend_kind
519 }
520
521 pub fn execution_mode(&self) -> CpuExecutionMode {
532 self.execution_mode
533 }
534
535 pub fn requested_placement(&self) -> CpuPlacement {
544 self.requested_placement
545 }
546
547 pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
556 self.resolved_placement.as_ref()
557 }
558
559 pub fn topology(&self) -> &CpuTopology {
568 &self.topology
569 }
570
571 pub fn domain_id(&self) -> CpuDomainId {
580 self.domain_id
581 }
582
583 pub fn domain_cpus(&self) -> &CpuSet {
592 &self.domain_cpus
593 }
594
595 pub fn worker_count(&self) -> usize {
604 self.worker_count
605 }
606
607 pub fn thread_budget(&self) -> usize {
620 self.thread_budget
621 }
622
623 pub fn placement_guarantee(&self) -> CpuPlacementGuarantee {
634 self.placement_guarantee
635 }
636
637 pub fn domain_ownership(&self) -> CpuDomainOwnership {
648 self.domain_ownership
649 }
650
651 pub fn executor_affinity(&self) -> CpuExecutorAffinity {
662 self.executor_affinity
663 }
664
665 pub fn executor_shutdown(&self) -> CpuExecutorShutdown {
676 self.executor_shutdown
677 }
678
679 pub fn provider_diagnostic(&self) -> &'static str {
692 self.provider_diagnostic
693 }
694}
695
696fn provider_diagnostic(kind: CpuBackendKind, ownership: CpuDomainOwnership) -> &'static str {
697 if ownership == CpuDomainOwnership::ExternalManaged {
698 return match kind {
699 CpuBackendKind::Faer => "faer (externally managed CPU executor)",
700 CpuBackendKind::Blas => "BLAS/LAPACK (externally managed CPU executor)",
701 };
702 }
703 match kind {
704 CpuBackendKind::Faer => "faer (tenferro-managed Rayon affinity)",
705 CpuBackendKind::Blas => {
706 #[cfg(feature = "blas-openblas")]
707 return "OpenBLAS (external worker affinity)";
708 #[cfg(feature = "blas-mkl")]
709 return "Intel MKL (external worker affinity)";
710 #[cfg(feature = "blas-accelerate")]
711 return "Apple Accelerate (external worker affinity)";
712 #[cfg(feature = "provider-inject")]
713 return "runtime-injected BLAS/LAPACK (external worker affinity)";
714 #[cfg(not(any(
715 feature = "blas-openblas",
716 feature = "blas-mkl",
717 feature = "blas-accelerate",
718 feature = "provider-inject"
719 )))]
720 return "linked BLAS/LAPACK provider (identity unknown; external worker affinity)";
721 }
722 }
723}
724
725fn ensure_cpu_backend_kind_available(kind: CpuBackendKind, op: &'static str) -> crate::Result<()> {
726 let _ = op;
727 match kind {
728 CpuBackendKind::Faer => {
729 #[cfg(feature = "cpu-faer")]
730 {
731 Ok(())
732 }
733 #[cfg(not(feature = "cpu-faer"))]
734 {
735 Err(crate::Error::invalid_argument(
736 op,
737 "configuration",
738 "CpuBackendKind::Faer requires the cpu-faer feature".to_string(),
739 ))
740 }
741 }
742 CpuBackendKind::Blas => {
743 #[cfg(feature = "cpu-blas")]
744 {
745 Ok(())
746 }
747 #[cfg(not(feature = "cpu-blas"))]
748 {
749 Err(crate::Error::invalid_argument(
750 op,
751 "configuration",
752 "CpuBackendKind::Blas requires the cpu-blas feature".to_string(),
753 ))
754 }
755 }
756 }
757}
758
759fn constructor_tensor_error(op: &'static str, error: crate::Error) -> CpuBackendError {
760 CpuBackendError::Tensor(match error {
761 crate::Error::Validation { source, .. } => crate::Error::validation(op, source),
762 error => error,
763 })
764}
765
766#[allow(dead_code)]
769pub(super) fn unavailable_cpu_backend_kind(kind: CpuBackendKind, op: &'static str) -> crate::Error {
770 crate::Error::invalid_argument(
771 op,
772 "configuration",
773 format!("CPU backend kind {} is not compiled in", kind.name()),
774 )
775}
776
777struct ManagedEngineRegistry {
778 node_engines: Mutex<BTreeMap<NumaNodeId, Arc<CpuEngine>>>,
779 node_domain_ids: BTreeMap<NumaNodeId, CpuDomainId>,
780 all_allowed: OnceLock<Arc<CpuEngine>>,
781 all_allowed_build: Mutex<()>,
782 base_engine: Arc<CpuEngine>,
783 thread_budget: usize,
784}
785
786struct ExternalEngineRegistry {
787 by_id: BTreeMap<CpuDomainId, Arc<CpuEngine>>,
788 by_node: BTreeMap<NumaNodeId, Arc<CpuEngine>>,
789 all_allowed: Option<Arc<CpuEngine>>,
790 default_domain: CpuDomainId,
791}
792
793enum CpuEngineRegistry {
794 ManagedLazy(ManagedEngineRegistry),
795 ExternalPrebuilt(ExternalEngineRegistry),
796}
797
798struct CpuBackendState {
799 topology: CpuTopology,
800 engines: CpuEngineRegistry,
801 arbiter: ResourceArbiter,
802 kind: CpuBackendKind,
803 buffer_limit: AtomicUsize,
804 indexed_plan_cache_limits: Mutex<IndexedPlanCacheLimits>,
805}
806
807impl CpuBackendState {
808 fn managed_engine_for(
809 &self,
810 placement: &ResolvedCpuPlacement,
811 requested: CpuPlacement,
812 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
813 let cache_configuration = self.indexed_plan_cache_limits.lock().map_err(|_| {
817 CpuPlacementError::InternalState {
818 requested,
819 backend: self.kind,
820 message: "CPU indexed-plan cache configuration lock is poisoned",
821 }
822 })?;
823 let cache_limits = *cache_configuration;
824 let CpuEngineRegistry::ManagedLazy(registry) = &self.engines else {
825 return Err(CpuPlacementError::InternalState {
826 requested,
827 backend: self.kind,
828 message: "managed placement requested from an external engine registry",
829 });
830 };
831 match placement {
832 ResolvedCpuPlacement::NumaNode { id, .. } => {
833 let mut engines = registry
834 .node_engines
835 .lock()
836 .unwrap_or_else(std::sync::PoisonError::into_inner);
837 if let Some(engine) = engines.get(id) {
838 return Ok(Arc::clone(engine));
839 }
840 let Some(domain_id) = registry.node_domain_ids.get(id).copied() else {
841 return Err(CpuPlacementError::InternalState {
842 requested,
843 backend: self.kind,
844 message: "managed NUMA node has no coordinator-stable domain ID",
845 });
846 };
847 let engine = Arc::new(
848 CpuEngine::new_managed(
849 domain_id,
850 placement.clone(),
851 registry.thread_budget,
852 self.buffer_limit.load(Ordering::Relaxed),
853 )
854 .map_err(|error| {
855 CpuPlacementError::EngineConstruction {
856 requested,
857 backend: self.kind,
858 source: CpuEngineConstructionError::Context(error),
859 }
860 })?,
861 );
862 self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
863 engines.insert(*id, Arc::clone(&engine));
864 Ok(engine)
865 }
866 ResolvedCpuPlacement::AllAllowed { .. } => {
867 if let Some(engine) = registry.all_allowed.get() {
868 return Ok(Arc::clone(engine));
869 }
870 let _build = registry
871 .all_allowed_build
872 .lock()
873 .unwrap_or_else(std::sync::PoisonError::into_inner);
874 if let Some(engine) = registry.all_allowed.get() {
875 return Ok(Arc::clone(engine));
876 }
877 let engine = Arc::new(
878 CpuEngine::new_managed(
879 CpuDomainId::new(0),
880 placement.clone(),
881 registry.thread_budget,
882 self.buffer_limit.load(Ordering::Relaxed),
883 )
884 .map_err(|error| {
885 CpuPlacementError::EngineConstruction {
886 requested,
887 backend: self.kind,
888 source: CpuEngineConstructionError::Context(error),
889 }
890 })?,
891 );
892 self.configure_new_indexed_plan_cache(&engine, requested, cache_limits)?;
893 let _ = registry.all_allowed.set(Arc::clone(&engine));
894 Ok(engine)
895 }
896 }
897 }
898
899 fn configure_new_indexed_plan_cache(
900 &self,
901 engine: &CpuEngine,
902 requested: CpuPlacement,
903 limits: IndexedPlanCacheLimits,
904 ) -> Result<(), CpuPlacementError> {
905 let mut resources =
906 engine
907 .resources
908 .lock()
909 .map_err(|_| CpuPlacementError::InternalState {
910 requested,
911 backend: self.kind,
912 message: "new CPU engine indexed-plan cache lock is poisoned",
913 })?;
914 resources.indexed_plan_cache.set_limits(limits);
915 Ok(())
916 }
917
918 fn managed_base_engine(
919 &self,
920 requested: CpuPlacement,
921 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
922 match &self.engines {
923 CpuEngineRegistry::ManagedLazy(registry) => Ok(Arc::clone(®istry.base_engine)),
924 CpuEngineRegistry::ExternalPrebuilt(_) => Err(CpuPlacementError::InternalState {
925 requested,
926 backend: self.kind,
927 message: "managed compatibility placement requested from an external registry",
928 }),
929 }
930 }
931
932 fn external_engine_for(
933 &self,
934 requested: CpuPlacement,
935 ) -> Result<Arc<CpuEngine>, CpuPlacementError> {
936 let CpuEngineRegistry::ExternalPrebuilt(registry) = &self.engines else {
937 return Err(CpuPlacementError::InternalState {
938 requested,
939 backend: self.kind,
940 message: "external placement requested from a managed engine registry",
941 });
942 };
943 let engine = match requested {
944 CpuPlacement::Auto => registry.by_id.get(®istry.default_domain),
945 CpuPlacement::NumaNode(id) => registry.by_node.get(&id),
946 CpuPlacement::AllAllowed => registry.all_allowed.as_ref(),
947 };
948 engine
949 .cloned()
950 .ok_or(CpuPlacementError::UnregisteredExternalPlacement { requested })
951 }
952
953 fn is_external(&self) -> bool {
954 matches!(&self.engines, CpuEngineRegistry::ExternalPrebuilt(_))
955 }
956
957 fn initialized_engines(&self, op: &'static str) -> crate::Result<Vec<Arc<CpuEngine>>> {
958 let mut engines = match &self.engines {
959 CpuEngineRegistry::ManagedLazy(registry) => {
960 let mut engines = vec![Arc::clone(®istry.base_engine)];
961 if let Some(engine) = registry.all_allowed.get() {
962 engines.push(Arc::clone(engine));
963 }
964 engines.extend(
965 registry
966 .node_engines
967 .lock()
968 .map_err(|_| poisoned_cpu_lock(op, "CPU engine registry"))?
969 .values()
970 .cloned(),
971 );
972 engines
973 }
974 CpuEngineRegistry::ExternalPrebuilt(registry) => {
975 registry.by_id.values().cloned().collect()
976 }
977 };
978 if engines.len() > 1 {
979 engines.sort_unstable_by_key(|engine| Arc::as_ptr(engine) as usize);
980 engines.dedup_by(|left, right| Arc::ptr_eq(left, right));
981 }
982 Ok(engines)
983 }
984}
985
986fn poisoned_cpu_lock(op: &'static str, lock: &'static str) -> crate::Error {
987 crate::Error::runtime_state(op, format!("{lock} lock poisoned"))
988}
989
990fn lock_engine_resources<'a>(
991 engine: &'a CpuEngine,
992 op: &'static str,
993) -> crate::Result<std::sync::MutexGuard<'a, EngineResources>> {
994 engine
995 .resources
996 .lock()
997 .map_err(|_| poisoned_cpu_lock(op, "CPU engine resources"))
998}
999
1000fn saturating_add_tensor_cache_stats(total: &mut CacheStats, value: CacheStats) {
1001 total.entries = total.entries.saturating_add(value.entries);
1002 total.retained_bytes = total.retained_bytes.saturating_add(value.retained_bytes);
1003 total.hits = total.hits.saturating_add(value.hits);
1004 total.misses = total.misses.saturating_add(value.misses);
1005 total.evictions = total.evictions.saturating_add(value.evictions);
1006 total.clears = total.clears.saturating_add(value.clears);
1007}
1008
1009#[derive(Clone)]
1024pub struct CpuBackend {
1025 runtime_identity: CpuRuntimeIdentity,
1026 shared: Arc<CpuBackendState>,
1027 requested: CpuPlacement,
1028 resolved: ResolvedCpuExecution,
1029 engine: Arc<CpuEngine>,
1030 provider_bundle: CpuProviderBundle,
1031 allocation_domain: Option<Arc<dyn SharedTensorAllocationDomain>>,
1032}
1033
1034#[derive(Clone, Debug)]
1050pub struct CpuRuntimeIdentity {
1051 marker: Arc<()>,
1052}
1053
1054impl CpuRuntimeIdentity {
1055 fn fresh() -> Self {
1056 Self {
1057 marker: Arc::new(()),
1058 }
1059 }
1060}
1061
1062impl PartialEq for CpuRuntimeIdentity {
1063 fn eq(&self, other: &Self) -> bool {
1064 Arc::ptr_eq(&self.marker, &other.marker)
1065 }
1066}
1067
1068impl Eq for CpuRuntimeIdentity {}
1069
1070fn resolve_discovered_topology(
1071 kind: CpuBackendKind,
1072 topology: Result<CpuTopology, CpuTopologyError>,
1073) -> Result<CpuTopology, CpuPlacementError> {
1074 topology.map_err(|source| CpuPlacementError::TopologyDiscovery {
1075 requested: CpuPlacement::Auto,
1076 backend: kind,
1077 source,
1078 })
1079}
1080
1081fn coordinator_node_domain_ids(topology: &CpuTopology) -> BTreeMap<NumaNodeId, CpuDomainId> {
1082 topology
1083 .nodes()
1084 .iter()
1085 .enumerate()
1086 .filter_map(|(index, node)| {
1087 u64::try_from(index)
1088 .ok()
1089 .and_then(|index| index.checked_add(1))
1090 .map(|id| (node.id(), CpuDomainId::new(id)))
1091 })
1092 .collect()
1093}
1094
1095impl fmt::Debug for CpuBackend {
1096 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1097 f.debug_struct("CpuBackend")
1098 .field("kind", &self.kind())
1099 .field("provider_bundle", &self.provider_bundle)
1100 .field("requested_placement", &self.requested)
1101 .field("resolved_execution", &self.resolved)
1102 .field("engine_placement", &self.engine.placement())
1103 .field("num_threads", &self.num_threads())
1104 .field("allocation_domain", &self.allocation_domain())
1105 .field("buffer_pool_cache_stats", &self.buffer_pool_cache_stats())
1106 .field("buffer_pool_limit_bytes", &self.buffer_pool_limit_bytes())
1107 .finish_non_exhaustive()
1108 }
1109}
1110
1111impl CpuBackend {
1112 fn from_thread_budget_and_kind(
1113 thread_budget: usize,
1114 kind: CpuBackendKind,
1115 max_retained_capacity_bytes: usize,
1116 ) -> Result<Self, CpuPlacementError> {
1117 let topology = resolve_discovered_topology(kind, discover_cpu_topology())?;
1118 let resolved = resolve_placement(kind, CpuPlacement::Auto, &topology)?;
1119 #[cfg(not(any(target_os = "linux", target_os = "android")))]
1120 {
1121 let context = CpuContext::with_threads(thread_budget).map_err(|error| {
1122 CpuPlacementError::EngineConstruction {
1123 requested: CpuPlacement::Auto,
1124 backend: kind,
1125 source: CpuEngineConstructionError::Tensor(error),
1126 }
1127 })?;
1128 Ok(Self::compatibility_with_topology(
1129 Arc::new(context),
1130 max_retained_capacity_bytes,
1131 kind,
1132 topology,
1133 resolved,
1134 ))
1135 }
1136 #[cfg(any(target_os = "linux", target_os = "android"))]
1137 {
1138 let engine_placement = ResolvedCpuPlacement::AllAllowed {
1139 cpus: topology.allowed_cpus().clone(),
1140 };
1141 let engine = Arc::new(
1142 CpuEngine::new_managed(
1143 CpuDomainId::new(0),
1144 engine_placement,
1145 thread_budget,
1146 max_retained_capacity_bytes,
1147 )
1148 .map_err(|error| CpuPlacementError::EngineConstruction {
1149 requested: CpuPlacement::Auto,
1150 backend: kind,
1151 source: CpuEngineConstructionError::Context(error),
1152 })?,
1153 );
1154 let all_allowed = OnceLock::new();
1155 let _ = all_allowed.set(Arc::clone(&engine));
1156 Ok(Self {
1157 shared: Arc::new(CpuBackendState {
1158 engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
1159 node_engines: Mutex::new(BTreeMap::new()),
1160 node_domain_ids: coordinator_node_domain_ids(&topology),
1161 all_allowed,
1162 all_allowed_build: Mutex::new(()),
1163 base_engine: Arc::clone(&engine),
1164 thread_budget,
1165 }),
1166 topology,
1167 arbiter: ResourceArbiter::global(),
1168 kind,
1169 buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
1170 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1171 }),
1172 runtime_identity: CpuRuntimeIdentity::fresh(),
1173 requested: CpuPlacement::Auto,
1174 resolved,
1175 engine,
1176 provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
1177 allocation_domain: None,
1178 })
1179 }
1180 }
1181
1182 fn compatibility(
1183 ctx: Arc<CpuContext>,
1184 max_retained_capacity_bytes: usize,
1185 kind: CpuBackendKind,
1186 ) -> Self {
1187 let topology = discover_cpu_topology().unwrap_or_else(|_| {
1188 let allowed = crate::process_cpu_affinity().unwrap_or_else(|| {
1189 CpuSet::new((0..crate::available_parallelism()).map(CpuId::new))
1190 .unwrap_or_else(|_| CpuSet::singleton(CpuId::new(0)))
1191 });
1192 CpuTopology::all_allowed(allowed)
1193 });
1194 let resolved = if kind == CpuBackendKind::Blas {
1195 ResolvedCpuExecution::ProviderDefaultExclusive
1196 } else {
1197 ResolvedCpuExecution::Compatibility
1198 };
1199 Self::compatibility_with_topology(
1200 ctx,
1201 max_retained_capacity_bytes,
1202 kind,
1203 topology,
1204 resolved,
1205 )
1206 }
1207
1208 fn compatibility_with_topology(
1209 ctx: Arc<CpuContext>,
1210 max_retained_capacity_bytes: usize,
1211 kind: CpuBackendKind,
1212 topology: CpuTopology,
1213 resolved: ResolvedCpuExecution,
1214 ) -> Self {
1215 let placement = ResolvedCpuPlacement::AllAllowed {
1216 cpus: topology.allowed_cpus().clone(),
1217 };
1218 let base_engine = Arc::new(CpuEngine::from_context(
1219 CpuDomainId::new(0),
1220 placement,
1221 ctx,
1222 max_retained_capacity_bytes,
1223 ));
1224 Self {
1225 shared: Arc::new(CpuBackendState {
1226 engines: CpuEngineRegistry::ManagedLazy(ManagedEngineRegistry {
1227 node_engines: Mutex::new(BTreeMap::new()),
1228 node_domain_ids: coordinator_node_domain_ids(&topology),
1229 all_allowed: OnceLock::new(),
1230 all_allowed_build: Mutex::new(()),
1231 base_engine: Arc::clone(&base_engine),
1232 thread_budget: base_engine.domain().thread_budget().get(),
1233 }),
1234 topology,
1235 arbiter: ResourceArbiter::global(),
1236 kind,
1237 buffer_limit: AtomicUsize::new(max_retained_capacity_bytes),
1238 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1239 }),
1240 runtime_identity: CpuRuntimeIdentity::fresh(),
1241 requested: CpuPlacement::Auto,
1242 resolved,
1243 engine: base_engine,
1244 provider_bundle: CpuProviderBundle::standard(kind, kind == CpuBackendKind::Blas),
1245 allocation_domain: None,
1246 }
1247 }
1248
1249 pub fn new() -> Self {
1259 let context = Arc::new(CpuContext::from_env());
1260 Self::from_thread_budget_and_kind(
1261 context.num_threads(),
1262 CpuBackendKind::default_compiled(),
1263 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1264 )
1265 .unwrap_or_else(|error| {
1266 eprintln!(
1267 "tenferro_cpu: using the unpinned compatibility context after placement error: {error}"
1268 );
1269 Self::from_context(context)
1270 })
1271 }
1272
1273 pub fn from_external_managed_domains(
1331 default_domain: CpuDomainId,
1332 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1333 ) -> Result<Self, CpuBackendError> {
1334 let op = "CpuBackend::from_external_managed_domains";
1335 let kind = CpuBackendKind::default_compiled();
1336 let topology = resolve_discovered_topology(kind, discover_cpu_topology())
1337 .map_err(|source| CpuBackendError::placement(op, source))?;
1338 Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1339 default_domain,
1340 domains,
1341 topology,
1342 ResourceArbiter::global(),
1343 CpuProviderBundle::standard(kind, false),
1344 )
1345 }
1346
1347 pub fn from_external_managed_domains_with_provider_bundle(
1402 default_domain: CpuDomainId,
1403 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1404 provider_bundle: CpuProviderBundle,
1405 ) -> Result<Self, CpuBackendError> {
1406 let op = "CpuBackend::from_external_managed_domains_with_provider_bundle";
1407 let kind = CpuBackendKind::default_compiled();
1408 let topology = resolve_discovered_topology(kind, discover_cpu_topology())
1409 .map_err(|source| CpuBackendError::placement(op, source))?;
1410 Self::from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1411 default_domain,
1412 domains,
1413 topology,
1414 ResourceArbiter::global(),
1415 provider_bundle,
1416 )
1417 }
1418
1419 fn from_external_managed_domains_with_topology_arbiter_and_provider_bundle(
1420 default_domain: CpuDomainId,
1421 domains: impl IntoIterator<Item = ExternalCpuDomain>,
1422 topology: CpuTopology,
1423 arbiter: ResourceArbiter,
1424 provider_bundle: CpuProviderBundle,
1425 ) -> Result<Self, CpuBackendError> {
1426 let domains: Vec<_> = domains.into_iter().collect();
1427 if domains.is_empty() {
1428 return Err(ExternalCpuDomainRegistryError::EmptyRegistry.into());
1429 }
1430
1431 let mut domain_ids = BTreeSet::new();
1432 let mut node_ids = BTreeSet::new();
1433 let mut has_all_allowed = false;
1434 for domain in &domains {
1435 if !domain_ids.insert(domain.id()) {
1436 return Err(
1437 ExternalCpuDomainRegistryError::DuplicateDomainId { id: domain.id() }.into(),
1438 );
1439 }
1440 match domain.placement() {
1441 ResolvedCpuPlacement::NumaNode { id, .. } => {
1442 if !node_ids.insert(*id) {
1443 return Err(ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
1444 placement: CpuPlacement::NumaNode(*id),
1445 }
1446 .into());
1447 }
1448 }
1449 ResolvedCpuPlacement::AllAllowed { cpus } => {
1450 if has_all_allowed {
1451 return Err(ExternalCpuDomainRegistryError::DuplicatePlacementIdentity {
1452 placement: CpuPlacement::AllAllowed,
1453 }
1454 .into());
1455 }
1456 has_all_allowed = true;
1457 if domain.placement_guarantee() == CpuPlacementGuarantee::ExactDeclared
1458 && cpus != topology.allowed_cpus()
1459 {
1460 return Err(ExternalCpuDomainRegistryError::ExactAllAllowedMismatch {
1461 domain: domain.id(),
1462 declared: cpus.clone(),
1463 allowed: topology.allowed_cpus().clone(),
1464 }
1465 .into());
1466 }
1467 }
1468 }
1469 if let Some(cpu) = domain
1470 .cpus()
1471 .as_slice()
1472 .iter()
1473 .copied()
1474 .find(|cpu| !topology.allowed_cpus().contains(*cpu))
1475 {
1476 return Err(ExternalCpuDomainRegistryError::CpuOutsideAllowedSet {
1477 domain: domain.id(),
1478 cpu,
1479 }
1480 .into());
1481 }
1482 }
1483 if !domain_ids.contains(&default_domain) {
1484 return Err(
1485 ExternalCpuDomainRegistryError::MissingDefaultDomain { default_domain }.into(),
1486 );
1487 }
1488
1489 let buffer_limit = crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES;
1490 let mut by_id = BTreeMap::new();
1491 let mut by_node = BTreeMap::new();
1492 let mut all_allowed = None;
1493 for domain in domains {
1494 let id = domain.id();
1495 let placement = domain.placement().clone();
1496 let engine = Arc::new(CpuEngine::from_external(domain, buffer_limit));
1497 match placement {
1498 ResolvedCpuPlacement::NumaNode { id, .. } => {
1499 by_node.insert(id, Arc::clone(&engine));
1500 }
1501 ResolvedCpuPlacement::AllAllowed { .. } => {
1502 all_allowed = Some(Arc::clone(&engine));
1503 }
1504 }
1505 by_id.insert(id, engine);
1506 }
1507 let Some(engine) = by_id.get(&default_domain).cloned() else {
1508 return Err(
1509 ExternalCpuDomainRegistryError::MissingDefaultDomain { default_domain }.into(),
1510 );
1511 };
1512 let resolved = ResolvedCpuExecution::ExternalManaged(engine.placement().clone());
1513 let kind = CpuBackendKind::default_compiled();
1514 let backend = Self {
1515 runtime_identity: CpuRuntimeIdentity::fresh(),
1516 shared: Arc::new(CpuBackendState {
1517 topology,
1518 engines: CpuEngineRegistry::ExternalPrebuilt(ExternalEngineRegistry {
1519 by_id,
1520 by_node,
1521 all_allowed,
1522 default_domain,
1523 }),
1524 arbiter,
1525 kind,
1526 buffer_limit: AtomicUsize::new(buffer_limit),
1527 indexed_plan_cache_limits: Mutex::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS),
1528 }),
1529 requested: CpuPlacement::Auto,
1530 resolved,
1531 engine,
1532 provider_bundle,
1533 allocation_domain: None,
1534 };
1535 backend
1536 .validate_provider_bundle_for_domains(&backend.provider_bundle)
1537 .map_err(|source| {
1538 CpuBackendError::Tensor(crate::Error::backend_source(
1539 "CpuBackend ExternalManaged provider validation",
1540 source,
1541 ))
1542 })?;
1543 Ok(backend)
1544 }
1545
1546 pub fn with_kind(kind: CpuBackendKind) -> Result<Self, CpuBackendError> {
1563 let op = "CpuBackend::with_kind";
1564 ensure_cpu_backend_kind_available(kind, op)
1565 .map_err(|error| constructor_tensor_error(op, error))?;
1566 let context = CpuContext::from_env();
1567 Self::from_thread_budget_and_kind(
1568 context.num_threads(),
1569 kind,
1570 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1571 )
1572 .map_err(|error| CpuBackendError::placement(op, error))
1573 }
1574
1575 pub fn try_new() -> Result<Self, CpuBackendError> {
1594 let op = "CpuBackend::try_new";
1595 let context =
1596 CpuContext::try_from_env().map_err(|error| constructor_tensor_error(op, error))?;
1597 Self::from_thread_budget_and_kind(
1598 context.num_threads(),
1599 CpuBackendKind::default_compiled(),
1600 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1601 )
1602 .map_err(|error| CpuBackendError::placement(op, error))
1603 }
1604
1605 pub fn from_context(ctx: Arc<CpuContext>) -> Self {
1618 Self::compatibility(
1619 ctx,
1620 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1621 CpuBackendKind::default_compiled(),
1622 )
1623 }
1624
1625 pub fn from_context_with_buffer_pool_limit(
1641 ctx: Arc<CpuContext>,
1642 max_retained_capacity_bytes: usize,
1643 ) -> Self {
1644 Self::from_context_with_buffer_pool_limit_and_kind(
1645 ctx,
1646 max_retained_capacity_bytes,
1647 CpuBackendKind::default_compiled(),
1648 )
1649 }
1650
1651 fn from_context_with_buffer_pool_limit_and_kind(
1652 ctx: Arc<CpuContext>,
1653 max_retained_capacity_bytes: usize,
1654 kind: CpuBackendKind,
1655 ) -> Self {
1656 Self::compatibility(ctx, max_retained_capacity_bytes, kind)
1657 }
1658
1659 pub fn with_threads(num_threads: usize) -> Result<Self, CpuBackendError> {
1676 let op = "CpuBackend::with_threads";
1677 let context = CpuContext::with_threads(num_threads)
1678 .map_err(|error| constructor_tensor_error(op, error))?;
1679 Self::from_thread_budget_and_kind(
1680 context.num_threads(),
1681 CpuBackendKind::default_compiled(),
1682 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1683 )
1684 .map_err(|error| CpuBackendError::placement(op, error))
1685 }
1686
1687 pub fn with_threads_and_kind(
1708 num_threads: usize,
1709 kind: CpuBackendKind,
1710 ) -> Result<Self, CpuBackendError> {
1711 let op = "CpuBackend::with_threads_and_kind";
1712 ensure_cpu_backend_kind_available(kind, op)
1713 .map_err(|error| constructor_tensor_error(op, error))?;
1714 let context = CpuContext::with_threads(num_threads)
1715 .map_err(|error| constructor_tensor_error(op, error))?;
1716 Self::from_thread_budget_and_kind(
1717 context.num_threads(),
1718 kind,
1719 crate::buffer_pool::DEFAULT_MAX_RETAINED_CAPACITY_BYTES,
1720 )
1721 .map_err(|error| CpuBackendError::placement(op, error))
1722 }
1723
1724 pub fn for_placement(&self, requested: CpuPlacement) -> Result<Self, CpuPlacementError> {
1748 self.for_placement_with_affinity(
1749 requested,
1750 cfg!(any(target_os = "linux", target_os = "android")),
1751 )
1752 }
1753
1754 fn for_placement_with_affinity(
1755 &self,
1756 requested: CpuPlacement,
1757 managed_affinity_available: bool,
1758 ) -> Result<Self, CpuPlacementError> {
1759 if self.shared.is_external() {
1760 let engine = self.shared.external_engine_for(requested)?;
1761 return Ok(Self {
1762 runtime_identity: CpuRuntimeIdentity::fresh(),
1763 shared: Arc::clone(&self.shared),
1764 requested,
1765 resolved: ResolvedCpuExecution::ExternalManaged(engine.placement().clone()),
1766 engine,
1767 provider_bundle: self.provider_bundle.clone(),
1768 allocation_domain: self.allocation_domain.clone(),
1769 });
1770 }
1771 let resolved = resolve_placement_with_affinity(
1772 self.kind(),
1773 requested,
1774 &self.shared.topology,
1775 managed_affinity_available,
1776 )?;
1777 if requested == CpuPlacement::Auto && !managed_affinity_available {
1778 return Ok(Self {
1779 runtime_identity: CpuRuntimeIdentity::fresh(),
1780 shared: Arc::clone(&self.shared),
1781 requested,
1782 resolved,
1783 engine: self.shared.managed_base_engine(requested)?,
1784 provider_bundle: self.provider_bundle.clone(),
1785 allocation_domain: self.allocation_domain.clone(),
1786 });
1787 }
1788 let engine_placement = match &resolved {
1789 ResolvedCpuExecution::Managed(placement) => placement.clone(),
1790 ResolvedCpuExecution::ExternalManaged(_) => {
1791 return Err(CpuPlacementError::InternalState {
1792 requested,
1793 backend: self.kind(),
1794 message: "managed resolver returned an external execution mode",
1795 });
1796 }
1797 ResolvedCpuExecution::ProviderDefaultExclusive => ResolvedCpuPlacement::AllAllowed {
1798 cpus: self.shared.topology.allowed_cpus().clone(),
1799 },
1800 ResolvedCpuExecution::Compatibility => {
1801 return Err(CpuPlacementError::InternalState {
1802 requested,
1803 backend: self.kind(),
1804 message: "placement resolution returned an internal compatibility mode",
1805 });
1806 }
1807 };
1808 let engine = self
1809 .shared
1810 .managed_engine_for(&engine_placement, requested)?;
1811 Ok(Self {
1812 runtime_identity: CpuRuntimeIdentity::fresh(),
1813 shared: Arc::clone(&self.shared),
1814 requested,
1815 resolved,
1816 engine,
1817 provider_bundle: self.provider_bundle.clone(),
1818 allocation_domain: self.allocation_domain.clone(),
1819 })
1820 }
1821
1822 pub fn placement(&self) -> CpuPlacement {
1832 self.requested
1833 }
1834
1835 pub fn resolved_placement(&self) -> Option<&ResolvedCpuPlacement> {
1854 match &self.resolved {
1855 ResolvedCpuExecution::Managed(placement)
1856 | ResolvedCpuExecution::ExternalManaged(placement) => Some(placement),
1857 ResolvedCpuExecution::Compatibility
1858 | ResolvedCpuExecution::ProviderDefaultExclusive => None,
1859 }
1860 }
1861
1862 pub fn topology(&self) -> &CpuTopology {
1872 &self.shared.topology
1873 }
1874
1875 pub fn supports_placement(&self, placement: CpuPlacement) -> bool {
1885 if self.shared.is_external() {
1886 self.shared.external_engine_for(placement).is_ok()
1887 } else {
1888 resolve_placement(self.kind(), placement, &self.shared.topology).is_ok()
1889 }
1890 }
1891
1892 pub fn execution_info(&self) -> CpuExecutionInfo {
1901 let domain = self.engine.domain();
1902 let capabilities = domain.executor_capabilities();
1903 let (executor_affinity, executor_shutdown) =
1904 if domain.ownership() == CpuDomainOwnership::ExternalManaged {
1905 (
1906 CpuExecutorAffinity::CallerDeclaredUnverified,
1907 CpuExecutorShutdown::CallerOwned,
1908 )
1909 } else {
1910 (capabilities.affinity, capabilities.shutdown)
1911 };
1912 CpuExecutionInfo {
1913 backend_kind: self.kind(),
1914 execution_mode: match &self.resolved {
1915 ResolvedCpuExecution::Managed(_) => CpuExecutionMode::Managed,
1916 ResolvedCpuExecution::ExternalManaged(_) => CpuExecutionMode::ExternalManaged,
1917 ResolvedCpuExecution::ProviderDefaultExclusive => {
1918 CpuExecutionMode::ProviderDefaultExclusive
1919 }
1920 ResolvedCpuExecution::Compatibility => CpuExecutionMode::Compatibility,
1921 },
1922 requested_placement: self.requested,
1923 resolved_placement: self.resolved_placement().cloned(),
1924 topology: self.shared.topology.clone(),
1925 domain_id: domain.id(),
1926 domain_cpus: domain.cpus().clone(),
1927 worker_count: capabilities.worker_count.get(),
1928 thread_budget: domain.thread_budget().get(),
1929 placement_guarantee: domain.placement_guarantee(),
1930 domain_ownership: domain.ownership(),
1931 executor_affinity,
1932 executor_shutdown,
1933 provider_diagnostic: provider_diagnostic(self.kind(), domain.ownership()),
1934 }
1935 }
1936
1937 #[cfg(all(
1938 test,
1939 feature = "cpu-faer",
1940 any(target_os = "linux", target_os = "android")
1941 ))]
1942 fn coordinator_id_for_test(&self) -> usize {
1943 Arc::as_ptr(&self.shared) as usize
1944 }
1945
1946 #[cfg(test)]
1947 pub(crate) fn context_id_for_test(&self) -> usize {
1948 Arc::as_ptr(self.engine.domain().executor()) as *const () as usize
1949 }
1950
1951 pub fn kind(&self) -> CpuBackendKind {
1962 self.shared.kind
1963 }
1964
1965 pub fn provider_bundle(&self) -> &CpuProviderBundle {
1967 &self.provider_bundle
1968 }
1969
1970 pub fn runtime_identity(&self) -> CpuRuntimeIdentity {
1977 self.runtime_identity.clone()
1978 }
1979
1980 pub fn with_provider_bundle(
2000 mut self,
2001 bundle: CpuProviderBundle,
2002 ) -> Result<Self, CpuProviderBundleInstallError> {
2003 self.validate_provider_bundle_for_domains(&bundle)?;
2004 self.provider_bundle = bundle;
2005 self.runtime_identity = CpuRuntimeIdentity::fresh();
2006 Ok(self)
2007 }
2008
2009 fn validate_provider_bundle_for_domains(
2010 &self,
2011 bundle: &CpuProviderBundle,
2012 ) -> Result<(), CpuProviderBundleInstallError> {
2013 let allowed = self.shared.topology.allowed_cpus();
2014 let validate_engine = |engine: &CpuEngine| {
2015 let domain = engine.domain();
2016 bundle.validate_for_domain(
2017 domain.id(),
2018 domain.thread_budget(),
2019 domain.placement_guarantee(),
2020 domain.cpus(),
2021 allowed,
2022 )
2023 };
2024
2025 match &self.shared.engines {
2026 CpuEngineRegistry::ExternalPrebuilt(registry) => {
2027 for engine in registry.by_id.values() {
2028 validate_engine(engine)?;
2029 }
2030 }
2031 CpuEngineRegistry::ManagedLazy(registry) => {
2032 validate_engine(®istry.base_engine)?;
2033
2034 #[cfg(any(target_os = "linux", target_os = "android"))]
2039 for node in self.shared.topology.nodes() {
2040 let Some(domain_id) = registry.node_domain_ids.get(&node.id()).copied() else {
2041 continue;
2042 };
2043 let budget =
2044 std::num::NonZeroUsize::new(registry.thread_budget.min(node.cpus().len()))
2045 .expect("usable topology nodes have non-empty CPU sets");
2046 bundle.validate_for_domain(
2047 domain_id,
2048 budget,
2049 CpuPlacementGuarantee::ExactDeclared,
2050 node.cpus(),
2051 allowed,
2052 )?;
2053 }
2054 }
2055 }
2056 Ok(())
2057 }
2058
2059 pub fn num_threads(&self) -> usize {
2070 self.engine.domain().thread_budget().get()
2071 }
2072
2073 pub fn buffer_pool_len(&self) -> crate::Result<usize> {
2090 self.shared
2091 .initialized_engines("CpuBackend::buffer_pool_len")?
2092 .iter()
2093 .try_fold(0, |total, engine| {
2094 Ok(total
2095 + lock_engine_resources(engine, "CpuBackend::buffer_pool_len")?
2096 .buffers
2097 .len())
2098 })
2099 }
2100
2101 pub fn buffer_pool_stats(&self) -> crate::Result<BufferPoolStats> {
2120 self.shared
2121 .initialized_engines("CpuBackend::buffer_pool_stats")?
2122 .iter()
2123 .try_fold(BufferPoolStats::default(), |mut total, engine| {
2124 let stats = lock_engine_resources(engine, "CpuBackend::buffer_pool_stats")?
2125 .buffers
2126 .stats();
2127 total.buffers += stats.buffers;
2128 total.capacity_bytes += stats.capacity_bytes;
2129 Ok(total)
2130 })
2131 }
2132
2133 pub fn buffer_pool_cache_stats(&self) -> crate::Result<CacheStats> {
2152 let stats = self.buffer_pool_stats()?;
2153 Ok(CacheStats {
2154 entries: stats.buffers,
2155 retained_bytes: stats.capacity_bytes,
2156 hits: 0,
2157 misses: 0,
2158 evictions: 0,
2159 clears: 0,
2160 })
2161 }
2162
2163 pub fn indexed_plan_cache_limits(&self) -> crate::Result<IndexedPlanCacheLimits> {
2180 self.shared
2181 .indexed_plan_cache_limits
2182 .lock()
2183 .map(|limits| *limits)
2184 .map_err(|_| {
2185 poisoned_cpu_lock(
2186 "CpuBackend::indexed_plan_cache_limits",
2187 "CPU indexed-plan cache configuration",
2188 )
2189 })
2190 }
2191
2192 pub fn set_indexed_plan_cache_limits(
2213 &mut self,
2214 limits: IndexedPlanCacheLimits,
2215 ) -> crate::Result<()> {
2216 let mut configured_limits = self.shared.indexed_plan_cache_limits.lock().map_err(|_| {
2220 poisoned_cpu_lock(
2221 "CpuBackend::set_indexed_plan_cache_limits",
2222 "CPU indexed-plan cache configuration",
2223 )
2224 })?;
2225 let engines = self
2226 .shared
2227 .initialized_engines("CpuBackend::set_indexed_plan_cache_limits")?;
2228 let mut resources = engines
2229 .iter()
2230 .map(|engine| {
2231 lock_engine_resources(engine, "CpuBackend::set_indexed_plan_cache_limits")
2232 })
2233 .collect::<crate::Result<Vec<_>>>()?;
2234 *configured_limits = limits;
2235 for resource in &mut resources {
2236 resource.indexed_plan_cache.set_limits(limits);
2237 }
2238 Ok(())
2239 }
2240
2241 pub fn indexed_plan_cache_stats(&self) -> crate::Result<CacheStats> {
2258 self.shared
2259 .initialized_engines("CpuBackend::indexed_plan_cache_stats")?
2260 .iter()
2261 .try_fold(CacheStats::default(), |mut total, engine| {
2262 let stats = lock_engine_resources(engine, "CpuBackend::indexed_plan_cache_stats")?
2263 .indexed_plan_cache
2264 .stats();
2265 saturating_add_tensor_cache_stats(&mut total, stats);
2266 Ok(total)
2267 })
2268 }
2269
2270 pub fn clear_indexed_plan_cache(&mut self) -> crate::Result<()> {
2288 let engines = self
2289 .shared
2290 .initialized_engines("CpuBackend::clear_indexed_plan_cache")?;
2291 let mut resources = engines
2292 .iter()
2293 .map(|engine| lock_engine_resources(engine, "CpuBackend::clear_indexed_plan_cache"))
2294 .collect::<crate::Result<Vec<_>>>()?;
2295 for resource in &mut resources {
2296 resource.indexed_plan_cache.clear();
2297 }
2298 Ok(())
2299 }
2300
2301 pub fn buffer_pool_limit_bytes(&self) -> usize {
2316 self.shared.buffer_limit.load(Ordering::Relaxed)
2317 }
2318
2319 pub fn set_buffer_pool_limit_bytes(
2342 &mut self,
2343 max_retained_capacity_bytes: usize,
2344 ) -> crate::Result<()> {
2345 let engines = self
2346 .shared
2347 .initialized_engines("CpuBackend::set_buffer_pool_limit_bytes")?;
2348 let mut resources = engines
2349 .iter()
2350 .map(|engine| lock_engine_resources(engine, "CpuBackend::set_buffer_pool_limit_bytes"))
2351 .collect::<crate::Result<Vec<_>>>()?;
2352 self.shared
2353 .buffer_limit
2354 .store(max_retained_capacity_bytes, Ordering::Relaxed);
2355 for resource in &mut resources {
2356 resource
2357 .buffers
2358 .set_max_retained_capacity_bytes(max_retained_capacity_bytes);
2359 }
2360 Ok(())
2361 }
2362
2363 pub fn reset_buffer_pool(&mut self) -> crate::Result<()> {
2386 let engines = self
2387 .shared
2388 .initialized_engines("CpuBackend::reset_buffer_pool")?;
2389 let mut resources = engines
2390 .iter()
2391 .map(|engine| lock_engine_resources(engine, "CpuBackend::reset_buffer_pool"))
2392 .collect::<crate::Result<Vec<_>>>()?;
2393 for resource in &mut resources {
2394 resource.buffers.clear();
2395 }
2396 Ok(())
2397 }
2398
2399 pub(crate) fn runtime_cache_stats(
2400 &self,
2401 ) -> crate::Result<tenferro_runtime::runtime::CacheStats> {
2402 let resources = lock_engine_resources(&self.engine, "CpuBackend::runtime_cache_stats")?;
2403 let buffers = resources.buffers.cache_stats();
2404 let gemm = tenferro_tensor::RuntimeCacheControl::stats(&resources.gemm_analysis_cache);
2405 let indexed = resources.indexed_plan_cache.stats();
2406 Ok(tenferro_runtime::runtime::CacheStats {
2407 entries: buffers
2408 .entries
2409 .saturating_add(gemm.entries)
2410 .saturating_add(indexed.entries),
2411 retained_bytes: buffers
2412 .retained_bytes
2413 .saturating_add(gemm.retained_bytes)
2414 .saturating_add(indexed.retained_bytes),
2415 hits: indexed.hits,
2416 misses: indexed.misses,
2417 evictions: indexed.evictions,
2418 clears: indexed.clears,
2419 })
2420 }
2421
2422 pub(crate) fn clear_runtime_caches(&self) -> crate::Result<()> {
2423 let mut resources =
2424 lock_engine_resources(&self.engine, "CpuBackend::clear_runtime_caches")?;
2425 resources.buffers.clear();
2426 tenferro_tensor::RuntimeCacheControl::clear(&mut resources.gemm_analysis_cache);
2427 resources.indexed_plan_cache.clear();
2428 Ok(())
2429 }
2430
2431 pub fn install<R: Send>(&self, op: impl FnOnce() -> R + Send) -> R {
2452 let owner = inherited_or_new_execution_owner();
2453 let permit = self.acquire_execution_permit(owner);
2454 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
2455 match entry.enter(ParallelMode::Sequential, |_| op()) {
2456 Ok(result) => result,
2457 Err(error) => panic!("CpuBackend::install executor failed: {error}"),
2458 }
2459 }
2460
2461 fn try_install<R: Send>(
2462 &self,
2463 op: impl FnOnce() -> crate::Result<R> + Send,
2464 ) -> crate::Result<R> {
2465 let owner = inherited_or_new_execution_owner();
2466 let permit = self.acquire_execution_permit(owner);
2467 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
2468 let mode = entry.preferred_engine_mode();
2469 entry
2470 .enter(mode, |context| context.with_native_parallelism(op))
2471 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2472 }
2473
2474 fn try_install_with_context<R: Send>(
2475 &self,
2476 op: impl FnOnce(&CpuExecutionContext<'_>) -> crate::Result<R> + Send,
2477 ) -> crate::Result<R> {
2478 let owner = inherited_or_new_execution_owner();
2479 let permit = self.acquire_execution_permit(owner);
2480 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
2481 let mode = entry.preferred_engine_mode();
2482 entry
2483 .enter(mode, |context| {
2484 context.with_native_parallelism(|| op(context))
2485 })
2486 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2487 }
2488
2489 fn try_install_fresh<R: FreshCpuOutput + Send>(
2490 &self,
2491 op: impl FnOnce() -> crate::Result<R> + Send,
2492 ) -> crate::Result<R> {
2493 let domain = self.engine.domain().id();
2494 let mut output = self.try_install(op)?;
2495 output.tag_fresh(domain);
2496 Ok(output)
2497 }
2498
2499 fn try_install_fresh_with_context<R: FreshCpuOutput + Send>(
2500 &self,
2501 op: impl FnOnce(&CpuExecutionContext<'_>) -> crate::Result<R> + Send,
2502 ) -> crate::Result<R> {
2503 let domain = self.engine.domain().id();
2504 let mut output = self.try_install_with_context(op)?;
2505 output.tag_fresh(domain);
2506 Ok(output)
2507 }
2508
2509 fn install_with_pool_unmarked<R: Send>(
2510 &mut self,
2511 op: impl FnOnce(&mut BufferPool) -> crate::Result<R> + Send,
2512 ) -> crate::Result<R> {
2513 let owner = inherited_or_new_execution_owner();
2514 let permit = self.acquire_execution_permit(owner);
2515 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
2516 let mode = entry.preferred_engine_mode();
2517 entry
2518 .enter(mode, |context| {
2519 context.with_native_parallelism(|| {
2520 self.with_execution_resources(&permit, |resources| {
2521 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
2522 op(buffers.get_mut())
2523 })
2524 })
2525 })
2526 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2527 }
2528
2529 fn install_with_pool_context_unmarked<R: Send>(
2530 &mut self,
2531 op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
2532 ) -> crate::Result<R> {
2533 let owner = inherited_or_new_execution_owner();
2534 let permit = self.acquire_execution_permit(owner);
2535 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
2536 let mode = entry.preferred_engine_mode();
2537 entry
2538 .enter(mode, |context| {
2539 context.with_native_parallelism(|| {
2540 self.with_execution_resources(&permit, |resources| {
2541 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
2542 op(context, buffers.get_mut())
2543 })
2544 })
2545 })
2546 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2547 }
2548
2549 fn install_with_indexed_pool_context_unmarked<R: Send>(
2550 &mut self,
2551 op: impl FnOnce(
2552 &CpuExecutionContext<'_>,
2553 &mut BufferPool,
2554 &mut IndexedPlanCache,
2555 ) -> crate::Result<R>
2556 + Send,
2557 ) -> crate::Result<R> {
2558 let owner = inherited_or_new_execution_owner();
2559 let permit = self.acquire_execution_permit(owner);
2560 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
2561 let mode = entry.preferred_engine_mode();
2562 entry
2563 .enter(mode, |context| {
2564 context.with_native_parallelism(|| {
2565 self.with_execution_resources(&permit, |resources| {
2566 let EngineResources {
2567 buffers,
2568 indexed_plan_cache,
2569 ..
2570 } = resources;
2571 let mut buffers = BufferPoolLoan::new(buffers);
2572 op(context, buffers.get_mut(), indexed_plan_cache)
2573 })
2574 })
2575 })
2576 .map_err(|error| crate::Error::backend_source("CPU tensor execution", error))?
2577 }
2578
2579 fn install_with_pool<R: FreshCpuOutput + Send>(
2580 &mut self,
2581 op: impl FnOnce(&mut BufferPool) -> crate::Result<R> + Send,
2582 ) -> crate::Result<R> {
2583 let domain = self.engine.domain().id();
2584 let mut output = self.install_with_pool_unmarked(op)?;
2585 output.tag_fresh(domain);
2586 Ok(output)
2587 }
2588
2589 fn install_with_pool_context<R: FreshCpuOutput + Send>(
2590 &mut self,
2591 op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
2592 ) -> crate::Result<R> {
2593 let domain = self.engine.domain().id();
2594 let mut output = self.install_with_pool_context_unmarked(op)?;
2595 output.tag_fresh(domain);
2596 Ok(output)
2597 }
2598
2599 fn install_with_indexed_pool_context<R: FreshCpuOutput + Send>(
2600 &mut self,
2601 op: impl FnOnce(
2602 &CpuExecutionContext<'_>,
2603 &mut BufferPool,
2604 &mut IndexedPlanCache,
2605 ) -> crate::Result<R>
2606 + Send,
2607 ) -> crate::Result<R> {
2608 let domain = self.engine.domain().id();
2609 let mut output = self.install_with_indexed_pool_context_unmarked(op)?;
2610 output.tag_fresh(domain);
2611 Ok(output)
2612 }
2613
2614 #[doc(hidden)]
2639 pub fn with_linalg_pool<R: Send>(
2640 &mut self,
2641 op: impl FnOnce(&CpuExecutionContext<'_>, &mut BufferPool) -> crate::Result<R> + Send,
2642 ) -> crate::Result<R> {
2643 let owner = inherited_or_new_execution_owner();
2644 let permit = self.acquire_execution_permit(owner);
2645 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
2646 let mode = entry.preferred_linalg_mode(self.kind());
2647 entry
2648 .enter(mode, |context| {
2649 context.with_native_parallelism(|| {
2650 self.with_execution_resources(&permit, |resources| {
2651 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
2652 op(context, buffers.get_mut())
2653 })
2654 })
2655 })
2656 .map_err(|error| crate::Error::backend_source("CPU linalg execution", error))?
2657 }
2658
2659 fn with_execution_resources<R>(
2660 &self,
2661 permit: &ResourcePermit,
2662 op: impl FnOnce(&mut EngineResources) -> R,
2663 ) -> R {
2664 if permit.is_reentrant() {
2665 let mut resources =
2666 EngineResources::new(self.shared.buffer_limit.load(Ordering::Relaxed));
2667 return op(&mut resources);
2668 }
2669 let mut resources = self
2670 .engine
2671 .resources
2672 .lock()
2673 .unwrap_or_else(std::sync::PoisonError::into_inner);
2674 op(&mut resources)
2675 }
2676
2677 fn acquire_execution_permit(&self, owner: ResourceOwner) -> ResourcePermit {
2678 match &self.resolved {
2679 ResolvedCpuExecution::Managed(placement)
2680 | ResolvedCpuExecution::ExternalManaged(placement) => self
2681 .shared
2682 .arbiter
2683 .acquire_recovering(placement.cpus().clone(), owner),
2684 ResolvedCpuExecution::Compatibility => self
2685 .shared
2686 .arbiter
2687 .acquire_recovering(self.shared.topology.allowed_cpus().clone(), owner),
2688 ResolvedCpuExecution::ProviderDefaultExclusive => self
2689 .shared
2690 .arbiter
2691 .acquire_provider_exclusive_recovering(owner),
2692 }
2693 }
2694
2695 #[cfg(test)]
2696 fn try_acquire_execution_permit_for_test(
2697 &self,
2698 ) -> Result<Option<ResourcePermit>, crate::arbiter::ResourceArbiterError> {
2699 match &self.resolved {
2700 ResolvedCpuExecution::Managed(placement)
2701 | ResolvedCpuExecution::ExternalManaged(placement) => {
2702 self.shared.arbiter.try_acquire(placement.cpus().clone())
2703 }
2704 ResolvedCpuExecution::Compatibility => self
2705 .shared
2706 .arbiter
2707 .try_acquire(self.shared.topology.allowed_cpus().clone()),
2708 ResolvedCpuExecution::ProviderDefaultExclusive => {
2709 self.shared.arbiter.try_acquire_provider_exclusive()
2710 }
2711 }
2712 }
2713}
2714
2715impl BackendRuntimeCache for CpuBackend {
2716 type RuntimeCache = gemm::GemmAnalysisCache;
2717}
2718
2719impl TensorElementwise for CpuBackend {
2720 fn elementwise_read_into(
2721 &mut self,
2722 op: ElementwiseReadOp,
2723 inputs: &[TensorRead<'_>],
2724 out: TensorWrite<'_>,
2725 ) -> crate::Result<()> {
2726 self.install_with_pool_context_unmarked(|context, buffers| {
2727 let exec_context = context.strided_exec_context();
2728 tenferro_tensor::backend::elementwise_read_into_with_context(
2729 op,
2730 inputs,
2731 out,
2732 &exec_context,
2733 |inputs, out| elementwise_read_into_fallback_with_pool(buffers, op, inputs, out),
2734 )
2735 })
2736 }
2737
2738 fn add(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2739 self.install_with_pool(|buffers| elementwise::add_with_pool(buffers, lhs, rhs))
2740 }
2741
2742 fn add_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2743 self.install_with_pool(|buffers| elementwise::add_read_with_pool(buffers, lhs, rhs))
2744 }
2745
2746 fn sub(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2747 self.install_with_pool(|buffers| elementwise::sub_with_pool(buffers, lhs, rhs))
2748 }
2749
2750 fn sub_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2751 self.install_with_pool(|buffers| elementwise::sub_read_with_pool(buffers, lhs, rhs))
2752 }
2753
2754 fn mul(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2755 self.install_with_pool(|buffers| elementwise::mul_with_pool(buffers, lhs, rhs))
2756 }
2757
2758 fn mul_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2759 self.install_with_pool(|buffers| elementwise::mul_read_with_pool(buffers, lhs, rhs))
2760 }
2761
2762 fn neg(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2763 self.install_with_pool(|buffers| elementwise::neg_with_pool(buffers, input))
2764 }
2765
2766 fn neg_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2767 self.install_with_pool(|buffers| elementwise::neg_read_with_pool(buffers, input))
2768 }
2769
2770 fn conj(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2771 self.install_with_pool(|buffers| elementwise::conj_with_pool(buffers, input))
2772 }
2773
2774 fn conj_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2775 self.install_with_pool(|buffers| elementwise::conj_read_with_pool(buffers, input))
2776 }
2777
2778 fn div(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2779 self.install_with_pool(|buffers| elementwise::div_with_pool(buffers, lhs, rhs))
2780 }
2781
2782 fn div_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2783 self.install_with_pool(|buffers| elementwise::div_read_with_pool(buffers, lhs, rhs))
2784 }
2785
2786 fn rem(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2787 self.install_with_pool(|buffers| elementwise::rem_with_pool(buffers, lhs, rhs))
2788 }
2789
2790 fn rem_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2791 self.install_with_pool(|buffers| elementwise::rem_read_with_pool(buffers, lhs, rhs))
2792 }
2793
2794 fn abs(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2795 self.install_with_pool(|buffers| elementwise::abs_with_pool(buffers, input))
2796 }
2797
2798 fn abs_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2799 self.install_with_pool(|buffers| elementwise::abs_read_with_pool(buffers, input))
2800 }
2801
2802 fn sign(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2803 self.install_with_pool(|buffers| elementwise::sign_with_pool(buffers, input))
2804 }
2805
2806 fn sign_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2807 self.install_with_pool(|buffers| elementwise::sign_read_with_pool(buffers, input))
2808 }
2809
2810 fn maximum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2811 self.install_with_pool(|buffers| elementwise::maximum_with_pool(buffers, lhs, rhs))
2812 }
2813
2814 fn maximum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2815 self.install_with_pool(|buffers| elementwise::maximum_read_with_pool(buffers, lhs, rhs))
2816 }
2817
2818 fn minimum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2819 self.install_with_pool(|buffers| elementwise::minimum_with_pool(buffers, lhs, rhs))
2820 }
2821
2822 fn minimum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2823 self.install_with_pool(|buffers| elementwise::minimum_read_with_pool(buffers, lhs, rhs))
2824 }
2825
2826 fn compare(&mut self, lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor> {
2827 self.install_with_pool(|buffers| elementwise::compare_with_pool(buffers, lhs, rhs, dir))
2828 }
2829
2830 fn compare_read(
2831 &mut self,
2832 lhs: TensorRead<'_>,
2833 rhs: TensorRead<'_>,
2834 dir: &CompareDir,
2835 ) -> crate::Result<Tensor> {
2836 self.install_with_pool(|buffers| {
2837 elementwise::compare_read_with_pool(buffers, lhs, rhs, dir)
2838 })
2839 }
2840
2841 fn select(
2842 &mut self,
2843 pred: &Tensor,
2844 on_true: &Tensor,
2845 on_false: &Tensor,
2846 ) -> crate::Result<Tensor> {
2847 self.install_with_pool(|buffers| {
2848 elementwise::select_with_pool(buffers, pred, on_true, on_false)
2849 })
2850 }
2851
2852 fn select_read(
2853 &mut self,
2854 pred: TensorRead<'_>,
2855 on_true: TensorRead<'_>,
2856 on_false: TensorRead<'_>,
2857 ) -> crate::Result<Tensor> {
2858 self.install_with_pool(|buffers| {
2859 elementwise::select_read_with_pool(buffers, pred, on_true, on_false)
2860 })
2861 }
2862
2863 fn clamp(&mut self, input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor> {
2864 self.install_with_pool(|buffers| elementwise::clamp_with_pool(buffers, input, lower, upper))
2865 }
2866
2867 fn clamp_read(
2868 &mut self,
2869 input: TensorRead<'_>,
2870 lower: TensorRead<'_>,
2871 upper: TensorRead<'_>,
2872 ) -> crate::Result<Tensor> {
2873 self.install_with_pool(|buffers| {
2874 elementwise::clamp_read_with_pool(buffers, input, lower, upper)
2875 })
2876 }
2877}
2878
2879impl TensorAnalytic for CpuBackend {
2880 fn exp(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2881 self.install_with_pool(|buffers| analytic::exp_with_pool(buffers, input))
2882 }
2883
2884 fn exp_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2885 self.install_with_pool(|buffers| analytic::exp_read_with_pool(buffers, input))
2886 }
2887
2888 fn log(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2889 self.install_with_pool(|buffers| analytic::log_with_pool(buffers, input))
2890 }
2891
2892 fn log_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2893 self.install_with_pool(|buffers| analytic::log_read_with_pool(buffers, input))
2894 }
2895
2896 fn sin(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2897 self.install_with_pool(|buffers| analytic::sin_with_pool(buffers, input))
2898 }
2899
2900 fn sin_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2901 self.install_with_pool(|buffers| analytic::sin_read_with_pool(buffers, input))
2902 }
2903
2904 fn cos(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2905 self.install_with_pool(|buffers| analytic::cos_with_pool(buffers, input))
2906 }
2907
2908 fn cos_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2909 self.install_with_pool(|buffers| analytic::cos_read_with_pool(buffers, input))
2910 }
2911
2912 fn tanh(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2913 self.install_with_pool(|buffers| analytic::tanh_with_pool(buffers, input))
2914 }
2915
2916 fn tanh_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2917 self.install_with_pool(|buffers| analytic::tanh_read_with_pool(buffers, input))
2918 }
2919
2920 fn sqrt(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2921 self.install_with_pool(|buffers| analytic::sqrt_with_pool(buffers, input))
2922 }
2923
2924 fn sqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2925 self.install_with_pool(|buffers| analytic::sqrt_read_with_pool(buffers, input))
2926 }
2927
2928 fn rsqrt(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2929 self.install_with_pool(|buffers| analytic::rsqrt_with_pool(buffers, input))
2930 }
2931
2932 fn rsqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2933 self.install_with_pool(|buffers| analytic::rsqrt_read_with_pool(buffers, input))
2934 }
2935
2936 fn pow(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2937 self.install_with_pool(|buffers| analytic::pow_with_pool(buffers, lhs, rhs))
2938 }
2939
2940 fn pow_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
2941 self.install_with_pool(|buffers| analytic::pow_read_with_pool(buffers, lhs, rhs))
2942 }
2943
2944 fn expm1(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2945 self.install_with_pool(|buffers| analytic::expm1_with_pool(buffers, input))
2946 }
2947
2948 fn expm1_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2949 self.install_with_pool(|buffers| analytic::expm1_read_with_pool(buffers, input))
2950 }
2951
2952 fn log1p(&mut self, input: &Tensor) -> crate::Result<Tensor> {
2953 self.install_with_pool(|buffers| analytic::log1p_with_pool(buffers, input))
2954 }
2955
2956 fn log1p_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2957 self.install_with_pool(|buffers| analytic::log1p_read_with_pool(buffers, input))
2958 }
2959}
2960
2961impl TensorStructural for CpuBackend {
2962 fn to_contiguous_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
2963 self.install_with_pool(|buffers| {
2964 materialize_tensor_read(buffers, "CpuBackend::to_contiguous_read", input)
2965 })
2966 }
2967
2968 fn copy_read_into(&mut self, src: TensorRead<'_>, dst: TensorWrite<'_>) -> crate::Result<()> {
2969 self.try_install(|| copy_tensor_read_into("CpuBackend::copy_read_into", src, dst))
2970 }
2971
2972 fn transpose(&mut self, input: &Tensor, perm: &[usize]) -> crate::Result<Tensor> {
2973 self.install_with_pool(|buffers| structural::transpose_with_pool(buffers, input, perm))
2974 }
2975
2976 fn transpose_read(&mut self, input: TensorRead<'_>, perm: &[usize]) -> crate::Result<Tensor> {
2977 self.install_with_pool(|buffers| structural::transpose_read_with_pool(buffers, input, perm))
2978 }
2979
2980 fn reshape(&mut self, input: &Tensor, shape: &[usize]) -> crate::Result<Tensor> {
2981 self.try_install(|| structural::reshape(input, shape))
2982 }
2983
2984 fn reshape_read(&mut self, input: TensorRead<'_>, shape: &[usize]) -> crate::Result<Tensor> {
2985 let materializes = matches!(&input, TensorRead::View(_));
2986 if materializes {
2987 self.install_with_pool(|buffers| {
2988 structural::reshape_read_with_pool(buffers, input, shape)
2989 })
2990 } else {
2991 self.install_with_pool_unmarked(|buffers| {
2992 structural::reshape_read_with_pool(buffers, input, shape)
2993 })
2994 }
2995 }
2996
2997 fn broadcast_in_dim(
2998 &mut self,
2999 input: &Tensor,
3000 shape: &[usize],
3001 dims: &[usize],
3002 ) -> crate::Result<Tensor> {
3003 self.install_with_pool(|buffers| {
3004 structural::broadcast_in_dim_with_pool(buffers, input, shape, dims)
3005 })
3006 }
3007
3008 fn broadcast_in_dim_read(
3009 &mut self,
3010 input: TensorRead<'_>,
3011 shape: &[usize],
3012 dims: &[usize],
3013 ) -> crate::Result<Tensor> {
3014 self.install_with_pool(|buffers| {
3015 structural::broadcast_in_dim_read_with_pool(buffers, input, shape, dims)
3016 })
3017 }
3018
3019 fn cast(&mut self, input: &Tensor, to: crate::DType) -> crate::Result<Tensor> {
3020 self.install_with_pool(|buffers| structural::cast_with_pool(buffers, input, to))
3021 }
3022
3023 fn extract_diagonal(
3024 &mut self,
3025 input: &Tensor,
3026 axis_a: usize,
3027 axis_b: usize,
3028 ) -> crate::Result<Tensor> {
3029 self.install_with_pool(|buffers| {
3030 structural::extract_diagonal_with_pool(buffers, input, axis_a, axis_b)
3031 })
3032 }
3033
3034 fn embed_diagonal(
3035 &mut self,
3036 input: &Tensor,
3037 axis_a: usize,
3038 axis_b: usize,
3039 ) -> crate::Result<Tensor> {
3040 self.install_with_pool(|buffers| {
3041 structural::embed_diagonal_with_pool(buffers, input, axis_a, axis_b)
3042 })
3043 }
3044
3045 fn tril(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor> {
3046 self.install_with_pool(|buffers| structural::tril_with_pool(buffers, input, k))
3047 }
3048
3049 fn triu(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor> {
3050 self.install_with_pool(|buffers| structural::triu_with_pool(buffers, input, k))
3051 }
3052}
3053
3054impl TensorReduction for CpuBackend {
3055 fn reduce_sum(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3056 self.try_install_fresh_with_context(|context| {
3057 let exec_context = context.strided_exec_context();
3058 reduction::reduce_sum(input, axes, &exec_context)
3059 })
3060 }
3061
3062 fn reduce_sum_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3063 self.install_with_pool_context(|context, buffers| {
3064 let exec_context = context.strided_exec_context();
3065 reduction::reduce_sum_read(buffers, input, axes, &exec_context)
3066 })
3067 }
3068
3069 fn reduce_sum_squares_read(
3070 &mut self,
3071 input: TensorRead<'_>,
3072 axes: &[usize],
3073 ) -> crate::Result<Tensor> {
3074 self.install_with_pool_context(|context, buffers| {
3075 let exec_context = context.strided_exec_context();
3076 reduction::reduce_sum_squares_read(buffers, input, axes, &exec_context)
3077 })
3078 }
3079
3080 fn reduce_prod(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3081 self.try_install_fresh_with_context(|context| {
3082 let exec_context = context.strided_exec_context();
3083 reduction::reduce_prod(input, axes, &exec_context)
3084 })
3085 }
3086
3087 fn reduce_prod_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3088 self.install_with_pool_context(|context, buffers| {
3089 let exec_context = context.strided_exec_context();
3090 reduction::reduce_prod_read(buffers, input, axes, &exec_context)
3091 })
3092 }
3093
3094 fn reduce_max(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3095 self.try_install_fresh(|| reduction::reduce_max(input, axes))
3096 }
3097
3098 fn reduce_max_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3099 self.install_with_pool(|buffers| reduction::reduce_max_read(buffers, input, axes))
3100 }
3101
3102 fn reduce_min(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3103 self.try_install_fresh(|| reduction::reduce_min(input, axes))
3104 }
3105
3106 fn reduce_min_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
3107 self.install_with_pool(|buffers| reduction::reduce_min_read(buffers, input, axes))
3108 }
3109}
3110
3111impl TensorDot for CpuBackend {
3112 fn dot_general(
3113 &mut self,
3114 lhs: &Tensor,
3115 rhs: &Tensor,
3116 config: &DotGeneralConfig,
3117 ) -> crate::Result<Tensor> {
3118 self.run_backend_session_cached(None, move |session| session.dot_general(lhs, rhs, config))
3119 }
3120
3121 fn dot_general_read(
3122 &mut self,
3123 lhs: TensorRead<'_>,
3124 rhs: TensorRead<'_>,
3125 config: &DotGeneralConfig,
3126 ) -> crate::Result<Tensor> {
3127 self.run_backend_session_cached(None, move |session| {
3128 session.dot_general_read(lhs, rhs, config)
3129 })
3130 }
3131
3132 fn dot_general_read_into(
3133 &mut self,
3134 lhs: TensorRead<'_>,
3135 rhs: TensorRead<'_>,
3136 config: &DotGeneralConfig,
3137 out: TensorWrite<'_>,
3138 ) -> crate::Result<()> {
3139 self.run_backend_session_cached(None, move |session| {
3140 session.dot_general_read_into(lhs, rhs, config, out)
3141 })
3142 }
3143
3144 fn dot_general_read_into_accum(
3145 &mut self,
3146 lhs: TensorRead<'_>,
3147 rhs: TensorRead<'_>,
3148 config: &DotGeneralConfig,
3149 accumulation: DotGeneralAccumulation,
3150 out: TensorWrite<'_>,
3151 ) -> crate::Result<()> {
3152 self.run_backend_session_cached(None, move |session| {
3153 session.dot_general_read_into_accum(lhs, rhs, config, accumulation, out)
3154 })
3155 }
3156
3157 fn dot_general_with_conj(
3158 &mut self,
3159 lhs: &Tensor,
3160 rhs: &Tensor,
3161 config: &DotGeneralConfig,
3162 lhs_conj: bool,
3163 rhs_conj: bool,
3164 ) -> crate::Result<Tensor> {
3165 self.run_backend_session_cached(None, move |session| {
3166 session.dot_general_with_conj(lhs, rhs, config, lhs_conj, rhs_conj)
3167 })
3168 }
3169}
3170
3171impl BackendCachedDot for CpuBackend {
3172 fn dot_general_cached(
3173 &mut self,
3174 cache: &mut Self::RuntimeCache,
3175 cache_slot: Option<usize>,
3176 lhs: &Tensor,
3177 rhs: &Tensor,
3178 config: &DotGeneralConfig,
3179 ) -> crate::Result<Tensor> {
3180 self.run_backend_session_cached(Some(cache), move |session| {
3181 session.dot_general_cached(cache_slot, lhs, rhs, config)
3182 })
3183 }
3184
3185 fn dot_general_with_conj_cached(
3186 &mut self,
3187 cache: &mut Self::RuntimeCache,
3188 cache_slot: Option<usize>,
3189 lhs: &Tensor,
3190 rhs: &Tensor,
3191 config: &DotGeneralConfig,
3192 lhs_conj: bool,
3193 rhs_conj: bool,
3194 ) -> crate::Result<Tensor> {
3195 self.run_backend_session_cached(Some(cache), move |session| {
3196 session.dot_general_with_conj_cached(cache_slot, lhs, rhs, config, lhs_conj, rhs_conj)
3197 })
3198 }
3199
3200 fn dot_general_read_into_accum_cached(
3201 &mut self,
3202 cache: &mut Self::RuntimeCache,
3203 cache_slot: Option<usize>,
3204 lhs: TensorRead<'_>,
3205 rhs: TensorRead<'_>,
3206 config: &DotGeneralConfig,
3207 accumulation: DotGeneralAccumulation,
3208 out: TensorWrite<'_>,
3209 ) -> crate::Result<()> {
3210 self.run_backend_session_cached(Some(cache), move |session| {
3211 session.dot_general_read_into_accum_cached(
3212 cache_slot,
3213 lhs,
3214 rhs,
3215 config,
3216 accumulation,
3217 out,
3218 )
3219 })
3220 }
3221
3222 fn grouped_gemm_cached(
3223 &mut self,
3224 cache: &mut Self::RuntimeCache,
3225 cache_slot: Option<usize>,
3226 lhs: TensorRead<'_>,
3227 rhs: TensorRead<'_>,
3228 config: &GroupedGemmConfig<'_>,
3229 out: TensorWrite<'_>,
3230 ) -> crate::Result<()> {
3231 self.run_backend_session_cached(Some(cache), move |session| {
3232 session.grouped_gemm_cached(cache_slot, lhs, rhs, config, out)
3233 })
3234 }
3235}
3236
3237impl TensorIndexing for CpuBackend {
3238 fn gather(
3239 &mut self,
3240 operand: &Tensor,
3241 start_indices: &Tensor,
3242 config: &GatherConfig,
3243 ) -> crate::Result<Tensor> {
3244 self.install_with_indexed_pool_context(|context, buffers, cache| {
3245 let exec_context = context.strided_exec_context();
3246 indexing::gather_with_pool(
3247 buffers,
3248 cache,
3249 &exec_context,
3250 operand,
3251 start_indices,
3252 config,
3253 )
3254 })
3255 }
3256
3257 fn scatter(
3258 &mut self,
3259 operand: &Tensor,
3260 scatter_indices: &Tensor,
3261 updates: &Tensor,
3262 config: &ScatterConfig,
3263 ) -> crate::Result<Tensor> {
3264 self.install_with_indexed_pool_context(|context, buffers, cache| {
3265 let exec_context = context.strided_exec_context();
3266 indexing::scatter_with_pool(
3267 buffers,
3268 cache,
3269 &exec_context,
3270 operand,
3271 scatter_indices,
3272 updates,
3273 config,
3274 )
3275 })
3276 }
3277
3278 fn slice(&mut self, input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor> {
3279 self.install_with_pool_context(|context, buffers| {
3280 let exec_context = context.strided_exec_context();
3281 indexing::try_slice_with_pool(buffers, &exec_context, input, config)
3282 })
3283 }
3284
3285 fn dynamic_slice(
3286 &mut self,
3287 input: &Tensor,
3288 starts: &Tensor,
3289 slice_sizes: &[usize],
3290 ) -> crate::Result<Tensor> {
3291 self.install_with_indexed_pool_context(|context, buffers, cache| {
3292 let exec_context = context.strided_exec_context();
3293 indexing::dynamic_slice_with_pool(
3294 buffers,
3295 cache,
3296 &exec_context,
3297 input,
3298 starts,
3299 slice_sizes,
3300 )
3301 })
3302 }
3303
3304 fn dynamic_update_slice(
3305 &mut self,
3306 operand: &Tensor,
3307 update: &Tensor,
3308 starts: &Tensor,
3309 ) -> crate::Result<Tensor> {
3310 self.install_with_indexed_pool_context(|context, buffers, cache| {
3311 let exec_context = context.strided_exec_context();
3312 indexing::dynamic_update_slice_with_pool(
3313 buffers,
3314 cache,
3315 &exec_context,
3316 operand,
3317 update,
3318 starts,
3319 )
3320 })
3321 }
3322
3323 fn pad(&mut self, input: &Tensor, config: &PadConfig) -> crate::Result<Tensor> {
3324 self.install_with_pool_context(|context, buffers| {
3325 let exec_context = context.strided_exec_context();
3326 indexing::try_pad_with_pool(buffers, &exec_context, input, config)
3327 })
3328 }
3329
3330 fn concatenate(&mut self, inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor> {
3331 self.install_with_pool_context(|context, buffers| {
3332 let exec_context = context.strided_exec_context();
3333 indexing::try_concatenate_with_pool(buffers, &exec_context, inputs, axis)
3334 })
3335 }
3336
3337 fn reverse(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor> {
3338 self.install_with_pool_context(|context, buffers| {
3339 let exec_context = context.strided_exec_context();
3340 indexing::reverse_with_pool(buffers, &exec_context, input, axes)
3341 })
3342 }
3343}
3344
3345impl CpuBackend {
3346 pub fn with_allocation_domain(mut self, domain: Arc<dyn SharedTensorAllocationDomain>) -> Self {
3371 self.allocation_domain = Some(domain);
3372 self.runtime_identity = CpuRuntimeIdentity::fresh();
3373 self
3374 }
3375
3376 pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
3386 self.allocation_domain.as_ref().map(|domain| domain.id())
3387 }
3388
3389 pub fn shared_allocation_domain(&self) -> Option<&Arc<dyn SharedTensorAllocationDomain>> {
3399 self.allocation_domain.as_ref()
3400 }
3401
3402 fn run_backend_session_cached<R: Send>(
3403 &mut self,
3404 cache: Option<&mut gemm::GemmAnalysisCache>,
3405 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
3406 ) -> R {
3407 let providers = self.provider_bundle.clone();
3408 let owner = inherited_or_new_execution_owner();
3409 let permit = self.acquire_execution_permit(owner);
3410 let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
3411 let enter_managed_session = entry.supports_infallible_session_entry()
3412 && !matches!(
3413 &self.resolved,
3414 ResolvedCpuExecution::ProviderDefaultExclusive
3415 );
3416 let run = |entered| {
3417 self.with_execution_resources(&permit, |resources| {
3418 let mut buffers = BufferPoolLoan::new(&mut resources.buffers);
3419 let cache = cache.unwrap_or(&mut resources.gemm_analysis_cache);
3420 let session_started = Instant::now();
3421 let mut session = CpuExecSession {
3422 entry,
3423 entered,
3424 buffers: buffers.get_mut(),
3425 gemm_analysis_cache: cache,
3426 indexed_plan_cache: &mut resources.indexed_plan_cache,
3427 providers: &providers,
3428 backend_kind: self.kind(),
3429 allocation_domain: self.allocation_domain.as_ref(),
3430 };
3431 record_cpu_session_profile(
3432 "with_backend_session_cached.session_construct",
3433 session_started.elapsed(),
3434 );
3435 let exec_started = Instant::now();
3436 let result = f(&mut session);
3437 record_cpu_session_profile(
3438 "with_backend_session_cached.exec_body",
3439 exec_started.elapsed(),
3440 );
3441 result
3442 })
3443 };
3444 if enter_managed_session {
3445 entry.enter_managed_session(|context| run(Some(context)))
3446 } else {
3447 with_execution_owner(owner, || run(None))
3448 }
3449 }
3450}
3451
3452impl BackendSessionHost for CpuBackend {
3453 fn with_backend_session<R: Send>(
3454 &mut self,
3455 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
3456 ) -> R {
3457 self.run_backend_session_cached(None, f)
3458 }
3459
3460 fn with_backend_session_cached<R: Send>(
3461 &mut self,
3462 cache: &mut Self::RuntimeCache,
3463 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
3464 ) -> R {
3465 if !cpu_session_profile_enabled() {
3466 return self.run_backend_session_cached(Some(cache), f);
3467 }
3468 let total_started = Instant::now();
3469 let result =
3470 profile_cpu_session_section("with_backend_session_cached.exec_session", || {
3471 self.run_backend_session_cached(Some(cache), f)
3472 });
3473 record_cpu_session_profile("with_backend_session_cached.total", total_started.elapsed());
3474 maybe_print_cpu_session_profile();
3475 result
3476 }
3477}
3478
3479impl TensorBuffer for CpuBackend {
3480 fn reclaim_buffer(&mut self, tensor: Tensor) {
3481 let owner = inherited_or_new_execution_owner();
3482 with_execution_owner(owner, || {
3483 let permit = self.acquire_execution_permit(owner);
3484 self.with_execution_resources(&permit, |resources| {
3485 let buffers = &mut resources.buffers;
3486 match tensor {
3487 Tensor::F32(t) => reclaim_typed(buffers, t),
3488 Tensor::F64(t) => reclaim_typed(buffers, t),
3489 Tensor::I32(t) => reclaim_typed(buffers, t),
3490 Tensor::I64(t) => reclaim_typed(buffers, t),
3491 Tensor::Bool(t) => reclaim_typed(buffers, t),
3492 Tensor::C32(t) => reclaim_typed(buffers, t),
3493 Tensor::C64(t) => reclaim_typed(buffers, t),
3494 }
3495 })
3496 })
3497 }
3498}
3499
3500impl<T, R> TensorViewCanonicalization<T, R> for CpuBackend
3501where
3502 T: TensorScalar + PoolScalar,
3503 R: TensorRank,
3504 R::Shape: Send + Sync,
3505 R::Strides: Send + Sync,
3506{
3507 fn to_contiguous(
3508 &mut self,
3509 view: &TypedTensorView<'_, T, R>,
3510 ) -> crate::Result<TypedTensor<T, R>> {
3511 self.install_with_pool(|buffers| {
3512 structural::typed_materialize_view_with_pool(buffers, view, "CpuBackend::to_contiguous")
3513 })
3514 }
3515
3516 fn copy_into(
3517 &mut self,
3518 src: &TypedTensorView<'_, T, R>,
3519 dst: &mut TypedTensorViewMut<'_, T, R>,
3520 ) -> crate::Result<()> {
3521 self.try_install(|| structural::typed_copy_view_into(src, dst, "CpuBackend::copy_into"))
3522 }
3523}
3524
3525impl TensorFusion for CpuBackend {
3526 fn execute_elementwise_fusion(
3527 &mut self,
3528 inputs: &[&Tensor],
3529 plan: &ElementwiseFusionPlan,
3530 ) -> crate::Result<Option<Vec<Tensor>>> {
3531 self.install_with_pool_context(|context, buffers| {
3532 let exec_context = context.strided_exec_context();
3533 elementwise::elementwise_fusion_with_pool(buffers, &exec_context, inputs, plan)
3534 })
3535 }
3536
3537 fn execute_broadcast_multiply(
3538 &mut self,
3539 lhs: TensorRead<'_>,
3540 lhs_shape: &[usize],
3541 lhs_dims: &[usize],
3542 rhs: TensorRead<'_>,
3543 rhs_shape: &[usize],
3544 rhs_dims: &[usize],
3545 ) -> crate::Result<Option<Tensor>> {
3546 self.install_with_pool(|buffers| {
3547 elementwise::broadcast_multiply_read_with_pool(
3548 buffers, lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims,
3549 )
3550 })
3551 }
3552
3553 fn execute_broadcast_multiply_value(
3554 &mut self,
3555 lhs: TensorRead<'_>,
3556 lhs_shape: &[usize],
3557 lhs_dims: &[usize],
3558 rhs: TensorRead<'_>,
3559 rhs_shape: &[usize],
3560 rhs_dims: &[usize],
3561 ) -> crate::Result<Option<TensorValue>> {
3562 let domain = self.engine.domain().id();
3563 self.install_with_pool_unmarked(|buffers| {
3564 elementwise::broadcast_multiply_value_with_pool_and_tag(
3565 buffers,
3566 lhs,
3567 lhs_shape,
3568 lhs_dims,
3569 rhs,
3570 rhs_shape,
3571 rhs_dims,
3572 |tensor| tag_fresh_output(tensor, domain),
3573 )
3574 })
3575 }
3576}
3577
3578impl TensorDeviceTransfer for CpuBackend {
3579 fn download_to_host(&mut self, tensor: &Tensor) -> crate::Result<Tensor> {
3580 if tensor.is_backend_buffer() {
3581 return Err(crate::Error::runtime_state(
3582 "CpuBackend::download_to_host",
3583 "CPU backend received a backend buffer; download the tensor to host with its owning backend before CPU execution",
3584 ));
3585 }
3586 Ok(tensor.clone())
3587 }
3588
3589 fn upload_host_tensor(&mut self, tensor: &Tensor) -> crate::Result<Tensor> {
3590 if tensor.is_backend_buffer() {
3591 return Err(crate::Error::runtime_state(
3592 "CpuBackend::upload_host_tensor",
3593 "CPU backend upload_host_tensor expects a host tensor; download backend buffers to host before CPU execution",
3594 ));
3595 }
3596 Ok(tensor.clone())
3597 }
3598}
3599
3600impl TensorBackend for CpuBackend {}
3601
3602pub(crate) fn reclaim_typed<T: PoolScalar>(pool: &mut BufferPool, typed: TypedTensor<T>) {
3603 let (buffer, _, _) = typed.into_parts();
3604 match buffer {
3605 Buffer::Host(data) => T::pool_release(pool, data),
3606 Buffer::Backend(_) => {}
3607 }
3608}
3609
3610impl Default for CpuBackend {
3611 fn default() -> Self {
3612 Self::new()
3613 }
3614}
3615
3616#[cfg(test)]
3617mod tests;