cubecl_runtime/throughput/
cache.rs

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

cubecl_runtime/throughput/
cache.rs
