cubecl_environment/stream/
policy.rs

1use crate::sync::{AtomicU8, Ordering};
2
3/// How the current stream is derived when no explicit override is active.
4///
5/// Loaded from the `[streaming]` section of `cubecl.toml` by `cubecl-runtime`,
6/// or set programmatically with [`set_policy`]. An explicit [`set_policy`] call
7/// always wins over the configuration file.
8#[non_exhaustive]
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
10#[serde(rename_all = "kebab-case")]
11pub enum StreamPolicy {
12    /// One stream per OS thread (the historical behavior and the default).
13    ///
14    /// Correct for applications that submit work from plain threads. Under an
15    /// async executor with work stealing, a logical task hopping between
16    /// worker threads changes streams, causing spurious synchronization.
17    #[default]
18    PerThread,
19    /// One stream per logical task.
20    ///
21    /// With the `tokio` feature enabled, a task keeps a stable stream across
22    /// `.await` points even when the executor moves it between worker
23    /// threads. Outside a task, or without the `tokio` feature, this behaves
24    /// like [`PerThread`](StreamPolicy::PerThread).
25    PerTask,
26    /// Everything runs on stream `0`.
27    ///
28    /// The effective behavior on wasm and no-std targets today.
29    Single,
30}
31
32const SET_BY_USER: u8 = 0b1000_0000;
33
34static POLICY: AtomicU8 = AtomicU8::new(0);
35
36fn encode(policy: StreamPolicy) -> u8 {
37    match policy {
38        StreamPolicy::PerThread => 0,
39        StreamPolicy::PerTask => 1,
40        StreamPolicy::Single => 2,
41    }
42}
43
44fn decode(bits: u8) -> StreamPolicy {
45    match bits & !SET_BY_USER {
46        1 => StreamPolicy::PerTask,
47        2 => StreamPolicy::Single,
48        _ => StreamPolicy::PerThread,
49    }
50}
51
52/// Sets the active stream policy.
53///
54/// Takes precedence over any policy loaded from configuration files, no matter
55/// the call order. Should be called before the first kernel submissions:
56/// already-resolved stream ids are not revisited.
57pub fn set_policy(policy: StreamPolicy) {
58    POLICY.store(encode(policy) | SET_BY_USER, Ordering::Relaxed);
59}
60
61/// Returns the active stream policy.
62pub fn policy() -> StreamPolicy {
63    decode(POLICY.load(Ordering::Relaxed))
64}
65
66/// Sets the stream policy from a configuration file.
67///
68/// A no-op if [`set_policy`] was already called: the user's explicit choice
69/// wins over configuration.
70#[doc(hidden)]
71pub fn set_policy_from_config(policy: StreamPolicy) {
72    let _ = POLICY.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
73        if current & SET_BY_USER != 0 {
74            None
75        } else {
76            Some(encode(policy))
77        }
78    });
79}
80
81/// Resets the process-global policy so policy-mutating tests don't leak state.
82///
83/// Gated like `tests_policy_lock`: its only callers are the tests that take
84/// that lock, which need std.
85#[cfg(all(test, feature = "std"))]
86pub(crate) fn tests_reset_policy() {
87    POLICY.store(0, Ordering::Relaxed);
88}
89
90/// Serializes tests that read or mutate the process-global policy.
91#[cfg(all(test, feature = "std"))]
92pub(crate) fn tests_policy_lock() -> std::sync::MutexGuard<'static, ()> {
93    static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
94    LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
95}
96
97// Gated like `tests_policy_lock`, which these tests use to serialize.
98#[cfg(all(test, feature = "std"))]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn explicit_policy_wins_over_config() {
104        let _guard = tests_policy_lock();
105
106        set_policy_from_config(StreamPolicy::Single);
107        assert_eq!(policy(), StreamPolicy::Single);
108
109        set_policy(StreamPolicy::PerTask);
110        set_policy_from_config(StreamPolicy::Single);
111        assert_eq!(policy(), StreamPolicy::PerTask);
112
113        tests_reset_policy();
114    }
115}
cubecl_environment/stream/policy.rs

cubecl_environment/stream/
policy.rs
