use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use computegraph::graph::Graph;
use computegraph::{LocalValueId, OperationRole, ValueKey, ValueRef};
use tenferro_ops::input_key::TensorInputKey;
use tenferro_ops::std_tensor_op::StdTensorOp;
use tenferro_ops::{ShapeGuardContext, TensorMeta};
use tenferro_tensor::{DType, Tensor, TensorBackend, TypedTensor};
use tidu::eager::BackwardExecutor;
use tidu::{ADRuleError, ADRuleKind, ADRuleResult, LinearizedGraph, PrimitiveGraph};
use crate::eager_builder::EagerPrimitiveBuilder;
use crate::eager_exec::{exec_op_on_tensors, exec_op_on_tensors_with_extension_executor};
use crate::extension_runtime::ExtensionExecutor;
use crate::metadata::{
push_metadata_scope, register_scoped_live_graph_metadata, tensor_meta_from_tensor,
GlobalMetadataScope,
};
use super::zero_like_tensor;
pub(crate) struct TenferroBackwardCallbacks<'a, B: TensorBackend + 'static> {
backend: &'a mut B,
extension_executor: Option<&'a mut ExtensionExecutor<B>>,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
deferred_error: Option<ADRuleError>,
}
impl<'a, B: TensorBackend + 'static> TenferroBackwardCallbacks<'a, B> {
pub(crate) fn new(
backend: &'a mut B,
extension_executor: Option<&'a mut ExtensionExecutor<B>>,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
) -> Self {
Self {
backend,
extension_executor,
metadata_scopes,
deferred_error: None,
}
}
pub(crate) fn take_error(&mut self) -> Option<ADRuleError> {
self.deferred_error.take()
}
fn record_error(&mut self, err: ADRuleError) {
if self.deferred_error.is_none() {
self.deferred_error = Some(err);
}
}
}
pub(super) fn missing_tangent_base_key(
key: &ValueKey<StdTensorOp>,
) -> Option<ValueKey<StdTensorOp>> {
let ValueKey::Input(tangent_key) = key else {
return None;
};
let TensorInputKey::Tangent { of, .. } = tangent_key else {
return None;
};
Some(ValueKey::Input((**of).clone()))
}
pub(super) fn eager_forward_input_metadata(
key: &ValueKey<StdTensorOp>,
initial_data: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
) -> ADRuleResult<TensorMeta> {
if let Some(value) = initial_data.get(key) {
return Ok(tensor_meta_from_tensor(value.as_ref()));
}
let base_key = missing_tangent_base_key(key).ok_or_else(|| {
eager_ad_invalid_input(format!("missing concrete eager value for {key:?}"))
})?;
let base = initial_data.get(&base_key).ok_or_else(|| {
eager_ad_invalid_input(format!("missing base eager value for {base_key:?}"))
})?;
Ok(tensor_meta_from_tensor(base.as_ref()))
}
pub(super) fn eager_forward_value<B: TensorBackend>(
all_values: &mut HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
key: &ValueKey<StdTensorOp>,
initial_data: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
backend: &mut B,
) -> ADRuleResult<Arc<Tensor>> {
if let Some(value) = all_values.get(key) {
return Ok(Arc::clone(value));
}
let base_key = missing_tangent_base_key(key).ok_or_else(|| {
eager_ad_invalid_input(format!("missing concrete eager value for {key:?}"))
})?;
let base = initial_data.get(&base_key).ok_or_else(|| {
eager_ad_invalid_input(format!("missing base eager value for {base_key:?}"))
})?;
let value = Arc::new(zero_like_tensor(base.as_ref(), backend).map_err(|err| {
eager_ad_invalid_input(format!("failed to create eager tangent zero: {err}"))
})?);
all_values.insert(key.clone(), Arc::clone(&value));
Ok(value)
}
fn live_graph_values(graph: &Graph<StdTensorOp>) -> HashSet<LocalValueId> {
let mut producers = HashMap::new();
for (op_index, op_node) in graph.operations().iter().enumerate() {
for &output_id in &op_node.outputs {
producers.insert(output_id, op_index);
}
}
let mut live = HashSet::new();
let mut stack = graph.outputs().to_vec();
while let Some(local_id) = stack.pop() {
if !live.insert(local_id) {
continue;
}
let Some(&op_index) = producers.get(&local_id) else {
continue;
};
for input in &graph.operations()[op_index].inputs {
if let ValueRef::Local(input_id) = input {
stack.push(*input_id);
}
}
}
live
}
fn linear_op_depends_on_tangents(mode: &OperationRole) -> bool {
matches!(mode, OperationRole::Linearized { active_mask } if active_mask.iter().any(|is_active| *is_active))
}
pub(super) fn zero_from_exact_metadata<B: TensorBackend>(
meta: &TensorMeta,
backend: &mut B,
) -> ADRuleResult<Option<Tensor>> {
let Some(shape) = meta.exact_shape() else {
return Ok(None);
};
let shape = shape
.into_iter()
.map(|dim| dim.constant_value())
.collect::<Option<Vec<_>>>();
let Some(shape) = shape else {
return Ok(None);
};
let host =
match meta.dtype {
DType::F32 => Tensor::F32(TypedTensor::zeros(shape).map_err(|err| {
eager_ad_invalid_input(format!("failed to create F32 zero: {err}"))
})?),
DType::F64 => Tensor::F64(TypedTensor::zeros(shape).map_err(|err| {
eager_ad_invalid_input(format!("failed to create F64 zero: {err}"))
})?),
DType::I32 => Tensor::I32(TypedTensor::zeros(shape).map_err(|err| {
eager_ad_invalid_input(format!("failed to create I32 zero: {err}"))
})?),
DType::I64 => Tensor::I64(TypedTensor::zeros(shape).map_err(|err| {
eager_ad_invalid_input(format!("failed to create I64 zero: {err}"))
})?),
DType::Bool => {
let len = checked_zero_element_count(&shape)?;
Tensor::Bool(
TypedTensor::from_vec_col_major(shape, vec![false; len]).map_err(|err| {
eager_ad_invalid_input(format!("failed to create bool zero: {err}"))
})?,
)
}
DType::C32 => Tensor::C32(TypedTensor::zeros(shape).map_err(|err| {
eager_ad_invalid_input(format!("failed to create C32 zero: {err}"))
})?),
DType::C64 => Tensor::C64(TypedTensor::zeros(shape).map_err(|err| {
eager_ad_invalid_input(format!("failed to create C64 zero: {err}"))
})?),
};
Ok(Some(backend.upload_host_tensor(&host).map_err(|err| {
eager_ad_invalid_input(format!("failed to upload eager zero tensor: {err}"))
})?))
}
fn prefill_missing_linear_zero_values<B: TensorBackend>(
linear: &LinearizedGraph<StdTensorOp>,
external_data: &mut HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
ctx: &mut ShapeGuardContext,
backend: &mut B,
) -> ADRuleResult<()> {
for value in linear.as_graph().values() {
if external_data.contains_key(&value.key) {
continue;
}
let Some(meta) = ctx
.metadata_if_available(&ValueRef::External(value.key.clone()))
.cloned()
else {
continue;
};
let Some(zero) = zero_from_exact_metadata(&meta, backend)? else {
continue;
};
external_data.insert(value.key.clone(), Arc::new(zero));
}
Ok(())
}
fn checked_zero_element_count(shape: &[usize]) -> ADRuleResult<usize> {
shape.iter().try_fold(1usize, |acc, &dim| {
acc.checked_mul(dim).ok_or_else(|| {
eager_ad_invalid_input(format!(
"zero tensor shape product overflows for shape {shape:?}"
))
})
})
}
fn eager_ad_invalid_input(message: impl Into<String>) -> ADRuleError {
ADRuleError::invalid_input("tenferro-ad.eager", ADRuleKind::Transpose, message)
}
impl<B: TensorBackend + 'static> BackwardExecutor<StdTensorOp>
for TenferroBackwardCallbacks<'_, B>
{
fn execute_forward(
&mut self,
graph: PrimitiveGraph<'_, StdTensorOp>,
initial_data: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
) -> HashMap<ValueKey<StdTensorOp>, Arc<Tensor>> {
if self.deferred_error.is_some() {
return initial_data.clone();
}
let graph = graph.as_graph();
let mut all_values = initial_data.clone();
let live_values = live_graph_values(graph);
let mut input_metadata = Vec::with_capacity(graph.inputs().len());
for &input_id in graph.inputs() {
let key = graph.values()[input_id].key.clone();
match eager_forward_input_metadata(&key, initial_data) {
Ok(meta) => input_metadata.push((key, meta)),
Err(err) => {
self.record_error(err);
return all_values;
}
}
}
for op_node in graph.operations() {
if linear_op_depends_on_tangents(&op_node.role) {
continue;
}
if !op_node
.outputs
.iter()
.any(|output_id| live_values.contains(output_id))
{
continue;
}
let mut resolved_values = Vec::with_capacity(op_node.inputs.len());
for input in &op_node.inputs {
let resolved = match input {
ValueRef::Local(local_id) => {
let key = &graph.values()[*local_id].key;
eager_forward_value(&mut all_values, key, initial_data, self.backend)
}
ValueRef::External(key) => {
eager_forward_value(&mut all_values, key, initial_data, self.backend)
}
};
match resolved {
Ok(value) => resolved_values.push(value),
Err(err) => {
self.record_error(err);
return all_values;
}
}
}
let resolved_inputs: Vec<&Tensor> =
resolved_values.iter().map(|value| value.as_ref()).collect();
let outputs_result =
if let Some(extension_executor) = self.extension_executor.as_deref_mut() {
exec_op_on_tensors_with_extension_executor(
&op_node.operation,
&resolved_inputs,
self.backend,
Some(extension_executor),
)
} else {
exec_op_on_tensors(&op_node.operation, &resolved_inputs, self.backend)
};
let outputs = match outputs_result {
Ok(outputs) => outputs,
Err(err) => {
self.record_error(eager_ad_invalid_input(format!(
"eager forward exec failed for {:?}: {err}",
op_node.operation
)));
return all_values;
}
};
for (output_id, output) in op_node.outputs.iter().zip(outputs) {
let key = graph.values()[*output_id].key.clone();
all_values.insert(key, Arc::new(output));
}
}
let metadata_scope =
match register_scoped_live_graph_metadata(graph, &live_values, input_metadata) {
Ok(scope) => scope,
Err(err) => {
self.record_error(eager_ad_invalid_input(format!(
"eager replay metadata registration failed: {err}"
)));
return all_values;
}
};
push_metadata_scope(&mut self.metadata_scopes, Arc::new(metadata_scope));
all_values
}
fn run_transposed_linear(
&mut self,
linear: &LinearizedGraph<StdTensorOp>,
cotangent_out: &[Option<Arc<Tensor>>],
external_data: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
ctx: &mut ShapeGuardContext,
) -> tidu::ADRuleResult<Vec<Option<Arc<Tensor>>>> {
if let Some(err) = self.take_error() {
return Err(err);
}
let mut external_data = external_data.clone();
ctx.refresh_global_metadata();
prefill_missing_linear_zero_values(linear, &mut external_data, ctx, self.backend)?;
let mut builder = if let Some(extension_executor) = self.extension_executor.as_deref_mut() {
EagerPrimitiveBuilder::with_extension_executor(self.backend, extension_executor)
} else {
EagerPrimitiveBuilder::new(self.backend)
};
builder.external_data = external_data;
let cotangent_seed_ids = cotangent_out
.iter()
.map(|maybe_seed| {
maybe_seed
.as_ref()
.map(|seed| builder.push_tensor(Arc::clone(seed)))
})
.collect::<Vec<_>>();
let transpose_result =
tidu::linear_transpose_with_builder(linear, &mut builder, &cotangent_seed_ids, ctx);
if let Some(err) = builder.take_error() {
return Err(err);
}
let cotangent_ids = transpose_result?;
cotangent_ids
.into_iter()
.map(|maybe_id| maybe_id.map(|id| builder.tensor(id)).transpose())
.collect()
}
fn add_operands(&mut self, a: &Arc<Tensor>, b: &Arc<Tensor>) -> Arc<Tensor> {
if self.deferred_error.is_some() {
return Arc::clone(a);
}
match self.backend.add(a.as_ref(), b.as_ref()) {
Ok(sum) => Arc::new(sum),
Err(err) => {
self.record_error(eager_ad_invalid_input(format!(
"eager cotangent add failed: {err}"
)));
Arc::clone(a)
}
}
}
}