Add shared state map (#934)

Author: nathanielsimard

crates/cubecl-common/src/bytes/default_controller.rs

@@ -2,6 +2,7 @@
 
 use crate::bytes::{AllocationController, AllocationError};
 use alloc::alloc::Layout;
+use alloc::vec::Vec;
 use bytemuck::Contiguous;
 use core::{alloc::LayoutError, marker::PhantomData, mem::MaybeUninit, ptr::NonNull};
 

crates/cubecl-common/src/lib.rs

@@ -15,6 +15,9 @@ pub mod rand;
 /// Device module.
 pub mod device;
 
+/// Map utilities and implementations.
+pub mod map;
+
 /// Utilities module to manipulate bytes.
 #[cfg(feature = "serde")]
 pub mod bytes;

crates/cubecl-common/src/map.rs

@@ -0,0 +1,102 @@
+use crate::stub::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
+use alloc::sync::Arc;
+use hashbrown::HashMap;
+
+/// A thread-safe map that allows concurrent access to values using read-write locks.
+pub struct SharedStateMap<K, V> {
+    state: Mutex<Option<HashMap<K, Arc<RwLock<V>>>>>,
+}
+
+/// A value in the [SharedStateMap] that provides read and write access.
+pub struct SharedState<V> {
+    val: Arc<RwLock<V>>,
+}
+
+impl<V> SharedState<V> {
+    /// Acquires a read lock on the value, returning a read guard.
+    pub fn read(&self) -> RwLockReadGuard<'_, V> {
+        self.val.read().unwrap()
+    }
+
+    /// Acquires a write lock on the value, returning a write guard.
+    pub fn write(&self) -> RwLockWriteGuard<'_, V> {
+        self.val.write().unwrap()
+    }
+}
+
+impl<K, V> Default for SharedStateMap<K, V>
+where
+    K: core::hash::Hash + core::cmp::PartialEq + core::cmp::Eq,
+{
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<K, V> SharedStateMap<K, V>
+where
+    K: core::hash::Hash + core::cmp::PartialEq + core::cmp::Eq,
+{
+    /// Creates a new, empty `SharedStateMap`.
+    pub const fn new() -> Self {
+        Self {
+            state: Mutex::new(None),
+        }
+    }
+
+    /// Retrieves a value associated with the given key, if it exists.
+    pub fn get(&self, k: &K) -> Option<SharedState<V>> {
+        let mut state = self.state.lock().unwrap();
+        let map = get_or_init(&mut state);
+
+        match map.get(k) {
+            Some(val) => Some(SharedState { val: val.clone() }),
+            None => None,
+        }
+    }
+
+    /// Retrieves a value associated with the given key, or inserts a new value using the provided
+    /// initializer function if the key does not exist.
+    pub fn get_or_init<Fn: FnMut(&K) -> V>(&self, k: &K, mut init: Fn) -> SharedState<V>
+    where
+        K: Clone,
+    {
+        let mut state = self.state.lock().unwrap();
+        let map = get_or_init(&mut state);
+
+        match map.get(k) {
+            Some(val) => SharedState { val: val.clone() },
+            None => {
+                let val = init(k);
+                let val = Arc::new(RwLock::new(val));
+                map.insert(k.clone(), val.clone());
+                SharedState { val: val.clone() }
+            }
+        }
+    }
+
+    /// Inserts a key-value pair into the map.
+    pub fn insert(&self, k: K, v: V) {
+        let mut state = self.state.lock().unwrap();
+        let map = get_or_init(&mut state);
+
+        map.insert(k, Arc::new(RwLock::new(v)));
+    }
+
+    /// Clears the map, removing all key-value pairs.
+    pub fn clear(&self) {
+        let mut state = self.state.lock().unwrap();
+        let map = get_or_init(&mut state);
+        map.clear();
+    }
+}
+
+fn get_or_init<T: Default>(state: &mut Option<T>) -> &mut T {
+    match state {
+        Some(state) => state,
+        None => {
+            *state = Some(T::default());
+            state.as_mut().unwrap()
+        }
+    }
+}

crates/cubecl-common/src/stub.rs

@@ -1,14 +1,15 @@
 #[cfg(not(feature = "std"))]
-use spin::{
-    Mutex as MutexImported, MutexGuard, Once as OnceImported, RwLock as RwLockImported,
-    RwLockReadGuard, RwLockWriteGuard,
-};
+use spin::{Mutex as MutexImported, MutexGuard, Once as OnceImported, RwLock as RwLockImported};
 #[cfg(feature = "std")]
 use std::sync::{
     Mutex as MutexImported, MutexGuard, OnceLock as OnceImported, RwLock as RwLockImported,
-    RwLockReadGuard, RwLockWriteGuard,
 };
 
+#[cfg(not(feature = "std"))]
+pub use spin::{RwLockReadGuard, RwLockWriteGuard};
+#[cfg(feature = "std")]
+pub use std::sync::{RwLockReadGuard, RwLockWriteGuard};
+
 /// A mutual exclusion primitive useful for protecting shared data
 ///
 /// This mutex will block threads waiting for the lock to become available. The

crates/cubecl-runtime/src/tune/local.rs

@@ -9,6 +9,7 @@ use core::{
     fmt::Display,
     hash::Hash,
 };
+use cubecl_common::map::SharedStateMap;
 use hashbrown::HashMap;
 
 #[cfg(not(feature = "std"))]
@@ -17,7 +18,7 @@ use alloc::string::ToString;
 /// A local tuner allows to create a tuner for a specific key that can be different from the server
 /// key.
 pub struct LocalTuner<AK: AutotuneKey, ID> {
-    state: spin::Mutex<Option<HashMap<ID, Tuner<AK>>>>,
+    state: SharedStateMap<ID, Tuner<AK>>,
     name: &'static str,
     sets: spin::RwLock<Option<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>>,
 }
@@ -45,7 +46,7 @@ where
     /// Create a new local tuner.
     pub const fn new(name: &'static str) -> Self {
         Self {
-            state: spin::Mutex::new(None),
+            state: SharedStateMap::new(),
             name,
             sets: spin::RwLock::new(None),
         }
@@ -94,8 +95,7 @@ where
 
     /// Clear the autotune state.
     pub fn clear(&self) {
-        let mut state = self.state.lock();
-        *state = None;
+        self.state.clear()
     }
 
     #[cfg(feature = "autotune-checks")]
@@ -131,20 +131,16 @@ where
 
         // If this is cached and ready, use the operation.
         let autotune_job = {
-            let mut state = self.state.lock();
-            let map = state.get_or_insert_with(Default::default);
-
-            let tuner = match map.get_mut(id) {
-                Some(val) => val,
-                None => map.entry(id.clone()).or_insert_with(move || {
-                    let name = self.name.replace("::", "-");
-                    Tuner::new(&name, &id.to_string())
-                }),
-            };
+            let tuner_state = self.state.get_or_init(id, move |id| {
+                let name = self.name.replace("::", "-");
+                Tuner::new(&name, &id.to_string())
+            });
+            let tuner = tuner_state.read();
 
-            match tuner.fastest(&key) {
+            let mut tuner = match tuner.fastest(&key) {
                 TuneCacheResult::Hit { fastest_index } => {
-                    core::mem::drop(state);
+                    core::mem::drop(tuner);
+                    core::mem::drop(tuner_state);
 
                     #[cfg(feature = "autotune-checks")]
                     self.checks(&operations, &inputs);
@@ -156,7 +152,8 @@ where
                     return result;
                 }
                 TuneCacheResult::Pending => {
-                    core::mem::drop(state);
+                    core::mem::drop(tuner);
+                    core::mem::drop(tuner_state);
 
                     let op = operations.fastest(0);
                     let result = op
@@ -166,26 +163,38 @@ where
                 }
                 #[cfg(std_io)]
                 TuneCacheResult::Unchecked => {
+                    core::mem::drop(tuner);
+                    let mut tuner = tuner_state.write();
+
                     // If the cache checksum hasn't been checked, do so now, and retry.
                     let checksum = operations.compute_checksum();
                     tuner.validate_checksum(&key, &checksum);
 
                     // Check if with validation we can use its result
                     if let TuneCacheResult::Hit { fastest_index } = tuner.fastest(&key) {
-                        core::mem::drop(state);
+                        core::mem::drop(tuner);
+                        core::mem::drop(tuner_state);
 
                         let op = operations.fastest(fastest_index);
                         let result = op
                             .execute(inputs)
                             .expect("Should run when selected by autotune.");
                         return result;
                     }
+
+                    tuner
                 }
 
                 #[cfg(not(std_io))]
-                TuneCacheResult::Unchecked => (),
-                TuneCacheResult::Miss => (),
-            }
+                TuneCacheResult::Unchecked => {
+                    core::mem::drop(tuner);
+                    tuner_state.write()
+                }
+                TuneCacheResult::Miss => {
+                    core::mem::drop(tuner);
+                    tuner_state.write()
+                }
+            };
 
             if tuner.autotuning.contains(&key) {
                 Box::new(move || {})
@@ -198,9 +207,8 @@ where
         autotune_job();
 
         let index_to_run = {
-            let mut state = self.state.lock();
-            let map = state.as_mut().unwrap();
-            let tuner = map.get_mut(id).unwrap();
+            let tuner_state = self.state.get(id).unwrap();
+            let mut tuner = tuner_state.write();
 
             tuner.handle_results();