use std::any::Any;
use std::fmt::Debug;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use computegraph::graph::GraphBuilder;
#[cfg(not(feature = "autodiff"))]
use computegraph::types::ValueRef;
#[cfg(feature = "autodiff")]
use computegraph::types::{LocalValueId, OperationRole, ValueKey, ValueRef};
use tenferro_tensor::{DType, Tensor};
#[cfg(feature = "autodiff")]
use tidu::{ADRuleError, ADRuleKind, ADRuleResult};
#[cfg(feature = "autodiff")]
use crate::ad::context::ShapeGuardContext;
#[cfg(feature = "autodiff")]
use crate::ad::PrimitiveRuleBuilder;
use crate::std_tensor_op::StdTensorOp;
use crate::sym_dim::SymDim;
#[cfg(feature = "autodiff")]
use std::collections::HashMap;
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("{message}")]
pub struct ExtensionLoweringError {
message: String,
}
impl ExtensionLoweringError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
pub type ExtensionLoweringResult =
std::result::Result<Option<Vec<ValueRef<StdTensorOp>>>, ExtensionLoweringError>;
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 infer_output_meta(
&self,
input_dtypes: &[DType],
input_shapes: &[&[SymDim]],
) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>>;
fn eager_execute(&self, inputs: &[&Tensor]) -> tenferro_tensor::Result<Vec<Tensor>>;
fn lower_to_standard_ops(
&self,
_builder: &mut GraphBuilder<StdTensorOp>,
_inputs: &[ValueRef<StdTensorOp>],
_input_dtypes: &[DType],
_input_shapes: &[&[SymDim]],
) -> ExtensionLoweringResult {
Ok(None)
}
}
#[cfg(feature = "autodiff")]
pub trait ExtensionAdRule: Debug + Send + Sync + 'static {
fn family_id(&self) -> &'static str;
fn linearize(
&self,
op: &dyn ExtensionOp,
builder: &mut dyn PrimitiveRuleBuilder,
primal_in: &[ValueKey<StdTensorOp>],
primal_out: &[ValueKey<StdTensorOp>],
tangent_in: &[Option<LocalValueId>],
ctx: &mut ShapeGuardContext,
) -> ADRuleResult<Vec<Option<LocalValueId>>>;
fn transpose_rule(
&self,
op: &dyn ExtensionOp,
builder: &mut dyn PrimitiveRuleBuilder,
cotangent_out: &[Option<LocalValueId>],
inputs: &[ValueRef<StdTensorOp>],
mode: &OperationRole,
ctx: &mut ShapeGuardContext,
) -> ADRuleResult<Vec<Option<LocalValueId>>>;
}
#[cfg(feature = "autodiff")]
#[derive(Debug, thiserror::Error)]
pub enum ExtensionRegistryError {
#[error("AD rule for family_id {family_id:?} already registered")]
DuplicateRule { family_id: &'static str },
#[error("family_id {family_id:?} does not match the namespaced format")]
MalformedFamilyId { family_id: &'static str },
}
#[cfg(feature = "autodiff")]
type RuleMap = HashMap<&'static str, Arc<dyn ExtensionAdRule>>;
#[cfg(feature = "autodiff")]
#[derive(Clone, Default)]
pub struct ExtensionRuleSet {
rules: Arc<RuleMap>,
}
#[cfg(feature = "autodiff")]
impl std::fmt::Debug for ExtensionRuleSet {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut families: Vec<_> = self.rules.keys().copied().collect();
families.sort_unstable();
f.debug_struct("ExtensionRuleSet")
.field("families", &families)
.finish()
}
}
#[cfg(feature = "autodiff")]
impl ExtensionRuleSet {
pub fn new() -> Self {
Self::default()
}
pub fn register_rule(
&mut self,
rule: Arc<dyn ExtensionAdRule>,
) -> Result<(), ExtensionRegistryError> {
let family_id = rule.family_id();
validate_rule_insert(&self.rules, family_id)?;
let rules = Arc::make_mut(&mut self.rules);
rules.insert(family_id, rule);
Ok(())
}
pub fn with_rule(
mut self,
rule: Arc<dyn ExtensionAdRule>,
) -> Result<Self, ExtensionRegistryError> {
self.register_rule(rule)?;
Ok(self)
}
pub fn merge(&mut self, other: ExtensionRuleSet) -> Result<(), ExtensionRegistryError> {
let mut rules = (*self.rules).clone();
for rule in other.rules.values() {
insert_rule(&mut rules, Arc::clone(rule))?;
}
self.rules = Arc::new(rules);
Ok(())
}
pub fn lookup_rule(&self, family_id: &str) -> Option<Arc<dyn ExtensionAdRule>> {
self.rules.get(family_id).cloned()
}
pub fn is_rule_registered(&self, family_id: &str) -> bool {
self.rules.contains_key(family_id)
}
}
#[cfg(feature = "autodiff")]
fn is_valid_family_id(family_id: &str) -> bool {
let mut parts = family_id.rsplitn(2, '.');
let Some(version_part) = parts.next() else {
return false;
};
let Some(prefix) = parts.next() else {
return false;
};
if !version_part.starts_with('v') {
return false;
}
let digits = &version_part[1..];
if digits.is_empty() || !digits.chars().all(|c| c.is_ascii_digit()) {
return false;
}
let Some((crate_name, op_name)) = prefix.split_once('.') else {
return false;
};
if crate_name.is_empty() || op_name.is_empty() {
return false;
}
let any_invalid = |s: &str| s.chars().any(|c| c.is_whitespace() || !c.is_ascii());
if any_invalid(crate_name) || any_invalid(op_name) {
return false;
}
true
}
#[cfg(feature = "autodiff")]
fn insert_rule(
rules: &mut RuleMap,
rule: Arc<dyn ExtensionAdRule>,
) -> Result<(), ExtensionRegistryError> {
let family_id = rule.family_id();
validate_rule_insert(rules, family_id)?;
rules.insert(family_id, rule);
Ok(())
}
#[cfg(feature = "autodiff")]
fn validate_rule_insert(
rules: &RuleMap,
family_id: &'static str,
) -> Result<(), ExtensionRegistryError> {
if !is_valid_family_id(family_id) {
return Err(ExtensionRegistryError::MalformedFamilyId { family_id });
}
if rules.contains_key(family_id) {
return Err(ExtensionRegistryError::DuplicateRule { family_id });
}
Ok(())
}
#[cfg(feature = "autodiff")]
pub fn linearize_extension_rule(
op: &dyn ExtensionOp,
builder: &mut dyn PrimitiveRuleBuilder,
primal_in: &[ValueKey<StdTensorOp>],
primal_out: &[ValueKey<StdTensorOp>],
tangent_in: &[Option<LocalValueId>],
ctx: &mut ShapeGuardContext,
) -> ADRuleResult<Vec<Option<LocalValueId>>> {
match ctx.extension_rule_for(op.family_id()) {
Some(rule) => rule.linearize(op, builder, primal_in, primal_out, tangent_in, ctx),
None => Err(ADRuleError::unsupported(op.family_id(), ADRuleKind::Jvp)),
}
}
#[cfg(feature = "autodiff")]
pub fn transpose_extension_rule(
op: &dyn ExtensionOp,
builder: &mut dyn PrimitiveRuleBuilder,
cotangent_out: &[Option<LocalValueId>],
inputs: &[ValueRef<StdTensorOp>],
mode: &OperationRole,
ctx: &mut ShapeGuardContext,
) -> ADRuleResult<Vec<Option<LocalValueId>>> {
match ctx.extension_rule_for(op.family_id()) {
Some(rule) => rule.transpose_rule(op, builder, cotangent_out, inputs, mode, ctx),
None => Err(ADRuleError::unsupported(
op.family_id(),
ADRuleKind::Transpose,
)),
}
}
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)
}