From fd772accc16c735d3bba6002bbe859c0e5e4b67d Mon Sep 17 00:00:00 2001
From: nathaniel <nathaniel.simard.42@gmail.com>
Date: Thu, 2 May 2024 18:51:58 -0400
Subject: [PATCH] Cleanup

---
 .../src/runtime/memory_management.rs          | 51 ++++++++++---------
 1 file changed, 28 insertions(+), 23 deletions(-)

diff --git a/crates/burn-autodiff/src/runtime/memory_management.rs b/crates/burn-autodiff/src/runtime/memory_management.rs
index 93ba899ef4..84a36816e7 100644
--- a/crates/burn-autodiff/src/runtime/memory_management.rs
+++ b/crates/burn-autodiff/src/runtime/memory_management.rs
@@ -1,4 +1,4 @@
-use crate::{graph::Node, tensor::NodeRefCount, NodeID};
+use crate::{tensor::NodeRefCount, NodeID};
 use std::{
     collections::{HashMap, HashSet},
     mem,
@@ -19,28 +19,6 @@ enum NodeMemoryStatus {
     Unknown,
 }
 
-// Wrapper over hash set for fast popping of any node
-#[derive(new, Default)]
-struct PopNodeSet {
-    hash_set: HashSet<NodeID>,
-}
-
-impl PopNodeSet {
-    fn pop(&mut self) -> Option<NodeID> {
-        self.hash_set
-            .iter()
-            .next()
-            .copied()
-            .and_then(|node_id| self.hash_set.take(&node_id))
-    }
-    fn contains(&self, node_id: &NodeID) -> bool {
-        self.hash_set.contains(node_id)
-    }
-    fn insert(&mut self, node_id: NodeID) {
-        self.hash_set.insert(node_id);
-    }
-}
-
 impl GraphMemoryManagement {
     /// Register a new node with its parent.
     pub fn register(&mut self, node: NodeRefCount, parents: Vec<NodeID>) {
@@ -245,3 +223,30 @@ impl GraphMemoryManagement {
         }
     }
 }
+
+/// Wrapper over hash set for fast popping of any node
+#[derive(new, Default)]
+struct PopNodeSet {
+    hash_set: HashSet<NodeID>,
+}
+
+impl PopNodeSet {
+    #[inline(always)]
+    fn pop(&mut self) -> Option<NodeID> {
+        self.hash_set
+            .iter()
+            .next()
+            .copied()
+            .and_then(|node_id| self.hash_set.take(&node_id))
+    }
+
+    #[inline(always)]
+    fn contains(&self, node_id: &NodeID) -> bool {
+        self.hash_set.contains(node_id)
+    }
+
+    #[inline(always)]
+    fn insert(&mut self, node_id: NodeID) {
+        self.hash_set.insert(node_id);
+    }
+}