1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4use sha2::{Digest, Sha256};
5
6use super::super::{
7 BufferUsage, CanonicalRational, CapabilityId, ContractVersion, OperationId, SemanticValue,
8 VNextError,
9};
10use super::foundation::invalid_operation;
11use super::{
12 AliasPolicy, AttributeId, AttributeSchema, DimensionConstraint, ElementType, LayoutConstraint,
13 ResolvedTensorLayout, ResolvedTensorSpec, ResolvedValueBinding, ResolvedValueRole,
14 ResolvedValueStorage, StrideConstraint, TensorAccess, TensorContract,
15};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum ResourcePresenceRequirement {
20 Forbidden,
21 Optional,
22 Required,
23}
24
25impl ResourcePresenceRequirement {
26 pub const fn accepts(self, present: bool) -> bool {
27 matches!(
28 (self, present),
29 (Self::Forbidden, false) | (Self::Optional, _) | (Self::Required, true)
30 )
31 }
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
38#[serde(deny_unknown_fields)]
39pub struct ResourceRequirements {
40 pub minimum_value_alignment_bytes: u64,
41 pub scratch: ResourcePresenceRequirement,
42 pub binding: ResourcePresenceRequirement,
45 pub persistent: ResourcePresenceRequirement,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49#[serde(rename_all = "snake_case")]
50pub enum OracleSpec {
51 Exact,
52 AbsoluteTolerance {
53 tolerance: CanonicalRational,
54 },
55 RelativeTolerance {
56 tolerance: CanonicalRational,
57 },
58 ReferenceOperation {
59 operation_id: OperationId,
60 version: ContractVersion,
61 },
62}
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
65#[serde(rename_all = "snake_case")]
66pub enum ProfilePhase {
67 Load,
68 Prepare,
69 Forward,
73 Prefill,
74 Decode,
75 Transfer,
76 Synchronize,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
80pub struct ProviderRequirement {
81 pub minimum_version: ContractVersion,
82 pub required_capabilities: BTreeSet<CapabilityId>,
83}
84
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86pub struct OperationDescriptor {
87 pub id: OperationId,
88 pub version: ContractVersion,
89 pub inputs: Vec<TensorContract>,
90 pub outputs: Vec<TensorContract>,
91 pub attributes: AttributeSchema,
92 pub resources: ResourceRequirements,
93 pub oracle: OracleSpec,
94 pub provider: ProviderRequirement,
95 pub profile_phase: ProfilePhase,
96}
97
98impl OperationDescriptor {
99 pub fn validate(&self) -> Result<(), VNextError> {
100 if self.version.major == 0 {
101 return Err(VNextError::InvalidExecutionPlan {
102 reason: format!("operation `{}` has an unstable zero major version", self.id),
103 });
104 }
105 if self.outputs.is_empty() {
106 return Err(VNextError::InvalidExecutionPlan {
107 reason: format!("operation `{}` has no outputs", self.id),
108 });
109 }
110 for (index, input) in self.inputs.iter().enumerate() {
111 input.validate(&format!("operation.{}.inputs[{index}]", self.id))?;
112 if !matches!(input.access(), TensorAccess::Read | TensorAccess::ReadWrite)
113 || !matches!(input.alias(), AliasPolicy::NoAlias)
114 {
115 return Err(VNextError::InvalidExecutionPlan {
116 reason: format!(
117 "operation `{}` input {index} has invalid access or alias semantics",
118 self.id
119 ),
120 });
121 }
122 }
123 for (index, output) in self.outputs.iter().enumerate() {
124 output.validate(&format!("operation.{}.outputs[{index}]", self.id))?;
125 if !matches!(
126 output.access(),
127 TensorAccess::Write | TensorAccess::ReadWrite
128 ) {
129 return Err(VNextError::InvalidExecutionPlan {
130 reason: format!("operation `{}` output {index} is not writable", self.id),
131 });
132 }
133 if let AliasPolicy::MayAlias { tensor_index }
134 | AliasPolicy::MustAlias { tensor_index } = output.alias()
135 {
136 if *tensor_index as usize >= self.inputs.len() {
137 return Err(VNextError::InvalidExecutionPlan {
138 reason: format!("operation `{}` output {index} aliases no input", self.id),
139 });
140 }
141 }
142 }
143 if self.resources.minimum_value_alignment_bytes == 0
144 || !self
145 .resources
146 .minimum_value_alignment_bytes
147 .is_power_of_two()
148 {
149 return Err(VNextError::InvalidExecutionPlan {
150 reason: format!("operation `{}` has invalid resource requirements", self.id),
151 });
152 }
153 match self.oracle {
154 OracleSpec::AbsoluteTolerance { tolerance }
155 | OracleSpec::RelativeTolerance { tolerance }
156 if tolerance.numerator() < 0 =>
157 {
158 return Err(VNextError::InvalidExecutionPlan {
159 reason: format!("operation `{}` has a negative oracle tolerance", self.id),
160 });
161 }
162 OracleSpec::AbsoluteTolerance { .. } | OracleSpec::RelativeTolerance { .. }
163 if self
164 .outputs
165 .iter()
166 .any(|output| output.element_types().contains(&ElementType::Bool)) =>
167 {
168 return Err(VNextError::InvalidExecutionPlan {
169 reason: format!(
170 "operation `{}` applies numeric oracle tolerance to a possible boolean output",
171 self.id
172 ),
173 });
174 }
175 _ => {}
176 }
177 if self.provider.minimum_version.major == 0 {
178 return Err(VNextError::InvalidExecutionPlan {
179 reason: format!("operation `{}` has a zero provider major version", self.id),
180 });
181 }
182 if self.provider.minimum_version.major != self.version.major {
183 return Err(VNextError::InvalidExecutionPlan {
184 reason: format!(
185 "operation `{}` version {} and provider minimum version {} have incompatible major versions",
186 self.id, self.version, self.provider.minimum_version
187 ),
188 });
189 }
190 Ok(())
191 }
192
193 pub fn fingerprint(&self) -> Result<String, VNextError> {
194 self.validate()?;
195 let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
196 context: "serialize operation descriptor",
197 message: error.to_string(),
198 })?;
199 Ok(format!("{:x}", Sha256::digest(bytes)))
200 }
201
202 pub fn validate_attributes(
203 &self,
204 values: &BTreeMap<AttributeId, SemanticValue>,
205 ) -> Result<(), VNextError> {
206 self.attributes
207 .validate_values(values, &format!("operation.{}.attributes", self.id))
208 }
209
210 pub fn validate_resolved_bindings(
211 &self,
212 bindings: &[ResolvedValueBinding],
213 ) -> Result<(), VNextError> {
214 self.validate()?;
215 if bindings.len() != self.inputs.len() + self.outputs.len() {
216 return Err(invalid_operation(format!(
217 "operation `{}` expects {} value bindings, received {}",
218 self.id,
219 self.inputs.len() + self.outputs.len(),
220 bindings.len()
221 )));
222 }
223
224 let mut dimensions = BTreeMap::<String, u64>::new();
225 let mut strides = BTreeMap::<String, u64>::new();
226 let mut positions = BTreeSet::new();
227 for (index, binding) in bindings.iter().enumerate() {
228 let expected_position = if index < self.inputs.len() {
229 (ResolvedValueRole::Input, index as u32)
230 } else {
231 (
232 ResolvedValueRole::Output,
233 (index - self.inputs.len()) as u32,
234 )
235 };
236 if (binding.role(), binding.ordinal()) != expected_position {
237 return Err(invalid_operation(format!(
238 "operation `{}` bindings are not in canonical input/output ordinal order",
239 self.id
240 )));
241 }
242 if !positions.insert((binding.role(), binding.ordinal())) {
243 return Err(invalid_operation(format!(
244 "operation `{}` contains duplicate ordinal bindings",
245 self.id
246 )));
247 }
248 if let Some(previous) = bindings[..index]
249 .iter()
250 .find(|previous| previous.value_id() == binding.value_id())
251 {
252 let repeated_readonly_input = previous.role() == ResolvedValueRole::Input
253 && binding.role() == ResolvedValueRole::Input
254 && previous.access() == TensorAccess::Read
255 && binding.access() == TensorAccess::Read
256 && previous.tensor() == binding.tensor()
257 && previous.storage() == binding.storage()
258 && previous.usage() == binding.usage();
259 if !repeated_readonly_input {
260 return Err(invalid_operation(format!(
261 "operation `{}` repeats a value outside identical read-only input slots",
262 self.id
263 )));
264 }
265 }
266 let contract = match binding.role() {
267 ResolvedValueRole::Input => self.inputs.get(binding.ordinal() as usize),
268 ResolvedValueRole::Output => self.outputs.get(binding.ordinal() as usize),
269 }
270 .ok_or_else(|| {
271 invalid_operation(format!(
272 "operation `{}` binding ordinal is out of range",
273 self.id
274 ))
275 })?;
276 if binding.access() != contract.access() || binding.alias() != contract.alias() {
277 return Err(invalid_operation(format!(
278 "operation `{}` binding access or alias differs from its contract",
279 self.id
280 )));
281 }
282 Self::validate_resolved_tensor(
283 &self.id,
284 contract,
285 binding.tensor(),
286 &mut dimensions,
287 &mut strides,
288 )?;
289 }
290 let inputs = &bindings[..self.inputs.len()];
291 let outputs = &bindings[self.inputs.len()..];
292 for (index, input) in inputs.iter().enumerate() {
293 for previous in &inputs[..index] {
294 if storage_overlaps(input.storage(), previous.storage())
295 && (input.value_id() != previous.value_id()
296 || input.access() != TensorAccess::Read
297 || previous.access() != TensorAccess::Read)
298 && !shares_only_readonly_transform_signs(input, previous)
299 {
300 return Err(invalid_operation(format!(
301 "operation `{}` shares input storage between different or writable values",
302 self.id
303 )));
304 }
305 }
306 }
307 for (index, output) in outputs.iter().enumerate() {
308 let aliased_inputs = inputs
309 .iter()
310 .enumerate()
311 .filter(|(_, input)| storage_overlaps(output.storage(), input.storage()))
312 .map(|(ordinal, _)| ordinal as u32)
313 .collect::<Vec<_>>();
314 match output.alias() {
315 AliasPolicy::NoAlias if !aliased_inputs.is_empty() => {
316 return Err(invalid_operation(format!(
317 "operation `{}` output {index} aliases despite a no-alias contract",
318 self.id
319 )));
320 }
321 AliasPolicy::MayAlias { tensor_index } => {
322 if aliased_inputs.iter().any(|ordinal| ordinal != tensor_index)
323 || (aliased_inputs.contains(tensor_index)
324 && output.storage() != inputs[*tensor_index as usize].storage())
325 {
326 return Err(invalid_operation(format!(
327 "operation `{}` output {index} partially aliases or aliases the wrong input",
328 self.id
329 )));
330 }
331 }
332 AliasPolicy::MustAlias { tensor_index }
333 if aliased_inputs != [*tensor_index]
334 || output.storage() != inputs[*tensor_index as usize].storage() =>
335 {
336 return Err(invalid_operation(format!(
337 "operation `{}` output {index} does not exactly alias its declared input",
338 self.id
339 )));
340 }
341 _ => {}
342 }
343 if outputs[..index]
344 .iter()
345 .any(|previous| storage_overlaps(output.storage(), previous.storage()))
346 {
347 return Err(invalid_operation(format!(
348 "operation `{}` output resources overlap",
349 self.id
350 )));
351 }
352 }
353 Ok(())
354 }
355
356 fn validate_resolved_tensor(
357 operation_id: &OperationId,
358 contract: &TensorContract,
359 tensor: &ResolvedTensorSpec,
360 dimensions: &mut BTreeMap<String, u64>,
361 strides: &mut BTreeMap<String, u64>,
362 ) -> Result<(), VNextError> {
363 if tensor.dimensions().len() != contract.dimensions().len()
364 || !contract.element_types().contains(&tensor.element_type())
365 {
366 return Err(invalid_operation(format!(
367 "operation `{operation_id}` resolved tensor rank or element type is incompatible"
368 )));
369 }
370 for (constraint, extent) in contract.dimensions().iter().zip(tensor.dimensions()) {
371 let compatible = match constraint {
372 DimensionConstraint::Exact(expected) => expected == extent,
373 DimensionConstraint::Range { minimum, maximum } => {
374 minimum <= extent && extent <= maximum
375 }
376 DimensionConstraint::Symbol(symbol) => match dimensions.get(symbol) {
377 Some(expected) => expected == extent,
378 None => {
379 dimensions.insert(symbol.clone(), *extent);
380 true
381 }
382 },
383 };
384 if !compatible {
385 return Err(invalid_operation(format!(
386 "operation `{operation_id}` resolved tensor violates a dimension constraint"
387 )));
388 }
389 }
390
391 let mut matched_strides = None;
392 let layout_matches =
393 contract
394 .layouts()
395 .iter()
396 .any(|layout| match (layout, tensor.layout()) {
397 (LayoutConstraint::Contiguous, ResolvedTensorLayout::Contiguous) => true,
398 (
399 LayoutConstraint::Blocked {
400 block: expected_block,
401 axis_order: expected_axis_order,
402 },
403 ResolvedTensorLayout::Blocked {
404 block: actual_block,
405 axis_order: actual_axis_order,
406 ..
407 },
408 ) => expected_block == actual_block && expected_axis_order == actual_axis_order,
409 (
410 LayoutConstraint::Strided {
411 strides: constraints,
412 },
413 ResolvedTensorLayout::Strided { byte_strides },
414 ) if constraints.len() == byte_strides.len() => {
415 let mut candidate = strides.clone();
416 let matches =
417 constraints
418 .iter()
419 .zip(byte_strides)
420 .all(|(constraint, actual)| match constraint {
421 StrideConstraint::ExactBytes(expected) => expected == actual,
422 StrideConstraint::Symbol(symbol) => match candidate.get(symbol)
423 {
424 Some(expected) => expected == actual,
425 None => {
426 candidate.insert(symbol.clone(), *actual);
427 true
428 }
429 },
430 });
431 if matches {
432 matched_strides = Some(candidate);
433 }
434 matches
435 }
436 _ => false,
437 });
438 if !layout_matches {
439 return Err(invalid_operation(format!(
440 "operation `{operation_id}` resolved tensor layout is incompatible"
441 )));
442 }
443 if let Some(candidate) = matched_strides {
444 *strides = candidate;
445 }
446 Ok(())
447 }
448}
449
450fn shares_only_readonly_transform_signs(
451 left: &ResolvedValueBinding,
452 right: &ResolvedValueBinding,
453) -> bool {
454 if left.access() != TensorAccess::Read
455 || right.access() != TensorAccess::Read
456 || left.usage() != BufferUsage::Weights
457 || right.usage() != BufferUsage::Weights
458 {
459 return false;
460 }
461 let (Some(left_weight), Some(right_weight)) = (left.weight(), right.weight()) else {
462 return false;
463 };
464 left.storage().components().iter().all(|left_component| {
465 right.storage().components().iter().all(|right_component| {
466 let overlaps = left_component.resource_id() == right_component.resource_id()
467 && left_component.offset_bytes()
468 < right_component
469 .offset_bytes()
470 .saturating_add(right_component.length_bytes())
471 && right_component.offset_bytes()
472 < left_component
473 .offset_bytes()
474 .saturating_add(left_component.length_bytes());
475 if !overlaps {
476 return true;
477 }
478 super::same_shared_transform_sign_component(
479 left_weight,
480 left_component,
481 right_weight,
482 right_component,
483 )
484 })
485 })
486}
487
488fn storage_overlaps(left: &ResolvedValueStorage, right: &ResolvedValueStorage) -> bool {
489 left.components().iter().any(|left| {
490 right.components().iter().any(|right| {
491 left.resource_id() == right.resource_id()
492 && left.offset_bytes() < right.offset_bytes().saturating_add(right.length_bytes())
493 && right.offset_bytes() < left.offset_bytes().saturating_add(left.length_bytes())
494 })
495 })
496}
497
498pub trait OperationContract: Send + Sync {
500 fn descriptor(&self) -> &OperationDescriptor;
501
502 fn validate_signature(
503 &self,
504 inputs: &[TensorContract],
505 outputs: &[TensorContract],
506 ) -> Result<(), VNextError>;
507}