1use std::cell::Cell;
4use std::sync::{Arc, Mutex, OnceLock};
5
6use tenferro::{CompiledGraph, GraphCompiler, Runtime, Tensor, TracedGraph};
7use tenferro_ad::{AdContext, EagerRuntime};
8use tenferro_cpu::{BufferPoolStats, CpuBackend};
9use tenferro_tensor::{BackendSession, BackendSessionHost};
10
11#[derive(Clone, Debug)]
16pub enum ExecutionContext {
17 Cpu(Arc<CpuExecutionContext>),
19 #[cfg(feature = "tenferro-cuda")]
21 Cuda(Arc<crate::cuda::CudaExecutionContext>),
22}
23
24impl ExecutionContext {
25 pub fn is_global_default_cpu(&self) -> bool {
32 #[cfg(feature = "global-defaults")]
33 {
34 match self {
35 ExecutionContext::Cpu(context) => {
36 let own = context.eager_runtime().map(|runtime| runtime.id());
37 let global = defaults::default_eager_ctx().map(|runtime| runtime.id());
38 matches!((own, global), (Ok(a), Ok(b)) if a == b)
39 }
40 #[cfg(feature = "tenferro-cuda")]
41 ExecutionContext::Cuda(_) => false,
42 }
43 }
44 #[cfg(not(feature = "global-defaults"))]
45 {
46 let _ = self;
47 false
48 }
49 }
50}
51
52#[derive(Debug, Clone, thiserror::Error)]
70pub enum CpuExecutionContextError {
71 #[error("failed to initialize {component}: {source}")]
73 Initialization {
74 component: &'static str,
76 #[source]
78 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
79 },
80 #[error("CPU graph {operation} failed: {source}")]
82 Graph {
83 operation: &'static str,
85 #[source]
87 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
88 },
89}
90
91const CANONICAL_SESSION_REENTRY_MESSAGE: &str = "recursive tensorbackend canonical session entry";
92
93thread_local! {
94 static CANONICAL_SESSION_ACTIVE: Cell<bool> = const { Cell::new(false) };
95}
96
97struct CanonicalSessionGuard {
98 previous: bool,
99}
100
101impl CanonicalSessionGuard {
102 fn assert_inactive() {
103 CANONICAL_SESSION_ACTIVE.with(|active| {
104 assert!(!active.get(), "{CANONICAL_SESSION_REENTRY_MESSAGE}");
105 });
106 }
107
108 fn enter() -> Self {
109 Self::assert_inactive();
110 CANONICAL_SESSION_ACTIVE.with(|active| Self {
111 previous: active.replace(true),
112 })
113 }
114}
115
116impl Drop for CanonicalSessionGuard {
117 fn drop(&mut self) {
118 CANONICAL_SESSION_ACTIVE.with(|active| active.set(self.previous));
119 }
120}
121
122impl CpuExecutionContextError {
123 fn initialization(
124 component: &'static str,
125 source: impl std::error::Error + Send + Sync + 'static,
126 ) -> Self {
127 Self::Initialization {
128 component,
129 source: Arc::new(source),
130 }
131 }
132
133 fn graph(
134 operation: &'static str,
135 source: impl std::error::Error + Send + Sync + 'static,
136 ) -> Self {
137 Self::Graph {
138 operation,
139 source: Arc::new(source),
140 }
141 }
142}
143
144struct GraphState {
145 compiler: GraphCompiler,
146 runtime: Runtime,
147 backend: CpuBackend,
148}
149
150pub struct CpuExecutionContext {
170 backend: Mutex<CpuBackend>,
171 graph: OnceLock<Result<Mutex<GraphState>, CpuExecutionContextError>>,
172 eager: OnceLock<Result<Arc<EagerRuntime>, CpuExecutionContextError>>,
173}
174
175impl std::fmt::Debug for CpuExecutionContext {
176 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177 f.debug_struct("CpuExecutionContext")
178 .field("graph_initialized", &self.graph.get().is_some())
179 .field("eager_initialized", &self.eager.get().is_some())
180 .finish_non_exhaustive()
181 }
182}
183
184impl CpuExecutionContext {
185 pub fn from_backend(backend: CpuBackend) -> Self {
190 Self {
191 backend: Mutex::new(backend),
192 graph: OnceLock::new(),
193 eager: OnceLock::new(),
194 }
195 }
196
197 pub fn with_backend<R>(&self, f: impl FnOnce(&mut CpuBackend) -> R) -> R {
202 let mut backend = match self.backend.lock() {
203 Ok(guard) => guard,
204 Err(poisoned) => poisoned.into_inner(),
205 };
206 f(&mut backend)
207 }
208
209 pub(crate) fn with_session<R: Send>(
210 &self,
211 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
212 ) -> R {
213 CanonicalSessionGuard::assert_inactive();
214 let mut backend = match self.backend.lock() {
215 Ok(guard) => guard,
216 Err(poisoned) => poisoned.into_inner(),
217 };
218 backend.with_backend_session(|session| {
219 let _guard = CanonicalSessionGuard::enter();
220 f(session)
221 })
222 }
223
224 fn backend_clone(&self) -> CpuBackend {
225 self.with_backend(|backend| backend.clone())
226 }
227
228 fn graph_state(&self) -> Result<&Mutex<GraphState>, CpuExecutionContextError> {
229 self.graph
230 .get_or_init(|| {
231 let backend = self.backend_clone();
232 build_graph_runtime(&backend).map(|runtime| {
233 Mutex::new(GraphState {
234 compiler: GraphCompiler::new(),
235 runtime,
236 backend,
237 })
238 })
239 })
240 .as_ref()
241 .map_err(Clone::clone)
242 }
243
244 fn with_graph_state<R>(
245 &self,
246 f: impl FnOnce(&mut GraphCompiler, &mut Runtime, &mut CpuBackend) -> R,
247 ) -> Result<R, CpuExecutionContextError> {
248 let mut graph = match self.graph_state()?.lock() {
249 Ok(guard) => guard,
250 Err(poisoned) => poisoned.into_inner(),
251 };
252 let GraphState {
253 compiler,
254 runtime,
255 backend,
256 } = &mut *graph;
257 Ok(f(compiler, runtime, backend))
258 }
259
260 pub fn compile_graph(
267 &self,
268 graph: &TracedGraph,
269 ) -> Result<CompiledGraph, CpuExecutionContextError> {
270 self.with_graph_state(|compiler, _, _| compiler.compile_traced_graph(graph))?
271 .map_err(|source| CpuExecutionContextError::graph("compilation", source))
272 }
273
274 pub fn run_graph(
284 &self,
285 graph: &CompiledGraph,
286 inputs: &[&Tensor],
287 ) -> Result<Vec<Tensor>, CpuExecutionContextError> {
288 self.with_graph_state(|_, runtime, _| runtime.run_compiled(graph, inputs))?
289 .map_err(|source| CpuExecutionContextError::graph("execution", source))
290 }
291
292 pub fn eager_runtime(&self) -> Result<Arc<EagerRuntime>, CpuExecutionContextError> {
302 self.eager
303 .get_or_init(|| build_eager_runtime(self.backend_clone()))
304 .as_ref()
305 .map(Arc::clone)
306 .map_err(Clone::clone)
307 }
308
309 pub fn graph_cache_stats(
316 &self,
317 ) -> Result<tenferro::RuntimeCacheStats, CpuExecutionContextError> {
318 self.with_graph_state(|_, runtime, _| runtime.cache_stats())?
319 .map_err(|source| CpuExecutionContextError::graph("cache statistics", source))
320 }
321
322 pub fn graph_buffer_pool_stats(&self) -> Result<BufferPoolStats, CpuExecutionContextError> {
329 self.with_graph_state(|_, _, backend| backend.buffer_pool_stats())?
330 .map_err(|source| CpuExecutionContextError::graph("buffer-pool statistics", source))
331 }
332
333 pub fn reset_graph_buffer_pool(&self) -> Result<(), CpuExecutionContextError> {
340 self.with_graph_state(|_, _, backend| backend.reset_buffer_pool())?
341 .map_err(|source| CpuExecutionContextError::graph("buffer-pool reset", source))
342 }
343
344 pub fn reset_graph_runtime(&self) -> Result<(), CpuExecutionContextError> {
351 self.with_graph_state(|compiler, runtime, backend| {
352 let replacement = build_graph_runtime(backend)?;
353 *compiler = GraphCompiler::new();
354 let old = std::mem::replace(runtime, replacement);
355 drop(old);
356 backend
357 .reset_buffer_pool()
358 .map_err(|source| CpuExecutionContextError::graph("buffer-pool reset", source))
359 })??;
360 Ok(())
361 }
362}
363
364fn build_graph_runtime(backend: &CpuBackend) -> Result<Runtime, CpuExecutionContextError> {
365 let mut builder = Runtime::builder();
366 builder
367 .register_engine(
368 tenferro_cpu::runtime_engine_registration(backend).map_err(|source| {
369 CpuExecutionContextError::initialization("graph CPU engine", source)
370 })?,
371 )
372 .map_err(|source| CpuExecutionContextError::initialization("graph CPU engine", source))?;
373 builder
374 .install_extension_module(
375 tenferro_einsum::extension_module::<CpuBackend>(
376 tenferro_cpu::runtime_engine_id().map_err(|source| {
377 CpuExecutionContextError::initialization("einsum extension", source)
378 })?,
379 )
380 .map_err(|source| {
381 CpuExecutionContextError::initialization("einsum extension", source)
382 })?,
383 )
384 .map_err(|source| CpuExecutionContextError::initialization("einsum extension", source))?;
385 builder
386 .build()
387 .map_err(|source| CpuExecutionContextError::initialization("graph runtime", source))
388}
389
390fn build_eager_runtime(backend: CpuBackend) -> Result<Arc<EagerRuntime>, CpuExecutionContextError> {
391 let ad_context = AdContext::builder()
392 .with_semantic_extension_rules(tenferro_linalg::semantic_ad_rules().map_err(|source| {
393 CpuExecutionContextError::initialization("linalg AD rules", source)
394 })?)
395 .map_err(|source| CpuExecutionContextError::initialization("linalg AD rules", source))?
396 .build()
397 .map_err(|source| CpuExecutionContextError::initialization("AD context", source))?;
398 let runtime = EagerRuntime::with_cpu_backend_and_ad_context(backend, &ad_context)
399 .map_err(|source| CpuExecutionContextError::initialization("eager runtime", source))?;
400 let engine_id = tenferro_cpu::runtime_engine_id()
405 .map_err(|source| CpuExecutionContextError::initialization("CPU runtime engine", source))?;
406 let einsum_module = tenferro_einsum::extension_module::<CpuBackend>(engine_id.clone())
407 .map_err(|source| CpuExecutionContextError::initialization("einsum extension", source))?;
408 runtime
409 .install_extension_module(einsum_module)
410 .map_err(|source| CpuExecutionContextError::initialization("einsum runtime", source))?;
411 let linalg_module = tenferro_linalg::extension_module::<CpuBackend>(engine_id)
412 .map_err(|source| CpuExecutionContextError::initialization("linalg extension", source))?;
413 runtime
414 .install_extension_module(linalg_module)
415 .map_err(|source| CpuExecutionContextError::initialization("linalg runtime", source))?;
416 Ok(runtime)
417}
418
419#[cfg(feature = "global-defaults")]
420mod defaults {
421 use super::*;
422 use tenferro_cpu::CpuContext;
423
424 static DEFAULT_CONTEXT: OnceLock<Arc<CpuExecutionContext>> = OnceLock::new();
425
426 #[cfg(test)]
427 thread_local! {
428 static FORCE_EAGER_CONTEXT_FAILURE: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
429 }
430
431 #[cfg(test)]
432 static DEFAULT_CONTEXT_HITS: std::sync::atomic::AtomicUsize =
433 std::sync::atomic::AtomicUsize::new(0);
434
435 fn default_context() -> &'static Arc<CpuExecutionContext> {
436 DEFAULT_CONTEXT.get_or_init(|| {
437 #[cfg(test)]
438 DEFAULT_CONTEXT_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
439 Arc::new(CpuExecutionContext::from_backend(CpuBackend::from_context(
440 Arc::new(CpuContext::from_env()),
441 )))
442 })
443 }
444
445 #[derive(Debug, Clone, thiserror::Error)]
460 pub enum EagerContextError {
461 #[error("failed to register tenferro linalg AD rule: {source}")]
463 Registration {
464 #[source]
466 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
467 },
468 }
469
470 pub fn with_default_backend<R>(f: impl FnOnce(&mut CpuBackend) -> R) -> R {
472 default_context().with_backend(f)
473 }
474
475 pub(crate) fn with_default_session<R: Send>(
476 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
477 ) -> R {
478 default_context().with_session(f)
479 }
480
481 pub(crate) fn with_default_graph_runtime<R>(
482 f: impl FnOnce(&mut GraphCompiler, &Runtime, &mut CpuBackend) -> R,
483 ) -> anyhow::Result<R> {
484 default_context()
485 .with_graph_state(|compiler, runtime, backend| f(compiler, runtime, backend))
486 .map_err(anyhow::Error::new)
487 }
488
489 pub(crate) fn default_engine_buffer_pool_stats() -> anyhow::Result<BufferPoolStats> {
490 default_context()
491 .graph_buffer_pool_stats()
492 .map_err(anyhow::Error::new)
493 }
494
495 pub(crate) fn reset_default_engine_buffer_pool() -> anyhow::Result<()> {
496 default_context()
497 .reset_graph_buffer_pool()
498 .map_err(anyhow::Error::new)
499 }
500
501 pub(crate) fn reset_default_engine() -> anyhow::Result<()> {
502 default_context()
503 .reset_graph_runtime()
504 .map_err(anyhow::Error::new)
505 }
506
507 pub fn default_eager_ctx() -> Result<Arc<EagerRuntime>, EagerContextError> {
525 #[cfg(test)]
526 if FORCE_EAGER_CONTEXT_FAILURE.with(std::cell::Cell::get) {
527 return Err(EagerContextError::Registration {
528 source: Arc::new(std::io::Error::other(
529 "forced default eager context registration failure",
530 )),
531 });
532 }
533 default_context()
534 .eager_runtime()
535 .map_err(|source| EagerContextError::Registration {
536 source: Arc::new(source),
537 })
538 }
539
540 pub fn default_cpu_execution_context() -> Arc<CpuExecutionContext> {
556 Arc::clone(default_context())
557 }
558
559 #[cfg(test)]
560 pub(crate) fn default_context_hits() -> usize {
561 DEFAULT_CONTEXT_HITS.load(std::sync::atomic::Ordering::Relaxed)
562 }
563
564 #[cfg(test)]
565 pub(crate) fn with_forced_eager_context_failure<T>(f: impl FnOnce() -> T) -> T {
566 let previous = FORCE_EAGER_CONTEXT_FAILURE.with(|failure| failure.replace(true));
567 let result = f();
568 FORCE_EAGER_CONTEXT_FAILURE.with(|failure| failure.set(previous));
569 result
570 }
571}
572
573#[cfg(all(test, feature = "global-defaults"))]
574pub(crate) use defaults::with_forced_eager_context_failure;
575#[cfg(feature = "global-defaults")]
576pub use defaults::{
577 default_cpu_execution_context, default_eager_ctx, with_default_backend, EagerContextError,
578};
579#[cfg(feature = "global-defaults")]
580pub(crate) use defaults::{
581 default_engine_buffer_pool_stats, reset_default_engine, reset_default_engine_buffer_pool,
582 with_default_graph_runtime, with_default_session,
583};
584
585#[cfg(test)]
586mod tests {
587 use std::num::NonZeroUsize;
588 use std::sync::mpsc;
589 use std::time::Duration;
590
591 use super::*;
592 use tenferro::program::{CoreSemanticOp, ProgramInputSpec};
593 use tenferro::{DType, TensorSessionOpsExt, TraceContext};
594 use tenferro_ad::EagerTensor;
595 use tenferro_cpu::{CpuContext, ExternalCpuDomain};
596 use tenferro_tensor::CpuDomainId;
597
598 fn context() -> CpuExecutionContext {
599 CpuExecutionContext::from_backend(CpuBackend::with_threads(1).unwrap())
600 }
601
602 #[test]
603 fn explicit_session_runs_a_concrete_operation() {
604 let context = context();
605 let lhs = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
606 let rhs = Tensor::from_vec_col_major(vec![2, 1], vec![5.0_f64, 6.0]).unwrap();
607 let result = context
608 .with_session(|session| lhs.matmul(&rhs, session))
609 .unwrap();
610
611 assert_eq!(result.as_slice::<f64>().unwrap(), &[23.0, 34.0]);
612 }
613
614 #[test]
615 fn recursive_session_entry_fails_before_lock_and_restores_guard() {
616 let context = context();
617 let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
618 context.with_session(|_| context.with_session(|_| ()))
619 }))
620 .expect_err("recursive canonical session entry should panic");
621 let message = panic
622 .downcast_ref::<&str>()
623 .copied()
624 .or_else(|| panic.downcast_ref::<String>().map(String::as_str))
625 .expect("recursive entry panic should contain a string message");
626 assert_eq!(message, CANONICAL_SESSION_REENTRY_MESSAGE);
627
628 assert_eq!(context.with_session(|_| 7usize), 7);
629 }
630
631 #[test]
632 fn explicit_plain_graph_and_eager_paths_share_only_the_supplied_backend() {
633 let context = context();
634 assert!(format!("{context:?}").contains("graph_initialized: false"));
635 assert_eq!(context.with_backend(|backend| backend.num_threads()), 1);
636
637 let mut trace = TraceContext::new();
638 let input = trace
639 .input(ProgramInputSpec::new(DType::F64, [2_usize.into()]))
640 .unwrap();
641 let output = trace.add_op(CoreSemanticOp::Neg, &[input]).unwrap()[0];
642 let graph = trace.finish(&[output]).unwrap();
643 let compiled = context.compile_graph(&graph).unwrap();
644 let input = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0]).unwrap();
645 let output = context.run_graph(&compiled, &[&input]).unwrap();
646 assert_eq!(output[0].as_slice::<f64>().unwrap(), &[-1.0, 2.0]);
647 context.run_graph(&compiled, &[&input]).unwrap();
648 let cached = context.graph_cache_stats().unwrap().prepared_plans;
649 assert!(cached.entries > 0);
650 assert!(cached.hits > 0);
651 context.reset_graph_runtime().unwrap();
652 assert_eq!(
653 context.graph_cache_stats().unwrap().prepared_plans.entries,
654 0
655 );
656
657 let eager = context.eager_runtime().unwrap();
658 assert!(Arc::ptr_eq(&eager, &context.eager_runtime().unwrap()));
659 }
660
661 #[test]
662 fn separate_eager_contexts_reject_cross_context_operations() {
663 let first = context().eager_runtime().unwrap();
664 let second = context().eager_runtime().unwrap();
665 let a = EagerTensor::from_tensor_in(
666 Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
667 first,
668 )
669 .unwrap();
670 let b = EagerTensor::from_tensor_in(
671 Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
672 second,
673 )
674 .unwrap();
675 assert!(matches!(
676 a.add(&b),
677 Err(tenferro_ad::Error::ContextMismatch { .. })
678 ));
679 }
680
681 #[test]
682 fn caller_managed_backend_remains_caller_owned_after_context_drop() {
683 let executor = Arc::new(CpuContext::with_threads(1).unwrap());
684 let id = CpuDomainId::new(7);
685 let domain =
686 ExternalCpuDomain::new_caller_managed(id, executor.clone(), NonZeroUsize::MIN).unwrap();
687 let backend = CpuBackend::from_external_managed_domains(id, [domain]).unwrap();
688 let context = CpuExecutionContext::from_backend(backend);
689 assert_eq!(context.with_backend(|backend| backend.num_threads()), 1);
690 drop(context);
691 assert_eq!(executor.num_threads(), 1);
692 }
693
694 #[test]
695 fn independent_contexts_do_not_share_a_backend_mutex() {
696 let first = Arc::new(context());
697 let second = Arc::new(context());
698 let (entered_tx, entered_rx) = mpsc::channel();
699 let (release_tx, release_rx) = mpsc::channel();
700 let release_rx = Arc::new(Mutex::new(release_rx));
701 let handles = [first, second].map(|context| {
702 let entered_tx = entered_tx.clone();
703 let release_rx = Arc::clone(&release_rx);
704 std::thread::spawn(move || {
705 context.with_backend(|_| {
706 entered_tx.send(()).unwrap();
707 release_rx.lock().unwrap().recv().unwrap();
708 });
709 })
710 });
711 entered_rx.recv_timeout(Duration::from_secs(2)).unwrap();
712 entered_rx.recv_timeout(Duration::from_secs(2)).unwrap();
713 release_tx.send(()).unwrap();
714 release_tx.send(()).unwrap();
715 for handle in handles {
716 handle.join().unwrap();
717 }
718 }
719
720 #[cfg(feature = "global-defaults")]
721 #[test]
722 fn explicit_paths_do_not_initialize_the_default_context() {
723 let before = defaults::default_context_hits();
724 let context = context();
725 context.with_backend(|backend| assert_eq!(backend.num_threads(), 1));
726 context.eager_runtime().unwrap();
727 assert_eq!(defaults::default_context_hits(), before);
728 }
729}