Impl Hash for Id

Author: elftausend

src/cache/location_hasher.rs

@@ -1,21 +1,21 @@
 use core::{fmt::Display, ops::BitXor, panic::Location};
 
 #[derive(Default)]
-pub struct LocationHasher {
+pub struct FxHasher {
     hash: u64,
 }
 
 const K: u64 = 0x517cc1b727220a95;
 
-impl core::hash::Hasher for LocationHasher {
+impl core::hash::Hasher for FxHasher {
     #[inline]
     fn finish(&self) -> u64 {
         self.hash
     }
 
     #[inline]
     fn write(&mut self, _bytes: &[u8]) {
-        unimplemented!("LocationHasher only hashes u64, (u32 and usize as u64 cast).")
+        unimplemented!("(this) FxHasher only hashes u64, (u32 and usize as u64 cast).")
     }
 
     #[inline]

src/cache/owned_cache.rs

@@ -3,12 +3,12 @@ use std::collections::HashMap;
 
 use std::rc::Rc;
 
-use crate::{flag::AllocFlag, Alloc, Device, HashLocation, LocationHasher, ShallowCopy, Shape};
+use crate::{flag::AllocFlag, Alloc, Device, HashLocation, FxHasher, ShallowCopy, Shape};
 
 #[derive(Debug, Clone)]
 pub struct Cache {
     pub nodes:
-        HashMap<HashLocation<'static>, Rc<dyn core::any::Any>, BuildHasherDefault<LocationHasher>>,
+        HashMap<HashLocation<'static>, Rc<dyn core::any::Any>, BuildHasherDefault<FxHasher>>,
 }
 
 impl Default for Cache {

src/devices/opencl/unified.rs

@@ -6,7 +6,7 @@ use crate::{AllocFlag, DeviceError};
 
 use super::CLPtr;
 use crate::{
-    Base, Buffer, CachedCPU, CachedModule, Device, HashLocation, LocationHasher, OnDropBuffer,
+    Base, Buffer, CachedCPU, CachedModule, Device, HashLocation, FxHasher, OnDropBuffer,
     OpenCL, Shape, UnifiedMemChain, CPU,
 };
 use min_cl::api::{create_buffer, MemFlags};
@@ -62,7 +62,7 @@ pub unsafe fn to_cached_unified<OclMods, CpuMods, T, S>(
     cache: &mut HashMap<
         HashLocation<'static>,
         Rc<dyn core::any::Any>,
-        BuildHasherDefault<LocationHasher>,
+        BuildHasherDefault<FxHasher>,
     >,
     location: HashLocation<'static>,
 ) -> crate::Result<*mut c_void>
@@ -126,7 +126,7 @@ pub fn construct_buffer<'a, OclMods: OnDropBuffer, CpuMods: OnDropBuffer, T: 'st
     cache: &mut HashMap<
         HashLocation<'static>,
         Rc<dyn core::any::Any>,
-        BuildHasherDefault<LocationHasher>,
+        BuildHasherDefault<FxHasher>,
     >,
     location: HashLocation<'static>,
 ) -> crate::Result<Buffer<'a, T, OpenCL<OclMods>, S>> {

src/id.rs

@@ -20,6 +20,13 @@ pub struct Id {
     pub len: usize,
 }
 
+impl core::hash::Hash for Id {
+    #[inline]
+    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
+        self.id.hash(state);
+    }
+}
+
 impl Deref for Id {
     type Target = u64;
 

src/modules/autograd/tape.rs

@@ -2,7 +2,7 @@ use core::{any::Any, fmt::Debug, hash::BuildHasherDefault, panic::Location};
 use std::collections::{HashMap, HashSet};
 
 use crate::{
-    Alloc, Buffer, Buffers, HasId, HashLocation, LazyGraph, LocationHasher, Parents, Shape,
+    Alloc, Buffer, Buffers, HasId, HashLocation, LazyGraph, FxHasher, Parents, Shape,
     TapeActions, UpdateArgs, WriteBuf,
 };
 
@@ -16,7 +16,7 @@ pub struct Tape {
     // Caches gradients for each [`Buffer`]'s id ([`Ident`]).
     // pub grads: Gradients,
     grad_fns: Vec<GradFn>,
-    grad_fns_loc: HashMap<HashLocation<'static>, GradFn, BuildHasherDefault<LocationHasher>>,
+    grad_fns_loc: HashMap<HashLocation<'static>, GradFn, BuildHasherDefault<FxHasher>>,
     grad_fn_order: Vec<HashLocation<'static>>,
 
     unconsumed_locations: HashSet<HashLocation<'static>>,

src/parents.rs

@@ -1,8 +1,17 @@
-use crate::{HasId, Id, UpdateArg};
+use core::hash::Hasher;
+
+use crate::{FxHasher, HasId, Id, UpdateArg};
 
 pub trait Parents<const N: usize>: AllParents {
     fn ids(&self) -> [Id; N];
     fn maybe_ids(&self) -> [Option<Id>; N];
+
+    #[inline]
+    fn hash(&self) -> u64 {
+        let mut hasher = FxHasher::default();
+        core::hash::Hash::hash(&self.ids(), &mut hasher);
+        hasher.finish()
+    }
 }
 
 impl Parents<0> for () {
@@ -103,3 +112,67 @@ impl<T: HasId + Copy, const N: usize> Parents<N> for [T; N] {
 impl<T: HasId + Copy, const N: usize> AllParents for [T; N] {}
 
 pub trait AllParents {}
+
+#[cfg(test)]
+mod tests {
+
+    #[cfg(feature = "std")]
+    #[ignore = "slow"]
+    #[test]
+    fn test_collisions() {
+        use std::collections::HashSet;
+        use crate::{Id, Parents};
+
+        let handle = std::thread::spawn(|| {
+            let mut hashes = HashSet::new();
+            for i in 20000..30000u16 {
+                for j in 20000..30000 {
+                    let i = Id {
+                        id: i as u64,
+                        len: 0,
+                    };
+                    let j = Id {
+                        id: j,
+                        len: 0,
+                    };
+                    let parents = (i, j);
+                    let hash = parents.hash();
+                    if hashes.contains(&(hash)) {
+                        panic!("collision {}, {}, hash: {hash}", i.id, j.id,);
+                    }
+                    hashes.insert(hash);
+                }
+                if i % 1000 == 0 {
+                    println!("i: {}", i);
+                }
+            }
+            hashes
+        });
+        let mut hashes = HashSet::new();
+
+        for i in 10000..20000 {
+            for j in 10000..20000 {
+                let i = Id {
+                    id: i,
+                    len: 0,
+                };
+                let j = Id {
+                    id: j,
+                    len: 0,
+                };
+                let parents = (i, j);
+                let hash = parents.hash();
+                if hashes.contains(&(hash)) {
+                    panic!("collision");
+                }
+                hashes.insert(hash);
+            }
+            if i % 1000 == 0 {
+                println!("i: {}", i);
+            }
+        }
+
+        let other_hashes = handle.join().unwrap();
+        assert_eq!(hashes.intersection(&other_hashes).count(), 0);
+    }
+}