burn_autodiff/ops/
base.rs

1use super::Backward;
2use crate::{
3    checkpoint::{
4        base::Checkpointer,
5        builder::{ActionType, CheckpointerBuilder},
6        retro_forward::RetroForward,
7        strategy::CheckpointStrategy,
8    },
9    grads::Gradients,
10    graph::{ComputingProperty, NodeId, NodeRef, Parent, Requirement, Step},
11    tensor::AutodiffTensor,
12};
13use alloc::boxed::Box;
14use burn_backend::{Backend, TensorMetadata, tensor::FloatTensor};
15use burn_std::Shape;
16use core::marker::PhantomData;
17
18#[cfg(feature = "distributed")]
19use burn_backend::distributed::DistributedParams;
20
21/// Operation in preparation.
22///
23/// Each mode has its own set of functions to minimize cloning for unused backward states.
24#[derive(new)]
25pub struct OpsPrep<Backward, B, S, C, const N: usize, Mode = Init> {
26    nodes: [NodeRef; N],
27    requirement: Requirement,
28    backward: Backward,
29    compute_property: ComputingProperty,
30    checkpointer_builder: CheckpointerBuilder,
31    checkpoint_strategy: PhantomData<C>,
32    phantom_backend: PhantomData<B>,
33    phantom_state: PhantomData<S>,
34    marker: PhantomData<Mode>,
35}
36
37/// Operation is initialized
38pub struct Init;
39/// Operation has been tagged as memory bound
40pub struct MemoryBound;
41/// Memory bound operation has received its RetroForward
42pub struct MemoryBoundRetroForward;
43/// Operation's compute property is fixed
44pub struct ComputePropertyDone;
45/// Tracked operation tag.
46pub struct Tracked;
47/// Untracked operation tag.
48pub struct UnTracked;
49
50impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, Init>
51where
52    B: Backend,
53    BO: Backward<B, N, State = S>,
54{
55    /// Indicates that the operation is compute bound, meaning its computation
56    /// is heavy and should not be recomputed
57    pub fn compute_bound(self) -> OpsPrep<BO, B, S, C, N, ComputePropertyDone> {
58        OpsPrep::new(
59            self.nodes,
60            self.requirement,
61            self.backward,
62            ComputingProperty::ComputeBound,
63            self.checkpointer_builder,
64        )
65    }
66
67    /// Indicates that the operation is memory bound, meaning its computation
68    /// is light and can be recomputed
69    pub fn memory_bound(self) -> OpsPrep<BO, B, S, C, N, MemoryBound> {
70        OpsPrep::new(
71            self.nodes,
72            self.requirement,
73            self.backward,
74            self.compute_property,
75            self.checkpointer_builder,
76        )
77    }
78}
79
80impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, MemoryBound>
81where
82    B: Backend,
83    BO: Backward<B, N, State = S>,
84    C: CheckpointStrategy,
85{
86    /// Registers the retro forward, if needed
87    pub fn retro_forward<R: RetroForward>(
88        self,
89        retro_forward: R,
90    ) -> OpsPrep<BO, B, S, C, N, MemoryBoundRetroForward> {
91        OpsPrep::new(
92            self.nodes,
93            self.requirement,
94            self.backward,
95            C::compute_property(retro_forward),
96            self.checkpointer_builder,
97        )
98    }
99}
100
101impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, MemoryBoundRetroForward>
102where
103    B: Backend,
104    BO: Backward<B, N, State = S>,
105    C: CheckpointStrategy,
106{
107    /// Checkpoints the parents, if needed
108    pub fn parents<'a, B2, A>(mut self, parents: A) -> OpsPrep<BO, B, S, C, N, ComputePropertyDone>
109    where
110        B2: Backend,
111        A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
112    {
113        let compute_property = match C::checkpoint_parents(parents, &mut self.checkpointer_builder)
114        {
115            Ok(..) => self.compute_property,
116            Err(..) => ComputingProperty::ComputeBound,
117        };
118
119        OpsPrep::new(
120            self.nodes,
121            self.requirement,
122            self.backward,
123            compute_property,
124            self.checkpointer_builder,
125        )
126    }
127}
128
129impl<BO, B, C, const N: usize> OpsPrep<BO, B, (), C, N, ComputePropertyDone>
130where
131    B: Backend,
132    BO: Backward<B, N, State = ()>,
133{
134    /// Prepare a stateless operation.
135    pub fn stateless(self, output: FloatTensor<B>) -> AutodiffTensor<B> {
136        match self.stateful() {
137            OpsKind::Tracked(prep) => prep.finish((), output),
138            OpsKind::UnTracked(prep) => prep.finish(output),
139        }
140    }
141}
142
143impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, ComputePropertyDone>
144where
145    B: Backend,
146    S: Clone + Send + core::fmt::Debug + 'static,
147    BO: Backward<B, N, State = S>,
148{
149    /// Prepare an operation that requires a state during the backward pass.
150    pub fn stateful(self) -> OpsKind<BO, B, S, C, N> {
151        match self.requirement.is_none() {
152            false => OpsKind::Tracked(OpsPrep::new(
153                self.nodes,
154                self.requirement,
155                self.backward,
156                self.compute_property,
157                self.checkpointer_builder,
158            )),
159            true => OpsKind::UnTracked(OpsPrep::new(
160                self.nodes,
161                self.requirement,
162                self.backward,
163                self.compute_property,
164                self.checkpointer_builder,
165            )),
166        }
167    }
168}
169
170impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, UnTracked>
171where
172    B: Backend,
173    S: Clone + Send + core::fmt::Debug + 'static,
174    BO: Backward<B, N, State = S>,
175{
176    /// Finish the preparation of an untracked operation and returns the output tensor.
177    pub fn finish(self, output: FloatTensor<B>) -> AutodiffTensor<B> {
178        let output = AutodiffTensor::from_parents(
179            output,
180            &self.nodes,
181            self.requirement,
182            self.compute_property,
183        );
184        let parents = self.nodes.map(|node| node.clone_if_require_grad());
185        let ops = Ops::new(parents, output.node.clone(), ());
186
187        // We register the ops in the graph even if untracked, otherwise memory bound operations
188        // that have an untracked parent would not be able to retrieve it
189        output.register_step(UntrackedOpsStep::new(ops), self.checkpointer_builder)
190    }
191}
192
193impl<BO, B, S, C, const N: usize> OpsPrep<BO, B, S, C, N, Tracked>
194where
195    B: Backend,
196    S: Clone + Send + core::fmt::Debug + 'static,
197    BO: Backward<B, N, State = S>,
198{
199    /// Finish the preparation of a tracked operation and returns the output tensor.
200    pub fn finish(self, state: S, output: FloatTensor<B>) -> AutodiffTensor<B> {
201        let output = AutodiffTensor::from_parents(
202            output,
203            &self.nodes,
204            self.requirement,
205            self.compute_property,
206        );
207        let parents = self.nodes.map(|node| node.clone_if_require_grad());
208        let ops = Ops::new(parents, output.node.clone(), state);
209
210        output.register_step(OpsStep::new(ops, self.backward), self.checkpointer_builder)
211    }
212
213    /// Checkpoints the tensor
214    pub fn checkpoint(&mut self, tensor: &AutodiffTensor<B>) -> NodeId {
215        self.checkpointer_builder
216            .checkpoint(tensor, ActionType::Explicit);
217
218        tensor.node.id
219    }
220}
221
222/// Enum used before finishing tracked and untracked operations.
223pub enum OpsKind<BO, B, S, C, const N: usize> {
224    /// Tracked operation preparation.
225    Tracked(OpsPrep<BO, B, S, C, N, Tracked>),
226    /// Untracked operation preparation.
227    UnTracked(OpsPrep<BO, B, S, C, N, UnTracked>),
228}
229
230/// Operation containing its parent nodes, its own node and the backward step state.
231#[derive(new, Debug)]
232pub struct Ops<S, const N: usize> {
233    /// Parents nodes.
234    pub parents: [Option<NodeRef>; N],
235    /// The node.
236    pub node: NodeRef,
237    /// The state.
238    pub state: S,
239}
240
241/// Operation implementing backward [step](Step) with type erasing.
242#[derive(new, Debug)]
243struct OpsStep<B, T, SB, const N: usize>
244where
245    B: Backend,
246    T: Backward<B, N, State = SB>,
247    SB: Clone + Send + core::fmt::Debug + 'static,
248{
249    ops: Ops<SB, N>,
250    backward: T,
251    phantom: PhantomData<B>,
252}
253
254impl<B, T, SB, const N: usize> Step for OpsStep<B, T, SB, N>
255where
256    B: Backend,
257    T: Backward<B, N, State = SB>,
258    SB: Clone + Send + core::fmt::Debug + 'static,
259{
260    fn step(self: Box<Self>, grads: &mut Gradients, checkpointer: &mut Checkpointer) {
261        self.backward.backward(self.ops, grads, checkpointer);
262    }
263
264    fn node(&self) -> NodeId {
265        self.ops.node.id
266    }
267
268    fn parents(&self) -> &[Parent] {
269        &self.ops.node.parents
270    }
271
272    fn depth(&self) -> usize {
273        self.ops.node.order
274    }
275
276    #[cfg(feature = "distributed")]
277    fn distributed_params(&self) -> Option<DistributedParams> {
278        self.ops.node.distributed_params.clone()
279    }
280}
281
282#[derive(new, Debug)]
283struct UntrackedOpsStep<const N: usize> {
284    ops: Ops<(), N>,
285}
286
287impl<const N: usize> Step for UntrackedOpsStep<N> {
288    fn step(self: Box<Self>, _grads: &mut Gradients, _checkpointer: &mut Checkpointer) {
289        // Nothing to do
290    }
291
292    fn node(&self) -> NodeId {
293        self.ops.node.id
294    }
295
296    fn parents(&self) -> &[Parent] {
297        &self.ops.node.parents
298    }
299    fn depth(&self) -> usize {
300        self.ops.node.order
301    }
302
303    #[cfg(feature = "distributed")]
304    fn distributed_params(&self) -> Option<DistributedParams> {
305        self.ops.node.distributed_params.clone()
306    }
307}
308
309/// Make sure the grad tensor has the given shape.
310///
311/// If broadcasting happened during the forward pass, the gradients will be sum along the
312/// broadcasted dimension.
313pub fn broadcast_shape<B: Backend>(mut grad: FloatTensor<B>, shape: &Shape) -> FloatTensor<B> {
314    let shape_grad = grad.shape();
315    let ndims = shape_grad.num_dims();
316
317    for i in 0..ndims {
318        if shape_grad[i] != shape[i] {
319            if shape[i] != 1 {
320                panic!(
321                    "Invalid broadcast shapes: Next grad shape {:?}, Previous grad shape {:?}. {}",
322                    shape, shape_grad, "Expected the shape of the next grad to be 1."
323                );
324            }
325            grad = B::float_sum_dim(grad, i);
326        }
327    }
328
329    grad
330}
burn_autodiff/ops/base.rs

burn_autodiff/ops/
base.rs
