1use std::collections::BTreeSet;
2
3use serde::{Deserialize, Deserializer, Serialize};
4
5use super::super::VNextError;
6use super::foundation::invalid_operation;
7use super::ElementType;
8
9fn is_axis_permutation(axis_order: &[u32], rank: usize) -> bool {
10 axis_order.len() == rank
11 && axis_order.iter().copied().collect::<BTreeSet<_>>()
12 == (0..rank as u32).collect::<BTreeSet<_>>()
13}
14
15#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
16#[serde(rename_all = "snake_case")]
17pub enum DimensionConstraint {
18 Exact(u64),
19 Symbol(String),
20 Range { minimum: u64, maximum: u64 },
21}
22
23#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
24#[serde(rename_all = "snake_case")]
25pub enum StrideConstraint {
26 ExactBytes(u64),
27 Symbol(String),
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
31#[serde(rename_all = "snake_case")]
32pub enum LayoutConstraint {
33 Contiguous,
34 Strided {
35 strides: Vec<StrideConstraint>,
36 },
37 Blocked {
38 block: Vec<u64>,
39 axis_order: Vec<u32>,
40 },
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(rename_all = "snake_case")]
45pub enum TensorAccess {
46 Read,
47 Write,
48 ReadWrite,
49}
50
51#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
52#[serde(rename_all = "snake_case")]
53pub enum AliasPolicy {
54 NoAlias,
55 MayAlias { tensor_index: u32 },
56 MustAlias { tensor_index: u32 },
57}
58
59#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
60pub struct TensorContract {
61 dimensions: Vec<DimensionConstraint>,
62 element_types: BTreeSet<ElementType>,
63 layouts: Vec<LayoutConstraint>,
64 access: TensorAccess,
65 alias: AliasPolicy,
66}
67
68#[derive(Deserialize)]
69#[serde(deny_unknown_fields)]
70struct TensorContractWire {
71 dimensions: Vec<DimensionConstraint>,
72 element_types: BTreeSet<ElementType>,
73 layouts: Vec<LayoutConstraint>,
74 access: TensorAccess,
75 alias: AliasPolicy,
76}
77
78impl<'de> Deserialize<'de> for TensorContract {
79 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
80 where
81 D: Deserializer<'de>,
82 {
83 let wire = TensorContractWire::deserialize(deserializer)?;
84 Self::new(
85 wire.dimensions,
86 wire.element_types,
87 wire.layouts,
88 wire.access,
89 wire.alias,
90 )
91 .map_err(serde::de::Error::custom)
92 }
93}
94
95impl TensorContract {
96 pub fn new(
97 dimensions: Vec<DimensionConstraint>,
98 element_types: BTreeSet<ElementType>,
99 mut layouts: Vec<LayoutConstraint>,
100 access: TensorAccess,
101 alias: AliasPolicy,
102 ) -> Result<Self, VNextError> {
103 layouts.sort();
104 layouts.dedup();
105 let contract = Self {
106 dimensions,
107 element_types,
108 layouts,
109 access,
110 alias,
111 };
112 contract.validate("tensor_contract")?;
113 Ok(contract)
114 }
115
116 pub fn dimensions(&self) -> &[DimensionConstraint] {
117 &self.dimensions
118 }
119
120 pub fn element_types(&self) -> &BTreeSet<ElementType> {
121 &self.element_types
122 }
123
124 pub fn layouts(&self) -> &[LayoutConstraint] {
125 &self.layouts
126 }
127
128 pub const fn access(&self) -> TensorAccess {
129 self.access
130 }
131
132 pub fn alias(&self) -> &AliasPolicy {
133 &self.alias
134 }
135
136 pub fn validate(&self, field: &str) -> Result<(), VNextError> {
137 if self.element_types.is_empty() {
138 return Err(VNextError::InvalidExecutionPlan {
139 reason: format!("{field} has no allowed element type"),
140 });
141 }
142 if self.layouts.is_empty() {
143 return Err(VNextError::InvalidExecutionPlan {
144 reason: format!("{field} has no allowed layout"),
145 });
146 }
147 for (index, dimension) in self.dimensions.iter().enumerate() {
148 match dimension {
149 DimensionConstraint::Exact(0) | DimensionConstraint::Range { minimum: 0, .. } => {
150 return Err(VNextError::InvalidExecutionPlan {
151 reason: format!("{field}.dimensions[{index}] permits a zero extent"),
152 });
153 }
154 DimensionConstraint::Range { minimum, maximum } if minimum > maximum => {
155 return Err(VNextError::InvalidExecutionPlan {
156 reason: format!("{field}.dimensions[{index}] has an inverted range"),
157 });
158 }
159 DimensionConstraint::Symbol(symbol) if symbol.trim().is_empty() => {
160 return Err(VNextError::InvalidExecutionPlan {
161 reason: format!("{field}.dimensions[{index}] has an empty symbol"),
162 });
163 }
164 _ => {}
165 }
166 }
167 for (index, layout) in self.layouts.iter().enumerate() {
168 match layout {
169 LayoutConstraint::Strided { strides }
170 if strides.len() != self.dimensions.len()
171 || strides.iter().any(|stride| match stride {
172 StrideConstraint::ExactBytes(bytes) => *bytes == 0,
173 StrideConstraint::Symbol(symbol) => symbol.trim().is_empty(),
174 }) =>
175 {
176 return Err(VNextError::InvalidExecutionPlan {
177 reason: format!("{field}.layouts[{index}] has invalid strides"),
178 });
179 }
180 LayoutConstraint::Blocked { block, axis_order }
181 if block.len() != self.dimensions.len()
182 || block.iter().any(|extent| *extent == 0)
183 || !is_axis_permutation(axis_order, self.dimensions.len()) =>
184 {
185 return Err(VNextError::InvalidExecutionPlan {
186 reason: format!("{field}.layouts[{index}] has an invalid block"),
187 });
188 }
189 _ => {}
190 }
191 }
192 Ok(())
193 }
194}
195
196#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
197#[serde(rename_all = "snake_case")]
198pub enum BlockedTensorPadding {
199 Exact,
200 ZeroFill { physical_dimensions: Vec<u64> },
201}
202
203#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
204#[serde(rename_all = "snake_case")]
205pub enum ResolvedTensorLayout {
206 Contiguous,
207 Strided {
208 byte_strides: Vec<u64>,
209 },
210 Blocked {
211 block: Vec<u64>,
212 axis_order: Vec<u32>,
213 padding: BlockedTensorPadding,
214 },
215}
216
217#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
220pub struct ResolvedTensorSpec {
221 dimensions: Vec<u64>,
222 element_type: ElementType,
223 layout: ResolvedTensorLayout,
224}
225
226#[derive(Deserialize)]
227#[serde(deny_unknown_fields)]
228struct ResolvedTensorSpecWire {
229 dimensions: Vec<u64>,
230 element_type: ElementType,
231 layout: ResolvedTensorLayout,
232}
233
234impl<'de> Deserialize<'de> for ResolvedTensorSpec {
235 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
236 where
237 D: Deserializer<'de>,
238 {
239 let wire = ResolvedTensorSpecWire::deserialize(deserializer)?;
240 Self::new(wire.dimensions, wire.element_type, wire.layout).map_err(serde::de::Error::custom)
241 }
242}
243
244impl ResolvedTensorSpec {
245 pub fn new(
246 dimensions: Vec<u64>,
247 element_type: ElementType,
248 layout: ResolvedTensorLayout,
249 ) -> Result<Self, VNextError> {
250 if dimensions.iter().any(|extent| *extent == 0) {
251 return Err(invalid_operation(
252 "resolved tensor dimensions must be non-zero",
253 ));
254 }
255 match &layout {
256 ResolvedTensorLayout::Strided { byte_strides }
257 if byte_strides.len() != dimensions.len()
258 || byte_strides.iter().any(|stride| *stride == 0) =>
259 {
260 return Err(invalid_operation(
261 "resolved tensor byte strides must match rank and be non-zero",
262 ));
263 }
264 ResolvedTensorLayout::Blocked {
265 block,
266 axis_order,
267 padding,
268 } => {
269 if block.len() != dimensions.len()
270 || block.iter().any(|extent| *extent == 0)
271 || !is_axis_permutation(axis_order, dimensions.len())
272 {
273 return Err(invalid_operation(
274 "resolved tensor block and axis order must form a non-zero ranked layout",
275 ));
276 }
277 match padding {
278 BlockedTensorPadding::Exact => {
279 if dimensions
280 .iter()
281 .zip(block)
282 .any(|(extent, block)| extent % block != 0)
283 {
284 return Err(invalid_operation(
285 "exact blocked tensors require every logical extent to be block-divisible",
286 ));
287 }
288 }
289 BlockedTensorPadding::ZeroFill {
290 physical_dimensions,
291 } => {
292 if physical_dimensions.len() != dimensions.len() {
293 return Err(invalid_operation(
294 "zero-filled blocked tensor padding must match tensor rank",
295 ));
296 }
297 let mut has_padding = false;
298 let mut padded_logical = Vec::with_capacity(dimensions.len());
299 for (logical, block) in dimensions.iter().zip(block) {
300 let expected = logical
301 .checked_add(block - 1)
302 .map(|extent| extent / block * block)
303 .ok_or_else(|| {
304 invalid_operation(
305 "zero-filled blocked tensor padding overflows u64",
306 )
307 })?;
308 padded_logical.push(expected);
309 has_padding |= expected != *logical;
310 }
311 let expected_physical = axis_order
312 .iter()
313 .map(|axis| padded_logical[*axis as usize])
314 .collect::<Vec<_>>();
315 if *physical_dimensions != expected_physical {
316 return Err(invalid_operation(
317 "zero-filled blocked tensor physical shape is not the minimal block-aligned axis permutation",
318 ));
319 }
320 if !has_padding {
321 return Err(invalid_operation(
322 "zero-filled blocked tensor layout must contain actual padding; use Exact otherwise",
323 ));
324 }
325 }
326 }
327 }
328 _ => {}
329 }
330 dimensions
331 .iter()
332 .try_fold(element_type.size_bytes(), |bytes, extent| {
333 bytes.checked_mul(*extent)
334 })
335 .ok_or_else(|| invalid_operation("resolved tensor byte size overflows u64"))?;
336 Ok(Self {
337 dimensions,
338 element_type,
339 layout,
340 })
341 }
342
343 pub fn dimensions(&self) -> &[u64] {
344 &self.dimensions
345 }
346
347 pub fn element_type(&self) -> ElementType {
348 self.element_type
349 }
350
351 pub fn layout(&self) -> &ResolvedTensorLayout {
352 &self.layout
353 }
354
355 pub fn minimum_storage_bytes(&self) -> Result<u64, VNextError> {
356 match &self.layout {
357 ResolvedTensorLayout::Contiguous => self
358 .dimensions
359 .iter()
360 .try_fold(self.element_type.size_bytes(), |bytes, extent| {
361 bytes.checked_mul(*extent)
362 })
363 .ok_or_else(|| invalid_operation("resolved tensor byte size overflows u64")),
364 ResolvedTensorLayout::Blocked { padding, .. } => {
365 let storage_dimensions = match padding {
366 BlockedTensorPadding::Exact => &self.dimensions,
367 BlockedTensorPadding::ZeroFill {
368 physical_dimensions,
369 } => physical_dimensions,
370 };
371 storage_dimensions
372 .iter()
373 .try_fold(self.element_type.size_bytes(), |bytes, extent| {
374 bytes.checked_mul(*extent)
375 })
376 .ok_or_else(|| {
377 invalid_operation("resolved blocked tensor byte size overflows u64")
378 })
379 }
380 ResolvedTensorLayout::Strided { byte_strides } => self
381 .dimensions
382 .iter()
383 .zip(byte_strides)
384 .try_fold(self.element_type.size_bytes(), |span, (extent, stride)| {
385 extent
386 .checked_sub(1)
387 .and_then(|steps| steps.checked_mul(*stride))
388 .and_then(|bytes| span.checked_add(bytes))
389 })
390 .ok_or_else(|| invalid_operation("resolved strided tensor span overflows u64")),
391 }
392 }
393}