Skip to main content

tenferro_cpu/backend/
execution_scope.rs

1//! Callback-lifetime admission for sequential high-level CPU operations.
2
3use std::cell::RefCell;
4use std::marker::PhantomData;
5use std::rc::Rc;
6use std::sync::Arc;
7
8use super::{CpuBackend, CpuRuntimeIdentity};
9use crate::arbiter::{has_active_execution, inherited_or_new_execution_owner, ResourcePermit};
10use crate::engine::CpuEngine;
11use crate::provider::CpuOperationEntry;
12use crate::resource_domain::CpuResourceDomain;
13use crate::CpuDomainOwnership;
14
15struct Scope {
16    identity: CpuRuntimeIdentity,
17    engine: Arc<CpuEngine>,
18    permit: Arc<ResourcePermit>,
19    operation_active: bool,
20}
21
22thread_local! {
23    static SCOPE: RefCell<Option<Scope>> = const { RefCell::new(None) };
24}
25
26struct ScopeGuard;
27
28impl Drop for ScopeGuard {
29    fn drop(&mut self) {
30        SCOPE.with(|slot| slot.borrow_mut().take());
31    }
32}
33
34// The operation loan belongs to this callback thread, never a child worker.
35pub(super) struct OperationGuard(PhantomData<Rc<()>>);
36
37impl Drop for OperationGuard {
38    fn drop(&mut self) {
39        SCOPE.with(|slot| {
40            if let Some(scope) = slot.borrow_mut().as_mut() {
41                scope.operation_active = false;
42            }
43        });
44    }
45}
46
47pub(super) enum ExecutionAdmission {
48    Standalone(ResourcePermit),
49    Shared(Arc<ResourcePermit>, OperationGuard),
50}
51
52impl ExecutionAdmission {
53    pub(super) fn permit(&self) -> &ResourcePermit {
54        match self {
55            Self::Standalone(permit) => permit,
56            Self::Shared(permit, _) => permit,
57        }
58    }
59}
60
61pub(crate) fn is_entered(domain: &CpuResourceDomain, permit: &ResourcePermit) -> bool {
62    SCOPE.with(|slot| {
63        slot.borrow().as_ref().is_some_and(|scope| {
64            scope.operation_active
65                && std::ptr::eq(scope.engine.domain(), domain)
66                && std::ptr::eq(scope.permit.as_ref(), permit)
67        })
68    })
69}
70
71impl CpuBackend {
72    /// Run sequential high-level CPU work in one entered execution scope.
73    ///
74    /// Clones of this immutable backend witness may execute ordinary tensor,
75    /// eager/AD and prepared trace operations in the callback without installing
76    /// the executor again. Construct the eager/traced runtime from such a clone.
77    /// Each operation still owns its usual exclusive buffer/cache borrow. The
78    /// scope holds the resource permit, including BLAS provider exclusion, until
79    /// return or unwind. Only Tenferro-managed CPU executors are supported.
80    ///
81    /// Enter the scope and prepare inputs before starting a steady-state timer.
82    /// This does not remove intrinsic output allocation or operation dispatch.
83    ///
84    /// # Examples
85    ///
86    /// ```
87    /// use tenferro_cpu::CpuBackend;
88    /// use tenferro_tensor::{Tensor, TensorElementwise};
89    ///
90    /// let owner = CpuBackend::with_threads(1)?;
91    /// let mut operations = owner.clone();
92    /// let x = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 3.0])?;
93    /// let y = owner.with_execution_scope(|| {
94    ///     let y = operations.add(&x, &x)?;
95    ///     operations.add(&y, &x)
96    /// })??;
97    /// assert_eq!(y.as_slice::<f64>()?, &[3.0, 9.0]);
98    /// # Ok::<(), Box<dyn std::error::Error>>(())
99    /// ```
100    ///
101    /// # Errors
102    ///
103    /// Returns [`crate::Error::RuntimeState`] if a scope or CPU execution is
104    /// already active, or [`crate::Error::Unsupported`] for an externally managed
105    /// executor. Executor admission errors retain their typed source in
106    /// [`crate::Error::BackendSource`]. A callback's return value, including its
107    /// own error result, is returned unchanged inside this method's result.
108    ///
109    /// # Panics
110    ///
111    /// Existing infallible backend-session APIs still panic on invalid nested
112    /// entry or a different backend witness. Do not enter backend operations from
113    /// inside an active borrowed session/provider operation or from other workers.
114    /// A panic in the callback propagates after releasing the scope and permit.
115    pub fn with_execution_scope<R: Send>(
116        &self,
117        operation: impl FnOnce() -> R + Send,
118    ) -> crate::Result<R> {
119        const OP: &str = "CpuBackend::with_execution_scope";
120        if has_active_execution() || SCOPE.with(|slot| slot.borrow().is_some()) {
121            return Err(crate::Error::runtime_state(
122                OP,
123                "CPU execution is already active; open the shared scope outside active scopes and backend sessions",
124            ));
125        }
126        if self.engine.domain().ownership() != CpuDomainOwnership::Managed {
127            return Err(crate::Error::unsupported(
128                OP,
129                "shared execution scopes require a Tenferro-managed CPU domain; use ordinary operation entry for external domains",
130            ));
131        }
132        let owner = inherited_or_new_execution_owner();
133        let permit = Arc::new(self.acquire_execution_permit(owner));
134        let entry = CpuOperationEntry::new(self.engine.domain(), &permit);
135        entry
136            .enter(entry.preferred_engine_mode(), |_| {
137                SCOPE.with(|slot| {
138                    *slot.borrow_mut() = Some(Scope {
139                        identity: self.runtime_identity.clone(),
140                        engine: Arc::clone(&self.engine),
141                        permit: Arc::clone(&permit),
142                        operation_active: false,
143                    });
144                });
145                let _guard = ScopeGuard;
146                operation()
147            })
148            .map_err(|error| crate::Error::backend_source(OP, error))
149    }
150
151    pub(super) fn execution_admission(&self) -> crate::Result<ExecutionAdmission> {
152        let shared = SCOPE.with(|slot| {
153            let mut slot = slot.borrow_mut();
154            let Some(scope) = slot.as_mut() else {
155                return Ok(None);
156            };
157            if scope.operation_active {
158                // Preserve the original backend/session nested-entry guard.
159                return Ok(None);
160            }
161            if scope.identity != self.runtime_identity {
162                return Err(crate::Error::runtime_state(
163                    "CPU execution scope",
164                    "operation backend does not match the scope; use a clone of its backend witness",
165                ));
166            }
167            scope.operation_active = true;
168            Ok(Some(ExecutionAdmission::Shared(
169                Arc::clone(&scope.permit),
170                OperationGuard(PhantomData),
171            )))
172        })?;
173        if let Some(shared) = shared {
174            return Ok(shared);
175        }
176        let owner = inherited_or_new_execution_owner();
177        Ok(ExecutionAdmission::Standalone(
178            self.acquire_execution_permit(owner),
179        ))
180    }
181
182    pub(super) fn infallible_execution_admission(&self) -> ExecutionAdmission {
183        // INVARIANT: BackendSessionHost and install have existing infallible
184        // callback contracts; invalid scope entry retains their panic boundary.
185        self.execution_admission()
186            .unwrap_or_else(|error| panic!("{error}"))
187    }
188}