cubecl_runtime/throughput/
cache.rs

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

cubecl_runtime/throughput/
cache.rs
