tidu/eager/record.rs
1use std::collections::HashMap;
2use std::error::Error;
3use std::fmt;
4use std::sync::Arc;
5
6use crate::{ADKey, ADRuleError, ADRuleKind, ADRuleResult, Primitive};
7use computegraph::graph::{Graph, GraphBuilder};
8use computegraph::resolve::resolve;
9use computegraph::{GraphOperation, OperationRole, ValueKey, ValueRef};
10
11use crate::LinearizedGraph;
12
13use super::trace::{Trace, TraceEdge, TraceNode};
14
15/// Error returned when eager graph recording receives inconsistent metadata.
16///
17/// These errors describe caller-visible recording contract violations. AD rule
18/// failures during linearization are still reported as [`crate::ADRuleError`].
19#[derive(Clone, Debug, PartialEq, Eq)]
20pub enum EagerRecordError {
21 /// A caller supplied a slice whose length must match graph metadata.
22 CountMismatch {
23 /// Name of the validated field.
24 field: &'static str,
25 /// Required length.
26 expected: usize,
27 /// Supplied length.
28 actual: usize,
29 },
30 /// A caller supplied keys in an order that does not match the graph.
31 KeyMismatch {
32 /// Name of the validated field.
33 field: &'static str,
34 /// Mismatching slot.
35 index: usize,
36 },
37 /// The graph has more outputs than eager tracing can address.
38 TooManyOutputs {
39 /// Supplied output count.
40 actual: usize,
41 /// Maximum accepted output count.
42 max: usize,
43 },
44}
45
46impl EagerRecordError {
47 pub(crate) fn count_mismatch(field: &'static str, expected: usize, actual: usize) -> Self {
48 Self::CountMismatch {
49 field,
50 expected,
51 actual,
52 }
53 }
54
55 pub(crate) fn key_mismatch(field: &'static str, index: usize) -> Self {
56 Self::KeyMismatch { field, index }
57 }
58
59 fn too_many_outputs(actual: usize) -> Self {
60 Self::TooManyOutputs {
61 actual,
62 max: u8::MAX as usize + 1,
63 }
64 }
65}
66
67impl fmt::Display for EagerRecordError {
68 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69 match self {
70 Self::CountMismatch {
71 field,
72 expected,
73 actual,
74 } => write!(f, "{field} expected {expected} entries, got {actual}"),
75 Self::KeyMismatch { field, index } => {
76 write!(f, "{field} does not match graph metadata at slot {index}")
77 }
78 Self::TooManyOutputs { actual, max } => write!(
79 f,
80 "eager recording supports at most {max} outputs, got {actual}"
81 ),
82 }
83 }
84}
85
86impl Error for EagerRecordError {}
87
88/// Result type used by eager graph recording APIs.
89pub type EagerRecordResult<T> = Result<T, EagerRecordError>;
90
91/// Graph invocation recorded as one eager reverse-mode trace node.
92///
93/// # Examples
94///
95/// ```
96/// use tidu::eager::RecordedGraph;
97/// use computegraph::GraphOperation;
98///
99/// #[derive(Clone, Debug, Hash, PartialEq, Eq)]
100/// enum Op { Add }
101///
102/// impl GraphOperation for Op {
103/// type Operand = f64;
104/// type Context = ();
105/// type InputKey = &'static str;
106///
107/// fn input_count(&self) -> usize { 2 }
108/// fn output_count(&self) -> usize { 1 }
109/// }
110///
111/// let recorded = RecordedGraph::from_primitive(Op::Add, vec!["x", "y"])?;
112/// assert_eq!(recorded.input_keys(), &["x", "y"]);
113/// assert_eq!(recorded.output_keys().len(), 1);
114/// # Ok::<(), tidu::eager::EagerRecordError>(())
115/// ```
116pub struct RecordedGraph<Op: GraphOperation> {
117 graph: Arc<Graph<Op>>,
118 input_keys: Vec<Op::InputKey>,
119 output_keys: Vec<ValueKey<Op>>,
120}
121
122impl<Op: GraphOperation> RecordedGraph<Op> {
123 /// Create a recorded graph from an already-built graph and aligned keys.
124 ///
125 /// # Examples
126 ///
127 /// ```
128 /// use std::sync::Arc;
129 /// use computegraph::graph::GraphBuilder;
130 /// use computegraph::{GraphOperation, OperationRole, ValueRef};
131 /// use tidu::eager::RecordedGraph;
132 ///
133 /// #[derive(Clone, Debug, Hash, PartialEq, Eq)]
134 /// enum Op { Id }
135 ///
136 /// impl GraphOperation for Op {
137 /// type Operand = f64;
138 /// type Context = ();
139 /// type InputKey = &'static str;
140 ///
141 /// fn input_count(&self) -> usize { 1 }
142 /// fn output_count(&self) -> usize { 1 }
143 /// }
144 ///
145 /// let mut builder = GraphBuilder::new();
146 /// let x = builder.add_input("x");
147 /// let y = builder.add_operation(Op::Id, vec![ValueRef::Local(x)], OperationRole::Primary);
148 /// builder.set_outputs(y.clone());
149 /// let graph = Arc::new(builder.build());
150 /// let output_keys = y.iter().map(|id| graph.values()[*id].key.clone()).collect();
151 /// let recorded = RecordedGraph::new(graph, vec!["x"], output_keys)?;
152 ///
153 /// assert_eq!(recorded.input_keys(), &["x"]);
154 /// # Ok::<(), tidu::eager::EagerRecordError>(())
155 /// ```
156 pub fn new(
157 graph: Arc<Graph<Op>>,
158 input_keys: Vec<Op::InputKey>,
159 output_keys: Vec<ValueKey<Op>>,
160 ) -> EagerRecordResult<Self> {
161 if graph.inputs().len() != input_keys.len() {
162 return Err(EagerRecordError::count_mismatch(
163 "RecordedGraph input keys",
164 graph.inputs().len(),
165 input_keys.len(),
166 ));
167 }
168 if graph.outputs().len() != output_keys.len() {
169 return Err(EagerRecordError::count_mismatch(
170 "RecordedGraph output keys",
171 graph.outputs().len(),
172 output_keys.len(),
173 ));
174 }
175 for (index, (&input_id, input_key)) in
176 graph.inputs().iter().zip(input_keys.iter()).enumerate()
177 {
178 if graph.values()[input_id].key != ValueKey::Input(input_key.clone()) {
179 return Err(EagerRecordError::key_mismatch(
180 "RecordedGraph input keys",
181 index,
182 ));
183 }
184 }
185 for (index, (&output_id, output_key)) in
186 graph.outputs().iter().zip(output_keys.iter()).enumerate()
187 {
188 if &graph.values()[output_id].key != output_key {
189 return Err(EagerRecordError::key_mismatch(
190 "RecordedGraph output keys",
191 index,
192 ));
193 }
194 }
195
196 Ok(Self {
197 graph,
198 input_keys,
199 output_keys,
200 })
201 }
202
203 /// Build a one-op recorded graph for an eager primitive invocation.
204 ///
205 /// # Examples
206 ///
207 /// ```
208 /// use tidu::eager::RecordedGraph;
209 /// use computegraph::GraphOperation;
210 ///
211 /// #[derive(Clone, Debug, Hash, PartialEq, Eq)]
212 /// enum Op { Add }
213 ///
214 /// impl GraphOperation for Op {
215 /// type Operand = f64;
216 /// type Context = ();
217 /// type InputKey = &'static str;
218 ///
219 /// fn input_count(&self) -> usize { 2 }
220 /// fn output_count(&self) -> usize { 1 }
221 /// }
222 ///
223 /// let recorded = RecordedGraph::from_primitive(Op::Add, vec!["x", "y"])?;
224 /// assert_eq!(recorded.as_graph().operations().len(), 1);
225 /// # Ok::<(), tidu::eager::EagerRecordError>(())
226 /// ```
227 pub fn from_primitive(op: Op, input_keys: Vec<Op::InputKey>) -> EagerRecordResult<Self> {
228 let mut builder = GraphBuilder::new();
229 let input_ids: Vec<_> = input_keys
230 .iter()
231 .cloned()
232 .map(|key| builder.add_input(key))
233 .collect();
234 let output_ids = builder.add_operation(
235 op,
236 input_ids
237 .iter()
238 .map(|local_id| ValueRef::Local(*local_id))
239 .collect(),
240 OperationRole::Primary,
241 );
242 builder.set_outputs(output_ids.clone());
243 let graph = Arc::new(builder.build());
244 let output_keys = output_ids
245 .iter()
246 .map(|output_id| graph.values()[*output_id].key.clone())
247 .collect();
248 Self::new(graph, input_keys, output_keys)
249 }
250
251 /// Borrow the recorded graph.
252 pub fn as_graph(&self) -> &Graph<Op> {
253 &self.graph
254 }
255
256 /// Graph input keys aligned with eager input edges.
257 pub fn input_keys(&self) -> &[Op::InputKey] {
258 &self.input_keys
259 }
260
261 /// Graph output keys aligned with eager output slots.
262 pub fn output_keys(&self) -> &[ValueKey<Op>] {
263 &self.output_keys
264 }
265}
266
267impl<Op: Primitive> RecordedGraph<Op>
268where
269 Op::InputKey: ADKey,
270{
271 /// Linearize selected output slots of this recorded eager graph.
272 ///
273 /// Eager walkers call this through their executor hook. Runtime owners can
274 /// also use it as the cache-miss fallback when memoizing transform results.
275 pub fn linearize(
276 &self,
277 output_slots: &[usize],
278 ctx: &mut Op::ADContext,
279 ) -> ADRuleResult<LinearizedGraph<Op>> {
280 let mut selected_outputs = Vec::with_capacity(output_slots.len());
281 for &slot in output_slots {
282 let output_key = match self.output_keys.get(slot).cloned() {
283 Some(output_key) => output_key,
284 None => {
285 return Err(ADRuleError::invalid_input(
286 "tidu::eager::RecordedGraph",
287 ADRuleKind::Jvp,
288 format!(
289 "requested output slot {}, but graph has {} outputs",
290 slot,
291 self.output_keys.len()
292 ),
293 ));
294 }
295 };
296 selected_outputs.push(output_key);
297 }
298 let view = resolve(vec![Arc::clone(&self.graph)]);
299 let aliases = HashMap::new();
300 crate::linearize(&view, &selected_outputs, &self.input_keys, 0, ctx, &aliases)
301 }
302}
303
304/// Input descriptor for recording one eager graph invocation.
305pub struct EagerInput<Op: GraphOperation> {
306 /// User-visible eager value key used for cotangent accumulation.
307 pub key: ValueKey<Op>,
308 /// Trace node that produced this value, or `None` for leaves.
309 pub trace: Option<Trace<Op>>,
310 /// Whether this value participates in reverse-mode propagation.
311 pub requires_grad: bool,
312 /// Concrete primal data for saved forward replay.
313 pub data: Arc<Op::Operand>,
314}
315
316/// Per-output trace metadata returned by [`Recorder::record_graph`].
317pub struct EagerOutput<Op: GraphOperation> {
318 /// User-visible eager output key.
319 pub key: ValueKey<Op>,
320 /// Shared trace node for all outputs when any input requires gradients.
321 pub trace: Option<Trace<Op>>,
322 /// Whether this output should be tracked by the downstream frontend.
323 pub requires_grad: bool,
324 /// Output slot within the recorded graph invocation.
325 pub output_slot: usize,
326}
327
328/// Caller-provided source of stable eager value keys.
329pub trait KeySource<Op: GraphOperation> {
330 /// Return a fresh input key that has not been used for another eager value.
331 fn fresh_input_key(&mut self) -> Op::InputKey;
332}
333
334/// Stateful eager operation recorder.
335pub struct Recorder<K> {
336 key_source: K,
337}
338
339impl<K> Recorder<K> {
340 /// Create a recorder from a downstream key source.
341 pub fn new(key_source: K) -> Self {
342 Self { key_source }
343 }
344
345 /// Borrow the underlying key source.
346 pub fn key_source_mut(&mut self) -> &mut K {
347 &mut self.key_source
348 }
349
350 /// Return the underlying key source.
351 pub fn into_key_source(self) -> K {
352 self.key_source
353 }
354
355 /// Return fresh graph input keys for one eager graph invocation.
356 ///
357 /// # Examples
358 ///
359 /// ```
360 /// use tidu::eager::{KeySource, Recorder};
361 /// use computegraph::GraphOperation;
362 ///
363 /// #[derive(Clone, Debug, Hash, PartialEq, Eq)]
364 /// enum Op { Id }
365 ///
366 /// impl GraphOperation for Op {
367 /// type Operand = f64;
368 /// type Context = ();
369 /// type InputKey = usize;
370 ///
371 /// fn input_count(&self) -> usize { 1 }
372 /// fn output_count(&self) -> usize { 1 }
373 /// }
374 ///
375 /// struct Keys(usize);
376 ///
377 /// impl KeySource<Op> for Keys {
378 /// fn fresh_input_key(&mut self) -> usize {
379 /// let key = self.0;
380 /// self.0 += 1;
381 /// key
382 /// }
383 /// }
384 ///
385 /// let mut recorder = Recorder::new(Keys(0));
386 /// assert_eq!(recorder.fresh_input_keys::<Op>(2), vec![0, 1]);
387 /// ```
388 pub fn fresh_input_keys<Op>(&mut self, count: usize) -> Vec<Op::InputKey>
389 where
390 Op: GraphOperation,
391 K: KeySource<Op>,
392 {
393 (0..count)
394 .map(|_| self.key_source.fresh_input_key())
395 .collect()
396 }
397
398 /// Record a concrete eager graph invocation for reverse-mode AD.
399 ///
400 /// # Examples
401 ///
402 /// ```
403 /// use std::collections::HashMap;
404 /// use std::sync::Arc;
405 /// use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey};
406 /// use tidu::{
407 /// ADKey, DiffPassId, Primitive, PrimitiveBuilder, PrimitiveValue,
408 /// };
409 /// use tidu::eager::{EagerInput, KeySource, RecordedGraph, Recorder};
410 ///
411 /// #[derive(Clone, Debug, Hash, PartialEq, Eq)]
412 /// enum Key {
413 /// User(&'static str),
414 /// Generated(usize),
415 /// Tangent(Box<Key>, DiffPassId),
416 /// }
417 ///
418 /// impl ADKey for Key {
419 /// fn tangent_of(&self, pass: DiffPassId) -> Self {
420 /// Self::Tangent(Box::new(self.clone()), pass)
421 /// }
422 /// }
423 ///
424 /// #[derive(Clone, Debug, Hash, PartialEq, Eq)]
425 /// enum Op { Id }
426 ///
427 /// impl GraphOperation for Op {
428 /// type Operand = f64;
429 /// type Context = ();
430 /// type InputKey = Key;
431 ///
432 /// fn input_count(&self) -> usize { 1 }
433 /// fn output_count(&self) -> usize { 1 }
434 /// }
435 ///
436 /// impl Primitive for Op {
437 /// type ADContext = ();
438 /// fn add() -> Self { Self::Id }
439 ///
440 /// fn jvp_rule(
441 /// &self,
442 /// _builder: &mut impl PrimitiveBuilder<Self>,
443 /// _primal_in: &[ValueKey<Self>],
444 /// _primal_out: &[ValueKey<Self>],
445 /// tangent_in: &[Option<LocalValueId>],
446 /// _ctx: &mut (),
447 /// ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
448 /// Ok(vec![tangent_in[0]])
449 /// }
450 ///
451 /// fn transpose_rule(
452 /// &self,
453 /// _builder: &mut impl PrimitiveBuilder<Self>,
454 /// cotangent_out: &[Option<LocalValueId>],
455 /// _inputs: &[PrimitiveValue<Self>],
456 /// _role: &OperationRole,
457 /// _ctx: &mut (),
458 /// ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
459 /// Ok(vec![cotangent_out[0]])
460 /// }
461 /// }
462 ///
463 /// struct Keys(usize);
464 /// impl KeySource<Op> for Keys {
465 /// fn fresh_input_key(&mut self) -> Key {
466 /// let key = Key::Generated(self.0);
467 /// self.0 += 1;
468 /// key
469 /// }
470 /// }
471 ///
472 /// let mut recorder = Recorder::new(Keys(0));
473 /// let graph_inputs = recorder.fresh_input_keys::<Op>(1);
474 /// let graph = RecordedGraph::from_primitive(Op::Id, graph_inputs)?;
475 /// let input = EagerInput {
476 /// key: ValueKey::Input(Key::User("x")),
477 /// trace: None,
478 /// requires_grad: true,
479 /// data: Arc::new(2.0),
480 /// };
481 ///
482 /// let outputs = recorder.record_graph(
483 /// graph,
484 /// &[input],
485 /// &[Arc::new(2.0)],
486 /// HashMap::new(),
487 /// )?;
488 /// assert!(outputs[0].trace.is_some());
489 /// # Ok::<(), Box<dyn std::error::Error>>(())
490 /// ```
491 pub fn record_graph<Op>(
492 &mut self,
493 graph: RecordedGraph<Op>,
494 inputs: &[EagerInput<Op>],
495 outputs: &[Arc<Op::Operand>],
496 retained_values: HashMap<ValueKey<Op>, Arc<Op::Operand>>,
497 ) -> EagerRecordResult<Vec<EagerOutput<Op>>>
498 where
499 Op: Primitive,
500 Op::InputKey: ADKey,
501 K: KeySource<Op>,
502 {
503 if inputs.len() != graph.input_keys().len() {
504 return Err(EagerRecordError::count_mismatch(
505 "Recorder::record_graph inputs",
506 graph.input_keys().len(),
507 inputs.len(),
508 ));
509 }
510 if outputs.len() != graph.output_keys().len() {
511 return Err(EagerRecordError::count_mismatch(
512 "Recorder::record_graph outputs",
513 graph.output_keys().len(),
514 outputs.len(),
515 ));
516 }
517 if outputs.len() > u8::MAX as usize + 1 {
518 return Err(EagerRecordError::too_many_outputs(outputs.len()));
519 }
520
521 let output_keys = fresh_value_keys(&mut self.key_source, outputs.len());
522 let requires_grad = inputs.iter().any(|input| input.requires_grad);
523
524 let trace = if requires_grad {
525 let saved_data = saved_graph_values(&graph, inputs, &retained_values);
526 Some(Trace::new(Arc::new(TraceNode::new(
527 graph,
528 output_keys.clone(),
529 saved_data,
530 inputs
531 .iter()
532 .map(|input| {
533 TraceEdge::new(
534 input.trace.as_ref().map(|trace| trace.node().clone()),
535 input.key.clone(),
536 input.requires_grad,
537 )
538 })
539 .collect(),
540 )?)))
541 } else {
542 None
543 };
544
545 Ok(output_keys
546 .into_iter()
547 .enumerate()
548 .map(|(output_slot, key)| EagerOutput {
549 key,
550 trace: trace.clone(),
551 requires_grad,
552 output_slot,
553 })
554 .collect())
555 }
556}
557
558fn saved_graph_values<Op: GraphOperation>(
559 graph: &RecordedGraph<Op>,
560 inputs: &[EagerInput<Op>],
561 retained_values: &HashMap<ValueKey<Op>, Arc<Op::Operand>>,
562) -> HashMap<ValueKey<Op>, Arc<Op::Operand>> {
563 let mut saved = HashMap::with_capacity(inputs.len() + retained_values.len());
564 for (input_key, input) in graph.input_keys().iter().zip(inputs.iter()) {
565 saved.insert(ValueKey::Input(input_key.clone()), input.data.clone());
566 }
567 for (key, value) in retained_values {
568 saved.insert(key.clone(), value.clone());
569 }
570 saved
571}
572
573fn fresh_value_keys<Op: GraphOperation>(
574 key_source: &mut impl KeySource<Op>,
575 count: usize,
576) -> Vec<ValueKey<Op>> {
577 (0..count)
578 .map(|_| ValueKey::Input(key_source.fresh_input_key()))
579 .collect()
580}
581
582#[cfg(test)]
583mod tests {
584 use super::*;
585 use crate::{DiffPassId, PrimitiveBuilder, PrimitiveValue};
586 use computegraph::LocalValueId;
587
588 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
589 enum Key {
590 User(&'static str),
591 Tangent { of: Box<Key>, pass: DiffPassId },
592 }
593
594 impl ADKey for Key {
595 fn tangent_of(&self, pass: DiffPassId) -> Self {
596 Self::Tangent {
597 of: Box::new(self.clone()),
598 pass,
599 }
600 }
601 }
602
603 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
604 enum Op {
605 Id,
606 }
607
608 impl GraphOperation for Op {
609 type Operand = f64;
610 type Context = ();
611 type InputKey = Key;
612
613 fn input_count(&self) -> usize {
614 1
615 }
616
617 fn output_count(&self) -> usize {
618 1
619 }
620 }
621
622 impl Primitive for Op {
623 type ADContext = ();
624
625 fn add() -> Self {
626 Self::Id
627 }
628
629 fn jvp_rule(
630 &self,
631 _builder: &mut impl PrimitiveBuilder<Self>,
632 _primal_in: &[ValueKey<Self>],
633 _primal_out: &[ValueKey<Self>],
634 tangent_in: &[Option<LocalValueId>],
635 _ctx: &mut (),
636 ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
637 Ok(tangent_in.to_vec())
638 }
639
640 fn transpose_rule(
641 &self,
642 _builder: &mut impl PrimitiveBuilder<Self>,
643 cotangent_out: &[Option<LocalValueId>],
644 _inputs: &[PrimitiveValue<Self>],
645 _role: &OperationRole,
646 _ctx: &mut (),
647 ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
648 Ok(cotangent_out.to_vec())
649 }
650 }
651
652 #[test]
653 fn too_many_outputs_error_formats_limit() {
654 let err = EagerRecordError::too_many_outputs(u8::MAX as usize + 2);
655
656 assert!(matches!(
657 err,
658 EagerRecordError::TooManyOutputs {
659 actual,
660 max: 256
661 } if actual == u8::MAX as usize + 2
662 ));
663 assert_eq!(
664 err.to_string(),
665 "eager recording supports at most 256 outputs, got 257"
666 );
667 }
668
669 #[test]
670 fn recorded_graph_linearize_rejects_unknown_output_slot() {
671 let recorded = RecordedGraph::from_primitive(Op::Id, vec![Key::User("x")])
672 .expect("test primitive graph metadata should be valid");
673
674 let err = match recorded.linearize(&[1], &mut ()) {
675 Ok(_) => panic!("RecordedGraph::linearize unexpectedly accepted bad output slot"),
676 Err(err) => err,
677 };
678
679 assert!(matches!(
680 err,
681 ADRuleError::InvalidInput {
682 op,
683 rule: ADRuleKind::Jvp,
684 message
685 } if op == "tidu::eager::RecordedGraph"
686 && message.contains("requested output slot 1")
687 ));
688 }
689}