tenferro_runtime/
extension.rs

1//! Public surface for out-of-tree extension primitives.
2//!
3//! This module exposes the Stage 6 `ExtensionOp` mechanism through the
4//! runtime crate. External crates implement
5//! [`tenferro_ops::ext_op::ExtensionOp`] and build traced graphs containing
6//! the extension via [`apply`].
7//!
8//! See `docs/spec/extension-op.md` for the normative contract.
9//!
10//! # Examples
11//!
12//! ```rust
13//! use tenferro_runtime::extension::{apply, ExtensionOp};
14//!
15//! // Construct an `Arc<dyn ExtensionOp>` and call `apply(op, &[input])`
16//! // to lower it into a `TracedTensor`.
17//! ```
18
19use std::sync::Arc;
20
21use computegraph::graph::GraphBuilder;
22use computegraph::types::{OperationRole, ValueRef};
23use tenferro_ops::std_tensor_op::StdTensorOp;
24use tenferro_ops::SymDim;
25use tenferro_tensor::{Tensor, TensorBackend};
26
27use crate::checkpoint::CheckpointNode;
28use crate::error::{Error, Result};
29use crate::metadata::{register_scoped_graph_metadata, MetadataScopeChain};
30use crate::traced::{merge_traced_inputs_map, next_traced_id, TracedTensor};
31
32pub use crate::compiler::CompilerOptions;
33#[doc(hidden)]
34pub use crate::compiler::{compile_std_to_exec, compile_std_to_exec_with_options};
35#[doc(hidden)]
36pub use crate::exec::{ExecInstruction, ExecOp, ExecOutputExtents, ExecOutputShapes, ExecProgram};
37#[doc(hidden)]
38pub use crate::shape_infer::{
39    infer_output_dtype, infer_output_extents, infer_output_shapes, promote_dtype,
40    promote_dtype_div_like, promote_dtype_for_binary_op, promote_dtypes,
41};
42pub use tenferro_ops::ext_op::{ExtensionOp, HostReference};
43pub use tenferro_ops::ExtensionFamilyId;
44
45pub use crate::extension_cache::{
46    ExtensionCacheKey, ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore,
47};
48pub use crate::extension_runtime::{
49    ExtensionExecutionContext, ExtensionExecutor, ExtensionRegistry, ExtensionRuntime,
50    ExtensionRuntimeRegistryError, HostReferenceRuntime,
51};
52
53/// Execute a lowered core program with caller-owned backend runtime cache state.
54///
55/// This owner-scoped hook is for operation-family runtimes that expand an
56/// extension into core tensor operations and need to run that lowered program
57/// while preserving the runtime cache owned by the outer graph executor.
58#[doc(hidden)]
59pub fn execute_lowered_program_with_backend_cache<B: TensorBackend + 'static>(
60    backend: &mut B,
61    program: &ExecProgram,
62    inputs: Vec<Tensor>,
63    backend_cache: &mut B::RuntimeCache,
64) -> Result<Vec<Tensor>> {
65    crate::exec::ensure_core_exec_program(
66        program,
67        "extension::execute_lowered_program_with_backend_cache",
68    )?;
69    crate::exec::eval_exec_ir_with_backend_cache(backend, program, inputs, backend_cache)
70}
71
72/// Apply an extension op in the traced graph.
73///
74/// The `op` value is cloned into a `StdTensorOp::Extension(Arc<dyn ExtensionOp>)`
75/// carrier. The returned vector contains one [`TracedTensor`] per declared
76/// output slot of the extension. Output shapes are inferred via
77/// [`ExtensionOp::infer_output_meta`] using the input shape hints.
78///
79/// `inputs.len()` must equal `op.input_count()`, and each input's
80/// `shape_hint` must be present (i.e. the extension must be used on
81/// tensors whose rank is known at graph-build time). For symbolic-shape
82/// composition, bind the placeholder tensors via
83/// [`crate::GraphExecutor::run_with_inputs`] at evaluation time.
84///
85/// # Examples
86///
87/// ```rust
88/// # use std::any::Any;
89/// use std::sync::Arc;
90/// use tenferro_runtime::extension::{apply, ExtensionOp};
91/// use tenferro_runtime::{DType, SymDim, TracedTensor};
92///
93/// # #[derive(Clone, Debug)]
94/// # struct IdentityExt;
95/// # impl ExtensionOp for IdentityExt {
96/// #     fn family_id(&self) -> &'static str { "example.identity.v1" }
97/// #     fn payload_hash(&self, _hasher: &mut dyn std::hash::Hasher) {}
98/// #     fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
99/// #         other.as_any().downcast_ref::<IdentityExt>().is_some()
100/// #     }
101/// #     fn clone_arc(&self) -> Arc<dyn ExtensionOp> { Arc::new(self.clone()) }
102/// #     fn as_any(&self) -> &dyn Any { self }
103/// #     fn input_count(&self) -> usize { 1 }
104/// #     fn output_count(&self) -> usize { 1 }
105/// #     fn infer_output_meta(
106/// #         &self,
107/// #         dtypes: &[DType],
108/// #         shapes: &[&[SymDim]],
109/// #     ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
110/// #         Ok(vec![(dtypes[0], shapes[0].to_vec())])
111/// #     }
112/// # }
113/// let op: Arc<dyn ExtensionOp> = Arc::new(IdentityExt);
114/// let a = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
115/// let outputs = apply(op, &[&a])?;
116/// assert_eq!(outputs.len(), 1);
117/// # Ok::<(), tenferro_runtime::Error>(())
118/// ```
119///
120/// # Errors
121///
122/// Returns [`Error::InvalidGraphBuild`] when the extension receives the wrong
123/// number of inputs or when [`ExtensionOp::infer_output_meta`] returns metadata
124/// whose count does not match [`ExtensionOp::output_count`].
125pub fn apply(op: Arc<dyn ExtensionOp>, inputs: &[&TracedTensor]) -> Result<Vec<TracedTensor>> {
126    if inputs.len() != op.input_count() {
127        return Err(Error::InvalidGraphBuild {
128            op: "extension::apply",
129            message: format!(
130                "op family {:?} expects {} inputs, got {}",
131                op.family_id(),
132                op.input_count(),
133                inputs.len()
134            ),
135        });
136    }
137
138    // Build the per-input dtype / shape slices the extension's
139    // `infer_output_meta` wants. Symbolic-shape inputs (shape_hint =
140    // None) use per-axis TensorAxis symbolic dims keyed by the input
141    // TracedTensor's id so downstream composition still resolves
142    // correctly via tenferro-internal-ops's SymDim API.
143    let input_dtypes: Vec<_> = inputs.iter().map(|t| t.dtype).collect();
144    let input_shape_storage: Vec<Vec<SymDim>> = inputs
145        .iter()
146        .map(|t| {
147            if let Some(hint) = t.shape_hint.clone() {
148                hint
149            } else {
150                (0..t.rank)
151                    .map(|axis| SymDim::tensor_axis(t.id, axis))
152                    .collect()
153            }
154        })
155        .collect();
156    let input_shape_refs: Vec<&[SymDim]> = input_shape_storage.iter().map(Vec::as_slice).collect();
157
158    let output_metas = op.infer_output_meta(&input_dtypes, &input_shape_refs)?;
159    if output_metas.len() != op.output_count() {
160        return Err(Error::InvalidGraphBuild {
161            op: "extension::apply",
162            message: format!(
163                "op family {:?}: infer_output_meta produced {} output metadata entries; op declared {} outputs",
164                op.family_id(),
165                output_metas.len(),
166                op.output_count()
167            ),
168        });
169    }
170
171    // Build the graph that carries the Extension op.
172    let mut builder = GraphBuilder::<StdTensorOp>::new();
173    for input in inputs {
174        builder.add_parent(input.graph.clone());
175    }
176    let op_inputs: Vec<ValueRef<StdTensorOp>> = inputs
177        .iter()
178        .map(|t| ValueRef::External(t.graph.values()[t.val].key.clone()))
179        .collect();
180    let carrier = StdTensorOp::Extension(op.clone());
181    let outputs = builder.add_operation(carrier, op_inputs, OperationRole::Primary);
182    builder.set_outputs(outputs.clone());
183    let graph = Arc::new(builder.build());
184    traced_outputs_from_graph(inputs, graph, &outputs, output_metas)
185}
186
187/// Apply an extension-provided lowering as ordinary traced graph operations.
188///
189/// This is for extension crates whose operation can be expanded at graph-build
190/// time. It preserves the same parent graph and metadata merging behavior as
191/// [`apply`], but does not insert a `StdTensorOp::Extension` carrier.
192pub fn apply_expanded_graph(
193    inputs: &[&TracedTensor],
194    output_metas: Vec<(tenferro_tensor::DType, Vec<SymDim>)>,
195    build: impl FnOnce(&mut GraphBuilder<StdTensorOp>, &[ValueRef<StdTensorOp>]) -> Result<Vec<usize>>,
196) -> Result<Vec<TracedTensor>> {
197    let mut builder = GraphBuilder::<StdTensorOp>::new();
198    for input in inputs {
199        builder.add_parent(input.graph.clone());
200    }
201    let op_inputs: Vec<ValueRef<StdTensorOp>> = inputs
202        .iter()
203        .map(|t| ValueRef::External(t.graph.values()[t.val].key.clone()))
204        .collect();
205    let outputs = build(&mut builder, &op_inputs)?;
206    if outputs.len() != output_metas.len() {
207        return Err(Error::InvalidGraphBuild {
208            op: "extension::apply_expanded_graph",
209            message: format!(
210                "extension expanded graph returned {} outputs for {} output metadata entries",
211                outputs.len(),
212                output_metas.len()
213            ),
214        });
215    }
216    builder.set_outputs(outputs.clone());
217    let graph = Arc::new(builder.build());
218    traced_outputs_from_graph(inputs, graph, &outputs, output_metas)
219}
220
221fn traced_outputs_from_graph(
222    inputs: &[&TracedTensor],
223    graph: Arc<computegraph::graph::Graph<StdTensorOp>>,
224    outputs: &[usize],
225    output_metas: Vec<(tenferro_tensor::DType, Vec<SymDim>)>,
226) -> Result<Vec<TracedTensor>> {
227    let metadata_scope = Arc::new(register_scoped_graph_metadata(
228        graph.as_ref(),
229        std::iter::empty(),
230    )?);
231
232    let merged_map = merge_traced_inputs_map(inputs.iter().copied());
233    let mut extra_roots = Vec::new();
234    let mut checkpoint_chain = None;
235    let metadata_scopes = MetadataScopeChain::with_scope(
236        Arc::clone(&metadata_scope),
237        inputs.iter().map(|input| &input.metadata_scopes),
238    );
239    for input in inputs {
240        extra_roots.extend(input.extra_roots.iter().cloned());
241        checkpoint_chain =
242            CheckpointNode::merge_chains(checkpoint_chain, input.checkpoint_chain.clone());
243    }
244    let all_inputs_concrete = inputs.iter().all(|t| t.shape_hint.is_some());
245    Ok(outputs
246        .iter()
247        .zip(output_metas)
248        .map(|(&val, (dtype, shape))| {
249            let shape_hint = if all_inputs_concrete {
250                Some(shape.clone())
251            } else {
252                None
253            };
254            TracedTensor {
255                id: next_traced_id(),
256                rank: shape.len(),
257                dtype,
258                graph: graph.clone(),
259                val,
260                data: None,
261                shape_hint,
262                inputs_map: merged_map.clone(),
263                extra_roots: extra_roots.clone(),
264                checkpoint_chain: checkpoint_chain.clone(),
265                metadata_scopes: metadata_scopes.clone(),
266            }
267        })
268        .collect())
269}
270
271#[cfg(test)]
272mod tests;
tenferro_runtime/extension.rs

tenferro_runtime/
extension.rs
