tidu/eager/
backward.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3
4use crate::{ADKey, ADRuleResult, Primitive};
5use computegraph::{GraphOperation, ValueKey};
6
7use crate::{LinearizedGraph, PrimitiveGraph};
8
9use super::record::RecordedGraph;
10use super::trace::{Trace, TraceNode};
11
12/// Downstream execution hooks for eager backward.
13pub trait BackwardExecutor<Op: Primitive>
14where
15    Op::InputKey: ADKey,
16{
17    /// Linearize one recorded eager graph node for the active output slots.
18    ///
19    /// The default implementation calls [`RecordedGraph::linearize`]. Runtime
20    /// owners may override this to memoize per-node transform results without
21    /// taking trace traversal ownership away from tidu.
22    fn linearize_recorded_graph(
23        &mut self,
24        graph: &RecordedGraph<Op>,
25        output_slots: &[usize],
26        ctx: &mut Op::ADContext,
27    ) -> ADRuleResult<Arc<LinearizedGraph<Op>>> {
28        graph.linearize(output_slots, ctx).map(Arc::new)
29    }
30
31    /// Replay a primitive graph and return any concrete values needed by
32    /// transpose execution.
33    fn execute_forward(
34        &mut self,
35        graph: PrimitiveGraph<'_, Op>,
36        initial_data: &HashMap<ValueKey<Op>, Arc<Op::Operand>>,
37    ) -> HashMap<ValueKey<Op>, Arc<Op::Operand>>;
38
39    /// Run a transposed linear graph with concrete cotangent seeds.
40    fn run_transposed_linear(
41        &mut self,
42        linear: &LinearizedGraph<Op>,
43        cotangent_out: &[Option<Arc<Op::Operand>>],
44        external_data: &HashMap<ValueKey<Op>, Arc<Op::Operand>>,
45        ctx: &mut Op::ADContext,
46    ) -> ADRuleResult<Vec<Option<Arc<Op::Operand>>>>;
47
48    /// Add two concrete operands for cotangent accumulation.
49    fn add_operands(&mut self, a: &Arc<Op::Operand>, b: &Arc<Op::Operand>) -> Arc<Op::Operand>;
50}
51
52/// Execute reverse-mode AD over an eager trace.
53pub fn backward<Op: Primitive>(
54    output_key: &ValueKey<Op>,
55    output_trace: Option<&Trace<Op>>,
56    seed: Arc<Op::Operand>,
57    executor: &mut impl BackwardExecutor<Op>,
58    ctx: &mut Op::ADContext,
59) -> ADRuleResult<HashMap<ValueKey<Op>, Arc<Op::Operand>>>
60where
61    Op::InputKey: ADKey,
62{
63    let sorted_nodes = topo_sort_trace(output_trace);
64    let mut cotangents: HashMap<ValueKey<Op>, Arc<Op::Operand>> = HashMap::new();
65    cotangents.insert(output_key.clone(), seed);
66
67    for node in sorted_nodes.iter().rev() {
68        let cotangent_out: Vec<Option<Arc<Op::Operand>>> = node
69            .primal_out_keys()
70            .iter()
71            .map(|key| cotangents.get(key).cloned())
72            .collect();
73        if cotangent_out.iter().all(Option::is_none) {
74            continue;
75        }
76
77        let mut active_output_slots = Vec::new();
78        let mut active_cotangent_out = Vec::new();
79        for (slot, maybe_cotangent) in cotangent_out.into_iter().enumerate() {
80            if let Some(cotangent) = maybe_cotangent {
81                active_output_slots.push(slot);
82                active_cotangent_out.push(Some(cotangent));
83            }
84        }
85
86        let linear =
87            executor.linearize_recorded_graph(node.computation(), &active_output_slots, ctx)?;
88        let replay_graph = PrimitiveGraph::new(linear.as_graph());
89        let all_values = executor.execute_forward(replay_graph, node.saved_data());
90        let cotangent_in = executor.run_transposed_linear(
91            linear.as_ref(),
92            &active_cotangent_out,
93            &all_values,
94            ctx,
95        )?;
96
97        for (edge, maybe_cotangent) in node.input_edges().iter().zip(cotangent_in) {
98            let cotangent = match maybe_cotangent {
99                Some(cotangent) => cotangent,
100                None => continue,
101            };
102            if !edge.requires_grad {
103                continue;
104            }
105
106            let accumulated = match cotangents.remove(&edge.key) {
107                Some(existing) => executor.add_operands(&existing, &cotangent),
108                None => cotangent,
109            };
110            cotangents.insert(edge.key.clone(), accumulated);
111        }
112    }
113
114    Ok(cotangents)
115}
116
117fn topo_sort_trace<Op: GraphOperation>(
118    output_trace: Option<&Trace<Op>>,
119) -> Vec<Arc<TraceNode<Op>>> {
120    fn visit<Op: GraphOperation>(
121        node: &Arc<TraceNode<Op>>,
122        visited: &mut HashSet<*const TraceNode<Op>>,
123        order: &mut Vec<Arc<TraceNode<Op>>>,
124    ) {
125        let ptr = Arc::as_ptr(node);
126        if !visited.insert(ptr) {
127            return;
128        }
129
130        for edge in node.input_edges() {
131            if let Some(parent) = &edge.node {
132                visit(parent, visited, order);
133            }
134        }
135
136        order.push(node.clone());
137    }
138
139    let mut visited = HashSet::new();
140    let mut order = Vec::new();
141    if let Some(trace) = output_trace {
142        visit(trace.node(), &mut visited, &mut order);
143    }
144    order
145}
tidu/eager/backward.rs

tidu/eager/
backward.rs
