use std::sync::Arc;
use tenferro_ops::dim_expr::DimExpr;
use tenferro_ops::ext_op::{
ExtensionAlias, ExtensionAliasDeclaration, ExtensionEffectAccess, ExtensionEffectDeclaration,
ExtensionOp,
};
use tenferro_ops::shape_extent::ShapeExtent;
use tenferro_tensor::Tensor;
use crate::checkpoint::RetainedValue;
use super::bindings::PendingBinding;
use super::identity::SemanticIdentity;
use super::metadata::SemanticProvenance;
use super::op::{SemanticOp, SemanticOperation};
use super::value::ProgramBuilderNonce;
use super::{
Alias, BindingKey, CoreSemanticOp, Effect, EffectAccess, EffectResource, FrozenProgram,
ImportedProgramValues, ProgramBindingError, ProgramBindings, ProgramBuildError,
ProgramFinishError, ProgramImport, ProgramInputSpec, ProgramShapeRelation,
ProgramStructuralError, ProgramValue, ProgramValueMetadata, SemanticPlacementConstraint,
SemanticProgram, ShapeGuard,
};
pub struct SemanticProgramBuilder {
owner: ProgramBuilderNonce,
inputs: Vec<ProgramValue>,
input_specs: Vec<ProgramInputSpec>,
values: Vec<ProgramValueMetadata>,
operations: Vec<SemanticOperation>,
bindings: Vec<PendingBinding>,
}
impl Default for SemanticProgramBuilder {
fn default() -> Self {
Self::new()
}
}
impl SemanticProgramBuilder {
pub fn new() -> Self {
Self {
owner: ProgramBuilderNonce::fresh(),
inputs: Vec::new(),
input_specs: Vec::new(),
values: Vec::new(),
operations: Vec::new(),
bindings: Vec::new(),
}
}
pub fn bind_input(
&mut self,
input: ProgramValue,
tensor: Tensor,
) -> Result<BindingKey, ProgramBuildError> {
self.validate_value(input)?;
if !self.inputs.contains(&input) {
return Err(ProgramBuildError::BindingTargetNotInput);
}
if self.bindings.iter().any(|binding| binding.input == input) {
return Err(ProgramBuildError::DuplicateBinding);
}
self.bind_input_retained(input, Arc::new(RetainedValue::from_tensor(tensor)))
}
#[doc(hidden)]
pub fn bind_input_retained(
&mut self,
input: ProgramValue,
tensor: Arc<RetainedValue>,
) -> Result<BindingKey, ProgramBuildError> {
self.validate_value(input)?;
if !self.inputs.contains(&input) {
return Err(ProgramBuildError::BindingTargetNotInput);
}
if self.bindings.iter().any(|binding| binding.input == input) {
return Err(ProgramBuildError::DuplicateBinding);
}
let key = BindingKey::new(input.slot, self.owner);
self.bindings.push(PendingBinding { key, input, tensor });
Ok(key)
}
pub fn input(&mut self, spec: ProgramInputSpec) -> Result<ProgramValue, ProgramBuildError> {
require_scalar_identity(spec.metadata(), "a program input")?;
let slot = self.next_value_slot()?;
let value = ProgramValue::new(slot, self.owner);
self.values.push(spec.metadata().clone());
self.inputs.push(value);
self.input_specs.push(spec);
Ok(value)
}
pub fn validate_value(&self, value: ProgramValue) -> Result<(), ProgramBuildError> {
if value.owner != self.owner || value.slot as usize >= self.values.len() {
return Err(ProgramBuildError::ForeignValue);
}
Ok(())
}
pub fn value_metadata(
&self,
value: ProgramValue,
) -> Result<&ProgramValueMetadata, ProgramBuildError> {
self.validate_value(value)?;
Ok(&self.values[value.slot as usize])
}
pub fn operation_count(&self) -> usize {
self.operations.len()
}
#[cfg(test)]
pub(crate) fn add_shape_guards_to_output(
&mut self,
output: ProgramValue,
guards: impl IntoIterator<Item = ShapeGuard>,
) -> Result<(), ProgramBuildError> {
self.validate_value(output)?;
let operation = self
.operations
.iter_mut()
.find(|operation| operation.outputs.contains(&output))
.ok_or(ProgramBuildError::GuardTargetNotOperationOutput)?;
let mut combined = operation.shape_guards.to_vec();
combined.extend(guards);
operation.shape_guards = combined.into_boxed_slice();
Ok(())
}
pub fn import(
&mut self,
request: ProgramImport<'_>,
) -> Result<ImportedProgramValues, ProgramBuildError> {
let transaction = ImportTransaction::prepare(self, request)?;
let roots = transaction.roots.clone();
self.inputs.extend(transaction.inputs);
self.input_specs.extend(transaction.input_specs);
self.values.extend(transaction.values);
self.operations.extend(transaction.operations);
self.bindings.extend(transaction.bindings);
Ok(ImportedProgramValues::new(roots))
}
pub(crate) fn snapshot_pruned(
&self,
outputs: &[ProgramValue],
) -> Result<(FrozenProgram, Vec<ProgramValue>), ProgramFinishError> {
self.snapshot_pruned_with_aliases_and_specs(outputs, &[], &[])
}
pub(crate) fn snapshot_pruned_with_input_specs(
&self,
outputs: &[ProgramValue],
input_specs: &[(usize, ProgramInputSpec)],
) -> Result<(FrozenProgram, Vec<ProgramValue>), ProgramFinishError> {
self.snapshot_pruned_with_aliases_and_specs(outputs, &[], input_specs)
}
pub(crate) fn snapshot_pruned_with_aliases(
&self,
outputs: &[ProgramValue],
aliases: &[crate::trace::CheckpointAlias],
) -> Result<(FrozenProgram, Vec<ProgramValue>), ProgramFinishError> {
self.snapshot_pruned_with_aliases_and_specs(outputs, aliases, &[])
}
fn snapshot_pruned_with_aliases_and_specs(
&self,
outputs: &[ProgramValue],
aliases: &[crate::trace::CheckpointAlias],
input_specs: &[(usize, ProgramInputSpec)],
) -> Result<(FrozenProgram, Vec<ProgramValue>), ProgramFinishError> {
let value_count = self.values.len();
let mut live = vec![false; value_count];
for output in outputs {
if output.owner != self.owner || output.slot as usize >= value_count {
return Err(ProgramFinishError::ForeignOutput);
}
live[output.slot as usize] = true;
}
let mut live_operation = vec![false; self.operations.len()];
for (index, operation) in self.operations.iter().enumerate() {
if operation.effects.is_empty() {
continue;
}
live_operation[index] = true;
for output in &operation.outputs {
live[output.slot as usize] = true;
}
}
for alias in aliases {
let slot = alias.input.slot as usize;
if live[slot] {
live[alias.target.slot as usize] = true;
}
}
let alias_of = |value: ProgramValue| {
aliases
.iter()
.find(|alias| alias.input == value)
.map(|alias| (alias.target, alias.bound_shape.clone()))
};
for _ in 0..=aliases.len() {
let mut changed = false;
for (index, operation) in self.operations.iter().enumerate().rev() {
if !live_operation[index]
&& !operation
.outputs
.iter()
.any(|output| live[output.slot as usize])
{
continue;
}
live_operation[index] = true;
for input in &operation.inputs {
if !live[input.slot as usize] {
live[input.slot as usize] = true;
changed = true;
}
}
}
for alias in aliases {
if live[alias.input.slot as usize] && !live[alias.target.slot as usize] {
live[alias.target.slot as usize] = true;
changed = true;
}
}
if !changed {
break;
}
}
loop {
let mut changed = false;
for operation in self.operations.iter().rev() {
let SemanticOp::Extension(extension) = &operation.op else {
continue;
};
if !operation.effects.is_empty() {
continue;
}
let live_outputs: Vec<bool> = operation
.outputs
.iter()
.map(|output| live[output.slot as usize])
.collect();
if !live_outputs.iter().any(|live| *live) || live_outputs.iter().all(|live| *live) {
continue;
}
if extension.prune_outputs(&live_outputs).is_some() {
continue;
}
for output in &operation.outputs {
live[output.slot as usize] = true;
}
changed = true;
}
if !changed {
break;
}
}
let mut creation_order: Vec<(usize, Option<ProgramInputSpec>)> =
Vec::with_capacity(self.inputs.len());
for (source_index, declared) in input_specs {
if self.input_specs.get(*source_index).is_some() {
creation_order.push((*source_index, Some(declared.clone())));
}
}
for index in 0..self.inputs.len() {
if input_specs.iter().any(|(declared, _)| *declared == index) {
continue;
}
creation_order.push((index, None));
}
let mut input_index = std::collections::HashMap::new();
for (source_index, _) in &creation_order {
let Some(input) = self.inputs.get(*source_index) else {
continue;
};
if !live[input.slot as usize] || alias_of(*input).is_some() {
continue;
}
input_index.insert(*source_index, input_index.len());
}
fn remap_expr(
expr: &DimExpr,
input_index: &std::collections::HashMap<usize, usize>,
) -> Option<DimExpr> {
Some(match expr {
DimExpr::Const(value) => DimExpr::Const(*value),
DimExpr::InputDim { input_idx, axis } => DimExpr::InputDim {
input_idx: *input_index.get(input_idx)?,
axis: *axis,
},
DimExpr::Add(lhs, rhs) => DimExpr::Add(
Box::new(remap_expr(lhs, input_index)?),
Box::new(remap_expr(rhs, input_index)?),
),
DimExpr::Sub(lhs, rhs) => DimExpr::Sub(
Box::new(remap_expr(lhs, input_index)?),
Box::new(remap_expr(rhs, input_index)?),
),
DimExpr::Mul(lhs, rhs) => DimExpr::Mul(
Box::new(remap_expr(lhs, input_index)?),
Box::new(remap_expr(rhs, input_index)?),
),
DimExpr::FloorDiv(lhs, rhs) => DimExpr::FloorDiv(
Box::new(remap_expr(lhs, input_index)?),
Box::new(remap_expr(rhs, input_index)?),
),
DimExpr::Min(lhs, rhs) => DimExpr::Min(
Box::new(remap_expr(lhs, input_index)?),
Box::new(remap_expr(rhs, input_index)?),
),
DimExpr::Max(lhs, rhs) => DimExpr::Max(
Box::new(remap_expr(lhs, input_index)?),
Box::new(remap_expr(rhs, input_index)?),
),
})
}
let mut rebuild = Self::new();
let mut kept_inputs = Vec::new();
let mut mapped: Vec<Option<ProgramValue>> = vec![None; value_count];
for (source_index, declared) in creation_order {
let (Some(input), Some(spec)) = (
self.inputs.get(source_index),
self.input_specs.get(source_index),
) else {
continue;
};
if !live[input.slot as usize] || alias_of(*input).is_some() {
continue;
}
kept_inputs.push(*input);
let built = match declared {
Some(declared) => declared,
None => {
let metadata = spec.metadata();
let extents: Vec<ShapeExtent<DimExpr>> = metadata
.shape()
.iter()
.map(
|extent| -> Result<ShapeExtent<DimExpr>, ProgramFinishError> {
let invalid = || ProgramFinishError::SnapshotRebuild {
source: ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidValueReference,
},
};
Ok(match extent {
ShapeExtent::Exact(expr) => ShapeExtent::Exact(
remap_expr(expr, &input_index).ok_or_else(invalid)?,
),
ShapeExtent::UpperBound(expr) => ShapeExtent::UpperBound(
remap_expr(expr, &input_index).ok_or_else(invalid)?,
),
ShapeExtent::Unknown => ShapeExtent::Unknown,
})
},
)
.collect::<Result<_, _>>()?;
let mut built = ProgramInputSpec::from_metadata(
ProgramValueMetadata::from_extents(metadata.dtype(), extents),
);
if let Some(identity) = metadata.scalar_identity() {
built = built.with_scalar_identity(identity);
}
built
}
};
let rebuilt = rebuild.input(built).map_err(snapshot_rebuild)?;
if let Some(binding) = self.bindings.iter().find(|binding| binding.input == *input) {
rebuild
.bind_input_retained(rebuilt, Arc::clone(&binding.tensor))
.map_err(snapshot_rebuild)?;
}
mapped[input.slot as usize] = Some(rebuilt);
}
let materialize_alias = |rebuild: &mut Self,
mapped: &mut Vec<Option<ProgramValue>>,
input: ProgramValue|
-> Result<(), ProgramFinishError> {
if mapped[input.slot as usize].is_some() {
return Ok(());
}
let Some((target, shape)) = alias_of(input) else {
return Err(ProgramFinishError::SnapshotRebuild {
source: ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidValueReference,
},
});
};
let Some(mapped_target) = mapped[target.slot as usize] else {
return Err(ProgramFinishError::SnapshotRebuild {
source: ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidSsaOrder,
},
});
};
let rank = shape.len();
let slice = CoreSemanticOp::Slice(tenferro_tensor::SliceConfig {
starts: vec![0; rank],
limits: shape,
strides: vec![1; rank],
});
let sliced = rebuild
.add_op(slice, &[mapped_target])
.map_err(snapshot_rebuild)?[0];
mapped[input.slot as usize] = Some(sliced);
Ok(())
};
for (index, operation) in self.operations.iter().enumerate() {
if !live_operation[index] {
continue;
}
let live_outputs: Vec<bool> = operation
.outputs
.iter()
.map(|output| live[output.slot as usize])
.collect();
for input in &operation.inputs {
if alias_of(*input).is_some() {
materialize_alias(&mut rebuild, &mut mapped, *input)?;
}
}
let inputs = operation
.inputs
.iter()
.map(|input| {
mapped[input.slot as usize].ok_or(ProgramFinishError::SnapshotRebuild {
source: ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidSsaOrder,
},
})
})
.collect::<Result<Vec<_>, _>>()?;
let rebuilt_outputs = match &operation.op {
SemanticOp::Core(core) => rebuild.add_op(core.clone(), &inputs),
SemanticOp::Extension(extension) => rebuild.add_extension(
extension
.prune_outputs(&live_outputs)
.unwrap_or_else(|| extension.clone_arc()),
&inputs,
),
}
.map_err(snapshot_rebuild)?;
discharge_static_guards(&mut rebuild);
let kept: Vec<ProgramValue> = operation
.outputs
.iter()
.zip(live_outputs.iter())
.filter(|(_, live)| **live)
.map(|(output, _)| *output)
.collect();
if rebuilt_outputs.len() != kept.len() {
return Err(ProgramFinishError::SnapshotRebuild {
source: ProgramBuildError::OutputMetadataCount {
expected: kept.len(),
actual: rebuilt_outputs.len(),
},
});
}
for (source, rebuilt) in kept.iter().zip(rebuilt_outputs.iter()) {
mapped[source.slot as usize] = Some(*rebuilt);
}
}
for output in outputs {
if alias_of(*output).is_some() {
materialize_alias(&mut rebuild, &mut mapped, *output)?;
}
}
let mapped_outputs = outputs
.iter()
.map(|output| {
mapped[output.slot as usize].ok_or(ProgramFinishError::SnapshotRebuild {
source: ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidValueReference,
},
})
})
.collect::<Result<Vec<_>, _>>()?;
Ok((rebuild.finish(&mapped_outputs)?, kept_inputs))
}
pub fn finish(self, outputs: &[ProgramValue]) -> Result<FrozenProgram, ProgramFinishError> {
if outputs
.iter()
.any(|output| output.owner != self.owner || output.slot as usize >= self.values.len())
{
return Err(ProgramFinishError::ForeignOutput);
}
validate_structure(
self.owner,
&self.inputs,
self.values.len(),
&self.operations,
)?;
validate_bindings(&self.inputs, &self.input_specs, &self.bindings)?;
let inputs = self.inputs.into_boxed_slice();
let outputs: Box<[ProgramValue]> = outputs.into();
let values = self.values.into_boxed_slice();
let operations = self.operations.into_boxed_slice();
let shape_guards: Box<[ShapeGuard]> = operations
.iter()
.flat_map(|operation| operation.shape_guards.iter().cloned())
.collect();
let identity =
SemanticIdentity::build(&inputs, &outputs, &values, &operations, &shape_guards);
let bindings = ProgramBindings::freeze(self.owner, self.bindings);
let program = SemanticProgram {
owner: self.owner,
inputs,
outputs,
values,
operations,
shape_guards,
identity,
};
Ok(FrozenProgram {
program: Arc::new(program),
bindings,
})
}
pub fn translate_input_shape_payload(
&self,
payload: &[DimExpr],
operation_inputs: &mut Vec<ProgramValue>,
) -> Vec<DimExpr> {
payload
.iter()
.map(|extent| self.translate_input_shape_extent(extent, operation_inputs))
.collect()
}
fn translate_input_shape_extent(
&self,
extent: &DimExpr,
operation_inputs: &mut Vec<ProgramValue>,
) -> DimExpr {
match extent {
DimExpr::Const(value) => DimExpr::Const(*value),
DimExpr::InputDim { input_idx, axis } => {
let Some(input) = self.inputs.get(*input_idx) else {
return extent.clone();
};
let positioned = match operation_inputs.iter().position(|value| value == input) {
Some(position) => position,
None => {
operation_inputs.push(*input);
operation_inputs.len() - 1
}
};
DimExpr::InputDim {
input_idx: positioned,
axis: *axis,
}
}
DimExpr::Add(lhs, rhs) => DimExpr::Add(
Box::new(self.translate_input_shape_extent(lhs, operation_inputs)),
Box::new(self.translate_input_shape_extent(rhs, operation_inputs)),
),
DimExpr::Sub(lhs, rhs) => DimExpr::Sub(
Box::new(self.translate_input_shape_extent(lhs, operation_inputs)),
Box::new(self.translate_input_shape_extent(rhs, operation_inputs)),
),
DimExpr::Mul(lhs, rhs) => DimExpr::Mul(
Box::new(self.translate_input_shape_extent(lhs, operation_inputs)),
Box::new(self.translate_input_shape_extent(rhs, operation_inputs)),
),
DimExpr::FloorDiv(lhs, rhs) => DimExpr::FloorDiv(
Box::new(self.translate_input_shape_extent(lhs, operation_inputs)),
Box::new(self.translate_input_shape_extent(rhs, operation_inputs)),
),
DimExpr::Min(lhs, rhs) => DimExpr::Min(
Box::new(self.translate_input_shape_extent(lhs, operation_inputs)),
Box::new(self.translate_input_shape_extent(rhs, operation_inputs)),
),
DimExpr::Max(lhs, rhs) => DimExpr::Max(
Box::new(self.translate_input_shape_extent(lhs, operation_inputs)),
Box::new(self.translate_input_shape_extent(rhs, operation_inputs)),
),
}
}
#[doc(hidden)]
pub fn inputs(&self) -> &[ProgramValue] {
&self.inputs
}
#[doc(hidden)]
pub fn resolved_extents(
&self,
value: ProgramValue,
) -> Result<Vec<ShapeExtent<DimExpr>>, ProgramBuildError> {
self.value_metadata(value)?
.shape()
.iter()
.enumerate()
.map(|(axis, _)| self.resolved_extent(value, axis))
.collect()
}
#[doc(hidden)]
pub fn resolved_extent(
&self,
value: ProgramValue,
axis: usize,
) -> Result<ShapeExtent<DimExpr>, ProgramBuildError> {
let metadata = self.value_metadata(value)?;
let extent =
metadata
.shape()
.get(axis)
.ok_or_else(|| ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidValueReference,
})?;
let resolve = |expr: &DimExpr| -> DimExpr {
if expr.max_input_idx().is_none() || self.bindings.is_empty() {
return expr.clone();
}
let mut resolved = expr.clone();
let mut positions = Vec::new();
collect_input_positions(expr, &mut positions);
positions.sort_unstable();
positions.dedup();
for position in positions {
let Some(input) = self.inputs.get(position) else {
continue;
};
let Some(bound) = self
.bindings
.iter()
.find(|binding| binding.input == *input)
.map(|binding| binding.tensor.shape())
else {
continue;
};
resolved = substitute_input_extents(&resolved, position, bound);
}
resolved
};
Ok(match extent {
ShapeExtent::Exact(expr) => ShapeExtent::Exact(resolve(expr)),
ShapeExtent::UpperBound(expr) => ShapeExtent::UpperBound(resolve(expr)),
ShapeExtent::Unknown => ShapeExtent::Unknown,
})
}
#[cfg(test)]
pub(crate) fn operation_views_for_test(
&self,
) -> impl ExactSizeIterator<Item = super::SemanticOperationView<'_>> + '_ {
self.operations
.iter()
.map(super::SemanticOperationView::new)
}
pub fn add_op(
&mut self,
op: CoreSemanticOp,
inputs: &[ProgramValue],
) -> Result<Box<[ProgramValue]>, ProgramBuildError> {
self.validate_inputs(inputs)?;
validate_arity(op.input_count(), inputs.len())?;
reject_core_external_dtype(&op)?;
let output_count = op.output_count();
let metadata = self.infer_core_metadata(&op, inputs)?;
validate_output_count(output_count, metadata.len())?;
let aliases = (0..output_count).map(Alias::fresh).collect();
self.append_operation(
SemanticOp::Core(op),
inputs,
metadata,
Vec::new(),
aliases,
Vec::new(),
)
}
pub fn add_extension(
&mut self,
op: Arc<dyn ExtensionOp>,
inputs: &[ProgramValue],
) -> Result<Box<[ProgramValue]>, ProgramBuildError> {
self.validate_inputs(inputs)?;
validate_arity(op.input_count(), inputs.len())?;
let effects = extension_effects(op.as_ref())?;
let aliases = extension_aliases(op.as_ref())?;
validate_aliases(&aliases, inputs.len(), op.output_count())?;
let (metadata, guards) = self.infer_extension_metadata(op.as_ref(), inputs)?;
validate_output_count(op.output_count(), metadata.len())?;
self.append_operation(
SemanticOp::Extension(op),
inputs,
metadata,
effects,
aliases,
guards,
)
}
fn next_value_slot(&self) -> Result<u32, ProgramBuildError> {
u32::try_from(self.values.len()).map_err(|_| ProgramBuildError::TooManyValues)
}
fn validate_inputs(&self, inputs: &[ProgramValue]) -> Result<(), ProgramBuildError> {
inputs
.iter()
.try_for_each(|&value| self.validate_value(value))
}
fn input_metadata(
&self,
inputs: &[ProgramValue],
) -> Result<Vec<&ProgramValueMetadata>, ProgramBuildError> {
inputs
.iter()
.map(|&value| self.value_metadata(value))
.collect()
}
fn infer_core_metadata(
&self,
op: &CoreSemanticOp,
inputs: &[ProgramValue],
) -> Result<Vec<ProgramValueMetadata>, ProgramBuildError> {
let input_metadata = self.input_metadata(inputs)?;
let precision = crate::shape_infer::extent_precision(
input_metadata.iter().flat_map(|metadata| metadata.shape()),
);
let input_dtypes: Vec<_> = input_metadata
.iter()
.map(|metadata| metadata.dtype())
.collect();
let input_shapes = inference_shapes(&input_metadata);
let input_shape_refs: Vec<_> = input_shapes.iter().map(Vec::as_slice).collect();
let standard = tenferro_ops::std_tensor_op::StdTensorOp::from(op);
let dtype = crate::shape_infer::infer_output_dtype(&standard, &input_dtypes)
.map_err(metadata_error)?;
if core_output_uses_local_shape_coordinates(op) {
let local_input_shapes: Vec<_> = input_metadata
.iter()
.enumerate()
.map(|(input_idx, metadata)| {
DimExpr::input_shape(input_idx, metadata.shape().len())
})
.collect();
let local_input_shape_refs: Vec<_> =
local_input_shapes.iter().map(Vec::as_slice).collect();
let output_extents =
crate::shape_infer::infer_output_extents(&standard, &local_input_shape_refs)
.map_err(metadata_error)?;
output_extents
.into_iter()
.map(|shape| {
resolve_inferred_extents(shape, precision, &input_shape_refs)
.map(|shape| ProgramValueMetadata::from_extents(dtype, shape))
})
.collect()
} else {
let output_extents =
crate::shape_infer::infer_output_extents(&standard, &input_shape_refs)
.map_err(metadata_error)?;
Ok(output_extents
.into_iter()
.map(|shape| {
ProgramValueMetadata::from_extents(
dtype,
crate::shape_infer::withhold_unproved_exactness(shape, precision)
.collect::<Vec<_>>(),
)
})
.collect())
}
}
fn infer_extension_metadata(
&self,
op: &dyn ExtensionOp,
inputs: &[ProgramValue],
) -> Result<(Vec<ProgramValueMetadata>, Vec<ShapeGuard>), ProgramBuildError> {
let input_metadata = self.input_metadata(inputs)?;
let precision = crate::shape_infer::extent_precision(
input_metadata.iter().flat_map(|metadata| metadata.shape()),
);
let input_dtypes: Vec<_> = input_metadata
.iter()
.map(|metadata| metadata.dtype())
.collect();
let input_shapes = inference_shapes(&input_metadata);
let input_shape_refs: Vec<_> = input_shapes.iter().map(Vec::as_slice).collect();
let inferred = crate::shape_infer::infer_extension_output_meta_with_constraints(
op,
&input_dtypes,
&input_shape_refs,
)
.map_err(metadata_error)?;
let metadata = inferred
.output_metas
.into_iter()
.map(|(dtype, shape)| {
ProgramValueMetadata::from_extents(
dtype,
crate::shape_infer::withhold_unproved_exactness(
shape.into_iter().map(ShapeExtent::Exact),
precision,
)
.collect::<Vec<_>>(),
)
})
.collect();
let guards = inferred
.constraints
.into_iter()
.map(|constraint| {
let relation = match constraint.relation {
tenferro_ops::ShapeRelation::Equal => ProgramShapeRelation::Equal,
};
ShapeGuard::new(relation, constraint.lhs, constraint.rhs)
})
.collect();
Ok((metadata, guards))
}
fn append_operation(
&mut self,
op: SemanticOp,
inputs: &[ProgramValue],
metadata: Vec<ProgramValueMetadata>,
effects: Vec<Effect>,
aliases: Vec<Alias>,
shape_guards: Vec<ShapeGuard>,
) -> Result<Box<[ProgramValue]>, ProgramBuildError> {
let identity = match &op {
SemanticOp::Core(_) => None,
SemanticOp::Extension(extension) => extension.scalar_identity(),
};
let mut metadata = metadata;
for value in &mut metadata {
if matches!(value.dtype(), tenferro_tensor::DType::External(_))
&& let Some(identity) = identity
{
*value = value.clone().with_scalar_identity(identity);
}
require_scalar_identity(value, "an operation output")?;
}
let provenance = match &op {
SemanticOp::Core(_) => SemanticProvenance::builder(None),
SemanticOp::Extension(extension) => {
SemanticProvenance::builder(Some(extension.family_id()))
}
};
let start = self.values.len();
let end = start
.checked_add(metadata.len())
.ok_or(ProgramBuildError::TooManyValues)?;
if end > u32::MAX as usize {
return Err(ProgramBuildError::TooManyValues);
}
let outputs: Box<[_]> = (start..end)
.map(|slot| ProgramValue::new(slot as u32, self.owner))
.collect();
self.values.extend(metadata);
self.operations.push(SemanticOperation {
op,
inputs: inputs.into(),
outputs: outputs.clone(),
effects: effects.into(),
aliases: aliases.into(),
shape_guards: shape_guards.into(),
placement: SemanticPlacementConstraint::any(),
provenance,
});
Ok(outputs)
}
}
fn discharge_static_guards(builder: &mut SemanticProgramBuilder) {
let Some(operation) = builder.operations.last_mut() else {
return;
};
let kept: Vec<ShapeGuard> = operation
.shape_guards
.iter()
.filter(|guard| {
let lhs = static_dim_value(guard.lhs());
let rhs = static_dim_value(guard.rhs());
match (lhs, rhs) {
(Some(lhs), Some(rhs)) => lhs != rhs,
_ => true,
}
})
.cloned()
.collect();
operation.shape_guards = kept.into_boxed_slice();
}
fn static_dim_value(expr: &DimExpr) -> Option<usize> {
Some(match expr {
DimExpr::Const(value) => *value,
DimExpr::InputDim { .. } => return None,
DimExpr::Add(lhs, rhs) => static_dim_value(lhs)?.checked_add(static_dim_value(rhs)?)?,
DimExpr::Sub(lhs, rhs) => static_dim_value(lhs)?.checked_sub(static_dim_value(rhs)?)?,
DimExpr::Mul(lhs, rhs) => static_dim_value(lhs)?.checked_mul(static_dim_value(rhs)?)?,
DimExpr::FloorDiv(lhs, rhs) => {
static_dim_value(lhs)?.checked_div(static_dim_value(rhs)?)?
}
DimExpr::Min(lhs, rhs) => static_dim_value(lhs)?.min(static_dim_value(rhs)?),
DimExpr::Max(lhs, rhs) => static_dim_value(lhs)?.max(static_dim_value(rhs)?),
})
}
fn snapshot_rebuild(source: ProgramBuildError) -> ProgramFinishError {
ProgramFinishError::SnapshotRebuild { source }
}
fn core_output_uses_local_shape_coordinates(op: &CoreSemanticOp) -> bool {
matches!(
op,
CoreSemanticOp::Reshape { .. }
| CoreSemanticOp::BroadcastInDim { .. }
| CoreSemanticOp::GatherDynamicSliceSizes { .. }
)
}
pub(super) fn reject_core_external_dtype(op: &CoreSemanticOp) -> Result<(), ProgramBuildError> {
let external = match op {
CoreSemanticOp::Convert { from, to } => [Some(*from), Some(*to)]
.into_iter()
.flatten()
.find(|dtype| matches!(dtype, tenferro_tensor::DType::External(_))),
CoreSemanticOp::Constant { dtype, .. } => {
matches!(dtype, tenferro_tensor::DType::External(_)).then_some(*dtype)
}
_ => None,
};
match external {
Some(dtype) => Err(ProgramBuildError::ExternalScalarWithoutIdentity {
dtype,
site: "a core operation",
}),
None => Ok(()),
}
}
pub(super) fn require_scalar_identity(
metadata: &ProgramValueMetadata,
site: &'static str,
) -> Result<(), ProgramBuildError> {
match (metadata.dtype(), metadata.scalar_identity()) {
(tenferro_tensor::DType::External(_), None) => {
Err(ProgramBuildError::ExternalScalarWithoutIdentity {
dtype: metadata.dtype(),
site,
})
}
_ => Ok(()),
}
}
fn keep_scalar_identity(
identity: &Option<&'static str>,
metadata: ProgramValueMetadata,
) -> ProgramValueMetadata {
match identity {
Some(identity) => metadata.with_scalar_identity(identity),
None => metadata,
}
}
fn validate_arity(expected: usize, actual: usize) -> Result<(), ProgramBuildError> {
if expected == actual {
Ok(())
} else {
Err(ProgramBuildError::Arity { expected, actual })
}
}
fn validate_output_count(expected: usize, actual: usize) -> Result<(), ProgramBuildError> {
if expected == actual {
Ok(())
} else {
Err(ProgramBuildError::OutputMetadataCount { expected, actual })
}
}
impl FrozenProgram {
#[doc(hidden)]
pub fn input_metadata_with_bound_shapes(&self) -> Box<[ProgramValueMetadata]> {
let bound_input_shapes: Vec<Option<Vec<DimExpr>>> = self
.program
.inputs
.iter()
.map(|input| {
self.bindings.tensor_for_input(*input).map(|tensor| {
tensor
.shape()
.iter()
.map(|&size| DimExpr::Const(size))
.collect()
})
})
.collect();
self.program
.inputs
.iter()
.map(|input| {
let metadata = self.program.values[input.slot as usize].clone();
ProgramValueMetadata::from_extents(
metadata.dtype(),
metadata.shape().iter().map(|extent| match extent {
ShapeExtent::Exact(expr) => ShapeExtent::Exact(
resolve_dim_expr_from_input_shapes(expr, &bound_input_shapes),
),
ShapeExtent::UpperBound(expr) => ShapeExtent::UpperBound(
resolve_dim_expr_from_input_shapes(expr, &bound_input_shapes),
),
ShapeExtent::Unknown => ShapeExtent::Unknown,
}),
)
})
.collect()
}
#[doc(hidden)]
pub fn with_input_prefix_bindings_from(
&self,
source: &FrozenProgram,
) -> Result<FrozenProgram, ProgramFinishError> {
if self.program.inputs.len() < source.program.inputs.len() {
return Err(ProgramFinishError::StructuralValidation {
source: ProgramStructuralError::InvalidValueReference,
});
}
let mut bindings = Vec::new();
for (source_input, destination_input) in
source.program.inputs.iter().zip(self.program.inputs.iter())
{
if let Some(tensor) = source.bindings.tensor_for_input(*source_input) {
bindings.push(PendingBinding {
key: BindingKey::new(destination_input.slot, self.program.owner),
input: *destination_input,
tensor,
});
}
}
let input_specs: Vec<_> = self
.program
.inputs
.iter()
.map(|input| {
ProgramInputSpec::from_metadata(self.program.values[input.slot as usize].clone())
})
.collect();
validate_bindings(&self.program.inputs, &input_specs, &bindings)?;
Ok(FrozenProgram {
program: Arc::clone(&self.program),
bindings: ProgramBindings::freeze(self.program.owner, bindings),
})
}
}
struct ImportTransaction {
inputs: Vec<ProgramValue>,
input_specs: Vec<ProgramInputSpec>,
values: Vec<ProgramValueMetadata>,
operations: Vec<SemanticOperation>,
bindings: Vec<PendingBinding>,
roots: Box<[ProgramValue]>,
}
impl ImportTransaction {
fn prepare(
destination: &SemanticProgramBuilder,
request: ProgramImport<'_>,
) -> Result<Self, ProgramBuildError> {
let source = request.program;
if let Some(bindings) = request.bindings
&& !bindings.belongs_to(source.owner)
{
return Err(ProgramBuildError::ForeignBindings);
}
for (mapped_source, mapped_destination) in request.input_map {
if mapped_source.owner != source.owner
|| mapped_source.slot as usize >= source.values.len()
{
return Err(ProgramBuildError::ForeignValue);
}
destination.validate_value(*mapped_destination)?;
}
if request
.roots
.iter()
.any(|root| root.owner != source.owner || root.slot as usize >= source.values.len())
{
return Err(ProgramBuildError::ForeignImportRoot);
}
let mut producer = vec![None; source.values.len()];
for (operation_index, operation) in source.operations.iter().enumerate() {
for output in &operation.outputs {
producer[output.slot as usize] = Some(operation_index);
}
}
let mut needed_values = vec![false; source.values.len()];
let mut needed_operations = vec![false; source.operations.len()];
let mut pending: Vec<_> = request
.roots
.iter()
.map(|root| root.slot as usize)
.collect();
if let Some(bindings) = request.bindings {
pending.extend(bindings.bound_inputs().map(|input| input.slot as usize));
}
for (operation_index, operation) in source.operations.iter().enumerate() {
if !operation.effects.is_empty() {
needed_operations[operation_index] = true;
for output in &operation.outputs {
needed_values[output.slot as usize] = true;
}
pending.extend(operation.inputs.iter().map(|input| input.slot as usize));
}
}
while let Some(slot) = pending.pop() {
if needed_values[slot] {
continue;
}
needed_values[slot] = true;
if let Some(operation_index) = producer[slot]
&& !needed_operations[operation_index]
{
needed_operations[operation_index] = true;
let operation = &source.operations[operation_index];
for output in &operation.outputs {
needed_values[output.slot as usize] = true;
}
pending.extend(operation.inputs.iter().map(|input| input.slot as usize));
}
}
loop {
let mut positions = Vec::new();
for input in &source.inputs {
if !needed_values[input.slot as usize] {
continue;
}
for extent in source.values[input.slot as usize].shape() {
if let ShapeExtent::Exact(expr) | ShapeExtent::UpperBound(expr) = extent {
collect_input_positions(expr, &mut positions);
}
}
}
for (operation, needed) in source.operations.iter().zip(&needed_operations) {
if !needed {
continue;
}
for guard in operation.shape_guards.iter() {
collect_input_positions(guard.lhs(), &mut positions);
collect_input_positions(guard.rhs(), &mut positions);
}
}
let mut changed = false;
for position in positions {
let Some(input) = source.inputs.get(position) else {
continue;
};
if !needed_values[input.slot as usize] {
needed_values[input.slot as usize] = true;
changed = true;
}
}
if !changed {
break;
}
}
let imported_input_count = source
.inputs
.iter()
.filter(|input| needed_values[input.slot as usize])
.count();
let imported_output_count: usize = source
.operations
.iter()
.zip(&needed_operations)
.filter(|(_, needed)| **needed)
.map(|(operation, _)| operation.outputs.len())
.sum();
let imported_value_count = imported_input_count
.checked_add(imported_output_count)
.ok_or(ProgramBuildError::TooManyValues)?;
let final_value_count = destination
.values
.len()
.checked_add(imported_value_count)
.ok_or(ProgramBuildError::TooManyValues)?;
if final_value_count > u32::MAX as usize {
return Err(ProgramBuildError::TooManyValues);
}
let mut transaction = Self {
inputs: Vec::with_capacity(imported_input_count),
input_specs: Vec::with_capacity(imported_input_count),
values: Vec::with_capacity(imported_value_count),
operations: Vec::with_capacity(
needed_operations.iter().filter(|needed| **needed).count(),
),
bindings: Vec::new(),
roots: Box::new([]),
};
let mapped_input: std::collections::HashMap<ProgramValue, ProgramValue> =
request.input_map.iter().copied().collect();
let mut remap = vec![None; source.values.len()];
let mut source_input_index = vec![None; source.inputs.len()];
let mut reserved = Vec::new();
for (position, &input) in source.inputs.iter().enumerate() {
if !needed_values[input.slot as usize] {
continue;
}
if let Some(&existing) = mapped_input.get(&input) {
remap[input.slot as usize] = Some(existing);
continue;
}
let slot = destination
.values
.len()
.checked_add(transaction.values.len())
.and_then(|slot| slot.checked_add(reserved.len()))
.ok_or(ProgramBuildError::TooManyValues)?;
let slot = u32::try_from(slot).map_err(|_| ProgramBuildError::TooManyValues)?;
let imported = ProgramValue::new(slot, destination.owner);
source_input_index[position] = Some(destination.inputs.len() + reserved.len());
let tensor = request
.bindings
.and_then(|bindings| bindings.tensor_for_input(input));
reserved.push((imported, &source.values[input.slot as usize], tensor));
remap[input.slot as usize] = Some(imported);
}
for (imported, source_metadata, tensor) in reserved {
let identity = source_metadata.scalar_identity();
let metadata = ProgramValueMetadata::from_extents(
source_metadata.dtype(),
source_metadata
.shape()
.iter()
.cloned()
.map(|extent| {
remap_imported_extent(
extent,
source,
&source_input_index,
&mapped_input,
destination,
)
})
.collect::<Vec<_>>(),
);
let metadata = keep_scalar_identity(&identity, metadata);
transaction.inputs.push(imported);
transaction
.input_specs
.push(ProgramInputSpec::from_metadata(metadata.clone()));
transaction.values.push(metadata);
if let Some(tensor) = tensor {
transaction.bindings.push(PendingBinding {
key: BindingKey::new(imported.slot, destination.owner),
input: imported,
tensor,
});
}
}
for (operation, needed) in source.operations.iter().zip(needed_operations) {
if !needed {
continue;
}
let inputs: Box<[_]> = operation
.inputs
.iter()
.map(|input| {
remap[input.slot as usize].ok_or(ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidSsaOrder,
})
})
.collect::<Result<_, _>>()?;
let mut outputs = Vec::with_capacity(operation.outputs.len());
for output in &operation.outputs {
let imported =
transaction.next_value(destination.values.len(), destination.owner)?;
let meta = source.values[output.slot as usize].clone();
let identity = meta.scalar_identity();
let resolved = ProgramValueMetadata::from_extents(
meta.dtype(),
meta.shape()
.iter()
.cloned()
.map(|extent| {
remap_imported_extent(
extent,
source,
&source_input_index,
&mapped_input,
destination,
)
})
.collect::<Vec<_>>(),
);
transaction
.values
.push(keep_scalar_identity(&identity, resolved));
remap[output.slot as usize] = Some(imported);
outputs.push(imported);
}
let op = match &operation.op {
SemanticOp::Core(op) => SemanticOp::Core(op.clone()),
SemanticOp::Extension(op) => SemanticOp::Extension(op.clone_arc()),
};
let guards = operation
.shape_guards
.iter()
.filter_map(|guard| {
let resolve = |expr: &DimExpr| -> Option<DimExpr> {
remap_imported_dim_expr(
expr,
source,
&source_input_index,
&mapped_input,
destination,
)
};
Some(ShapeGuard::new(
guard.relation(),
resolve(guard.lhs())?,
resolve(guard.rhs())?,
))
})
.collect::<Vec<_>>();
transaction.operations.push(SemanticOperation {
op,
inputs,
outputs: outputs.into(),
effects: operation.effects.clone(),
aliases: operation.aliases.clone(),
shape_guards: guards.into_boxed_slice(),
placement: operation.placement,
provenance: operation.provenance.clone(),
});
}
transaction.roots = request
.roots
.iter()
.map(|root| {
remap[root.slot as usize].ok_or(ProgramBuildError::InvalidImport {
source: ProgramStructuralError::InvalidValueReference,
})
})
.collect::<Result<_, _>>()?;
Ok(transaction)
}
fn next_value(
&self,
destination_value_count: usize,
owner: ProgramBuilderNonce,
) -> Result<ProgramValue, ProgramBuildError> {
let slot = destination_value_count
.checked_add(self.values.len())
.ok_or(ProgramBuildError::TooManyValues)?;
let slot = u32::try_from(slot).map_err(|_| ProgramBuildError::TooManyValues)?;
Ok(ProgramValue::new(slot, owner))
}
}
fn substitute_input_extents(expr: &DimExpr, position: usize, bound: &[usize]) -> DimExpr {
match expr {
DimExpr::Const(value) => DimExpr::Const(*value),
DimExpr::InputDim { input_idx, axis } if *input_idx == position => bound
.get(*axis)
.map_or_else(|| expr.clone(), |size| DimExpr::Const(*size)),
DimExpr::InputDim { .. } => expr.clone(),
DimExpr::Add(lhs, rhs) => DimExpr::Add(
Box::new(substitute_input_extents(lhs, position, bound)),
Box::new(substitute_input_extents(rhs, position, bound)),
),
DimExpr::Sub(lhs, rhs) => DimExpr::Sub(
Box::new(substitute_input_extents(lhs, position, bound)),
Box::new(substitute_input_extents(rhs, position, bound)),
),
DimExpr::Mul(lhs, rhs) => DimExpr::Mul(
Box::new(substitute_input_extents(lhs, position, bound)),
Box::new(substitute_input_extents(rhs, position, bound)),
),
DimExpr::FloorDiv(lhs, rhs) => DimExpr::FloorDiv(
Box::new(substitute_input_extents(lhs, position, bound)),
Box::new(substitute_input_extents(rhs, position, bound)),
),
DimExpr::Min(lhs, rhs) => DimExpr::Min(
Box::new(substitute_input_extents(lhs, position, bound)),
Box::new(substitute_input_extents(rhs, position, bound)),
),
DimExpr::Max(lhs, rhs) => DimExpr::Max(
Box::new(substitute_input_extents(lhs, position, bound)),
Box::new(substitute_input_extents(rhs, position, bound)),
),
}
}
fn collect_input_positions(expr: &DimExpr, positions: &mut Vec<usize>) {
match expr {
DimExpr::Const(_) => {}
DimExpr::InputDim { input_idx, .. } => positions.push(*input_idx),
DimExpr::Add(lhs, rhs)
| DimExpr::Sub(lhs, rhs)
| DimExpr::Mul(lhs, rhs)
| DimExpr::FloorDiv(lhs, rhs)
| DimExpr::Min(lhs, rhs)
| DimExpr::Max(lhs, rhs) => {
collect_input_positions(lhs, positions);
collect_input_positions(rhs, positions);
}
}
}
fn remap_imported_extent(
extent: ShapeExtent<DimExpr>,
source: &SemanticProgram,
source_input_index: &[Option<usize>],
mapped_input: &std::collections::HashMap<ProgramValue, ProgramValue>,
destination: &SemanticProgramBuilder,
) -> ShapeExtent<DimExpr> {
match extent {
ShapeExtent::Exact(expr) => {
remap_imported_dim_expr(&expr, source, source_input_index, mapped_input, destination)
.map(ShapeExtent::Exact)
.unwrap_or(ShapeExtent::Unknown)
}
ShapeExtent::UpperBound(expr) => {
remap_imported_dim_expr(&expr, source, source_input_index, mapped_input, destination)
.map(ShapeExtent::UpperBound)
.unwrap_or(ShapeExtent::Unknown)
}
ShapeExtent::Unknown => ShapeExtent::Unknown,
}
}
fn remap_imported_dim_expr(
expr: &DimExpr,
source: &SemanticProgram,
source_input_index: &[Option<usize>],
mapped_input: &std::collections::HashMap<ProgramValue, ProgramValue>,
destination: &SemanticProgramBuilder,
) -> Option<DimExpr> {
match expr {
DimExpr::Const(value) => Some(DimExpr::Const(*value)),
DimExpr::InputDim { input_idx, axis } => {
if let Some(index) = source_input_index.get(*input_idx).copied().flatten() {
return Some(DimExpr::InputDim {
input_idx: index,
axis: *axis,
});
}
let source_input = source.inputs.get(*input_idx)?;
let destination_value = mapped_input.get(source_input)?;
match destination
.values
.get(destination_value.slot as usize)?
.shape()
.get(*axis)?
{
ShapeExtent::Exact(expr) => Some(expr.clone()),
_ => None,
}
}
DimExpr::Add(lhs, rhs) => Some(DimExpr::Add(
Box::new(remap_imported_dim_expr(
lhs,
source,
source_input_index,
mapped_input,
destination,
)?),
Box::new(remap_imported_dim_expr(
rhs,
source,
source_input_index,
mapped_input,
destination,
)?),
)),
DimExpr::Sub(lhs, rhs) => Some(DimExpr::Sub(
Box::new(remap_imported_dim_expr(
lhs,
source,
source_input_index,
mapped_input,
destination,
)?),
Box::new(remap_imported_dim_expr(
rhs,
source,
source_input_index,
mapped_input,
destination,
)?),
)),
DimExpr::Mul(lhs, rhs) => Some(DimExpr::Mul(
Box::new(remap_imported_dim_expr(
lhs,
source,
source_input_index,
mapped_input,
destination,
)?),
Box::new(remap_imported_dim_expr(
rhs,
source,
source_input_index,
mapped_input,
destination,
)?),
)),
DimExpr::FloorDiv(lhs, rhs) => Some(DimExpr::FloorDiv(
Box::new(remap_imported_dim_expr(
lhs,
source,
source_input_index,
mapped_input,
destination,
)?),
Box::new(remap_imported_dim_expr(
rhs,
source,
source_input_index,
mapped_input,
destination,
)?),
)),
DimExpr::Min(lhs, rhs) => Some(DimExpr::Min(
Box::new(remap_imported_dim_expr(
lhs,
source,
source_input_index,
mapped_input,
destination,
)?),
Box::new(remap_imported_dim_expr(
rhs,
source,
source_input_index,
mapped_input,
destination,
)?),
)),
DimExpr::Max(lhs, rhs) => Some(DimExpr::Max(
Box::new(remap_imported_dim_expr(
lhs,
source,
source_input_index,
mapped_input,
destination,
)?),
Box::new(remap_imported_dim_expr(
rhs,
source,
source_input_index,
mapped_input,
destination,
)?),
)),
}
}
fn resolve_inferred_extents(
extents: impl IntoIterator<Item = ShapeExtent<DimExpr>>,
precision: crate::shape_infer::ExtentPrecision,
input_shapes: &[&[DimExpr]],
) -> Result<Vec<ShapeExtent<DimExpr>>, ProgramBuildError> {
extents
.into_iter()
.map(|extent| {
if matches!(precision, crate::shape_infer::ExtentPrecision::Unknown) {
return Ok(ShapeExtent::Unknown);
}
let resolved = match extent {
ShapeExtent::Exact(expression) => ShapeExtent::Exact(
crate::shape_infer::resolve_dim_expr_from_shapes(&expression, input_shapes)
.map_err(metadata_error)?,
),
ShapeExtent::UpperBound(expression) => ShapeExtent::UpperBound(
crate::shape_infer::resolve_dim_expr_from_shapes(&expression, input_shapes)
.map_err(metadata_error)?,
),
ShapeExtent::Unknown => ShapeExtent::Unknown,
};
Ok(match precision {
crate::shape_infer::ExtentPrecision::Exact => resolved,
crate::shape_infer::ExtentPrecision::Bounded => match resolved {
ShapeExtent::Exact(expression) | ShapeExtent::UpperBound(expression) => {
ShapeExtent::UpperBound(expression)
}
ShapeExtent::Unknown => ShapeExtent::Unknown,
},
crate::shape_infer::ExtentPrecision::Unknown => unreachable!("handled above"),
})
})
.collect()
}
fn inference_shapes(metadata: &[&ProgramValueMetadata]) -> Vec<Vec<DimExpr>> {
metadata
.iter()
.enumerate()
.map(|(input_idx, metadata)| {
metadata
.shape()
.iter()
.enumerate()
.map(|(axis, extent)| match extent {
ShapeExtent::Exact(expression) | ShapeExtent::UpperBound(expression) => {
expression.clone()
}
ShapeExtent::Unknown => DimExpr::InputDim { input_idx, axis },
})
.collect()
})
.collect()
}
fn metadata_error(source: crate::Error) -> ProgramBuildError {
ProgramBuildError::Metadata {
source: Box::new(source),
}
}
fn extension_effects(op: &dyn ExtensionOp) -> Result<Vec<Effect>, ProgramBuildError> {
let family = op.family_id();
let effects = match op.semantic_effects() {
ExtensionEffectDeclaration::Undeclared => {
return Err(ProgramBuildError::UndeclaredExtensionEffects { family })
}
ExtensionEffectDeclaration::Declared(effects) => effects,
};
effects
.iter()
.map(|effect| {
let resource = EffectResource::new(effect.family, effect.key)
.map_err(|source| ProgramBuildError::InvalidEffectResource { family, source })?;
let access = match effect.access {
ExtensionEffectAccess::Read => EffectAccess::Read,
ExtensionEffectAccess::Write => EffectAccess::Write,
};
Ok(Effect::new(resource, access))
})
.collect()
}
fn extension_aliases(op: &dyn ExtensionOp) -> Result<Vec<Alias>, ProgramBuildError> {
let family = op.family_id();
match op.semantic_aliases() {
ExtensionAliasDeclaration::Undeclared => {
Err(ProgramBuildError::UndeclaredExtensionAliases { family })
}
ExtensionAliasDeclaration::AllFresh => {
Ok((0..op.output_count()).map(Alias::fresh).collect())
}
ExtensionAliasDeclaration::Declared(aliases) => aliases
.iter()
.map(|alias| match *alias {
ExtensionAlias::Fresh { output } => Ok(Alias::fresh(output)),
ExtensionAlias::ViewOf { output, input } => Ok(Alias::view_of(output, input)),
ExtensionAlias::MustAlias { output, input } => Ok(Alias::must_alias(output, input)),
ExtensionAlias::ExternalAlias {
output,
family: resource_family,
key,
} => EffectResource::new(resource_family, key)
.map(|resource| Alias::external(output, resource))
.map_err(|source| ProgramBuildError::InvalidEffectResource { family, source }),
})
.collect(),
}
}
fn validate_aliases(
aliases: &[Alias],
input_count: usize,
output_count: usize,
) -> Result<(), ProgramBuildError> {
let mut seen = vec![false; output_count];
for &alias in aliases {
let output = alias.output();
let input = alias.input();
if output >= output_count || input.is_some_and(|input| input >= input_count) {
return Err(ProgramBuildError::AliasOutOfBounds {
output,
output_count,
input,
input_count,
});
}
if seen[output] {
return Err(ProgramBuildError::AliasCoverage {
expected: output_count,
actual: seen.iter().filter(|&&present| present).count(),
});
}
seen[output] = true;
}
let actual = seen.iter().filter(|&&present| present).count();
if actual != output_count {
return Err(ProgramBuildError::AliasCoverage {
expected: output_count,
actual,
});
}
Ok(())
}
fn validate_structure(
owner: ProgramBuilderNonce,
inputs: &[ProgramValue],
value_count: usize,
operations: &[SemanticOperation],
) -> Result<(), ProgramFinishError> {
let mut covered = vec![false; value_count];
for input in inputs {
if input.owner != owner
|| input.slot as usize >= value_count
|| covered[input.slot as usize]
{
return Err(ProgramFinishError::StructuralValidation {
source: ProgramStructuralError::InvalidValueReference,
});
}
covered[input.slot as usize] = true;
}
let mut previous_output = None;
for operation in operations {
let Some(first_output) = operation.outputs.first() else {
if operation.inputs.iter().any(|value| {
value.owner != owner
|| value.slot as usize >= value_count
|| !covered[value.slot as usize]
}) {
return Err(ProgramFinishError::StructuralValidation {
source: ProgramStructuralError::InvalidValueReference,
});
}
continue;
};
let output_start = first_output.slot as usize;
let valid_input = operation.inputs.iter().all(|value| {
value.owner == owner
&& (value.slot as usize) < output_start
&& (value.slot as usize) < value_count
&& covered[value.slot as usize]
});
let valid_output = operation.outputs.iter().enumerate().all(|(offset, value)| {
value.owner == owner
&& value.slot as usize == output_start + offset
&& (value.slot as usize) < value_count
&& !covered[value.slot as usize]
});
let ordered = previous_output.is_none_or(|previous| output_start > previous);
if !valid_input || !valid_output || !ordered {
let source = if operation
.inputs
.iter()
.chain(operation.outputs.iter())
.any(|value| value.owner != owner || value.slot as usize >= value_count)
{
ProgramStructuralError::InvalidValueReference
} else {
ProgramStructuralError::InvalidSsaOrder
};
return Err(ProgramFinishError::StructuralValidation { source });
}
for output in &operation.outputs {
covered[output.slot as usize] = true;
}
previous_output = operation.outputs.last().map(|value| value.slot as usize);
}
if covered.iter().any(|covered| !covered) {
return Err(ProgramFinishError::StructuralValidation {
source: ProgramStructuralError::InvalidSsaOrder,
});
}
Ok(())
}
fn validate_bindings(
inputs: &[ProgramValue],
input_specs: &[ProgramInputSpec],
bindings: &[PendingBinding],
) -> Result<(), ProgramFinishError> {
for binding in bindings {
let input_index = inputs
.iter()
.position(|input| *input == binding.input)
.ok_or(ProgramFinishError::BindingFinalization {
source: ProgramBindingError::InvalidTarget,
})?;
let spec = &input_specs[input_index];
let metadata = spec.metadata();
let actual_dtype = binding.tensor.dtype();
if actual_dtype != metadata.dtype() {
return Err(ProgramFinishError::BindingFinalization {
source: ProgramBindingError::DTypeMismatch {
expected: metadata.dtype(),
actual: actual_dtype,
},
});
}
let actual_shape = binding.tensor.shape();
if actual_shape.len() != metadata.shape().len() {
return Err(ProgramFinishError::BindingFinalization {
source: ProgramBindingError::RankMismatch {
expected: metadata.shape().len(),
actual: actual_shape.len(),
},
});
}
for (axis, (extent, &actual)) in
metadata.shape().iter().zip(actual_shape.iter()).enumerate()
{
match extent {
ShapeExtent::Exact(DimExpr::Const(expected)) if *expected != actual => {
return Err(ProgramFinishError::BindingFinalization {
source: ProgramBindingError::ExactExtentMismatch {
axis,
expected: *expected,
actual,
},
});
}
ShapeExtent::UpperBound(DimExpr::Const(bound)) if actual > *bound => {
return Err(ProgramFinishError::BindingFinalization {
source: ProgramBindingError::UpperBoundExceeded {
axis,
bound: *bound,
actual,
},
});
}
_ => {}
}
}
}
Ok(())
}
fn resolve_dim_expr_from_input_shapes(
expr: &DimExpr,
bound_input_shapes: &[Option<Vec<DimExpr>>],
) -> DimExpr {
match expr {
DimExpr::Const(_) => expr.clone(),
DimExpr::InputDim { input_idx, axis } => {
if let Some(Some(shape)) = bound_input_shapes.get(*input_idx)
&& let Some(dim) = shape.get(*axis)
{
return dim.clone();
}
expr.clone()
}
DimExpr::Add(a, b) => DimExpr::add(
resolve_dim_expr_from_input_shapes(a, bound_input_shapes),
resolve_dim_expr_from_input_shapes(b, bound_input_shapes),
),
DimExpr::Sub(a, b) => DimExpr::sub(
resolve_dim_expr_from_input_shapes(a, bound_input_shapes),
resolve_dim_expr_from_input_shapes(b, bound_input_shapes),
),
DimExpr::Mul(a, b) => DimExpr::mul(
resolve_dim_expr_from_input_shapes(a, bound_input_shapes),
resolve_dim_expr_from_input_shapes(b, bound_input_shapes),
),
DimExpr::FloorDiv(a, b) => DimExpr::floor_div(
resolve_dim_expr_from_input_shapes(a, bound_input_shapes),
resolve_dim_expr_from_input_shapes(b, bound_input_shapes),
),
DimExpr::Min(a, b) => DimExpr::min(
resolve_dim_expr_from_input_shapes(a, bound_input_shapes),
resolve_dim_expr_from_input_shapes(b, bound_input_shapes),
),
DimExpr::Max(a, b) => DimExpr::max(
resolve_dim_expr_from_input_shapes(a, bound_input_shapes),
resolve_dim_expr_from_input_shapes(b, bound_input_shapes),
),
}
}