Skip to main content

tenferro_runtime/program/
semantic.rs

1use std::fmt;
2use std::mem::size_of;
3use std::sync::Arc;
4
5use super::identity::SemanticIdentity;
6use super::op::{CoreSemanticOp, SemanticOp, SemanticOperation};
7use super::value::ProgramBuilderNonce;
8use super::{
9    Alias, Effect, ProgramBindings, ProgramQueryError, ProgramValue, ProgramValueMetadata,
10    SemanticFingerprint, SemanticOperationView, ShapeGuard,
11};
12
13/// Immutable backend-neutral semantic SSA program.
14pub struct SemanticProgram {
15    pub(crate) owner: ProgramBuilderNonce,
16    pub(crate) inputs: Box<[ProgramValue]>,
17    pub(crate) outputs: Box<[ProgramValue]>,
18    pub(crate) values: Box<[ProgramValueMetadata]>,
19    pub(crate) operations: Box<[SemanticOperation]>,
20    pub(crate) shape_guards: Box<[ShapeGuard]>,
21    pub(crate) identity: SemanticIdentity,
22}
23
24impl SemanticProgram {
25    /// Borrow ordered external inputs.
26    pub fn inputs(&self) -> &[ProgramValue] {
27        &self.inputs
28    }
29
30    /// Borrow ordered program outputs.
31    pub fn outputs(&self) -> &[ProgramValue] {
32        &self.outputs
33    }
34
35    /// Iterate over operations in semantic order.
36    pub fn operations(&self) -> impl ExactSizeIterator<Item = SemanticOperationView<'_>> + '_ {
37        self.operations.iter().map(SemanticOperationView::new)
38    }
39
40    /// Borrow metadata for a value owned by this program.
41    ///
42    /// # Errors
43    ///
44    /// Returns [`ProgramQueryError::ForeignValue`] for a foreign token.
45    pub fn value_metadata(
46        &self,
47        value: ProgramValue,
48    ) -> Result<&ProgramValueMetadata, ProgramQueryError> {
49        if value.owner != self.owner {
50            return Err(ProgramQueryError::ForeignValue);
51        }
52        self.values
53            .get(value.slot as usize)
54            .ok_or(ProgramQueryError::ForeignValue)
55    }
56
57    /// Borrow all symbolic guards in semantic operation order.
58    pub fn shape_guards(&self) -> &[ShapeGuard] {
59        &self.shape_guards
60    }
61
62    /// Return the cached normalized structural fingerprint.
63    pub const fn semantic_fingerprint(&self) -> SemanticFingerprint {
64        self.identity.fingerprint
65    }
66
67    /// Compare exact normalized semantics after the compact fingerprint check.
68    pub fn semantic_eq(&self, other: &Self) -> bool {
69        self.identity.exact_eq(self, &other.identity, other)
70    }
71
72    pub(crate) fn logical_retained_bytes(&self) -> Option<usize> {
73        checked_sum([
74            size_of::<SemanticProgram>(),
75            self.inputs.len().checked_mul(size_of::<ProgramValue>())?,
76            self.outputs.len().checked_mul(size_of::<ProgramValue>())?,
77            self.values
78                .len()
79                .checked_mul(size_of::<ProgramValueMetadata>())?,
80            checked_sum_options(
81                self.values
82                    .iter()
83                    .map(ProgramValueMetadata::logical_retained_bytes),
84            )?,
85            self.operations
86                .len()
87                .checked_mul(size_of::<SemanticOperation>())?,
88            checked_sum_options(
89                self.operations
90                    .iter()
91                    .map(semantic_operation_retained_bytes),
92            )?,
93            self.shape_guards
94                .len()
95                .checked_mul(size_of::<ShapeGuard>())?,
96            checked_sum_options(
97                self.shape_guards
98                    .iter()
99                    .map(ShapeGuard::logical_retained_bytes),
100            )?,
101            self.identity.ordinals_retained_bytes()?,
102        ])
103    }
104
105    #[cfg(test)]
106    pub(crate) fn set_fingerprint_for_test(&mut self, fingerprint: SemanticFingerprint) {
107        self.identity.fingerprint = fingerprint;
108    }
109
110    #[cfg(test)]
111    pub(crate) fn set_first_provenance_for_test(&mut self, label: &str) {
112        if let Some(operation) = self.operations.first_mut() {
113            operation.provenance = super::metadata::SemanticProvenance::builder(Some(label));
114        }
115    }
116
117    #[cfg(test)]
118    pub(crate) fn fingerprint_computations_for_test(&self) -> usize {
119        self.identity.fingerprint_computations
120    }
121}
122
123fn semantic_operation_retained_bytes(operation: &SemanticOperation) -> Option<usize> {
124    checked_sum([
125        semantic_op_retained_bytes(&operation.op)?,
126        operation
127            .inputs
128            .len()
129            .checked_mul(size_of::<ProgramValue>())?,
130        operation
131            .outputs
132            .len()
133            .checked_mul(size_of::<ProgramValue>())?,
134        operation.effects.len().checked_mul(size_of::<Effect>())?,
135        operation.aliases.len().checked_mul(size_of::<Alias>())?,
136        operation
137            .shape_guards
138            .len()
139            .checked_mul(size_of::<ShapeGuard>())?,
140        checked_sum_options(
141            operation
142                .shape_guards
143                .iter()
144                .map(ShapeGuard::logical_retained_bytes),
145        )?,
146        operation.provenance.view().label().map_or(0, str::len),
147    ])
148}
149
150fn semantic_op_retained_bytes(op: &SemanticOp) -> Option<usize> {
151    match op {
152        SemanticOp::Core(op) => core_semantic_op_retained_bytes(op),
153        SemanticOp::Extension(_) => Some(0),
154    }
155}
156
157fn core_semantic_op_retained_bytes(op: &CoreSemanticOp) -> Option<usize> {
158    match op {
159        CoreSemanticOp::DotGeneral { config } => dot_general_config_retained_bytes(config),
160        CoreSemanticOp::Transpose { perm } => vec_bytes::<usize>(perm.len()),
161        CoreSemanticOp::Reshape { to_shape } => dim_expr_vec_bytes(to_shape),
162        CoreSemanticOp::BroadcastInDim { shape, dims } => {
163            checked_sum([dim_expr_vec_bytes(shape)?, vec_bytes::<usize>(dims.len())?])
164        }
165        CoreSemanticOp::Constant { bytes, .. } => vec_bytes::<u8>(bytes.len()),
166        CoreSemanticOp::ReduceSum { axes }
167        | CoreSemanticOp::ReduceSumSquares { axes }
168        | CoreSemanticOp::ReduceProd { axes }
169        | CoreSemanticOp::ReduceMax { axes }
170        | CoreSemanticOp::ReduceMin { axes }
171        | CoreSemanticOp::Reverse { axes } => vec_bytes::<usize>(axes.len()),
172        CoreSemanticOp::Gather(config) => gather_config_retained_bytes(config),
173        CoreSemanticOp::GatherDynamicSliceSizes {
174            offset_dims,
175            collapsed_slice_dims,
176            start_index_map,
177            slice_sizes,
178            ..
179        } => checked_sum([
180            vec_bytes::<usize>(offset_dims.len())?,
181            vec_bytes::<usize>(collapsed_slice_dims.len())?,
182            vec_bytes::<usize>(start_index_map.len())?,
183            dim_expr_vec_bytes(slice_sizes)?,
184        ]),
185        CoreSemanticOp::Scatter(config) => scatter_config_retained_bytes(config),
186        CoreSemanticOp::Slice(config) => slice_config_retained_bytes(config),
187        CoreSemanticOp::DynamicSlice { slice_sizes } => vec_bytes::<usize>(slice_sizes.len()),
188        CoreSemanticOp::Pad(config) => pad_config_retained_bytes(config),
189        CoreSemanticOp::Add
190        | CoreSemanticOp::Sub
191        | CoreSemanticOp::Mul
192        | CoreSemanticOp::Neg
193        | CoreSemanticOp::Conj
194        | CoreSemanticOp::Convert { .. }
195        | CoreSemanticOp::Div
196        | CoreSemanticOp::Rem
197        | CoreSemanticOp::Abs
198        | CoreSemanticOp::Sign
199        | CoreSemanticOp::Maximum
200        | CoreSemanticOp::Minimum
201        | CoreSemanticOp::Compare(_)
202        | CoreSemanticOp::Select
203        | CoreSemanticOp::Clamp
204        | CoreSemanticOp::Exp
205        | CoreSemanticOp::Log
206        | CoreSemanticOp::Sin
207        | CoreSemanticOp::Cos
208        | CoreSemanticOp::Tanh
209        | CoreSemanticOp::Sqrt
210        | CoreSemanticOp::Rsqrt
211        | CoreSemanticOp::Pow
212        | CoreSemanticOp::Expm1
213        | CoreSemanticOp::Log1p
214        | CoreSemanticOp::ExtractDiag { .. }
215        | CoreSemanticOp::EmbedDiag { .. }
216        | CoreSemanticOp::Tril { .. }
217        | CoreSemanticOp::Triu { .. }
218        | CoreSemanticOp::DynamicUpdateSlice
219        | CoreSemanticOp::Concatenate { .. }
220        | CoreSemanticOp::ShapeOf { .. }
221        | CoreSemanticOp::DynamicTruncate { .. }
222        | CoreSemanticOp::PadToMatch { .. } => Some(0),
223    }
224}
225
226fn dot_general_config_retained_bytes(config: &tenferro_tensor::DotGeneralConfig) -> Option<usize> {
227    checked_sum([
228        vec_bytes::<usize>(config.lhs_contracting_dims.len())?,
229        vec_bytes::<usize>(config.rhs_contracting_dims.len())?,
230        vec_bytes::<usize>(config.lhs_batch_dims.len())?,
231        vec_bytes::<usize>(config.rhs_batch_dims.len())?,
232    ])
233}
234
235fn gather_config_retained_bytes(config: &tenferro_tensor::GatherConfig) -> Option<usize> {
236    checked_sum([
237        vec_bytes::<usize>(config.offset_dims.len())?,
238        vec_bytes::<usize>(config.collapsed_slice_dims.len())?,
239        vec_bytes::<usize>(config.start_index_map.len())?,
240        vec_bytes::<usize>(config.slice_sizes.len())?,
241    ])
242}
243
244fn scatter_config_retained_bytes(config: &tenferro_tensor::ScatterConfig) -> Option<usize> {
245    checked_sum([
246        vec_bytes::<usize>(config.update_window_dims.len())?,
247        vec_bytes::<usize>(config.inserted_window_dims.len())?,
248        vec_bytes::<usize>(config.scatter_dims_to_operand_dims.len())?,
249    ])
250}
251
252fn slice_config_retained_bytes(config: &tenferro_tensor::SliceConfig) -> Option<usize> {
253    checked_sum([
254        vec_bytes::<usize>(config.starts.len())?,
255        vec_bytes::<usize>(config.limits.len())?,
256        vec_bytes::<usize>(config.strides.len())?,
257    ])
258}
259
260fn pad_config_retained_bytes(config: &tenferro_tensor::PadConfig) -> Option<usize> {
261    checked_sum([
262        vec_bytes::<i64>(config.edge_padding_low.len())?,
263        vec_bytes::<i64>(config.edge_padding_high.len())?,
264        vec_bytes::<i64>(config.interior_padding.len())?,
265    ])
266}
267
268fn dim_expr_vec_bytes(values: &[tenferro_ops::dim_expr::DimExpr]) -> Option<usize> {
269    checked_sum([
270        vec_bytes::<tenferro_ops::dim_expr::DimExpr>(values.len())?,
271        checked_sum_options(values.iter().map(dim_expr_logical_retained_bytes))?,
272    ])
273}
274
275fn dim_expr_logical_retained_bytes(expression: &tenferro_ops::dim_expr::DimExpr) -> Option<usize> {
276    use tenferro_ops::dim_expr::DimExpr;
277
278    match expression {
279        DimExpr::Const(_) | DimExpr::InputDim { .. } => Some(0),
280        DimExpr::Add(left, right)
281        | DimExpr::Sub(left, right)
282        | DimExpr::Mul(left, right)
283        | DimExpr::FloorDiv(left, right)
284        | DimExpr::Min(left, right)
285        | DimExpr::Max(left, right) => checked_sum([
286            2usize.checked_mul(size_of::<DimExpr>())?,
287            dim_expr_logical_retained_bytes(left)?,
288            dim_expr_logical_retained_bytes(right)?,
289        ]),
290    }
291}
292
293fn vec_bytes<T>(len: usize) -> Option<usize> {
294    len.checked_mul(size_of::<T>())
295}
296
297fn checked_sum(values: impl IntoIterator<Item = usize>) -> Option<usize> {
298    values
299        .into_iter()
300        .try_fold(0usize, |sum, value| sum.checked_add(value))
301}
302
303fn checked_sum_options(values: impl IntoIterator<Item = Option<usize>>) -> Option<usize> {
304    values
305        .into_iter()
306        .try_fold(0usize, |sum, value| sum.checked_add(value?))
307}
308
309impl fmt::Debug for SemanticProgram {
310    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
311        formatter
312            .debug_struct("SemanticProgram")
313            .field("inputs", &self.inputs.len())
314            .field("outputs", &self.outputs.len())
315            .field("values", &self.values.len())
316            .field("operations", &self.operations.len())
317            .field("shape_guards", &self.shape_guards.len())
318            .finish()
319    }
320}
321
322/// One atomically frozen semantic program and its separate tensor bindings.
323#[derive(Clone)]
324pub struct FrozenProgram {
325    /// Immutable backend-neutral semantic structure.
326    pub program: Arc<SemanticProgram>,
327    /// Process-local tensor defaults and large constants.
328    pub bindings: ProgramBindings,
329}
330
331impl fmt::Debug for FrozenProgram {
332    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
333        formatter
334            .debug_struct("FrozenProgram")
335            .field("program", &self.program)
336            .field("bindings", &self.bindings)
337            .finish()
338    }
339}