Use a u64 counter for autodiff NodeIDs (tracel-ai#843)

gmh5225 · Oct 4, 2023 · 28e2a99 · 28e2a99
1 parent ce120ea
commit 28e2a99
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/burn-autodiff/src/grads.rs b/burn-autodiff/src/grads.rs
@@ -6,7 +6,7 @@ use crate::{
 };
 
 /// Gradient identifier.
-pub type GradID = String;
+pub type GradID = u64;
 
 /// Gradients container used during the backward pass.
 pub struct Gradients {
@@ -80,13 +80,11 @@ impl Gradients {
         value: TensorPrimitive<B, D>,
     ) {
         if let Some(tensor_old) = self.container.remove::<B, D>(&node.id.value) {
-            self.container.register(
-                node.id.value.clone(),
-                Tensor::from_primitive(value).add(tensor_old),
-            );
+            self.container
+                .register(node.id.value, Tensor::from_primitive(value).add(tensor_old));
         } else {
             self.container
-                .register::<B, D>(node.id.value.clone(), Tensor::from_primitive(value));
+                .register::<B, D>(node.id.value, Tensor::from_primitive(value));
         }
     }
 }
diff --git a/burn-autodiff/src/graph/node.rs b/burn-autodiff/src/graph/node.rs
@@ -1,7 +1,6 @@
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
-use burn_common::id::IdGenerator;
-
 use super::Requirement;
 
 /// A node contains graph metadata and should be used wrapped in an Arc for cheap cloning.
@@ -27,15 +26,18 @@ impl Node {
 /// Unique identifier generated for each [node](Node).
 #[derive(Clone, Hash, PartialEq, Eq, Debug)]
 pub struct NodeID {
-    pub value: String,
+    pub value: u64,
 }
 
 impl NodeID {
     /// Create a unique [node id](NodeID).
     pub fn new() -> Self {
-        Self {
-            value: IdGenerator::generate(),
+        static COUNTER: AtomicU64 = AtomicU64::new(0);
+        let value = COUNTER.fetch_add(1, Ordering::Relaxed);
+        if value == u64::MAX {
+            panic!("NodeID overflowed");
         }
+        Self { value }
     }
 }