use std::any::Any;
use std::error::Error as StdError;
use std::fmt::Debug;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use computegraph::graph::GraphBuilder;
use computegraph::types::ValueRef;
use tenferro_tensor::{DType, ErrorKind, ValidationKind};
use crate::std_tensor_op::StdTensorOp;
use crate::sym_dim::SymDim;
use crate::ExtensionShapeContext;
#[doc(hidden)]
pub use crate::shape_constraint::ExtensionShapeConstraint;
#[doc(hidden)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ExtensionShapeInference {
pub output_metas: Vec<(DType, Vec<SymDim>)>,
pub constraints: Vec<ExtensionShapeConstraint>,
}
#[doc(hidden)]
pub fn invoke_extension_shape_inference(
op: &dyn ExtensionOp,
input_dtypes: &[DType],
input_shapes: &[&[SymDim]],
) -> tenferro_tensor::Result<ExtensionShapeInference> {
let expected_inputs = op.input_count();
if input_dtypes.len() != expected_inputs || input_shapes.len() != expected_inputs {
return Err(tenferro_tensor::Error::invalid_argument(
"extension",
"input metadata",
format!(
"family_id={:?}: infer_output_meta expects {expected_inputs} input metadata entries, got {} dtypes and {} shapes",
op.family_id(),
input_dtypes.len(),
input_shapes.len()
),
));
}
let mut ctx =
ExtensionShapeContext::new_for_inference(op.family_id(), input_dtypes, input_shapes);
let output_metas = op.infer_output_meta(&mut ctx)?;
if output_metas.len() != op.output_count() {
return Err(tenferro_tensor::Error::invalid_argument(
"extension",
"output metadata",
format!(
"family_id={:?}: infer_output_meta produced {} output metadata entries; op declared {} outputs",
op.family_id(),
output_metas.len(),
op.output_count()
),
));
}
Ok(ExtensionShapeInference {
output_metas,
constraints: ctx.into_constraints(),
})
}
#[derive(Debug, thiserror::Error)]
pub enum ExtensionLoweringError {
#[error("{message}")]
Message {
message: String,
kind: ErrorKind,
},
#[error("{source}")]
Source {
kind: ErrorKind,
#[source]
source: Box<dyn StdError + Send + Sync + 'static>,
},
}
impl ExtensionLoweringError {
pub fn new(message: impl Into<String>) -> Self {
Self::new_with_kind(
ErrorKind::Validation(ValidationKind::InvalidArgument),
message,
)
}
pub fn new_with_kind(kind: ErrorKind, message: impl Into<String>) -> Self {
Self::Message {
message: message.into(),
kind,
}
}
pub fn from_source<E>(source: E) -> Self
where
E: StdError + Send + Sync + 'static,
{
Self::from_source_with_kind(
ErrorKind::Validation(ValidationKind::InvalidArgument),
source,
)
}
pub fn from_source_with_kind<E>(kind: ErrorKind, source: E) -> Self
where
E: StdError + Send + Sync + 'static,
{
Self::Source {
kind,
source: Box::new(source),
}
}
pub fn kind(&self) -> ErrorKind {
match self {
Self::Message { kind, .. } | Self::Source { kind, .. } => *kind,
}
}
}
pub type ExtensionLoweringResult =
std::result::Result<ExtensionStandardLowering, ExtensionLoweringError>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ExtensionStandardLowering {
Lowered(Vec<ValueRef<StdTensorOp>>),
Unsupported,
}
pub trait ExtensionOp: Debug + Send + Sync + 'static {
fn family_id(&self) -> &'static str;
fn payload_hash(&self, hasher: &mut dyn Hasher);
fn payload_eq(&self, other: &dyn ExtensionOp) -> bool;
fn clone_arc(&self) -> Arc<dyn ExtensionOp>;
fn as_any(&self) -> &dyn Any;
fn input_count(&self) -> usize;
fn output_count(&self) -> usize;
fn semantic_effects(&self) -> ExtensionEffectDeclaration<'_> {
ExtensionEffectDeclaration::Undeclared
}
fn semantic_aliases(&self) -> ExtensionAliasDeclaration<'_> {
ExtensionAliasDeclaration::Undeclared
}
fn scalar_identity(&self) -> Option<&'static str> {
None
}
fn infer_output_meta(
&self,
ctx: &mut ExtensionShapeContext<'_>,
) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>>;
fn lower_to_standard_ops(
&self,
_builder: &mut GraphBuilder<StdTensorOp>,
_inputs: &[ValueRef<StdTensorOp>],
_input_dtypes: &[DType],
_input_shapes: &[&[SymDim]],
) -> ExtensionLoweringResult {
Ok(ExtensionStandardLowering::Unsupported)
}
fn prune_outputs(&self, _live_outputs: &[bool]) -> Option<Arc<dyn ExtensionOp>> {
None
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ExtensionEffectAccess {
Read,
Write,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ExtensionEffect {
pub family: &'static str,
pub key: u64,
pub access: ExtensionEffectAccess,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExtensionEffectDeclaration<'a> {
Undeclared,
Declared(&'a [ExtensionEffect]),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ExtensionAlias {
Fresh {
output: usize,
},
ViewOf {
output: usize,
input: usize,
},
MustAlias {
output: usize,
input: usize,
},
ExternalAlias {
output: usize,
family: &'static str,
key: u64,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExtensionAliasDeclaration<'a> {
Undeclared,
AllFresh,
Declared(&'a [ExtensionAlias]),
}
pub(crate) struct DynHasherProxy<'a, H: Hasher + ?Sized> {
inner: &'a mut H,
}
impl<'a, H: Hasher + ?Sized> DynHasherProxy<'a, H> {
pub(crate) fn new(inner: &'a mut H) -> Self {
Self { inner }
}
}
impl<H: Hasher + ?Sized> Hasher for DynHasherProxy<'_, H> {
fn finish(&self) -> u64 {
self.inner.finish()
}
fn write(&mut self, bytes: &[u8]) {
self.inner.write(bytes);
}
}
pub(crate) fn hash_extension<H: Hasher>(op: &(dyn ExtensionOp + '_), state: &mut H) {
op.family_id().as_bytes().hash(state);
op.payload_hash(&mut DynHasherProxy::new(state));
}
pub(crate) fn ext_op_eq(a: &dyn ExtensionOp, b: &dyn ExtensionOp) -> bool {
a.family_id() == b.family_id() && a.payload_eq(b)
}