cubecl_runtime/throughput/
cache.rs

1#[cfg(std_io)]
2use cubecl_environment::persistence::{Namespace, Store, StoreOptions};
3
4use crate::throughput::{ThroughputKey, ThroughputValue};
5use alloc::string::{String, ToString};
6use alloc::sync::Arc;
7use cubecl_environment::collections::HashMap;
8use cubecl_environment::sync::Mutex;
9
10static GLOBAL_CACHE: Mutex<Option<HashMap<String, Arc<Mutex<ThroughputCache>>>>> = Mutex::new(None);
11
12/// Caches the [`ThroughputValue`] for a given [`ThroughputKey`].
13///
14/// This cache is used to avoid recomputing throughput values for the same key.
15/// Stores on disk when std is available, otherwise stores in memory.
16pub struct ThroughputCache {
17    #[cfg(not(std_io))]
18    cache: HashMap<ThroughputKey, ThroughputValue>,
19    #[cfg(std_io)]
20    cache: Store<ThroughputKey, ThroughputValue>,
21}
22
23impl ThroughputCache {
24    /// Gets or creates a global `ThroughputCache` for the given device name.
25    pub fn get_for_device(name: &str) -> Arc<Mutex<Self>> {
26        let mut cache_map = GLOBAL_CACHE.lock();
27        let cache_map = cache_map.get_or_insert_with(HashMap::new);
28
29        cache_map
30            .entry(name.to_string())
31            .or_insert_with(|| Arc::new(Mutex::new(Self::new(name))))
32            .clone()
33    }
34
35    /// Creates a new `ThroughputCache` with the given name.
36    pub fn new(#[cfg_attr(not(std_io), allow(unused_variables))] name: &str) -> Self {
37        #[cfg(not(std_io))]
38        {
39            ThroughputCache {
40                cache: HashMap::new(),
41            }
42        }
43
44        #[cfg(std_io)]
45        {
46            let namespace = Namespace::scoped("throughput", name);
47
48            Self {
49                cache: Store::new(StoreOptions::new().storage(namespace)),
50            }
51        }
52    }
53
54    /// Inserts a new [`ThroughputValue`] into the cache for the given [`ThroughputKey`].
55    ///
56    /// Throughput measurements are nondeterministic, so a concurrent process (or an
57    /// earlier run) may have recorded a different value for the same key; the cache
58    /// keeps the existing value in that case rather than failing.
59    pub fn insert(&mut self, key: ThroughputKey, value: ThroughputValue) {
60        #[cfg(std_io)]
61        if let Err(err) = self.cache.insert(key, value) {
62            log::warn!("Concurrent throughput measurement, keeping the existing value: {err}");
63        }
64
65        #[cfg(not(std_io))]
66        self.cache.insert(key, value);
67    }
68
69    /// Returns the [`ThroughputValue`] for the given [`ThroughputKey`], if it exists in the cache.
70    pub fn get(&self, key: &ThroughputKey) -> Option<&ThroughputValue> {
71        self.cache.get(key)
72    }
73}
cubecl_runtime/throughput/cache.rs

cubecl_runtime/throughput/
cache.rs
