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