burn_autodiff/checkpoint/
base.rs

1use super::{
2    retro_forward::RetroForwards,
3    state::{BackwardStates, State},
4};
5use crate::collections::HashMap;
6use crate::graph::NodeId;
7
8use alloc::{format, vec, vec::Vec};
9use burn_std::config::{autodiff::AutodiffLogLevel, log_autodiff};
10
11#[derive(new, Debug)]
12/// Links a [NodeId] to its autodiff graph [NodeRef]
13pub(crate) struct NodeTree {
14    map: HashMap<NodeId, Vec<NodeId>>,
15}
16
17impl NodeTree {
18    /// Gives the parents of the node in the autodiff graph
19    pub(crate) fn parents(&self, node_id: &NodeId) -> Option<Vec<NodeId>> {
20        self.map.get(node_id).cloned()
21    }
22}
23
24#[derive(new, Debug)]
25/// Struct responsible of fetching the output for a node in the autodiff graph during a backward pass
26pub struct Checkpointer {
27    backward_states: BackwardStates,
28    retro_forwards: RetroForwards,
29    node_tree: NodeTree,
30}
31
32impl Checkpointer {
33    /// Gives the output of the given node, by recursively asking parents to compute themselves
34    /// or give their pre-computed tensors.
35    pub fn retrieve_node_output<T>(&mut self, node_id: NodeId) -> T
36    where
37        T: Clone + Send + 'static,
38    {
39        let sorted = self.topological_sort(node_id);
40        let num_nodes = sorted.len();
41        log_autodiff(AutodiffLogLevel::Basic, move || {
42            format!("retrieve_node_output {node_id:?}: {num_nodes} node(s) to compute")
43        });
44
45        sorted.into_iter().for_each(|node| {
46            log_autodiff(AutodiffLogLevel::Full, move || {
47                format!("execute_retro_forward {node:?}")
48            });
49            self.retro_forwards
50                .execute_retro_forward(node, &mut self.backward_states)
51        });
52
53        self.backward_states.get_state::<T>(&node_id)
54    }
55
56    /// Sorts the ancestors of NodeId in a way such that all parents come before their children
57    /// Useful to avoid recursivity later when mutating the states
58    ///
59    /// The sort on a compute bound state or a memory bound that is already computed is trivial.
60    /// The match on State::Computed also serves as a stopping criterion for the sort,
61    /// we don't need to look higher than that during recursivity.
62    fn topological_sort(&self, node_id: NodeId) -> Vec<NodeId> {
63        match self.backward_states.get_state_ref(&node_id) {
64            Some(state) => match state {
65                State::Recompute { n_required: _ } => {
66                    let mut sorted = Vec::new();
67                    let parents = self.node_tree.parents(&node_id).unwrap();
68                    for parent_node in parents {
69                        let parent_sorted = self.topological_sort(parent_node);
70                        for ps in parent_sorted {
71                            if !sorted.contains(&ps) {
72                                sorted.push(ps)
73                            }
74                        }
75                    }
76                    sorted.push(node_id);
77                    sorted
78                }
79                State::Computed {
80                    state_content: _,
81                    n_required: _,
82                } => vec![node_id],
83            },
84            None => panic!("Node {node_id:?} is not in the backward_states. "),
85        }
86    }
87
88    /// Checks if checkpointer has been drained adequately. Useful for testing
89    pub fn is_empty(&self) -> bool {
90        self.backward_states.is_empty() && self.retro_forwards.is_empty()
91    }
92}
burn_autodiff/checkpoint/base.rs

burn_autodiff/checkpoint/
base.rs
