Cleanup autodiff unused roots (tracel-ai#4039)

Author: laggui

Cargo.lock

@@ -64,6 +64,7 @@ indicatif = "0.18.0"
 js-sys = "0.3.77"
 libm = "0.2.15"
 log = { default-features = false, version = "0.4.28" }
+parking_lot = { version = "0.12.5", default-features = false }
 paste = "1"
 planus = { version = "=1.1" }
 polars = { version = "0.51.0", features = ["lazy"] }

Cargo.toml

@@ -17,7 +17,7 @@ workspace = true
 [features]
 default = ["std"]
 export_tests = ["burn-tensor-testgen"]
-std = []
+std = ["dep:parking_lot"]
 
 [dependencies]
 burn-common = { path = "../burn-common", version = "=0.20.0-pre.2", default-features = false }
@@ -26,11 +26,13 @@ burn-tensor-testgen = { path = "../burn-tensor-testgen", version = "=0.20.0-pre.
 
 derive-new = { workspace = true }
 spin = { workspace = true }
+parking_lot = { workspace = true, optional = true }
 log = { workspace = true }
 hashbrown = { workspace = true }
 num-traits = { workspace = true }
 portable-atomic = { workspace = true }
 
+
 [dev-dependencies]
 burn-tensor = { path = "../burn-tensor", version = "=0.20.0-pre.2", default-features = false, features = [
     "export_tests",

crates/burn-autodiff/Cargo.toml

@@ -10,15 +10,6 @@ use burn_tensor::backend::Backend;
 pub trait AutodiffClient: Send + Clone {
     /// Register a new step.
     fn register(&self, node_id: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder);
-    /// Register a new untracked step.
-    fn register_untracked(
-        &self,
-        node_id: NodeRefCount,
-        step: StepBoxed,
-        actions: CheckpointerBuilder,
-    ) {
-        self.register(node_id, step, actions);
-    }
     /// Call backpropagation from the given tensor.
     fn backward<B: Backend>(&self, tensor: AutodiffTensor<B>) -> Gradients;
 }

crates/burn-autodiff/src/runtime/client.rs

@@ -8,10 +8,16 @@ use crate::{
     tensor::{AutodiffTensor, NodeRefCount},
 };
 use alloc::sync::Arc;
-use burn_common::stub::Mutex;
+use alloc::vec::Vec;
 use burn_tensor::backend::Backend;
 use hashbrown::{HashMap, HashSet};
 
+#[cfg(feature = "std")]
+use parking_lot::{Mutex, MutexGuard};
+
+#[cfg(not(feature = "std"))]
+use spin::{Mutex, MutexGuard};
+
 /// A client for managing multiple graphs using mutex-based synchronization.
 ///
 /// The biggest benefit of using this client implementation is that each graph can modify its own
@@ -40,7 +46,6 @@ pub struct GraphLocator {
     /// This is to ensure that when merging graphs, we correctly move all previous graphs to
     /// the new merged one.
     keys: HashMap<NodeId, HashSet<NodeId>>,
-    untracked_nodes: HashSet<NodeId>,
 }
 
 /// Represents a single computation graph with a mutex-protected server.
@@ -65,7 +70,7 @@ impl core::fmt::Debug for Graph {
     }
 }
 
-static STATE: spin::Mutex<Option<GraphLocator>> = spin::Mutex::new(None);
+static STATE: Mutex<Option<GraphLocator>> = Mutex::new(None);
 
 impl GraphMutexClient {
     /// Retrieves or creates a graph for the given [NodeId] and parent dependencies.
@@ -95,52 +100,62 @@ impl AutodiffClient for GraphMutexClient {
     fn register(&self, node_id_ref: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder) {
         let node_id = *node_id_ref;
         let graph = GraphMutexClient::graph(node_id, step.parents());
-        let mut state = graph.state.lock().unwrap();
+        let mut state = graph.state.lock();
 
         state.server.register(node_id_ref, step, actions);
     }
 
-    fn register_untracked(
-        &self,
-        node_id_ref: NodeRefCount,
-        step: StepBoxed,
-        actions: CheckpointerBuilder,
-    ) {
-        let node_id = *node_id_ref;
-        // Register normally (might be needed by tracked children, required for checkpointing)
-        self.register(node_id_ref, step, actions);
-
-        // But mark this node for cleanup after backward
-        let mut state = STATE.lock();
-        if let Some(locator) = state.as_mut() {
-            locator.mark_untracked(node_id);
-        }
-    }
-
     fn backward<B: Backend>(&self, root: AutodiffTensor<B>) -> Gradients {
         let node_id = root.node.id;
         let graph = GraphMutexClient::graph(root.node.id, &[]);
 
         let grads = Gradients::new::<B>(root.node, root.primitive);
-        let mut state = graph.state.lock().unwrap();
-
-        let grads = state.server.backward::<GraphCleaner>(grads, node_id);
+        let grads = {
+            let mut state = graph.state.lock();
+            state.server.backward::<GraphCleaner>(grads, node_id)
+        }; // lock released
 
-        let mut cleaner = GraphCleaner::init();
-        cleaner.cleanup_orphaned_entries();
+        GraphCleaner::cleanup_orphaned_entries();
 
         grads
     }
 }
 
 struct GraphCleaner<'a> {
-    guard: spin::MutexGuard<'a, Option<GraphLocator>>,
+    guard: MutexGuard<'a, Option<GraphLocator>>,
 }
 
 impl<'a> GraphCleaner<'a> {
-    fn cleanup_orphaned_entries(&mut self) {
-        if let Some(state) = self.guard.as_mut() {
-            state.cleanup_untracked();
+    fn cleanup_orphaned_entries() {
+        let graphs = {
+            // Get the available graphs and release the lock
+            match STATE.lock().as_ref() {
+                Some(state) => state.graphs.clone(),
+                None => return,
+            }
+        };
+
+        let mut should_remove = Vec::new();
+        for graph in graphs.values() {
+            {
+                let mut guard = graph.state.lock();
+                // Double safety: in case it was marked as no longer useful, but other
+                // nodes are still relevant, we only check which nodes can safely be removed.
+                if !guard.server.maybe_useful() {
+                    guard
+                        .server
+                        .free_unused_roots(|node| should_remove.push(*node));
+                }
+            }
+        }
+
+        if !should_remove.is_empty() {
+            let mut state = STATE.lock();
+            if let Some(state) = state.as_mut() {
+                for node in should_remove {
+                    state.graphs.remove(&node);
+                }
+            }
         }
     }
 }
@@ -230,7 +245,7 @@ impl GraphLocator {
         let main = graphs.next().expect("At least one graph");
         self.register_key(main.origin, node);
 
-        let mut state = main.state.lock().unwrap();
+        let mut state = main.state.lock();
 
         for graph in graphs {
             self.merge_two(&mut state, &main, graph);
@@ -259,7 +274,7 @@ impl GraphLocator {
 
     /// Merges two graphs by combining their states and updating graph mappings.
     fn merge_two(&mut self, main_state: &mut GraphState, main: &Arc<Graph>, merged: Arc<Graph>) {
-        let mut locked = merged.state.lock().unwrap();
+        let mut locked = merged.state.lock();
         let mut state_old = GraphState::default();
         core::mem::swap(&mut state_old, &mut locked);
         main_state.server.extend(state_old.server);
@@ -292,18 +307,6 @@ impl GraphLocator {
         graph
     }
 
-    fn mark_untracked(&mut self, node_id: NodeId) {
-        self.untracked_nodes.insert(node_id);
-    }
-
-    /// Clean up untracked nodes.
-    fn cleanup_untracked(&mut self) {
-        let mut nodes = core::mem::take(&mut self.untracked_nodes);
-        for node in nodes.drain() {
-            self.remove_entry(&node);
-        }
-    }
-
     fn remove_entry(&mut self, node: &NodeId) {
         if let Some(graph) = self.graphs.remove(node) {
             let mut remove = false;

crates/burn-autodiff/src/runtime/graph.rs

@@ -89,7 +89,17 @@ impl GraphMemoryManagement {
         }
     }
 
-    fn clear_unused_roots(&mut self, to_delete: &mut Vec<NodeId>) {
+    pub(crate) fn free_unused_roots(&mut self, mut on_free_graph: impl FnMut(&NodeId)) {
+        let mut deletables = Vec::new();
+        self.clear_unused_roots(&mut deletables);
+
+        for node_id in deletables {
+            self.nodes.remove(&node_id);
+            on_free_graph(&node_id);
+        }
+    }
+
+    fn clear_unused_roots(&self, to_delete: &mut Vec<NodeId>) {
         for (id, parents) in self.nodes.iter() {
             let is_useful = matches!(
                 self.statuses.get(id.as_ref()),
@@ -250,6 +260,10 @@ impl GraphMemoryManagement {
             None => panic!("Node should be in the nodes map"),
         }
     }
+
+    pub(crate) fn maybe_useful(&self) -> bool {
+        self.nodes.keys().any(|node| Arc::strong_count(node) > 1)
+    }
 }
 
 /// Wrapper over hash set for fast popping of any node

crates/burn-autodiff/src/runtime/memory_management.rs

@@ -71,6 +71,14 @@ impl AutodiffServer {
         gradients
     }
 
+    pub(crate) fn free_unused_roots(&mut self, mut on_free_graph: impl FnMut(&NodeId)) {
+        self.memory_management.free_unused_roots(|node_id| {
+            self.steps.remove(node_id);
+            self.actions_builder.remove(node_id);
+            on_free_graph(node_id);
+        });
+    }
+
     fn build_tape(
         &mut self,
         node: NodeId,
@@ -128,4 +136,8 @@ impl AutodiffServer {
 
         grads
     }
+
+    pub(crate) fn maybe_useful(&self) -> bool {
+        self.memory_management.maybe_useful()
+    }
 }

crates/burn-autodiff/src/runtime/server.rs

@@ -156,19 +156,11 @@ impl<B: Backend> AutodiffTensor<B> {
         step_that_created_the_tensor: S,
         actions: CheckpointerBuilder,
     ) -> Self {
-        if self.is_tracked() {
-            self.node.client.register(
-                self.rc.clone(),
-                Box::new(step_that_created_the_tensor),
-                actions,
-            );
-        } else {
-            self.node.client.register_untracked(
-                self.rc.clone(),
-                Box::new(step_that_created_the_tensor),
-                actions,
-            );
-        }
+        self.node.client.register(
+            self.rc.clone(),
+            Box::new(step_that_created_the_tensor),
+            actions,
+        );
         self
     }