tenferro_runtime/
trace_ops.rs

1//! Typed operation surface over one [`TraceContext`](crate::TraceContext).
2//!
3//! Every standard operation and every traced constructor is written once here
4//! against a `(TraceContext, TraceValue)` pair. [`TracedTensor`] methods and the
5//! extension traits delegate to these helpers, so there is exactly one place that
6//! appends an operation and exactly one place that derives a tensor's cached
7//! metadata.
8
9use std::sync::Arc;
10
11use tenferro_ops::dim_expr::DimExpr;
12use tenferro_ops::ext_op::ExtensionOp;
13use tenferro_ops::shape_extent::ShapeExtent;
14use tenferro_ops::std_tensor_op::StdTensorOp;
15use tenferro_tensor::{DType, IntoShapeVec, Tensor, TensorScalar, TensorValue};
16
17use crate::checkpoint::RetainedValue;
18use crate::error::{Error, ErrorPhase, Result};
19use crate::program::{CoreSemanticOp, ProgramBuildError, ProgramInputSpec, ProgramValueMetadata};
20use crate::sym_dim::SymDim;
21use crate::trace::{TraceContext, TraceValue};
22use crate::traced::TracedTensor;
23
24/// Map a standard operation payload onto the closed core vocabulary.
25///
26/// # Errors
27///
28/// Returns [`Error::RuntimeState`] for the extension carrier, which must be
29/// appended through [`append_extension`] instead.
30pub(crate) fn core_op(op: &StdTensorOp) -> Result<CoreSemanticOp> {
31    CoreSemanticOp::try_from(op).map_err(|source| {
32        Error::runtime_state_source("semantic core operation", ErrorPhase::GraphBuild, source)
33    })
34}
35
36/// Append one standard operation to the context.
37pub(crate) fn append_core(
38    ctx: &TraceContext,
39    op: &StdTensorOp,
40    inputs: &[TraceValue],
41) -> Result<Box<[TraceValue]>> {
42    let core = core_op(op)?;
43    ctx.add_op(core, inputs)
44        .map_err(|source| build_error("semantic standard operation", source))
45}
46
47/// Append one extension operation to the context.
48pub(crate) fn append_extension(
49    ctx: &TraceContext,
50    op: Arc<dyn ExtensionOp>,
51    inputs: &[TraceValue],
52) -> Result<Box<[TraceValue]>> {
53    ctx.add_extension(op, inputs)
54        .map_err(|source| build_error("semantic extension operation", source))
55}
56
57/// Build a tensor for a value the context just appended.
58///
59/// `rank`/`dtype`/`hint` are the op-specific expectations; the returned cache is
60/// derived from the context metadata and cross-checked against them, so the cache
61/// can never disagree with the program.
62pub(crate) fn finish_value(
63    ctx: &TraceContext,
64    value: TraceValue,
65    rank: usize,
66    dtype: DType,
67    hint: Option<Vec<SymDim>>,
68) -> Result<TracedTensor> {
69    let metadata = ctx
70        .value_metadata(value)
71        .map_err(|source| build_error("semantic value metadata", source))?;
72    debug_assert_eq!(
73        rank,
74        metadata.shape().len(),
75        "cached rank must match the appended value"
76    );
77    debug_assert_eq!(
78        dtype,
79        metadata.dtype(),
80        "cached dtype must match the appended value"
81    );
82    let derived = metadata_sym_shape(&metadata);
83    let symbolic_id = next_symbolic_id();
84    // An extent the program only bounds must not resurface as a concrete shape, and a
85    // caller-supplied hint describes the shape the caller asked for, not the precision
86    // the program proved: drop it unless every extent is exact.
87    let hint = if crate::shape_infer::extent_precision(metadata.shape().iter())
88        == crate::shape_infer::ExtentPrecision::Exact
89    {
90        hint
91    } else {
92        None
93    };
94    Ok(TracedTensor {
95        ctx: ctx.clone(),
96        value,
97        rank: metadata.shape().len(),
98        dtype: metadata.dtype(),
99        shape_hint: restamp_hint(derived.or(hint), symbolic_id),
100        data: None,
101        symbolic_id,
102    })
103}
104
105/// Build a leaf tensor for a value the context just added as an input.
106fn finish_leaf(
107    ctx: &TraceContext,
108    value: TraceValue,
109    rank: usize,
110    dtype: DType,
111    hint: Option<Vec<SymDim>>,
112    data: Option<Arc<RetainedValue>>,
113) -> Result<TracedTensor> {
114    let mut tensor = finish_value(ctx, value, rank, dtype, hint)?;
115    tensor.data = data;
116    Ok(tensor)
117}
118
119/// Build a tensor purely from the value's context metadata.
120pub(crate) fn finish_derived(ctx: &TraceContext, value: TraceValue) -> Result<TracedTensor> {
121    let metadata = ctx
122        .value_metadata(value)
123        .map_err(|source| build_error("semantic value metadata", source))?;
124    let rank = metadata.shape().len();
125    let dtype = metadata.dtype();
126    let hint = metadata_sym_shape(&metadata);
127    finish_value(ctx, value, rank, dtype, hint)
128}
129
130/// Build a tensor for a value the context just bound to concrete data.
131pub(crate) fn finish_bound_value(
132    ctx: &TraceContext,
133    value: TraceValue,
134    data: Arc<RetainedValue>,
135) -> Result<TracedTensor> {
136    let metadata = ctx
137        .value_metadata(value)
138        .map_err(|source| build_error("semantic value metadata", source))?;
139    let rank = metadata.shape().len();
140    let dtype = metadata.dtype();
141    let hint = metadata_sym_shape(&metadata);
142    let mut tensor = finish_value(ctx, value, rank, dtype, hint)?;
143    tensor.data = Some(data);
144    Ok(tensor)
145}
146
147/// Rebuild a symbolic extent from a program dimension expression.
148///
149/// Program extents name program inputs by position, so the result names those
150/// positions as tensor ids; a consumer translates them into its own coordinates
151/// with an explicit tensor map.
152///
153/// Owner-scoped hook for the traced AD layer, which validates seed metadata
154/// against program extents.
155#[doc(hidden)]
156pub fn dim_expr_to_sym_dim_public(expr: &DimExpr) -> Option<SymDim> {
157    dim_expr_to_sym_dim(expr)
158}
159
160pub(crate) fn dim_expr_to_sym_dim(expr: &DimExpr) -> Option<SymDim> {
161    Some(match expr {
162        DimExpr::Const(value) => SymDim::from(*value),
163        DimExpr::InputDim { input_idx, axis } => SymDim::tensor_axis(*input_idx as u64, *axis),
164        DimExpr::Add(lhs, rhs) => dim_expr_to_sym_dim(lhs)? + dim_expr_to_sym_dim(rhs)?,
165        DimExpr::Sub(lhs, rhs) => dim_expr_to_sym_dim(lhs)? - dim_expr_to_sym_dim(rhs)?,
166        DimExpr::Mul(lhs, rhs) => dim_expr_to_sym_dim(lhs)? * dim_expr_to_sym_dim(rhs)?,
167        DimExpr::FloorDiv(lhs, rhs) => dim_expr_to_sym_dim(lhs)? / dim_expr_to_sym_dim(rhs)?,
168        DimExpr::Min(lhs, rhs) => dim_expr_to_sym_dim(lhs)?.min(dim_expr_to_sym_dim(rhs)?),
169        DimExpr::Max(lhs, rhs) => dim_expr_to_sym_dim(lhs)?.max(dim_expr_to_sym_dim(rhs)?),
170    })
171}
172
173/// Allocate the symbolic identity of one traced value.
174///
175/// The identity is unique per handle rather than derived from the program value,
176/// so no two tensors — in one context or across contexts — can name the same
177/// symbolic axis by accident.
178fn next_symbolic_id() -> u64 {
179    static NEXT_SYMBOLIC_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
180    NEXT_SYMBOLIC_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
181}
182
183/// Re-express a symbolic shape in terms of the value that owns it.
184///
185/// `sym_size`/`axis_sym_dim` must hand back expressions that a *consumer* of this
186/// value can convert into its own input coordinates, so every non-constant axis
187/// names this value's axis instead of the axes of the operation that produced it.
188fn restamp_hint(hint: Option<Vec<SymDim>>, own_id: u64) -> Option<Vec<SymDim>> {
189    hint.map(|shape| {
190        shape
191            .into_iter()
192            .enumerate()
193            .map(|(axis, dim)| {
194                if dim.constant_value().is_some() {
195                    dim
196                } else {
197                    SymDim::tensor_axis(own_id, axis)
198                }
199            })
200            .collect()
201    })
202}
203
204/// Derive the symbolic shape cache from value metadata.
205fn metadata_sym_shape(metadata: &ProgramValueMetadata) -> Option<Vec<SymDim>> {
206    metadata
207        .shape()
208        .iter()
209        .map(|extent| match extent {
210            ShapeExtent::Exact(expr) => dim_expr_to_sym_dim(expr),
211            _ => None,
212        })
213        .collect()
214}
215
216pub(crate) fn build_error(op: &'static str, source: ProgramBuildError) -> Error {
217    match source {
218        ProgramBuildError::Metadata { source } => *source,
219        other => Error::runtime_state_source(op, ErrorPhase::GraphBuild, other),
220    }
221}
222
223fn concrete_spec(
224    dtype: DType,
225    shape: &[usize],
226    identity: Option<&'static str>,
227) -> ProgramInputSpec {
228    let metadata = ProgramValueMetadata::new(dtype, shape.iter().map(|&size| DimExpr::Const(size)));
229    let spec = ProgramInputSpec::from_metadata(metadata);
230    match identity {
231        Some(identity) => spec.with_scalar_identity(identity),
232        None => spec,
233    }
234}
235
236fn concrete_hint(shape: &[usize]) -> Option<Vec<SymDim>> {
237    Some(shape.iter().map(|&size| SymDim::from(size)).collect())
238}
239
240fn symbolic_spec(dtype: DType, rank: usize, input_idx: usize) -> ProgramInputSpec {
241    ProgramInputSpec::new(
242        dtype,
243        (0..rank).map(|axis| DimExpr::InputDim { input_idx, axis }),
244    )
245}
246
247impl TraceContext {
248    /// Add a placeholder input with a fixed concrete shape.
249    ///
250    /// # Examples
251    ///
252    /// ```
253    /// use tenferro_runtime::TraceContext;
254    /// use tenferro_tensor::DType;
255    ///
256    /// let trace = TraceContext::new();
257    /// let x = trace.input_concrete_shape(DType::F64, &[2, 3])?;
258    /// assert_eq!(x.rank(), 2);
259    /// assert!(x.is_concrete_shape());
260    /// # Ok::<(), tenferro_runtime::Error>(())
261    /// ```
262    ///
263    /// # Errors
264    ///
265    /// Returns [`Error::RuntimeState`] when the trace cannot accept another input
266    /// or a binding already exists for it.
267    pub fn input_concrete_shape(&self, dtype: DType, shape: &[usize]) -> Result<TracedTensor> {
268        self.input_concrete_shape_with_identity(dtype, shape, None)
269    }
270
271    /// Add a placeholder input of concrete shape declaring an external scalar.
272    ///
273    /// # Errors
274    ///
275    /// Returns [`Error::RuntimeState`] when the trace cannot accept another input
276    /// or a binding already exists for it, and [`Error::RuntimeStateSource`]
277    /// wrapping `ProgramBuildError::ExternalScalarWithoutIdentity` when an
278    /// externally defined scalar dtype reaches the builder without the declared
279    /// identity.
280    pub fn input_concrete_shape_declaring_scalar(
281        &self,
282        dtype: DType,
283        shape: &[usize],
284        identity: &'static str,
285    ) -> Result<TracedTensor> {
286        self.input_concrete_shape_with_identity(dtype, shape, Some(identity))
287    }
288
289    fn input_concrete_shape_with_identity(
290        &self,
291        dtype: DType,
292        shape: &[usize],
293        identity: Option<&'static str>,
294    ) -> Result<TracedTensor> {
295        let spec = concrete_spec(dtype, shape, identity);
296        let value = self
297            .input(spec)
298            .map_err(|source| build_error("trace input", source))?;
299        finish_leaf(self, value, shape.len(), dtype, concrete_hint(shape), None)
300    }
301
302    /// Add a placeholder input with symbolic (rank-only) extents.
303    ///
304    /// # Examples
305    ///
306    /// ```
307    /// use tenferro_runtime::TraceContext;
308    /// use tenferro_tensor::DType;
309    ///
310    /// let trace = TraceContext::new();
311    /// let x = trace.input_symbolic_shape(DType::F64, 2)?;
312    /// assert_eq!(x.rank(), 2);
313    /// assert!(!x.is_concrete_shape());
314    /// # Ok::<(), tenferro_runtime::Error>(())
315    /// ```
316    ///
317    /// # Errors
318    ///
319    /// Returns [`Error::RuntimeState`] when the trace cannot accept another input.
320    pub fn input_symbolic_shape(&self, dtype: DType, rank: usize) -> Result<TracedTensor> {
321        let index = self.input_count();
322        let spec = symbolic_spec(dtype, rank, index);
323        let value = self
324            .input(spec)
325            .map_err(|source| build_error("trace input", source))?;
326        finish_leaf(self, value, rank, dtype, metadata_hint(self, value)?, None)
327    }
328
329    /// Add a constant leaf holding `tensor` with its concrete shape.
330    ///
331    /// # Examples
332    ///
333    /// ```
334    /// use tenferro_runtime::{Tensor, TraceContext};
335    ///
336    /// let trace = TraceContext::new();
337    /// let x = trace.from_tensor_concrete_shape(
338    ///     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?,
339    /// )?;
340    /// assert!(x.attached_value().is_some());
341    /// # Ok::<(), Box<dyn std::error::Error>>(())
342    /// ```
343    ///
344    /// # Errors
345    ///
346    /// Returns [`Error::Validation`] for a malformed tensor or
347    /// [`Error::RuntimeState`] when the trace cannot bind the constant.
348    pub fn from_tensor_concrete_shape(&self, tensor: Tensor) -> Result<TracedTensor> {
349        self.tensor_concrete_shape_with_identity(tensor, None)
350    }
351
352    /// Add a constant leaf of concrete shape declaring an external scalar.
353    ///
354    /// # Errors
355    ///
356    /// Returns [`Error::Validation`] for a malformed tensor, [`Error::RuntimeState`]
357    /// when the trace cannot bind the constant, and [`Error::RuntimeStateSource`]
358    /// wrapping `ProgramBuildError::ExternalScalarWithoutIdentity` when an
359    /// externally defined scalar dtype reaches the builder without the declared
360    /// identity.
361    pub fn from_tensor_concrete_shape_declaring_scalar(
362        &self,
363        tensor: Tensor,
364        identity: &'static str,
365    ) -> Result<TracedTensor> {
366        self.tensor_concrete_shape_with_identity(tensor, Some(identity))
367    }
368
369    fn tensor_concrete_shape_with_identity(
370        &self,
371        tensor: Tensor,
372        identity: Option<&'static str>,
373    ) -> Result<TracedTensor> {
374        let dtype = tensor.dtype();
375        let shape = tensor.shape().to_vec();
376        let spec = concrete_spec(dtype, &shape, identity);
377        let data = Arc::new(RetainedValue::from_tensor(tensor));
378        let value = self
379            .input_bound(spec, Arc::clone(&data))
380            .map_err(|source| build_error("trace constant", source))?;
381        finish_leaf(
382            self,
383            value,
384            shape.len(),
385            dtype,
386            concrete_hint(&shape),
387            Some(data),
388        )
389    }
390
391    /// Add a constant leaf holding `tensor` but advertising symbolic extents.
392    ///
393    /// # Errors
394    ///
395    /// Returns [`Error::RuntimeState`] when the trace cannot bind the constant.
396    pub fn from_tensor_symbolic_shape(&self, tensor: Tensor) -> Result<TracedTensor> {
397        let data = Arc::new(RetainedValue::from_tensor(tensor));
398        self.from_shared_tensor_value_symbolic_shape(data)
399    }
400
401    /// Add a constant leaf from a value, advertising symbolic extents.
402    ///
403    /// # Errors
404    ///
405    /// Returns [`Error::RuntimeState`] when the value cannot be retained or the
406    /// trace cannot bind it.
407    #[doc(hidden)]
408    pub fn from_tensor_value_symbolic_shape(&self, value: TensorValue) -> Result<TracedTensor> {
409        let data = Arc::new(RetainedValue::from_tensor_value(value)?);
410        self.from_shared_tensor_value_symbolic_shape(data)
411    }
412
413    /// Add a constant leaf from an already-retained value with symbolic extents.
414    ///
415    /// # Errors
416    ///
417    /// Returns [`Error::RuntimeState`] when the trace cannot bind the value.
418    #[doc(hidden)]
419    pub fn from_shared_tensor_value_symbolic_shape(
420        &self,
421        data: Arc<RetainedValue>,
422    ) -> Result<TracedTensor> {
423        let rank = data.shape().len();
424        let dtype = data.dtype();
425        let index = self.input_count();
426        let spec = symbolic_spec(dtype, rank, index);
427        let value = self
428            .input_bound(spec, Arc::clone(&data))
429            .map_err(|source| build_error("trace constant", source))?;
430        finish_leaf(
431            self,
432            value,
433            rank,
434            dtype,
435            metadata_hint(self, value)?,
436            Some(data),
437        )
438    }
439
440    /// Add a constant leaf from compact column-major data.
441    ///
442    /// # Errors
443    ///
444    /// Returns [`Error::TensorRuntime`] with a typed validation error when the
445    /// shape and data length disagree.
446    pub fn from_vec_col_major<T: TensorScalar>(
447        &self,
448        shape: impl IntoShapeVec,
449        data: Vec<T>,
450    ) -> Result<TracedTensor> {
451        let tensor = Tensor::from_vec_col_major(shape, data)?;
452        self.from_tensor_concrete_shape(tensor)
453    }
454}
455
456fn metadata_hint(ctx: &TraceContext, value: TraceValue) -> Result<Option<Vec<SymDim>>> {
457    let metadata = ctx
458        .value_metadata(value)
459        .map_err(|source| build_error("semantic value metadata", source))?;
460    Ok(metadata_sym_shape(&metadata))
461}
tenferro_runtime/trace_ops.rs

tenferro_runtime/
trace_ops.rs
