use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::Arc;
use tenferro_ops::ext_op::ExtensionOp;
use tenferro_runtime::program::{
ProgramBuildError, ProgramValue, SemanticOpRef, SemanticOperationView, SemanticProgramBuilder,
SemanticProvenanceView,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AdValue {
Absent,
Value(ProgramValue),
}
impl AdValue {
#[must_use]
pub const fn value(self) -> Option<ProgramValue> {
match self {
Self::Absent => None,
Self::Value(value) => Some(value),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum SemanticAdRuleRole {
Linearize,
LinearTranspose,
PrimalVjp,
}
#[derive(Debug, thiserror::Error)]
pub enum SemanticExtensionRegistryError {
#[error("semantic extension AD {role:?} rule for family {family_id:?} is already registered")]
DuplicateRule {
family_id: &'static str,
role: SemanticAdRuleRole,
},
#[error("semantic extension AD family {family_id:?} is not namespaced and versioned")]
MalformedFamilyId {
family_id: &'static str,
},
}
#[derive(Debug, thiserror::Error)]
pub enum SemanticAdError {
#[error("semantic AD extension dispatch received a core operation")]
CoreOperation,
#[error("semantic extension family {family_id:?} has observable effects")]
EffectfulExtension {
family_id: &'static str,
},
#[error("semantic extension family {family_id:?} has no {role:?} AD rule")]
MissingRule {
family_id: &'static str,
role: SemanticAdRuleRole,
},
#[error("semantic AD field {field} expects {expected} values, got {actual}")]
Arity {
field: &'static str,
expected: usize,
actual: usize,
},
#[error("semantic AD field {field}[{index}] does not belong to the destination builder")]
ForeignValue {
field: &'static str,
index: usize,
},
#[error("semantic extension family {family_id:?} does not support {role:?}: {message}")]
Unsupported {
family_id: &'static str,
role: SemanticAdRuleRole,
message: String,
},
#[error("semantic extension family {family_id:?} {role:?} rule failed: {source}")]
Rule {
family_id: &'static str,
role: SemanticAdRuleRole,
#[source]
source: Box<dyn std::error::Error + Send + Sync + 'static>,
},
#[error("semantic extension family {family_id:?} {role:?} invariant failed: {message}")]
Invariant {
family_id: &'static str,
role: SemanticAdRuleRole,
message: String,
},
#[error("semantic extension AD program construction failed: {0}")]
Build(#[from] ProgramBuildError),
}
#[derive(Clone, Copy)]
pub struct SemanticLinearizeRequest<'a> {
op: &'a dyn ExtensionOp,
primal_inputs: &'a [ProgramValue],
primal_outputs: &'a [ProgramValue],
tangent_inputs: &'a [AdValue],
active_outputs: &'a [bool],
provenance: SemanticProvenanceView<'a>,
}
impl<'a> SemanticLinearizeRequest<'a> {
pub const fn op(self) -> &'a dyn ExtensionOp {
self.op
}
pub const fn primal_inputs(self) -> &'a [ProgramValue] {
self.primal_inputs
}
pub const fn primal_outputs(self) -> &'a [ProgramValue] {
self.primal_outputs
}
pub const fn tangent_inputs(self) -> &'a [AdValue] {
self.tangent_inputs
}
pub const fn active_outputs(self) -> &'a [bool] {
self.active_outputs
}
pub const fn provenance(self) -> SemanticProvenanceView<'a> {
self.provenance
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SemanticLinearizeResult {
tangent_outputs: Box<[AdValue]>,
residuals: Box<[ProgramValue]>,
}
impl SemanticLinearizeResult {
#[must_use]
pub fn new(
tangent_outputs: impl IntoIterator<Item = AdValue>,
residuals: impl IntoIterator<Item = ProgramValue>,
) -> Self {
Self {
tangent_outputs: tangent_outputs.into_iter().collect(),
residuals: residuals.into_iter().collect(),
}
}
pub fn tangent_outputs(&self) -> &[AdValue] {
&self.tangent_outputs
}
pub fn residuals(&self) -> &[ProgramValue] {
&self.residuals
}
}
#[derive(Clone, Copy)]
pub struct SemanticLinearTransposeRequest<'a> {
op: &'a dyn ExtensionOp,
primal_inputs: &'a [ProgramValue],
primal_outputs: &'a [ProgramValue],
cotangent_outputs: &'a [AdValue],
active_inputs: &'a [bool],
residuals: &'a [ProgramValue],
provenance: SemanticProvenanceView<'a>,
}
impl<'a> SemanticLinearTransposeRequest<'a> {
pub const fn op(self) -> &'a dyn ExtensionOp {
self.op
}
pub const fn primal_inputs(self) -> &'a [ProgramValue] {
self.primal_inputs
}
pub const fn primal_outputs(self) -> &'a [ProgramValue] {
self.primal_outputs
}
pub const fn cotangent_outputs(self) -> &'a [AdValue] {
self.cotangent_outputs
}
pub const fn active_inputs(self) -> &'a [bool] {
self.active_inputs
}
pub const fn residuals(self) -> &'a [ProgramValue] {
self.residuals
}
pub const fn provenance(self) -> SemanticProvenanceView<'a> {
self.provenance
}
}
#[derive(Clone, Copy)]
pub struct SemanticPrimalVjpRequest<'a> {
op: &'a dyn ExtensionOp,
primal_inputs: &'a [ProgramValue],
primal_outputs: &'a [ProgramValue],
cotangent_outputs: &'a [AdValue],
active_inputs: &'a [bool],
provenance: SemanticProvenanceView<'a>,
}
impl<'a> SemanticPrimalVjpRequest<'a> {
pub const fn op(self) -> &'a dyn ExtensionOp {
self.op
}
pub const fn primal_inputs(self) -> &'a [ProgramValue] {
self.primal_inputs
}
pub const fn primal_outputs(self) -> &'a [ProgramValue] {
self.primal_outputs
}
pub const fn cotangent_outputs(self) -> &'a [AdValue] {
self.cotangent_outputs
}
pub const fn active_inputs(self) -> &'a [bool] {
self.active_inputs
}
pub const fn provenance(self) -> SemanticProvenanceView<'a> {
self.provenance
}
}
pub trait SemanticLinearizeRule: Debug + Send + Sync + 'static {
fn family_id(&self) -> &'static str;
fn linearize(
&self,
request: SemanticLinearizeRequest<'_>,
builder: &mut SemanticProgramBuilder,
) -> Result<SemanticLinearizeResult, SemanticAdError>;
}
pub trait SemanticLinearTransposeRule: Debug + Send + Sync + 'static {
fn family_id(&self) -> &'static str;
fn linear_transpose(
&self,
request: SemanticLinearTransposeRequest<'_>,
builder: &mut SemanticProgramBuilder,
) -> Result<Box<[AdValue]>, SemanticAdError>;
}
pub trait SemanticPrimalVjpRule: Debug + Send + Sync + 'static {
fn family_id(&self) -> &'static str;
fn primal_vjp(
&self,
request: SemanticPrimalVjpRequest<'_>,
builder: &mut SemanticProgramBuilder,
) -> Result<Box<[AdValue]>, SemanticAdError>;
}
type LinearizeMap = HashMap<&'static str, Arc<dyn SemanticLinearizeRule>>;
type LinearTransposeMap = HashMap<&'static str, Arc<dyn SemanticLinearTransposeRule>>;
type PrimalVjpMap = HashMap<&'static str, Arc<dyn SemanticPrimalVjpRule>>;
#[derive(Clone, Default)]
pub struct SemanticExtensionRuleSet {
linearize: Arc<LinearizeMap>,
linear_transpose: Arc<LinearTransposeMap>,
primal_vjp: Arc<PrimalVjpMap>,
}
impl Debug for SemanticExtensionRuleSet {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut linearize: Vec<_> = self.linearize.keys().copied().collect();
let mut linear_transpose: Vec<_> = self.linear_transpose.keys().copied().collect();
let mut primal_vjp: Vec<_> = self.primal_vjp.keys().copied().collect();
linearize.sort_unstable();
linear_transpose.sort_unstable();
primal_vjp.sort_unstable();
formatter
.debug_struct("SemanticExtensionRuleSet")
.field("linearize", &linearize)
.field("linear_transpose", &linear_transpose)
.field("primal_vjp", &primal_vjp)
.finish()
}
}
impl SemanticExtensionRuleSet {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn register_linearize(
&mut self,
rule: Arc<dyn SemanticLinearizeRule>,
) -> Result<(), SemanticExtensionRegistryError> {
validate_insert(
&self.linearize,
rule.family_id(),
SemanticAdRuleRole::Linearize,
)?;
Arc::make_mut(&mut self.linearize).insert(rule.family_id(), rule);
Ok(())
}
pub fn register_linear_transpose(
&mut self,
rule: Arc<dyn SemanticLinearTransposeRule>,
) -> Result<(), SemanticExtensionRegistryError> {
validate_insert(
&self.linear_transpose,
rule.family_id(),
SemanticAdRuleRole::LinearTranspose,
)?;
Arc::make_mut(&mut self.linear_transpose).insert(rule.family_id(), rule);
Ok(())
}
pub fn register_primal_vjp(
&mut self,
rule: Arc<dyn SemanticPrimalVjpRule>,
) -> Result<(), SemanticExtensionRegistryError> {
validate_insert(
&self.primal_vjp,
rule.family_id(),
SemanticAdRuleRole::PrimalVjp,
)?;
Arc::make_mut(&mut self.primal_vjp).insert(rule.family_id(), rule);
Ok(())
}
pub fn with_linearize(
mut self,
rule: Arc<dyn SemanticLinearizeRule>,
) -> Result<Self, SemanticExtensionRegistryError> {
self.register_linearize(rule)?;
Ok(self)
}
pub fn with_linear_transpose(
mut self,
rule: Arc<dyn SemanticLinearTransposeRule>,
) -> Result<Self, SemanticExtensionRegistryError> {
self.register_linear_transpose(rule)?;
Ok(self)
}
pub fn with_primal_vjp(
mut self,
rule: Arc<dyn SemanticPrimalVjpRule>,
) -> Result<Self, SemanticExtensionRegistryError> {
self.register_primal_vjp(rule)?;
Ok(self)
}
pub fn merge(&mut self, other: Self) -> Result<(), SemanticExtensionRegistryError> {
let mut candidate = self.clone();
for rule in other.linearize.values() {
candidate.register_linearize(Arc::clone(rule))?;
}
for rule in other.linear_transpose.values() {
candidate.register_linear_transpose(Arc::clone(rule))?;
}
for rule in other.primal_vjp.values() {
candidate.register_primal_vjp(Arc::clone(rule))?;
}
*self = candidate;
Ok(())
}
#[must_use]
pub fn lookup_linearize(&self, family_id: &str) -> Option<Arc<dyn SemanticLinearizeRule>> {
self.linearize.get(family_id).cloned()
}
#[must_use]
pub fn lookup_linear_transpose(
&self,
family_id: &str,
) -> Option<Arc<dyn SemanticLinearTransposeRule>> {
self.linear_transpose.get(family_id).cloned()
}
#[must_use]
pub fn lookup_primal_vjp(&self, family_id: &str) -> Option<Arc<dyn SemanticPrimalVjpRule>> {
self.primal_vjp.get(family_id).cloned()
}
#[allow(clippy::too_many_arguments)]
pub fn linearize_operation(
&self,
operation: SemanticOperationView<'_>,
primal_inputs: &[ProgramValue],
primal_outputs: &[ProgramValue],
tangent_inputs: &[AdValue],
active_outputs: &[bool],
builder: &mut SemanticProgramBuilder,
) -> Result<SemanticLinearizeResult, SemanticAdError> {
let op = extension_for_dispatch(operation)?;
validate_operation_inputs(operation, primal_inputs, primal_outputs, builder)?;
validate_len("tangent_inputs", op.input_count(), tangent_inputs.len())?;
validate_len("active_outputs", op.output_count(), active_outputs.len())?;
validate_ad_values("tangent_inputs", tangent_inputs, builder)?;
let rule = self
.lookup_linearize(op.family_id())
.ok_or(SemanticAdError::MissingRule {
family_id: op.family_id(),
role: SemanticAdRuleRole::Linearize,
})?;
let result = rule.linearize(
SemanticLinearizeRequest {
op,
primal_inputs,
primal_outputs,
tangent_inputs,
active_outputs,
provenance: operation.provenance(),
},
builder,
)?;
validate_len(
"tangent_outputs",
op.output_count(),
result.tangent_outputs.len(),
)?;
validate_ad_values("tangent_outputs", &result.tangent_outputs, builder)?;
validate_values("residuals", &result.residuals, builder)?;
Ok(result)
}
#[allow(clippy::too_many_arguments)]
pub fn linear_transpose_operation(
&self,
operation: SemanticOperationView<'_>,
primal_inputs: &[ProgramValue],
primal_outputs: &[ProgramValue],
cotangent_outputs: &[AdValue],
active_inputs: &[bool],
residuals: &[ProgramValue],
builder: &mut SemanticProgramBuilder,
) -> Result<Box<[AdValue]>, SemanticAdError> {
let op = extension_for_dispatch(operation)?;
validate_operation_inputs(operation, primal_inputs, primal_outputs, builder)?;
validate_len(
"cotangent_outputs",
op.output_count(),
cotangent_outputs.len(),
)?;
validate_len("active_inputs", op.input_count(), active_inputs.len())?;
validate_ad_values("cotangent_outputs", cotangent_outputs, builder)?;
validate_values("residuals", residuals, builder)?;
let rule =
self.lookup_linear_transpose(op.family_id())
.ok_or(SemanticAdError::MissingRule {
family_id: op.family_id(),
role: SemanticAdRuleRole::LinearTranspose,
})?;
let result = rule.linear_transpose(
SemanticLinearTransposeRequest {
op,
primal_inputs,
primal_outputs,
cotangent_outputs,
active_inputs,
residuals,
provenance: operation.provenance(),
},
builder,
)?;
validate_len("cotangent_inputs", op.input_count(), result.len())?;
validate_ad_values("cotangent_inputs", &result, builder)?;
Ok(result)
}
#[allow(clippy::too_many_arguments)]
pub fn primal_vjp_operation(
&self,
operation: SemanticOperationView<'_>,
primal_inputs: &[ProgramValue],
primal_outputs: &[ProgramValue],
cotangent_outputs: &[AdValue],
active_inputs: &[bool],
builder: &mut SemanticProgramBuilder,
) -> Result<Box<[AdValue]>, SemanticAdError> {
let op = extension_for_dispatch(operation)?;
validate_operation_inputs(operation, primal_inputs, primal_outputs, builder)?;
validate_len(
"cotangent_outputs",
op.output_count(),
cotangent_outputs.len(),
)?;
validate_len("active_inputs", op.input_count(), active_inputs.len())?;
validate_ad_values("cotangent_outputs", cotangent_outputs, builder)?;
let rule = self
.lookup_primal_vjp(op.family_id())
.ok_or(SemanticAdError::MissingRule {
family_id: op.family_id(),
role: SemanticAdRuleRole::PrimalVjp,
})?;
let result = rule.primal_vjp(
SemanticPrimalVjpRequest {
op,
primal_inputs,
primal_outputs,
cotangent_outputs,
active_inputs,
provenance: operation.provenance(),
},
builder,
)?;
validate_len("cotangent_inputs", op.input_count(), result.len())?;
validate_ad_values("cotangent_inputs", &result, builder)?;
Ok(result)
}
}
fn extension_for_dispatch(
operation: SemanticOperationView<'_>,
) -> Result<&dyn ExtensionOp, SemanticAdError> {
let SemanticOpRef::Extension(op) = operation.op() else {
return Err(SemanticAdError::CoreOperation);
};
if !operation.effects().is_empty() {
return Err(SemanticAdError::EffectfulExtension {
family_id: op.family_id(),
});
}
Ok(op)
}
fn validate_operation_inputs(
operation: SemanticOperationView<'_>,
primal_inputs: &[ProgramValue],
primal_outputs: &[ProgramValue],
builder: &SemanticProgramBuilder,
) -> Result<(), SemanticAdError> {
validate_len(
"primal_inputs",
operation.inputs().len(),
primal_inputs.len(),
)?;
validate_len(
"primal_outputs",
operation.outputs().len(),
primal_outputs.len(),
)?;
validate_values("primal_inputs", primal_inputs, builder)?;
validate_values("primal_outputs", primal_outputs, builder)
}
fn validate_values(
field: &'static str,
values: &[ProgramValue],
builder: &SemanticProgramBuilder,
) -> Result<(), SemanticAdError> {
for (index, value) in values.iter().copied().enumerate() {
if builder.validate_value(value).is_err() {
return Err(SemanticAdError::ForeignValue { field, index });
}
}
Ok(())
}
fn validate_ad_values(
field: &'static str,
values: &[AdValue],
builder: &SemanticProgramBuilder,
) -> Result<(), SemanticAdError> {
for (index, value) in values.iter().copied().enumerate() {
if let AdValue::Value(value) = value {
if builder.validate_value(value).is_err() {
return Err(SemanticAdError::ForeignValue { field, index });
}
}
}
Ok(())
}
fn validate_len(
field: &'static str,
expected: usize,
actual: usize,
) -> Result<(), SemanticAdError> {
if expected != actual {
return Err(SemanticAdError::Arity {
field,
expected,
actual,
});
}
Ok(())
}
fn validate_insert<T>(
map: &HashMap<&'static str, T>,
family_id: &'static str,
role: SemanticAdRuleRole,
) -> Result<(), SemanticExtensionRegistryError> {
if !is_valid_family_id(family_id) {
return Err(SemanticExtensionRegistryError::MalformedFamilyId { family_id });
}
if map.contains_key(family_id) {
return Err(SemanticExtensionRegistryError::DuplicateRule { family_id, role });
}
Ok(())
}
fn is_valid_family_id(family_id: &str) -> bool {
let Some((prefix, version)) = family_id.rsplit_once('.') else {
return false;
};
let Some(version) = version.strip_prefix('v') else {
return false;
};
let Some((crate_name, op_name)) = prefix.split_once('.') else {
return false;
};
!crate_name.is_empty()
&& !op_name.is_empty()
&& !version.is_empty()
&& version.bytes().all(|byte| byte.is_ascii_digit())
&& crate_name.is_ascii()
&& op_name.is_ascii()
&& !crate_name.chars().any(char::is_whitespace)
&& !op_name.chars().any(char::is_whitespace)
}