cubecl_environment/stream/
id.rs

1// `core`'s directly, not the `sync` shim: this is `stream_local`, which is
2// `std`, so 64-bit atomics are always there — see the note in `sync::base`.
3#[cfg(stream_local)]
4use core::cell::Cell;
5#[cfg(stream_local)]
6use core::sync::atomic::AtomicU64;
7
8#[cfg(stream_local)]
9use super::StreamPolicy;
10
11/// Unique identifier representing the stream on which work is submitted.
12///
13/// How the current stream is resolved depends on the active
14/// [`StreamPolicy`](super::StreamPolicy) and on any explicit override installed
15/// with [`StreamId::executes`] or [`Stream::enter`](super::Stream::enter).
16#[derive(
17    Debug, PartialEq, Eq, Clone, Copy, Hash, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
18)]
19pub struct StreamId {
20    /// The value representing the stream id.
21    pub value: u64,
22}
23
24#[cfg(stream_local)]
25static STREAM_COUNT: AtomicU64 = AtomicU64::new(0);
26
27#[cfg(stream_local)]
28std::thread_local! {
29    /// Explicitly scoped stream override, installed by [`StreamId::executes`].
30    static OVERRIDE: Cell<Option<u64>> = const { Cell::new(None) };
31    /// Lazily assigned per-thread default stream.
32    static DEFAULT: Cell<Option<u64>> = const { Cell::new(None) };
33}
34
35/// Replaces the current stream override, returning the previous one.
36///
37/// `None` means "no override": [`StreamId::current`] falls back to the active
38/// policy. Keeping the override separate from the per-thread default is what
39/// allows a scoped override to be fully undone, even on threads that never had
40/// a default assigned.
41#[cfg(stream_local)]
42pub(crate) fn set_override(value: Option<u64>) -> Option<u64> {
43    OVERRIDE.with(|cell| cell.replace(value))
44}
45
46impl StreamId {
47    /// Executes `f` on this stream, restoring the previous stream afterward.
48    ///
49    /// The previous state is saved before the call and restored on return —
50    /// including on unwind — so the caller never has to manage raw override
51    /// pairs. Restoring also works when no stream was active before the call.
52    pub fn executes<F, T>(self, f: F) -> T
53    where
54        F: FnOnce() -> T,
55    {
56        #[cfg(stream_local)]
57        {
58            struct Guard(Option<u64>);
59
60            impl Drop for Guard {
61                fn drop(&mut self) {
62                    set_override(self.0);
63                }
64            }
65
66            let _guard = Guard(set_override(Some(self.value)));
67            f()
68        }
69
70        #[cfg(not(stream_local))]
71        f()
72    }
73
74    /// Get the current stream id.
75    ///
76    /// Resolution order:
77    /// 1. An explicit override installed by [`StreamId::executes`].
78    /// 2. The active [`StreamPolicy`](super::StreamPolicy): a stable per-task
79    ///    id under [`PerTask`](super::StreamPolicy::PerTask), stream `0` under
80    ///    [`Single`](super::StreamPolicy::Single), or a lazily assigned
81    ///    per-thread id otherwise.
82    pub fn current() -> Self {
83        #[cfg(stream_local)]
84        {
85            if let Some(value) = OVERRIDE.with(|cell| cell.get()) {
86                return Self { value };
87            }
88
89            match super::policy() {
90                StreamPolicy::Single => Self { value: 0 },
91                StreamPolicy::PerTask => Self::per_task(),
92                StreamPolicy::PerThread => Self::per_thread(),
93            }
94        }
95
96        #[cfg(not(stream_local))]
97        Self { value: 0 }
98    }
99
100    /// Allocate a fresh stream id, distinct from every per-thread default.
101    ///
102    /// On no-std targets there is a single stream, so this returns id `0`.
103    pub fn allocate() -> Self {
104        #[cfg(stream_local)]
105        {
106            Self {
107                value: STREAM_COUNT.fetch_add(1, core::sync::atomic::Ordering::Relaxed),
108            }
109        }
110
111        #[cfg(not(stream_local))]
112        Self { value: 0 }
113    }
114
115    #[cfg(stream_local)]
116    fn per_thread() -> Self {
117        DEFAULT.with(|cell| match cell.get() {
118            Some(value) => Self { value },
119            None => {
120                let new = Self::allocate();
121                cell.set(Some(new.value));
122                new
123            }
124        })
125    }
126
127    #[cfg(all(stream_local, tokio_rt))]
128    fn per_task() -> Self {
129        match tokio::task::try_id() {
130            Some(id) => {
131                use core::hash::BuildHasher;
132
133                let hash = foldhash::fast::FixedState::default().hash_one(id);
134
135                // The high bit namespaces task-derived ids away from the
136                // counter-based thread and manual ids. A hash collision or a
137                // recycled tokio task id only merges two logical streams onto
138                // one backend stream, which is safe (FIFO ordering), while a
139                // single task always keeps one stable id across thread hops.
140                Self {
141                    value: hash | (1 << 63),
142                }
143            }
144            // Not inside a tokio task: behave like the per-thread policy.
145            None => Self::per_thread(),
146        }
147    }
148
149    #[cfg(all(stream_local, not(tokio_rt)))]
150    fn per_task() -> Self {
151        #[cfg(feature = "std")]
152        {
153            use std::sync::Once;
154
155            static WARN: Once = Once::new();
156            WARN.call_once(|| {
157                log::warn!(
158                    "Stream policy 'per-task' requires the 'tokio' feature of cubecl-environment; falling back to 'per-thread'."
159                );
160            });
161        }
162
163        Self::per_thread()
164    }
165}
166
167impl core::fmt::Display for StreamId {
168    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
169        f.write_fmt(format_args!("StreamId({:?})", self.value))
170    }
171}
172
173#[cfg(all(test, stream_local))]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn executes_restores_previous_override() {
179        let outer = StreamId { value: 1_000_000 };
180        let inner = StreamId { value: 2_000_000 };
181
182        outer.executes(|| {
183            assert_eq!(StreamId::current(), outer);
184            inner.executes(|| {
185                assert_eq!(StreamId::current(), inner);
186            });
187            assert_eq!(StreamId::current(), outer);
188        });
189    }
190
191    #[test]
192    fn executes_restores_no_override_state() {
193        // Regression: restoring after the outermost `executes` must return to
194        // "no override", not pin the resolved id onto the thread.
195        let scoped = StreamId { value: 500_000 };
196
197        scoped.executes(|| {
198            assert_eq!(StreamId::current(), scoped);
199        });
200
201        assert_eq!(OVERRIDE.with(|cell| cell.get()), None);
202        assert_ne!(StreamId::current(), scoped);
203    }
204
205    #[test]
206    fn current_is_stable_on_one_thread() {
207        let _guard = crate::stream::tests_policy_lock();
208
209        assert_eq!(StreamId::current(), StreamId::current());
210    }
211
212    #[test]
213    fn allocate_returns_distinct_ids() {
214        assert_ne!(StreamId::allocate(), StreamId::allocate());
215    }
216}
cubecl_environment/stream/id.rs

cubecl_environment/stream/
id.rs
