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::Erf
215        | CoreSemanticOp::ExtractDiag { .. }
216        | CoreSemanticOp::EmbedDiag { .. }
217        | CoreSemanticOp::Tril { .. }
218        | CoreSemanticOp::Triu { .. }
219        | CoreSemanticOp::DynamicUpdateSlice
220        | CoreSemanticOp::Concatenate { .. }
221        | CoreSemanticOp::ShapeOf { .. }
222        | CoreSemanticOp::DynamicTruncate { .. }
223        | CoreSemanticOp::PadToMatch { .. } => Some(0),
224    }
225}
226
227#[cfg(test)]
228mod tests;
229
230fn dot_general_config_retained_bytes(config: &tenferro_tensor::DotGeneralConfig) -> Option<usize> {
231    // Inline axes are part of the enclosing operation's size.
232    [
233        &config.lhs_contracting_dims,
234        &config.rhs_contracting_dims,
235        &config.lhs_batch_dims,
236        &config.rhs_batch_dims,
237    ]
238    .into_iter()
239    .filter(|axes| axes.spilled())
240    .try_fold(0usize, |sum, axes| {
241        sum.checked_add(vec_bytes::<usize>(axes.len())?)
242    })
243}
244
245fn gather_config_retained_bytes(config: &tenferro_tensor::GatherConfig) -> Option<usize> {
246    checked_sum([
247        vec_bytes::<usize>(config.offset_dims.len())?,
248        vec_bytes::<usize>(config.collapsed_slice_dims.len())?,
249        vec_bytes::<usize>(config.start_index_map.len())?,
250        vec_bytes::<usize>(config.slice_sizes.len())?,
251    ])
252}
253
254fn scatter_config_retained_bytes(config: &tenferro_tensor::ScatterConfig) -> Option<usize> {
255    checked_sum([
256        vec_bytes::<usize>(config.update_window_dims.len())?,
257        vec_bytes::<usize>(config.inserted_window_dims.len())?,
258        vec_bytes::<usize>(config.scatter_dims_to_operand_dims.len())?,
259    ])
260}
261
262fn slice_config_retained_bytes(config: &tenferro_tensor::SliceConfig) -> Option<usize> {
263    checked_sum([
264        vec_bytes::<usize>(config.starts.len())?,
265        vec_bytes::<usize>(config.limits.len())?,
266        vec_bytes::<usize>(config.strides.len())?,
267    ])
268}
269
270fn pad_config_retained_bytes(config: &tenferro_tensor::PadConfig) -> Option<usize> {
271    checked_sum([
272        vec_bytes::<i64>(config.edge_padding_low.len())?,
273        vec_bytes::<i64>(config.edge_padding_high.len())?,
274        vec_bytes::<i64>(config.interior_padding.len())?,
275    ])
276}
277
278fn dim_expr_vec_bytes(values: &[tenferro_ops::dim_expr::DimExpr]) -> Option<usize> {
279    checked_sum([
280        vec_bytes::<tenferro_ops::dim_expr::DimExpr>(values.len())?,
281        checked_sum_options(values.iter().map(dim_expr_logical_retained_bytes))?,
282    ])
283}
284
285fn dim_expr_logical_retained_bytes(expression: &tenferro_ops::dim_expr::DimExpr) -> Option<usize> {
286    use tenferro_ops::dim_expr::DimExpr;
287
288    match expression {
289        DimExpr::Const(_) | DimExpr::InputDim { .. } => Some(0),
290        DimExpr::Add(left, right)
291        | DimExpr::Sub(left, right)
292        | DimExpr::Mul(left, right)
293        | DimExpr::FloorDiv(left, right)
294        | DimExpr::Min(left, right)
295        | DimExpr::Max(left, right) => checked_sum([
296            2usize.checked_mul(size_of::<DimExpr>())?,
297            dim_expr_logical_retained_bytes(left)?,
298            dim_expr_logical_retained_bytes(right)?,
299        ]),
300    }
301}
302
303fn vec_bytes<T>(len: usize) -> Option<usize> {
304    len.checked_mul(size_of::<T>())
305}
306
307fn checked_sum(values: impl IntoIterator<Item = usize>) -> Option<usize> {
308    values
309        .into_iter()
310        .try_fold(0usize, |sum, value| sum.checked_add(value))
311}
312
313fn checked_sum_options(values: impl IntoIterator<Item = Option<usize>>) -> Option<usize> {
314    values
315        .into_iter()
316        .try_fold(0usize, |sum, value| sum.checked_add(value?))
317}
318
319impl fmt::Debug for SemanticProgram {
320    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
321        formatter
322            .debug_struct("SemanticProgram")
323            .field("inputs", &self.inputs.len())
324            .field("outputs", &self.outputs.len())
325            .field("values", &self.values.len())
326            .field("operations", &self.operations.len())
327            .field("shape_guards", &self.shape_guards.len())
328            .finish()
329    }
330}
331
332/// One atomically frozen semantic program and its separate tensor bindings.
333#[derive(Clone)]
334pub struct FrozenProgram {
335    /// Immutable backend-neutral semantic structure.
336    pub program: Arc<SemanticProgram>,
337    /// Process-local tensor defaults and large constants.
338    pub bindings: ProgramBindings,
339}
340
341impl fmt::Debug for FrozenProgram {
342    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
343        formatter
344            .debug_struct("FrozenProgram")
345            .field("program", &self.program)
346            .field("bindings", &self.bindings)
347            .finish()
348    }
349}
tenferro_runtime/program/semantic.rs

tenferro_runtime/program/
semantic.rs
