Skip to main content

tenferro_runtime/graph/
program.rs

1use std::sync::Arc;
2
3use tenferro_ops::dim_expr::DimExpr;
4use tenferro_ops::input_key::TensorInputKey;
5use tenferro_tensor::{DType, Tensor};
6
7use crate::exec::ExecProgram;
8use crate::graph::lowering_view::GraphProgramLoweringView;
9
10/// A compiled traced graph, independent of any execution backend.
11///
12/// # Examples
13///
14/// ```
15/// use tenferro_runtime::{GraphCompiler, TracedTensor};
16///
17/// let x = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
18/// let y = (&x + &x).unwrap();
19/// let mut compiler = GraphCompiler::new();
20/// let program = compiler.compile(&y).unwrap();
21/// assert_eq!(program.input_count(), 1);
22/// ```
23#[derive(Clone, Debug)]
24pub struct GraphProgram {
25    pub(crate) exec: ExecProgram,
26    pub(crate) inputs: Vec<GraphProgramInput>,
27}
28
29impl GraphProgram {
30    pub(crate) fn new(exec: ExecProgram, inputs: Vec<GraphProgramInput>) -> Self {
31        Self { exec, inputs }
32    }
33
34    /// Return the number of graph inputs expected by this program.
35    ///
36    /// # Examples
37    ///
38    /// ```
39    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
40    ///
41    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap();
42    /// let mut compiler = GraphCompiler::new();
43    /// let y = x.neg().unwrap();
44    /// let program = compiler.compile(&y).unwrap();
45    /// assert_eq!(program.input_count(), 1);
46    /// ```
47    #[inline(never)]
48    pub fn input_count(&self) -> usize {
49        self.inputs.len()
50    }
51
52    /// Return the number of graph outputs produced by this program.
53    ///
54    /// # Examples
55    ///
56    /// ```
57    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
58    ///
59    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap();
60    /// let mut compiler = GraphCompiler::new();
61    /// let y = x.neg().unwrap();
62    /// let program = compiler.compile(&y).unwrap();
63    /// assert_eq!(program.output_count(), 1);
64    /// ```
65    #[inline(never)]
66    pub fn output_count(&self) -> usize {
67        self.exec.output_slots.len()
68    }
69
70    /// Return the ordered input specs expected by this program.
71    ///
72    /// # Examples
73    ///
74    /// ```
75    /// use tenferro_runtime::{DType, GraphCompiler, TracedTensor};
76    ///
77    /// let x = TracedTensor::input_symbolic_shape(DType::F64, 1).unwrap();
78    /// let mut compiler = GraphCompiler::new();
79    /// let y = x.neg().unwrap();
80    /// let program = compiler
81    ///     .compile_with_input_specs(&y, &[(&x, DType::F64, &[4])])
82    ///     .unwrap();
83    /// assert_eq!(program.input_specs()[0].shape(), &[4]);
84    /// ```
85    #[inline(never)]
86    pub fn input_specs(&self) -> &[GraphProgramInput] {
87        &self.inputs
88    }
89
90    /// Return a read-only lowering view for peer executor integrations.
91    ///
92    /// The view exposes only immutable, lowering-oriented program metadata.
93    /// Native execution and mutation remain owned by [`GraphExecutor`](super::GraphExecutor).
94    ///
95    /// # Examples
96    ///
97    /// ```
98    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
99    ///
100    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
101    /// let mut compiler = GraphCompiler::new();
102    /// let y = x.neg().unwrap();
103    /// let program = compiler.compile(&y).unwrap();
104    /// assert_eq!(program.lowering_view().output_slots().len(), 1);
105    /// ```
106    #[inline(never)]
107    pub fn lowering_view(&self) -> GraphProgramLoweringView<'_> {
108        GraphProgramLoweringView::new(&self.exec)
109    }
110}
111
112/// A single ordered input required by a [`GraphProgram`].
113///
114/// # Examples
115///
116/// ```
117/// use tenferro_runtime::{GraphCompiler, TracedTensor};
118///
119/// let x = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
120/// let mut compiler = GraphCompiler::new();
121/// let y = x.neg().unwrap();
122/// let program = compiler.compile(&y).unwrap();
123/// let input = &program.input_specs()[0];
124/// assert_eq!(input.shape(), &[2]);
125/// ```
126#[derive(Clone, Debug)]
127pub struct GraphProgramInput {
128    pub(crate) key: TensorInputKey,
129    pub(crate) dtype: DType,
130    pub(crate) shape: Vec<usize>,
131    // Preserved for symbolic-shape diagnostics and future graph-input metadata
132    // without exposing `DimExpr` through the stable input-spec accessor.
133    #[allow(dead_code)]
134    pub(crate) dim_expr_shape: Vec<DimExpr>,
135    pub(crate) default_tensor: Option<Arc<Tensor>>,
136}
137
138impl GraphProgramInput {
139    pub(crate) fn new(
140        key: TensorInputKey,
141        dtype: DType,
142        shape: Vec<usize>,
143        dim_expr_shape: Vec<DimExpr>,
144        default_tensor: Option<Arc<Tensor>>,
145    ) -> Self {
146        Self {
147            key,
148            dtype,
149            shape,
150            dim_expr_shape,
151            default_tensor,
152        }
153    }
154
155    /// Return the dtype expected for this input.
156    ///
157    /// # Examples
158    ///
159    /// ```
160    /// use tenferro_runtime::{DType, GraphCompiler, TracedTensor};
161    ///
162    /// let x = TracedTensor::input_symbolic_shape(DType::F64, 1).unwrap();
163    /// let mut compiler = GraphCompiler::new();
164    /// let program = compiler
165    ///     .compile_with_input_specs(&x, &[(&x, DType::F64, &[2])])
166    ///     .unwrap();
167    /// assert_eq!(program.input_specs()[0].dtype(), DType::F64);
168    /// ```
169    #[inline(never)]
170    pub fn dtype(&self) -> DType {
171        self.dtype
172    }
173
174    /// Return the concrete shape expected for this input.
175    ///
176    /// # Examples
177    ///
178    /// ```
179    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
180    ///
181    /// let x = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
182    /// let mut compiler = GraphCompiler::new();
183    /// let program = compiler.compile(&x).unwrap();
184    /// assert_eq!(program.input_specs()[0].shape(), &[2]);
185    /// ```
186    #[inline(never)]
187    pub fn shape(&self) -> &[usize] {
188        &self.shape
189    }
190}