cubecl_runtime/execution/
stream_mode.rs

1use crate::client::Client;
2use alloc::vec::Vec;
3use cubecl_common::device::ServiceId;
4use cubecl_environment::stream::StreamId;
5use cubecl_environment::sync::{AtomicUsize, Mutex, Ordering};
6
7/// Whether a stream's launches run.
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum StreamMode {
10    /// They run.
11    Execute,
12    /// Their kernels compile — now under
13    /// [`CompileAndAutotune`](super::ProcessMode::CompileAndAutotune),
14    /// queued for a batch otherwise — and the launches are discarded.
15    Discard,
16}
17
18impl StreamMode {
19    /// Its slot in a [`StreamModeRegistry`]'s live counts.
20    fn index(self) -> usize {
21        match self {
22            Self::Execute => 0,
23            Self::Discard => 1,
24        }
25    }
26
27    /// The other mode.
28    fn opposite(self) -> Self {
29        match self {
30            Self::Execute => Self::Discard,
31            Self::Discard => Self::Execute,
32        }
33    }
34}
35
36/// One stream of one device's service: what a [`StreamModeOverride`] sets
37/// the mode of. Stream ids are the process's, not a device's, so a
38/// measurement on one device leaves the same stream of every other device
39/// alone.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub(crate) struct ServiceStream {
42    /// The device's service, as its clients reach it.
43    pub service: ServiceId,
44    /// The stream on it.
45    pub stream: StreamId,
46}
47
48/// Sets the mode of one client's stream on its device while it lives,
49/// restored on drop: what a measurement opens, so its launches run whatever
50/// the process mode discards.
51///
52/// Keyed on the device and the stream the client's launches go out on — a
53/// client bound to a stream of its own does not follow the thread's — so it
54/// holds wherever those launches are issued from, a task resumed on another
55/// thread included. Overrides of one stream nest, the newest deciding, and
56/// may drop in any order.
57///
58/// Only the client's stream follows it: a launch another client issues on
59/// another stream — a candidate holding a client bound elsewhere — keeps its
60/// own stream's mode, and under a process mode that discards is discarded,
61/// so a measurement launches through the client it switched.
62///
63/// Every launch on that stream of that device follows it, whichever thread
64/// issues it. Where threads share a stream — `StreamPolicy::Single`, or a
65/// build without per-thread streams — a launch another thread issues while a
66/// measurement runs executes too, rather than being discarded.
67#[derive(Debug)]
68pub struct StreamModeOverride {
69    id: usize,
70    mode: StreamMode,
71}
72
73/// One live [`StreamModeOverride`].
74#[derive(Debug)]
75struct StreamModeEntry {
76    id: usize,
77    stream: ServiceStream,
78    mode: StreamMode,
79}
80
81/// The live [`StreamModeOverride`]s: what sets each stream's mode, and how
82/// many set each mode, so a launch looks them up only when one sets the mode
83/// its process mode does not — a measurement under no override, which sets
84/// the mode every stream already has, costs other launches nothing.
85pub(crate) struct StreamModeRegistry {
86    /// The live overrides, oldest first.
87    entries: Mutex<Vec<StreamModeEntry>>,
88    /// How many live overrides set each mode, by [`StreamMode::index`].
89    live: [AtomicUsize; 2],
90    /// The id the next override takes.
91    next_id: AtomicUsize,
92}
93
94/// The process's stream modes.
95pub(crate) static STREAM_MODES: StreamModeRegistry = StreamModeRegistry {
96    entries: Mutex::new(Vec::new()),
97    live: [AtomicUsize::new(0), AtomicUsize::new(0)],
98    next_id: AtomicUsize::new(0),
99};
100
101impl StreamModeRegistry {
102    /// Put `stream` in `mode` until the guard it hands back drops.
103    pub(crate) fn set(
104        &'static self,
105        mode: StreamMode,
106        stream: ServiceStream,
107    ) -> StreamModeOverride {
108        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
109        let mut entries = self.entries.lock();
110        entries.push(StreamModeEntry { id, stream, mode });
111        self.live[mode.index()].fetch_add(1, Ordering::Release);
112        StreamModeOverride { id, mode }
113    }
114
115    /// `stream`'s mode now: what the newest live override sets for it, or
116    /// `default` when none does.
117    pub(crate) fn mode(&self, stream: ServiceStream, default: StreamMode) -> StreamMode {
118        // An override setting the default changes nothing, unless it is
119        // newer than one setting the other mode on the same stream, and then
120        // there is one of those to look for.
121        if self.live[default.opposite().index()].load(Ordering::Acquire) == 0 {
122            return default;
123        }
124        self.entries
125            .lock()
126            .iter()
127            .rev()
128            .find(|entry| entry.stream == stream)
129            .map_or(default, |entry| entry.mode)
130    }
131
132    fn unset(&self, id: usize, mode: StreamMode) {
133        let mut entries = self.entries.lock();
134        entries.retain(|entry| entry.id != id);
135        self.live[mode.index()].fetch_sub(1, Ordering::Release);
136    }
137}
138
139impl StreamModeOverride {
140    /// Put `client`'s stream on its device in `mode` until the guard drops.
141    pub fn new(mode: StreamMode, client: &Client) -> Self {
142        STREAM_MODES.set(mode, client.service_stream())
143    }
144}
145
146impl Drop for StreamModeOverride {
147    fn drop(&mut self) {
148        STREAM_MODES.unset(self.id, self.mode);
149    }
150}
cubecl_runtime/execution/stream_mode.rs

cubecl_runtime/execution/
stream_mode.rs
