tenferro_runtime/program/
op.rs

1use std::sync::Arc;
2
3use computegraph::GraphOperation;
4use tenferro_ops::dim_expr::DimExpr;
5use tenferro_ops::ext_op::ExtensionOp;
6use tenferro_ops::std_tensor_op::StdTensorOp;
7use tenferro_tensor::{
8    CompareDir, DType, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
9};
10
11use super::metadata::SemanticProvenance;
12use super::{
13    Alias, Effect, ProgramValue, SemanticPlacementConstraint, SemanticProvenanceView, ShapeGuard,
14};
15
16/// Closed backend-neutral vocabulary of core semantic tensor operations.
17#[derive(Clone, Debug, PartialEq)]
18#[non_exhaustive]
19pub enum CoreSemanticOp {
20    Add,
21    Sub,
22    Mul,
23    Neg,
24    Conj,
25    DotGeneral {
26        config: DotGeneralConfig,
27    },
28    Transpose {
29        perm: Vec<usize>,
30    },
31    Reshape {
32        to_shape: Vec<DimExpr>,
33    },
34    BroadcastInDim {
35        shape: Vec<DimExpr>,
36        dims: Vec<usize>,
37    },
38    Convert {
39        from: DType,
40        to: DType,
41    },
42    Constant {
43        dtype: DType,
44        bytes: Vec<u8>,
45    },
46    ReduceSum {
47        axes: Vec<usize>,
48    },
49    ReduceSumSquares {
50        axes: Vec<usize>,
51    },
52    Div,
53    Rem,
54    Abs,
55    Sign,
56    Maximum,
57    Minimum,
58    Compare(CompareDir),
59    Select,
60    Clamp,
61    Exp,
62    Log,
63    Sin,
64    Cos,
65    Tanh,
66    Sqrt,
67    Rsqrt,
68    Pow,
69    Expm1,
70    Log1p,
71    ExtractDiag {
72        axis_a: usize,
73        axis_b: usize,
74    },
75    EmbedDiag {
76        axis_a: usize,
77        axis_b: usize,
78    },
79    Tril {
80        k: i64,
81    },
82    Triu {
83        k: i64,
84    },
85    Gather(GatherConfig),
86    GatherDynamicSliceSizes {
87        offset_dims: Vec<usize>,
88        collapsed_slice_dims: Vec<usize>,
89        start_index_map: Vec<usize>,
90        index_vector_dim: usize,
91        slice_sizes: Vec<DimExpr>,
92    },
93    Scatter(ScatterConfig),
94    Slice(SliceConfig),
95    DynamicSlice {
96        slice_sizes: Vec<usize>,
97    },
98    DynamicUpdateSlice,
99    Pad(PadConfig),
100    Concatenate {
101        axis: usize,
102        input_count: usize,
103    },
104    Reverse {
105        axes: Vec<usize>,
106    },
107    ShapeOf {
108        axis: usize,
109    },
110    DynamicTruncate {
111        axis: usize,
112    },
113    PadToMatch {
114        axis: usize,
115    },
116    ReduceProd {
117        axes: Vec<usize>,
118    },
119    ReduceMax {
120        axes: Vec<usize>,
121    },
122    ReduceMin {
123        axes: Vec<usize>,
124    },
125}
126
127/// Failure to convert a standard-operation carrier into the closed core
128/// semantic vocabulary.
129#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
130pub enum CoreSemanticOpConversionError {
131    /// Extension payloads must enter through
132    /// [`SemanticProgramBuilder::add_extension`](super::SemanticProgramBuilder::add_extension).
133    #[error("extension operations are not core semantic operations")]
134    ExtensionCarrier,
135}
136
137impl CoreSemanticOp {
138    pub(crate) fn input_count(&self) -> usize {
139        let standard = StdTensorOp::from(self);
140        GraphOperation::input_count(&standard)
141    }
142
143    pub(crate) fn output_count(&self) -> usize {
144        let standard = StdTensorOp::from(self);
145        GraphOperation::output_count(&standard)
146    }
147}
148
149impl TryFrom<&StdTensorOp> for CoreSemanticOp {
150    type Error = CoreSemanticOpConversionError;
151
152    fn try_from(op: &StdTensorOp) -> Result<Self, Self::Error> {
153        Ok(match op {
154            StdTensorOp::Add => Self::Add,
155            StdTensorOp::Sub => Self::Sub,
156            StdTensorOp::Mul => Self::Mul,
157            StdTensorOp::Neg => Self::Neg,
158            StdTensorOp::Conj => Self::Conj,
159            StdTensorOp::DotGeneral { config } => Self::DotGeneral {
160                config: config.clone(),
161            },
162            StdTensorOp::Transpose { perm } => Self::Transpose { perm: perm.clone() },
163            StdTensorOp::Reshape { to_shape } => Self::Reshape {
164                to_shape: to_shape.clone(),
165            },
166            StdTensorOp::BroadcastInDim { shape, dims } => Self::BroadcastInDim {
167                shape: shape.clone(),
168                dims: dims.clone(),
169            },
170            StdTensorOp::Convert { from, to } => Self::Convert {
171                from: *from,
172                to: *to,
173            },
174            StdTensorOp::Constant { dtype, bytes } => Self::Constant {
175                dtype: *dtype,
176                bytes: bytes.clone(),
177            },
178            StdTensorOp::ReduceSum { axes } => Self::ReduceSum { axes: axes.clone() },
179            StdTensorOp::ReduceSumSquares { axes } => Self::ReduceSumSquares { axes: axes.clone() },
180            StdTensorOp::Div => Self::Div,
181            StdTensorOp::Rem => Self::Rem,
182            StdTensorOp::Abs => Self::Abs,
183            StdTensorOp::Sign => Self::Sign,
184            StdTensorOp::Maximum => Self::Maximum,
185            StdTensorOp::Minimum => Self::Minimum,
186            StdTensorOp::Compare(direction) => Self::Compare(direction.clone()),
187            StdTensorOp::Select => Self::Select,
188            StdTensorOp::Clamp => Self::Clamp,
189            StdTensorOp::Exp => Self::Exp,
190            StdTensorOp::Log => Self::Log,
191            StdTensorOp::Sin => Self::Sin,
192            StdTensorOp::Cos => Self::Cos,
193            StdTensorOp::Tanh => Self::Tanh,
194            StdTensorOp::Sqrt => Self::Sqrt,
195            StdTensorOp::Rsqrt => Self::Rsqrt,
196            StdTensorOp::Pow => Self::Pow,
197            StdTensorOp::Expm1 => Self::Expm1,
198            StdTensorOp::Log1p => Self::Log1p,
199            StdTensorOp::ExtractDiag { axis_a, axis_b } => Self::ExtractDiag {
200                axis_a: *axis_a,
201                axis_b: *axis_b,
202            },
203            StdTensorOp::EmbedDiag { axis_a, axis_b } => Self::EmbedDiag {
204                axis_a: *axis_a,
205                axis_b: *axis_b,
206            },
207            StdTensorOp::Tril { k } => Self::Tril { k: *k },
208            StdTensorOp::Triu { k } => Self::Triu { k: *k },
209            StdTensorOp::Gather(config) => Self::Gather(config.clone()),
210            StdTensorOp::GatherDynamicSliceSizes {
211                offset_dims,
212                collapsed_slice_dims,
213                start_index_map,
214                index_vector_dim,
215                slice_sizes,
216            } => Self::GatherDynamicSliceSizes {
217                offset_dims: offset_dims.clone(),
218                collapsed_slice_dims: collapsed_slice_dims.clone(),
219                start_index_map: start_index_map.clone(),
220                index_vector_dim: *index_vector_dim,
221                slice_sizes: slice_sizes.clone(),
222            },
223            StdTensorOp::Scatter(config) => Self::Scatter(config.clone()),
224            StdTensorOp::Slice(config) => Self::Slice(config.clone()),
225            StdTensorOp::DynamicSlice { slice_sizes } => Self::DynamicSlice {
226                slice_sizes: slice_sizes.clone(),
227            },
228            StdTensorOp::DynamicUpdateSlice => Self::DynamicUpdateSlice,
229            StdTensorOp::Pad(config) => Self::Pad(config.clone()),
230            StdTensorOp::Concatenate { axis, input_count } => Self::Concatenate {
231                axis: *axis,
232                input_count: *input_count,
233            },
234            StdTensorOp::Reverse { axes } => Self::Reverse { axes: axes.clone() },
235            StdTensorOp::ShapeOf { axis } => Self::ShapeOf { axis: *axis },
236            StdTensorOp::DynamicTruncate { axis } => Self::DynamicTruncate { axis: *axis },
237            StdTensorOp::PadToMatch { axis } => Self::PadToMatch { axis: *axis },
238            StdTensorOp::ReduceProd { axes } => Self::ReduceProd { axes: axes.clone() },
239            StdTensorOp::ReduceMax { axes } => Self::ReduceMax { axes: axes.clone() },
240            StdTensorOp::ReduceMin { axes } => Self::ReduceMin { axes: axes.clone() },
241            StdTensorOp::Extension(_) => {
242                return Err(CoreSemanticOpConversionError::ExtensionCarrier);
243            }
244        })
245    }
246}
247
248impl From<&CoreSemanticOp> for StdTensorOp {
249    fn from(op: &CoreSemanticOp) -> Self {
250        match op {
251            CoreSemanticOp::Add => Self::Add,
252            CoreSemanticOp::Sub => Self::Sub,
253            CoreSemanticOp::Mul => Self::Mul,
254            CoreSemanticOp::Neg => Self::Neg,
255            CoreSemanticOp::Conj => Self::Conj,
256            CoreSemanticOp::DotGeneral { config } => Self::DotGeneral {
257                config: config.clone(),
258            },
259            CoreSemanticOp::Transpose { perm } => Self::Transpose { perm: perm.clone() },
260            CoreSemanticOp::Reshape { to_shape } => Self::Reshape {
261                to_shape: to_shape.clone(),
262            },
263            CoreSemanticOp::BroadcastInDim { shape, dims } => Self::BroadcastInDim {
264                shape: shape.clone(),
265                dims: dims.clone(),
266            },
267            CoreSemanticOp::Convert { from, to } => Self::Convert {
268                from: *from,
269                to: *to,
270            },
271            CoreSemanticOp::Constant { dtype, bytes } => Self::Constant {
272                dtype: *dtype,
273                bytes: bytes.clone(),
274            },
275            CoreSemanticOp::ReduceSum { axes } => Self::ReduceSum { axes: axes.clone() },
276            CoreSemanticOp::ReduceSumSquares { axes } => {
277                Self::ReduceSumSquares { axes: axes.clone() }
278            }
279            CoreSemanticOp::Div => Self::Div,
280            CoreSemanticOp::Rem => Self::Rem,
281            CoreSemanticOp::Abs => Self::Abs,
282            CoreSemanticOp::Sign => Self::Sign,
283            CoreSemanticOp::Maximum => Self::Maximum,
284            CoreSemanticOp::Minimum => Self::Minimum,
285            CoreSemanticOp::Compare(direction) => Self::Compare(direction.clone()),
286            CoreSemanticOp::Select => Self::Select,
287            CoreSemanticOp::Clamp => Self::Clamp,
288            CoreSemanticOp::Exp => Self::Exp,
289            CoreSemanticOp::Log => Self::Log,
290            CoreSemanticOp::Sin => Self::Sin,
291            CoreSemanticOp::Cos => Self::Cos,
292            CoreSemanticOp::Tanh => Self::Tanh,
293            CoreSemanticOp::Sqrt => Self::Sqrt,
294            CoreSemanticOp::Rsqrt => Self::Rsqrt,
295            CoreSemanticOp::Pow => Self::Pow,
296            CoreSemanticOp::Expm1 => Self::Expm1,
297            CoreSemanticOp::Log1p => Self::Log1p,
298            CoreSemanticOp::ExtractDiag { axis_a, axis_b } => Self::ExtractDiag {
299                axis_a: *axis_a,
300                axis_b: *axis_b,
301            },
302            CoreSemanticOp::EmbedDiag { axis_a, axis_b } => Self::EmbedDiag {
303                axis_a: *axis_a,
304                axis_b: *axis_b,
305            },
306            CoreSemanticOp::Tril { k } => Self::Tril { k: *k },
307            CoreSemanticOp::Triu { k } => Self::Triu { k: *k },
308            CoreSemanticOp::Gather(config) => Self::Gather(config.clone()),
309            CoreSemanticOp::GatherDynamicSliceSizes {
310                offset_dims,
311                collapsed_slice_dims,
312                start_index_map,
313                index_vector_dim,
314                slice_sizes,
315            } => Self::GatherDynamicSliceSizes {
316                offset_dims: offset_dims.clone(),
317                collapsed_slice_dims: collapsed_slice_dims.clone(),
318                start_index_map: start_index_map.clone(),
319                index_vector_dim: *index_vector_dim,
320                slice_sizes: slice_sizes.clone(),
321            },
322            CoreSemanticOp::Scatter(config) => Self::Scatter(config.clone()),
323            CoreSemanticOp::Slice(config) => Self::Slice(config.clone()),
324            CoreSemanticOp::DynamicSlice { slice_sizes } => Self::DynamicSlice {
325                slice_sizes: slice_sizes.clone(),
326            },
327            CoreSemanticOp::DynamicUpdateSlice => Self::DynamicUpdateSlice,
328            CoreSemanticOp::Pad(config) => Self::Pad(config.clone()),
329            CoreSemanticOp::Concatenate { axis, input_count } => Self::Concatenate {
330                axis: *axis,
331                input_count: *input_count,
332            },
333            CoreSemanticOp::Reverse { axes } => Self::Reverse { axes: axes.clone() },
334            CoreSemanticOp::ShapeOf { axis } => Self::ShapeOf { axis: *axis },
335            CoreSemanticOp::DynamicTruncate { axis } => Self::DynamicTruncate { axis: *axis },
336            CoreSemanticOp::PadToMatch { axis } => Self::PadToMatch { axis: *axis },
337            CoreSemanticOp::ReduceProd { axes } => Self::ReduceProd { axes: axes.clone() },
338            CoreSemanticOp::ReduceMax { axes } => Self::ReduceMax { axes: axes.clone() },
339            CoreSemanticOp::ReduceMin { axes } => Self::ReduceMin { axes: axes.clone() },
340        }
341    }
342}
343
344pub(crate) enum SemanticOp {
345    Core(CoreSemanticOp),
346    Extension(Arc<dyn ExtensionOp>),
347}
348
349pub(crate) struct SemanticOperation {
350    pub(crate) op: SemanticOp,
351    pub(crate) inputs: Box<[ProgramValue]>,
352    pub(crate) outputs: Box<[ProgramValue]>,
353    pub(crate) effects: Box<[Effect]>,
354    pub(crate) aliases: Box<[Alias]>,
355    pub(crate) shape_guards: Box<[ShapeGuard]>,
356    pub(crate) placement: SemanticPlacementConstraint,
357    pub(crate) provenance: SemanticProvenance,
358}
359
360/// Borrowed semantic operation payload.
361#[derive(Clone, Copy)]
362#[non_exhaustive]
363pub enum SemanticOpRef<'a> {
364    /// Closed core operation.
365    Core(&'a CoreSemanticOp),
366    /// Extension semantic payload.
367    Extension(&'a dyn ExtensionOp),
368}
369
370impl std::fmt::Debug for SemanticOpRef<'_> {
371    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
372        match self {
373            Self::Core(_) => formatter.write_str("SemanticOpRef::Core(<bounded>)"),
374            Self::Extension(op) => formatter
375                .debug_tuple("SemanticOpRef::Extension")
376                .field(&op.family_id())
377                .finish(),
378        }
379    }
380}
381
382/// Allocation-free immutable view of one semantic operation.
383#[derive(Clone, Copy)]
384pub struct SemanticOperationView<'a> {
385    operation: &'a SemanticOperation,
386}
387
388impl<'a> SemanticOperationView<'a> {
389    pub(crate) const fn new(operation: &'a SemanticOperation) -> Self {
390        Self { operation }
391    }
392
393    /// Borrow the semantic operation payload.
394    pub fn op(self) -> SemanticOpRef<'a> {
395        match &self.operation.op {
396            SemanticOp::Core(op) => SemanticOpRef::Core(op),
397            SemanticOp::Extension(op) => SemanticOpRef::Extension(op.as_ref()),
398        }
399    }
400
401    /// Borrow ordered SSA inputs.
402    pub fn inputs(self) -> &'a [ProgramValue] {
403        &self.operation.inputs
404    }
405
406    /// Borrow ordered SSA outputs.
407    pub fn outputs(self) -> &'a [ProgramValue] {
408        &self.operation.outputs
409    }
410
411    /// Borrow ordered observable effects.
412    pub fn effects(self) -> &'a [Effect] {
413        &self.operation.effects
414    }
415
416    /// Borrow output alias declarations.
417    pub fn aliases(self) -> &'a [Alias] {
418        &self.operation.aliases
419    }
420
421    /// Borrow operation-local symbolic guards.
422    pub fn shape_guards(self) -> &'a [ShapeGuard] {
423        &self.operation.shape_guards
424    }
425
426    /// Return bounded diagnostic provenance without source identities.
427    pub fn provenance(self) -> SemanticProvenanceView<'a> {
428        self.operation.provenance.view()
429    }
430
431    /// Return the unresolved placement constraint.
432    pub fn placement(self) -> SemanticPlacementConstraint {
433        self.operation.placement
434    }
435}
436
437impl std::fmt::Debug for SemanticOperationView<'_> {
438    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
439        formatter
440            .debug_struct("SemanticOperationView")
441            .field("op", &self.op())
442            .field("inputs", &self.inputs().len())
443            .field("outputs", &self.outputs().len())
444            .field("effects", &self.effects().len())
445            .field("aliases", &self.aliases().len())
446            .field("shape_guards", &self.shape_guards().len())
447            .finish()
448    }
449}
tenferro_runtime/program/op.rs

tenferro_runtime/program/
op.rs
