Skip to main content

ferrum_interfaces/vnext/operation/
descriptor.rs

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/// Shape-independent resource contract. Concrete byte counts, scopes, and
35/// alignment are produced by the selected provider's versioned estimator and
36/// bound into the immutable execution plan.
37#[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    /// Small request-shaped control workspace whose contents are written in
43    /// the wave binding preamble and consumed by reusable compute.
44    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    /// Backend operation shared by prefill and decode. The exact request phase
70    /// is derived from the bound work shape rather than changing operation
71    /// identity or selecting another provider in the hot path.
72    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
498/// Object-safe semantic operation contract used while building a plan.
499pub 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}