cubecl_runtime/
launched.rs

1//! The kernels a workload launches, collected while it replays: what an
2//! environment shipped for that workload has to keep, and all it has to.
3//!
4//! A build compiles far more than it ships — every candidate an autotune race
5//! measured, of which only the winner is ever launched again. Replaying the
6//! workload the environment is shipped for, under a [`LaunchedKernels`], names
7//! exactly the kernels it runs; everything else in the compilation store can
8//! go. The replay may be a dry run: a dropped launch is still issued, and so
9//! still collected.
10//!
11//! Collected on the issuing thread, where every backend's launch passes, by
12//! the [stable hash](crate::id::KernelId::stable_hash) that also names a
13//! kernel's artifact in the compilation store. Outside a collection a launch
14//! pays one relaxed atomic load.
15
16use crate::id::KernelId;
17use cubecl_common::hash::StableHash;
18use cubecl_environment::collections::{HashMap, HashSet};
19use cubecl_environment::sync::{AtomicUsize, LazyLock, Mutex, Ordering};
20
21/// How many collections are open: what a launch reads before it takes the
22/// lock, so one outside every collection pays a relaxed load and nothing more.
23static COLLECTING: AtomicUsize = AtomicUsize::new(0);
24
25/// The collections open, each with what was launched since it opened.
26static OPEN: LazyLock<Mutex<Collections>> = LazyLock::new(|| Mutex::new(Collections::default()));
27
28#[derive(Default)]
29struct Collections {
30    next: u64,
31    open: HashMap<u64, HashSet<StableHash>>,
32}
33
34/// Collects every kernel launched, on every thread and device, from
35/// [`new`](Self::new) to [`finish`](Self::finish). Collections may overlap:
36/// each sees what was launched while it was open, and nothing before.
37#[must_use = "the kernels are collected until `finish`"]
38#[derive(Debug)]
39pub struct LaunchedKernels {
40    id: u64,
41}
42
43impl LaunchedKernels {
44    /// Start collecting.
45    #[allow(
46        clippy::new_without_default,
47        reason = "starting a collection is not a value"
48    )]
49    pub fn new() -> Self {
50        let mut collections = OPEN.lock();
51        let id = collections.next;
52        collections.next += 1;
53        collections.open.insert(id, HashSet::default());
54        // Raised under the lock: a launch that sees it finds the set there.
55        COLLECTING.fetch_add(1, Ordering::AcqRel);
56        Self { id }
57    }
58
59    /// Stop collecting, and hand back the stable hash of every kernel
60    /// launched since [`new`](Self::new).
61    pub fn finish(self) -> HashSet<StableHash> {
62        let launched = close(self.id);
63        core::mem::forget(self);
64        launched
65    }
66}
67
68impl Drop for LaunchedKernels {
69    fn drop(&mut self) {
70        close(self.id);
71    }
72}
73
74/// Close collection `id`, handing back what it collected.
75fn close(id: u64) -> HashSet<StableHash> {
76    let mut collections = OPEN.lock();
77    let launched = collections.open.remove(&id).unwrap_or_default();
78    COLLECTING.fetch_sub(1, Ordering::AcqRel);
79    launched
80}
81
82/// Note a launch in every open collection. The id is asked for only then:
83/// outside a collection a launch computes nothing here.
84pub(crate) fn note(kernel: impl FnOnce() -> KernelId) {
85    if COLLECTING.load(Ordering::Relaxed) == 0 {
86        return;
87    }
88    // Hashed before the lock, which every launching thread shares.
89    let hash = kernel().stable_hash();
90    let mut collections = OPEN.lock();
91    for launched in collections.open.values_mut() {
92        launched.insert(hash);
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    struct Launched;
101
102    #[test]
103    fn a_collection_names_what_was_launched_while_it_was_open() {
104        let before = KernelId::new::<Launched>().info(0u32);
105        let during = KernelId::new::<Launched>().info(1u32);
106        note(|| before.clone());
107
108        let collection = LaunchedKernels::new();
109        note(|| during.clone());
110        note(|| during.clone());
111        let launched = collection.finish();
112
113        // Other tests launch in parallel: only this test's kernels are known.
114        assert!(launched.contains(&during.stable_hash()));
115        assert!(!launched.contains(&before.stable_hash()));
116    }
117
118    /// An inner collection sees only what was launched while it was open,
119    /// and closing it leaves the outer one collecting.
120    #[test]
121    fn overlapping_collections_each_see_their_own_launches() {
122        let early = KernelId::new::<Launched>().info(10u32);
123        let late = KernelId::new::<Launched>().info(11u32);
124
125        let outer = LaunchedKernels::new();
126        note(|| early.clone());
127        let inner = LaunchedKernels::new();
128        note(|| late.clone());
129        let inner = inner.finish();
130        let outer = outer.finish();
131
132        assert!(inner.contains(&late.stable_hash()));
133        assert!(!inner.contains(&early.stable_hash()));
134        assert!(outer.contains(&early.stable_hash()) && outer.contains(&late.stable_hash()));
135    }
136}
cubecl_runtime/launched.rs

cubecl_runtime/
launched.rs
