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}