cubecl_runtime/throughput/
cache.rs

1#[cfg(persistence)]
2use cubecl_environment::persistence::{Namespace, Store, StoreOptions};
3
4use crate::throughput::{ThroughputKey, ThroughputValue};
5use alloc::format;
6use alloc::string::String;
7use alloc::sync::Arc;
8use cubecl_environment::collections::HashMap;
9use cubecl_environment::sync::Mutex;
10use cubecl_ir::{DeviceIdentity, DeviceProperties};
11
12/// The namespace segment naming which generation of probes wrote a value.
13#[cfg(persistence)]
14const GENERATION: &str = "probe-v";
15
16static GLOBAL_CACHE: Mutex<Option<HashMap<String, Arc<Mutex<ThroughputCache>>>>> = Mutex::new(None);
17
18/// Caches the [`ThroughputValue`] for a given [`ThroughputKey`].
19///
20/// This cache is used to avoid recomputing throughput values for the same key.
21/// Stores on disk when std is available, otherwise stores in memory.
22pub struct ThroughputCache {
23    #[cfg(not(persistence))]
24    cache: HashMap<ThroughputKey, ThroughputValue>,
25    #[cfg(persistence)]
26    cache: Store<ThroughputKey, ThroughputValue>,
27}
28
29impl ThroughputCache {
30    /// Gets or creates the global `ThroughputCache` for one device.
31    pub fn get_for_device(runtime: &str, properties: &DeviceProperties) -> Arc<Mutex<Self>> {
32        let hardware = &properties.hardware;
33        let name = device_key(
34            runtime,
35            &properties.identity,
36            properties.memory.max_page_size,
37            hardware
38                .num_cpu_cores
39                .or(hardware.num_streaming_multiprocessors)
40                .unwrap_or(0),
41        );
42        let mut cache_map = GLOBAL_CACHE.lock();
43        let cache_map = cache_map.get_or_insert_with(HashMap::new);
44
45        #[cfg(persistence)]
46        drop_earlier_generations();
47
48        cache_map
49            .entry(name.clone())
50            .or_insert_with(|| Arc::new(Mutex::new(Self::new(&name))))
51            .clone()
52    }
53
54    /// Creates a new `ThroughputCache` with the given name.
55    pub fn new(#[cfg_attr(not(persistence), allow(unused_variables))] name: &str) -> Self {
56        #[cfg(not(persistence))]
57        {
58            ThroughputCache {
59                cache: HashMap::new(),
60            }
61        }
62
63        #[cfg(persistence)]
64        {
65            Self {
66                cache: Store::new(StoreOptions::new().storage(namespace(name))),
67            }
68        }
69    }
70
71    /// Inserts a new [`ThroughputValue`] into the cache for the given [`ThroughputKey`].
72    ///
73    /// Throughput measurements are nondeterministic, so a concurrent process (or an
74    /// earlier run) may have recorded a different value for the same key; the cache
75    /// keeps the existing value in that case rather than failing.
76    pub fn insert(&mut self, key: ThroughputKey, value: ThroughputValue) {
77        #[cfg(persistence)]
78        if let Err(err) = self.cache.insert(key, value) {
79            log::warn!("Concurrent throughput measurement, keeping the existing value: {err}");
80        }
81
82        #[cfg(not(persistence))]
83        self.cache.insert(key, value);
84    }
85
86    /// Returns the [`ThroughputValue`] for the given [`ThroughputKey`], if it exists in the cache.
87    pub fn get(&self, key: &ThroughputKey) -> Option<&ThroughputValue> {
88        self.cache.get(key)
89    }
90}
91
92/// The part, never the device index, which `CUDA_VISIBLE_DEVICES` makes 0 for
93/// whichever card was pinned. Two cards this cannot separate are one part, and
94/// share a peak.
95fn device_key(runtime: &str, identity: &DeviceIdentity, capacity: u64, parallelism: u32) -> String {
96    let DeviceIdentity {
97        name,
98        fingerprint,
99        physical: _,
100    } = identity;
101    // A namespace is a path, and a runtime names itself `wgpu<spirv>`.
102    let segment = |text: &str| text.replace(|c: char| !c.is_ascii_alphanumeric(), "-");
103
104    format!(
105        "{}_{}_{}_mem{capacity}_par{parallelism}",
106        segment(runtime),
107        segment(fingerprint),
108        segment(name)
109    )
110}
111
112/// Namespaces this build wrote under an earlier [`PROBE_VERSION`].
113///
114/// Only this crate version's, and only earlier: another version may belong to a
115/// cubecl still in use, and a later generation to a build running right now. A
116/// namespace with no generation at all predates them.
117#[cfg(persistence)]
118fn is_earlier_generation(candidate: &str, scope: &str, current: u32) -> bool {
119    let Some(device) = candidate.strip_prefix(scope) else {
120        return false;
121    };
122
123    let generation = device
124        .strip_prefix(GENERATION)
125        .and_then(|device| device.split('/').next())
126        .and_then(|generation| generation.parse::<u32>().ok())
127        .unwrap_or(0);
128
129    generation < current
130}
131
132#[cfg(persistence)]
133fn drop_earlier_generations() {
134    use core::sync::atomic::{AtomicBool, Ordering};
135
136    static DROPPED: AtomicBool = AtomicBool::new(false);
137
138    if DROPPED.swap(true, Ordering::Relaxed) {
139        return;
140    }
141
142    let scope = String::from(Namespace::scoped("throughput", "").as_str());
143    let current = crate::throughput::PROBE_VERSION;
144
145    for candidate in cubecl_environment::persistence::namespaces() {
146        if is_earlier_generation(&candidate, &scope, current) {
147            cubecl_environment::persistence::open(&candidate).purge();
148        }
149    }
150}
151
152#[cfg(persistence)]
153fn namespace(device_key: &str) -> Namespace {
154    Namespace::scoped(
155        "throughput",
156        format!(
157            "{GENERATION}{}/{device_key}",
158            crate::throughput::PROBE_VERSION
159        ),
160    )
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166    use alloc::string::ToString;
167
168    fn identity(name: &str, fingerprint: &str) -> DeviceIdentity {
169        DeviceIdentity {
170            name: name.to_string(),
171            fingerprint: fingerprint.to_string(),
172            physical: None,
173        }
174    }
175
176    /// Two cards in one machine were served each other's peaks: pinning either
177    /// makes it index 0, and the index was all that told them apart.
178    #[test]
179    fn two_parts_of_one_architecture_do_not_share_an_entry() {
180        let turing = |name: &str| identity(name, "ptx_sm75");
181
182        assert_ne!(
183            device_key("cuda", &turing("NVIDIA GeForce RTX 2060"), 6 << 30, 30),
184            device_key(
185                "cuda",
186                &turing("NVIDIA GeForce GTX 1660 SUPER"),
187                6 << 30,
188                22
189            ),
190        );
191    }
192
193    #[test]
194    fn a_machine_that_gains_memory_does_not_reuse_its_ceiling() {
195        let xeon = identity("Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz", "cpu_x86_64");
196
197        assert_ne!(
198            device_key("cpu", &xeon, 32 << 30, 16),
199            device_key("cpu", &xeon, 64 << 30, 16),
200        );
201    }
202
203    #[test]
204    fn a_container_given_fewer_cores_does_not_reuse_the_host_ceiling() {
205        let xeon = identity("Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz", "cpu_x86_64");
206
207        assert_ne!(
208            device_key("cpu", &xeon, 64 << 30, 16),
209            device_key("cpu", &xeon, 64 << 30, 4),
210        );
211    }
212
213    /// A device key is one segment of a path, so a runtime that names itself
214    /// `wgpu<spirv>` must not put brackets or a separator in it.
215    #[test]
216    fn a_device_key_is_one_path_segment() {
217        let key = device_key(
218            "wgpu<spirv>",
219            &identity("Intel(R) Arc(tm) B390 (PTL)", "spirv_32902_45184"),
220            1 << 31,
221            0,
222        );
223
224        assert!(
225            key.chars()
226                .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_'),
227            "{key}"
228        );
229    }
230
231    /// A build drops what its own earlier probes wrote and nothing else:
232    /// another crate version may be in use, and a later probe version is
233    /// running right now.
234    #[cfg(persistence)]
235    #[test]
236    fn only_this_crate_version_s_earlier_probes_are_dropped() {
237        let stale = |ns: &str| is_earlier_generation(ns, "throughput/0.11.0/", 2);
238
239        assert!(stale("throughput/0.11.0/probe-v1/cuda_ptx_sm75_part"));
240        assert!(stale("throughput/0.11.0/cuda_dev0"));
241
242        assert!(!stale("throughput/0.11.0/probe-v2/cuda_ptx_sm75_part"));
243        assert!(!stale("throughput/0.11.0/probe-v3/cuda_ptx_sm75_part"));
244        assert!(!stale("throughput/0.10.0/probe-v1/cuda_ptx_sm75_part"));
245        assert!(!stale("autotune/0.11.0/device-0-0-cpu/matmul"));
246    }
247
248    /// The crate version in the namespace does not move when a probe changes
249    /// what it measures inside a release, so the probes carry their own.
250    #[cfg(persistence)]
251    #[test]
252    fn the_namespace_carries_the_probe_version() {
253        let generation = namespace("").as_str().to_string();
254
255        assert!(
256            namespace("cuda_ptx_sm75_part")
257                .as_str()
258                .starts_with(&format!("{generation}/"))
259        );
260    }
261}
cubecl_runtime/throughput/cache.rs

cubecl_runtime/throughput/
cache.rs
