tenferro_runtime/
extension_runtime.rs

1//! Backend-parametric runtime dispatch for extension ops.
2//!
3//! This module is intentionally generic: extension crates can register an
4//! executor for a family and keep runtime cache state outside both the
5//! semantic [`ExtensionOp`] payload and the
6//! tensor backend implementation.
7
8use std::collections::HashMap;
9use std::fmt::{self, Debug};
10use std::marker::PhantomData;
11use std::sync::Arc;
12
13use tenferro_ops::ext_op::ExtensionOp;
14use tenferro_tensor::{CacheStats, Tensor, TensorBackend, TensorRead};
15
16use crate::extension_cache::{ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore};
17
18/// Errors returned by backend-parametric extension runtime registries.
19#[derive(Debug, thiserror::Error)]
20pub enum ExtensionRuntimeRegistryError {
21    /// The `family_id` does not match the namespaced format
22    /// `"<crate-name>.<op-name>.v<major>"`.
23    #[error("family_id {family_id:?} does not match the namespaced format")]
24    MalformedFamilyId { family_id: &'static str },
25    /// A registry lock was poisoned by a panic in another thread.
26    #[error("{name} poisoned")]
27    PoisonedLock { name: &'static str },
28}
29
30/// Backend and cache state passed to one extension execution.
31pub struct ExtensionExecutionContext<'a, B: TensorBackend> {
32    backend: &'a mut B,
33    caches: &'a mut ExtensionCacheStore,
34}
35
36impl<B: TensorBackend> fmt::Debug for ExtensionExecutionContext<'_, B> {
37    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38        f.debug_struct("ExtensionExecutionContext")
39            .field("backend_type", &std::any::type_name::<B>())
40            .field("caches", &self.caches)
41            .finish_non_exhaustive()
42    }
43}
44
45impl<'a, B: TensorBackend> ExtensionExecutionContext<'a, B> {
46    /// Build a context from externally-owned backend and cache state.
47    pub fn new(backend: &'a mut B, caches: &'a mut ExtensionCacheStore) -> Self {
48        Self { backend, caches }
49    }
50
51    /// Borrow the backend for non-mutating inspection.
52    pub fn backend(&self) -> &B {
53        self.backend
54    }
55
56    /// Borrow the backend mutably for extension execution.
57    pub fn backend_mut(&mut self) -> &mut B {
58        self.backend
59    }
60
61    /// Borrow the extension runtime cache store.
62    pub fn caches(&self) -> &ExtensionCacheStore {
63        self.caches
64    }
65
66    /// Borrow the extension runtime cache store mutably.
67    pub fn caches_mut(&mut self) -> &mut ExtensionCacheStore {
68        self.caches
69    }
70
71    /// Execute a core-only execution program one instruction at a time.
72    ///
73    /// This is for extension runtimes that lower their own operation into a
74    /// temporary `ExecProgram` containing only core tensor ops. Nested
75    /// `ExecOp::Extension` instructions are rejected so extension dispatch
76    /// cannot bypass the owning runtime registry.
77    ///
78    /// # Examples
79    ///
80    /// ```
81    /// use tenferro_cpu::CpuBackend;
82    /// use tenferro_ops::dim_expr::DimExpr;
83    /// use tenferro_runtime::extension::{ExecInstruction, ExecOp, ExecProgram};
84    /// use tenferro_runtime::{DType, ExtensionCacheStore, ExtensionExecutionContext, Tensor};
85    ///
86    /// let program = ExecProgram {
87    ///     instructions: vec![ExecInstruction {
88    ///         op: ExecOp::Add,
89    ///         input_slots: vec![0, 1],
90    ///         output_slots: vec![2],
91    ///         dtype: DType::F64,
92    ///         output_shapes: vec![vec![]].into(),
93    ///         output_extents: vec![vec![]].into(),
94    ///         last_use: vec![true, true],
95    ///     }],
96    ///     input_slots: vec![0, 1],
97    ///     output_slots: vec![2],
98    ///     n_slots: 3,
99    /// };
100    /// let lhs = Tensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
101    /// let rhs = Tensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap();
102    ///
103    /// let mut backend = CpuBackend::new();
104    /// let mut caches = ExtensionCacheStore::new();
105    /// let mut ctx = ExtensionExecutionContext::new(&mut backend, &mut caches);
106    /// let outputs = ctx
107    ///     .execute_core_exec_program_unsegmented(&program, vec![lhs, rhs])
108    ///     .unwrap();
109    /// assert_eq!(outputs[0].as_slice::<f64>().unwrap(), &[3.0]);
110    /// ```
111    pub fn execute_core_exec_program_unsegmented(
112        &mut self,
113        program: &crate::extension::ExecProgram,
114        inputs: Vec<Tensor>,
115    ) -> crate::error::Result<Vec<Tensor>>
116    where
117        B: 'static,
118    {
119        crate::exec::ensure_core_exec_program(
120            program,
121            "ExtensionExecutionContext::execute_core_exec_program_unsegmented",
122        )?;
123        crate::exec::eval_exec_ir_unsegmented_with_cache(self.backend, program, inputs)
124    }
125
126    /// Borrow backend and extension cache store as disjoint mutable parts.
127    pub fn parts_mut(&mut self) -> (&mut B, &mut ExtensionCacheStore) {
128        (self.backend, self.caches)
129    }
130}
131
132/// A backend-specific runtime executor for one extension family.
133pub trait ExtensionRuntime<B: TensorBackend + 'static>: Debug + Send + Sync + 'static {
134    /// Extension family handled by this executor.
135    fn family_id(&self) -> &'static str;
136
137    /// Execute the extension op with backend and cache state supplied by core.
138    fn execute(
139        &self,
140        op: &dyn ExtensionOp,
141        inputs: &[&Tensor],
142        ctx: &mut ExtensionExecutionContext<'_, B>,
143    ) -> tenferro_tensor::Result<Vec<Tensor>>;
144
145    /// Execute the extension op on borrowed tensor reads.
146    ///
147    /// Implementations that need compact tensors must materialize inputs here
148    /// explicitly. Keeping this method required prevents implicit read-path
149    /// fallbacks from hiding backend or view handling bugs.
150    fn execute_reads(
151        &self,
152        op: &dyn ExtensionOp,
153        inputs: &[TensorRead<'_>],
154        ctx: &mut ExtensionExecutionContext<'_, B>,
155    ) -> tenferro_tensor::Result<Vec<Tensor>>;
156}
157
158/// Runtime adapter that delegates execution to an extension op's optional
159/// host/reference implementation.
160///
161/// Register one adapter per extension family. Backend-specific runtimes should
162/// implement [`ExtensionRuntime`] directly instead of using this adapter.
163///
164/// # Examples
165///
166/// ```rust
167/// use tenferro_cpu::CpuBackend;
168/// use tenferro_runtime::{ExtensionRuntime, HostReferenceRuntime};
169///
170/// let runtime = HostReferenceRuntime::<CpuBackend>::new("example.identity.v1");
171/// assert_eq!(runtime.family_id(), "example.identity.v1");
172/// ```
173#[derive(Clone, Copy)]
174pub struct HostReferenceRuntime<B: TensorBackend + 'static> {
175    family_id: &'static str,
176    _backend: PhantomData<fn() -> B>,
177}
178
179impl<B: TensorBackend + 'static> HostReferenceRuntime<B> {
180    /// Create a host-reference runtime for one extension family.
181    pub fn new(family_id: &'static str) -> Self {
182        Self {
183            family_id,
184            _backend: PhantomData,
185        }
186    }
187}
188
189impl<B: TensorBackend + 'static> Debug for HostReferenceRuntime<B> {
190    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
191        f.debug_struct("HostReferenceRuntime")
192            .field("backend_type", &std::any::type_name::<B>())
193            .field("family_id", &self.family_id)
194            .finish()
195    }
196}
197
198impl<B: TensorBackend + 'static> ExtensionRuntime<B> for HostReferenceRuntime<B> {
199    fn family_id(&self) -> &'static str {
200        self.family_id
201    }
202
203    fn execute(
204        &self,
205        op: &dyn ExtensionOp,
206        inputs: &[&Tensor],
207        _ctx: &mut ExtensionExecutionContext<'_, B>,
208    ) -> tenferro_tensor::Result<Vec<Tensor>> {
209        let host = op
210            .host_reference()
211            .ok_or(tenferro_tensor::Error::NoHostReference {
212                family_id: op.family_id(),
213            })?;
214        host.execute(inputs)
215    }
216
217    fn execute_reads(
218        &self,
219        op: &dyn ExtensionOp,
220        inputs: &[TensorRead<'_>],
221        ctx: &mut ExtensionExecutionContext<'_, B>,
222    ) -> tenferro_tensor::Result<Vec<Tensor>> {
223        let materialized_inputs: Vec<Tensor> = inputs
224            .iter()
225            .map(TensorRead::to_tensor)
226            .collect::<tenferro_tensor::Result<_>>()?;
227        let input_refs: Vec<&Tensor> = materialized_inputs.iter().collect();
228        self.execute(op, &input_refs, ctx)
229    }
230}
231
232fn validate_runtime_output_count(
233    op: &dyn ExtensionOp,
234    outputs: Vec<Tensor>,
235) -> tenferro_tensor::Result<Vec<Tensor>> {
236    let expected = op.output_count();
237    if outputs.len() != expected {
238        return Err(tenferro_tensor::Error::InvalidConfig {
239            op: "extension",
240            message: format!(
241                "family_id {:?}: runtime returned {} outputs but op declared {} outputs",
242                op.family_id(),
243                outputs.len(),
244                expected
245            ),
246        });
247    }
248    Ok(outputs)
249}
250
251fn validate_runtime_input_count(
252    op: &dyn ExtensionOp,
253    actual: usize,
254) -> tenferro_tensor::Result<()> {
255    let expected = op.input_count();
256    if actual != expected {
257        return Err(tenferro_tensor::Error::InvalidConfig {
258            op: "extension",
259            message: format!(
260                "family_id {:?}: op expects {} inputs, got {}",
261                op.family_id(),
262                expected,
263                actual
264            ),
265        });
266    }
267    Ok(())
268}
269
270/// Registry of backend-specific extension runtime executors.
271pub struct ExtensionRegistry<B: TensorBackend + 'static> {
272    executors: HashMap<&'static str, Arc<dyn ExtensionRuntime<B>>>,
273}
274
275impl<B: TensorBackend + 'static> fmt::Debug for ExtensionRegistry<B> {
276    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
277        let mut families = self.executors.keys().copied().collect::<Vec<_>>();
278        families.sort_unstable();
279        f.debug_struct("ExtensionRegistry")
280            .field("backend_type", &std::any::type_name::<B>())
281            .field("len", &self.executors.len())
282            .field("families", &families)
283            .finish_non_exhaustive()
284    }
285}
286
287impl<B: TensorBackend + 'static> ExtensionRegistry<B> {
288    /// Create an empty extension runtime registry.
289    ///
290    /// # Examples
291    ///
292    /// ```
293    /// use tenferro_runtime::ExtensionRegistry;
294    /// use tenferro_cpu::CpuBackend;
295    ///
296    /// let registry = ExtensionRegistry::<CpuBackend>::new();
297    /// assert!(!registry.contains("example.identity.v1"));
298    /// ```
299    pub fn new() -> Self {
300        Self {
301            executors: HashMap::new(),
302        }
303    }
304
305    /// Register one runtime executor.
306    ///
307    /// Registration is idempotent by family id: registering the same extension
308    /// family more than once succeeds and keeps the first runtime. This lets
309    /// extension crates register their own dependency extensions defensively.
310    pub fn register(
311        &mut self,
312        executor: Arc<dyn ExtensionRuntime<B>>,
313    ) -> Result<(), ExtensionRuntimeRegistryError> {
314        let family_id = executor.family_id();
315        if !is_valid_family_id(family_id) {
316            return Err(ExtensionRuntimeRegistryError::MalformedFamilyId { family_id });
317        }
318        if self.executors.contains_key(family_id) {
319            return Ok(());
320        }
321        self.executors.insert(family_id, executor);
322        Ok(())
323    }
324
325    /// Look up an executor by extension family id.
326    pub fn get(&self, family_id: &str) -> Option<Arc<dyn ExtensionRuntime<B>>> {
327        self.executors.get(family_id).cloned()
328    }
329
330    /// Return whether an executor is registered for `family_id`.
331    pub fn contains(&self, family_id: &str) -> bool {
332        self.executors.contains_key(family_id)
333    }
334
335    /// Number of registered runtime executors.
336    pub fn len(&self) -> usize {
337        self.executors.len()
338    }
339
340    /// Return whether no runtime executors are registered.
341    pub fn is_empty(&self) -> bool {
342        self.executors.is_empty()
343    }
344}
345
346impl<B: TensorBackend + 'static> Default for ExtensionRegistry<B> {
347    fn default() -> Self {
348        Self::new()
349    }
350}
351
352/// Runtime owner for backend-specific extension dispatch and caches.
353pub struct ExtensionExecutor<B: TensorBackend + 'static> {
354    registry: ExtensionRegistry<B>,
355    caches: ExtensionCacheStore,
356    _backend: PhantomData<fn() -> B>,
357}
358
359impl<B: TensorBackend + 'static> fmt::Debug for ExtensionExecutor<B> {
360    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
361        f.debug_struct("ExtensionExecutor")
362            .field("backend_type", &std::any::type_name::<B>())
363            .field("registry", &self.registry)
364            .field("caches", &self.caches)
365            .finish_non_exhaustive()
366    }
367}
368
369impl<B: TensorBackend + 'static> ExtensionExecutor<B> {
370    /// Create an executor with an empty registry and default cache limits.
371    ///
372    /// # Examples
373    ///
374    /// ```
375    /// use tenferro_runtime::ExtensionExecutor;
376    /// use tenferro_cpu::CpuBackend;
377    ///
378    /// let executor = ExtensionExecutor::<CpuBackend>::new();
379    /// assert_eq!(executor.cache_stats().entries, 0);
380    /// ```
381    pub fn new() -> Self {
382        Self {
383            registry: ExtensionRegistry::new(),
384            caches: ExtensionCacheStore::new(),
385            _backend: PhantomData,
386        }
387    }
388
389    /// Create an executor from explicit registry and cache store.
390    pub fn with_parts(registry: ExtensionRegistry<B>, caches: ExtensionCacheStore) -> Self {
391        Self {
392            registry,
393            caches,
394            _backend: PhantomData,
395        }
396    }
397
398    /// Borrow the runtime executor registry.
399    pub fn registry(&self) -> &ExtensionRegistry<B> {
400        &self.registry
401    }
402
403    /// Borrow the runtime executor registry mutably.
404    pub fn registry_mut(&mut self) -> &mut ExtensionRegistry<B> {
405        &mut self.registry
406    }
407
408    /// Borrow the extension cache store.
409    pub fn caches(&self) -> &ExtensionCacheStore {
410        &self.caches
411    }
412
413    /// Borrow the extension cache store mutably.
414    pub fn caches_mut(&mut self) -> &mut ExtensionCacheStore {
415        &mut self.caches
416    }
417
418    /// Execute an extension using a registered runtime executor.
419    pub fn execute(
420        &mut self,
421        backend: &mut B,
422        op: &dyn ExtensionOp,
423        inputs: &[&Tensor],
424    ) -> tenferro_tensor::Result<Vec<Tensor>> {
425        validate_runtime_input_count(op, inputs.len())?;
426        let Some(executor) = self.registry.get(op.family_id()) else {
427            return Err(tenferro_tensor::Error::InvalidConfig {
428                op: "extension",
429                message: format!(
430                    "missing runtime for family_id {:?}; register the extension on this runtime owner, for example `executor.register_extension(<extension_crate>::register_runtime)` or `eager_runtime.register_extension(<extension_crate>::register_runtime)`",
431                    op.family_id()
432                ),
433            });
434        };
435        let mut ctx = ExtensionExecutionContext::new(backend, &mut self.caches);
436        validate_runtime_output_count(op, executor.execute(op, inputs, &mut ctx)?)
437    }
438
439    /// Execute an extension using borrowed tensor reads.
440    ///
441    /// # Examples
442    ///
443    /// ```
444    /// use std::any::Any;
445    /// use std::sync::Arc;
446    ///
447    /// use tenferro_cpu::CpuBackend;
448    /// use tenferro_ops::{ext_op::{ExtensionOp, HostReference}, SymDim};
449    /// use tenferro_runtime::{DType, ExtensionExecutor, HostReferenceRuntime, Tensor};
450    /// use tenferro_tensor::TensorRead;
451    ///
452    /// #[derive(Clone, Debug)]
453    /// struct IdentityOp;
454    ///
455    /// impl ExtensionOp for IdentityOp {
456    ///     fn family_id(&self) -> &'static str {
457    ///         "example.identity.v1"
458    ///     }
459    ///
460    ///     fn payload_hash(&self, _hasher: &mut dyn std::hash::Hasher) {}
461    ///
462    ///     fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
463    ///         other.as_any().is::<IdentityOp>()
464    ///     }
465    ///
466    ///     fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
467    ///         Arc::new(self.clone())
468    ///     }
469    ///
470    ///     fn as_any(&self) -> &dyn Any {
471    ///         self
472    ///     }
473    ///
474    ///     fn input_count(&self) -> usize {
475    ///         1
476    ///     }
477    ///
478    ///     fn output_count(&self) -> usize {
479    ///         1
480    ///     }
481    ///
482    ///     fn infer_output_meta(
483    ///         &self,
484    ///         input_dtypes: &[DType],
485    ///         input_shapes: &[&[SymDim]],
486    ///     ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
487    ///         Ok(vec![(input_dtypes[0], input_shapes[0].to_vec())])
488    ///     }
489    ///
490    ///     fn host_reference(&self) -> Option<&dyn HostReference> {
491    ///         Some(self)
492    ///     }
493    /// }
494    ///
495    /// impl HostReference for IdentityOp {
496    ///     fn execute(&self, inputs: &[&Tensor]) -> tenferro_tensor::Result<Vec<Tensor>> {
497    ///         Ok(vec![inputs[0].clone()])
498    ///     }
499    /// }
500    ///
501    /// let mut executor = ExtensionExecutor::<CpuBackend>::new();
502    /// executor
503    ///     .registry_mut()
504    ///     .register(Arc::new(HostReferenceRuntime::<CpuBackend>::new(
505    ///         "example.identity.v1",
506    ///     )))?;
507    /// let input = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
508    /// let read = TensorRead::from_tensor(&input);
509    /// let mut backend = CpuBackend::new();
510    ///
511    /// let outputs = executor.execute_reads(&mut backend, &IdentityOp, &[read])?;
512    ///
513    /// assert_eq!(outputs[0].as_slice::<f64>().unwrap(), &[1.0, 2.0]);
514    /// # Ok::<(), Box<dyn std::error::Error>>(())
515    /// ```
516    pub fn execute_reads(
517        &mut self,
518        backend: &mut B,
519        op: &dyn ExtensionOp,
520        inputs: &[TensorRead<'_>],
521    ) -> tenferro_tensor::Result<Vec<Tensor>> {
522        validate_runtime_input_count(op, inputs.len())?;
523        let Some(executor) = self.registry.get(op.family_id()) else {
524            return Err(tenferro_tensor::Error::InvalidConfig {
525                op: "extension",
526                message: format!(
527                    "missing runtime for family_id {:?}; register the extension on this runtime owner, for example `executor.register_extension(<extension_crate>::register_runtime)` or `eager_runtime.register_extension(<extension_crate>::register_runtime)`",
528                    op.family_id()
529                ),
530            });
531        };
532        let mut ctx = ExtensionExecutionContext::new(backend, &mut self.caches);
533        validate_runtime_output_count(op, executor.execute_reads(op, inputs, &mut ctx)?)
534    }
535
536    /// Clear every runtime extension cache entry.
537    pub fn clear_caches(&mut self) {
538        self.caches.clear();
539    }
540
541    /// Return extension cache stats for all entries.
542    pub fn cache_stats(&self) -> CacheStats {
543        self.caches.stats(ExtensionCacheSelector::All)
544    }
545
546    /// Return the extension cache retention limits.
547    pub fn cache_limits(&self) -> ExtensionCacheLimits {
548        self.caches.limits()
549    }
550
551    /// Replace extension cache retention limits.
552    pub fn set_cache_limits(&mut self, limits: ExtensionCacheLimits) {
553        self.caches.set_limits(limits);
554    }
555}
556
557impl<B: TensorBackend + 'static> Default for ExtensionExecutor<B> {
558    fn default() -> Self {
559        Self::new()
560    }
561}
562
563#[cfg(test)]
564mod tests;
565
566fn is_valid_family_id(family_id: &str) -> bool {
567    let mut parts = family_id.rsplitn(2, '.');
568    let Some(version_part) = parts.next() else {
569        return false;
570    };
571    let Some(prefix) = parts.next() else {
572        return false;
573    };
574    if !version_part.starts_with('v') {
575        return false;
576    }
577    let digits = &version_part[1..];
578    if digits.is_empty() || !digits.chars().all(|c| c.is_ascii_digit()) {
579        return false;
580    }
581    let Some((crate_name, op_name)) = prefix.split_once('.') else {
582        return false;
583    };
584    if crate_name.is_empty() || op_name.is_empty() {
585        return false;
586    }
587    let any_invalid = |s: &str| s.chars().any(|c| c.is_whitespace() || !c.is_ascii());
588    !any_invalid(crate_name) && !any_invalid(op_name)
589}
tenferro_runtime/extension_runtime.rs

tenferro_runtime/
extension_runtime.rs
