tenferro_runtime/program/
semantic.rs1use 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
13pub 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 pub fn inputs(&self) -> &[ProgramValue] {
27 &self.inputs
28 }
29
30 pub fn outputs(&self) -> &[ProgramValue] {
32 &self.outputs
33 }
34
35 pub fn operations(&self) -> impl ExactSizeIterator<Item = SemanticOperationView<'_>> + '_ {
37 self.operations.iter().map(SemanticOperationView::new)
38 }
39
40 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 pub fn shape_guards(&self) -> &[ShapeGuard] {
59 &self.shape_guards
60 }
61
62 pub const fn semantic_fingerprint(&self) -> SemanticFingerprint {
64 self.identity.fingerprint
65 }
66
67 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#[derive(Clone)]
324pub struct FrozenProgram {
325 pub program: Arc<SemanticProgram>,
327 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}