Skip to main content

tensor4all_tensorbackend/
context.rs

1//! Explicit and optional process-global tenferro CPU execution contexts.
2
3use 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/// Caller-owned execution domain used by context-aware tensor algorithms.
12///
13/// Values are validated against the exact runtime represented by the selected
14/// context; no implicit host/device transfer is performed by this enum.
15#[derive(Clone, Debug)]
16pub enum ExecutionContext {
17    /// Host execution through one caller-owned CPU context.
18    Cpu(Arc<CpuExecutionContext>),
19    /// CUDA execution through one caller-owned CUDA context.
20    #[cfg(feature = "tenferro-cuda")]
21    Cuda(Arc<crate::cuda::CudaExecutionContext>),
22}
23
24impl ExecutionContext {
25    /// Check whether this is the process-global CPU context.
26    ///
27    /// Compatibility boundary: legacy CPU-global callers route through the
28    /// historical host code paths (bitwise-identical numerics) while explicit
29    /// contexts use the scoped primitives. Without the global-defaults
30    /// feature there is no global context, so this always reports false.
31    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/// Error returned by explicit CPU context graph or eager-runtime operations.
53///
54/// The original tenferro diagnostic is retained as the error source.
55///
56/// # Examples
57///
58/// ```
59/// use std::error::Error;
60/// use std::sync::Arc;
61/// use tensor4all_tensorbackend::CpuExecutionContextError;
62///
63/// let error = CpuExecutionContextError::Initialization {
64///     component: "graph runtime",
65///     source: Arc::new(std::io::Error::other("registration failed")),
66/// };
67/// assert!(error.source().is_some());
68/// ```
69#[derive(Debug, Clone, thiserror::Error)]
70pub enum CpuExecutionContextError {
71    /// A context-owned graph or eager runtime could not be initialized.
72    #[error("failed to initialize {component}: {source}")]
73    Initialization {
74        /// Context component being initialized.
75        component: &'static str,
76        /// Original tenferro diagnostic.
77        #[source]
78        source: Arc<dyn std::error::Error + Send + Sync + 'static>,
79    },
80    /// Graph compilation or execution failed.
81    #[error("CPU graph {operation} failed: {source}")]
82    Graph {
83        /// Graph operation that failed.
84        operation: &'static str,
85        /// Original tenferro diagnostic.
86        #[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
150/// Caller-owned CPU execution domain for plain, graph, and eager-AD work.
151///
152/// The supplied backend is the only source of CPU execution resources. Backend
153/// clones preserve its runtime identity; this constructor never uses
154/// `CpuBackend::new`, `CpuContext::from_env`, or a process-global fallback.
155/// Graph preparation caches and the eager runtime are owned by this context and
156/// are released when it is dropped.
157///
158/// # Examples
159///
160/// ```
161/// use tensor4all_tensorbackend::CpuExecutionContext;
162/// use tenferro_cpu::CpuBackend;
163///
164/// let context = CpuExecutionContext::from_backend(CpuBackend::with_threads(1)?);
165/// let threads = context.with_backend(|backend| backend.num_threads());
166/// assert_eq!(threads, 1);
167/// # Ok::<(), Box<dyn std::error::Error>>(())
168/// ```
169pub 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    /// Create an execution context from a caller-selected CPU backend.
186    ///
187    /// Runtime construction is lazy, so creating a context cannot fail and does
188    /// not allocate another executor or consult environment configuration.
189    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    /// Run a plain tensor operation with this context's backend.
198    ///
199    /// The closure runs while the context-local backend lock is held. A poisoned
200    /// lock is recovered because tenferro validates every new backend session.
201    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    /// Compile a backend-neutral traced graph using this context's compiler cache.
261    ///
262    /// # Errors
263    ///
264    /// Returns [`CpuExecutionContextError`] when graph-runtime initialization or
265    /// graph compilation fails.
266    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    /// Execute a compiled graph in this context's runtime and prepared-plan cache.
275    ///
276    /// `CompiledGraph` is backend-neutral. Backend-prepared executables and
277    /// workspaces never leave this context-owned runtime.
278    ///
279    /// # Errors
280    ///
281    /// Returns [`CpuExecutionContextError`] when runtime initialization,
282    /// preparation, or execution fails.
283    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    /// Return this context's eager reverse-AD runtime.
293    ///
294    /// Repeated calls return the same runtime and therefore the same eager
295    /// compilation cache.
296    ///
297    /// # Errors
298    ///
299    /// Returns [`CpuExecutionContextError`] when linalg AD-rule or eager-runtime
300    /// registration fails.
301    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    /// Return statistics for this context's runtime-owned graph caches.
310    ///
311    /// # Errors
312    ///
313    /// Returns [`CpuExecutionContextError`] when graph initialization or the
314    /// cache statistics query fails.
315    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    /// Return retained-buffer statistics for this context's graph backend.
323    ///
324    /// # Errors
325    ///
326    /// Returns [`CpuExecutionContextError`] when graph initialization or the
327    /// backend statistics query fails.
328    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    /// Release retained buffers owned by this context's graph backend.
334    ///
335    /// # Errors
336    ///
337    /// Returns [`CpuExecutionContextError`] when graph initialization or reset
338    /// fails.
339    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    /// Recreate this context's graph runtime and release its prepared caches.
345    ///
346    /// # Errors
347    ///
348    /// Returns [`CpuExecutionContextError`] when runtime reconstruction or
349    /// buffer release fails.
350    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    // [AI Supplied] Install the built-in extension modules before publishing
401    // the shared context. Lazy first use reconfigures the runtime and advances
402    // its epoch; doing that after tensors have prepared AD derivatives can
403    // invalidate those prepared programs under parallel first use.
404    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    /// Error returned when the process-global eager AD runtime cannot be initialized.
446    ///
447    /// # Examples
448    ///
449    /// ```
450    /// use std::error::Error;
451    /// use std::sync::Arc;
452    /// use tensor4all_tensorbackend::EagerContextError;
453    ///
454    /// let error = EagerContextError::Registration {
455    ///     source: Arc::new(std::io::Error::other("registration failed")),
456    /// };
457    /// assert!(error.source().is_some());
458    /// ```
459    #[derive(Debug, Clone, thiserror::Error)]
460    pub enum EagerContextError {
461        /// The tenferro linalg AD extension rule could not be registered.
462        #[error("failed to register tenferro linalg AD rule: {source}")]
463        Registration {
464            /// Original diagnostic returned by tenferro.
465            #[source]
466            source: Arc<dyn std::error::Error + Send + Sync + 'static>,
467        },
468    }
469
470    /// Run a closure against the optional process-global CPU backend.
471    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    /// Return the optional process-global eager context used by convenience APIs.
508    ///
509    /// # Errors
510    ///
511    /// Returns [`EagerContextError::Registration`] when eager runtime
512    /// initialization fails.
513    ///
514    /// # Examples
515    ///
516    /// ```
517    /// use std::sync::Arc;
518    /// use tensor4all_tensorbackend::default_eager_ctx;
519    ///
520    /// let first = default_eager_ctx().unwrap();
521    /// let second = default_eager_ctx().unwrap();
522    /// assert!(Arc::ptr_eq(&first, &second));
523    /// ```
524    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    /// Borrow the process-global CPU execution context.
541    ///
542    /// Compatibility entry for CPU-global convenience APIs (e.g. the legacy
543    /// context-free SRC entry): host tensors constructed through the global
544    /// default belong to this exact context, so they validate against it.
545    /// New code must take a caller-owned context instead of consulting this.
546    ///
547    /// # Examples
548    ///
549    /// ```
550    /// use tensor4all_tensorbackend::{default_cpu_execution_context, ExecutionContext};
551    ///
552    /// let context = ExecutionContext::Cpu(default_cpu_execution_context());
553    /// assert!(matches!(context, ExecutionContext::Cpu(_)));
554    /// ```
555    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}