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    #[allow(
73        deprecated,
74        reason = "portable_atomic lacks try_update on targets without native atomics"
75    )]
76    let _ = POLICY.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
77        if current & SET_BY_USER != 0 {
78            None
79        } else {
80            Some(encode(policy))
81        }
82    });
83}
84
85/// Resets the process-global policy so policy-mutating tests don't leak state.
86///
87/// Gated like `tests_policy_lock`: its only callers are the tests that take
88/// that lock, which need std.
89#[cfg(all(test, feature = "std"))]
90pub(crate) fn tests_reset_policy() {
91    POLICY.store(0, Ordering::Relaxed);
92}
93
94/// Serializes tests that read or mutate the process-global policy.
95#[cfg(all(test, feature = "std"))]
96pub(crate) fn tests_policy_lock() -> std::sync::MutexGuard<'static, ()> {
97    static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
98    LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
99}
100
101// Gated like `tests_policy_lock`, which these tests use to serialize.
102#[cfg(all(test, feature = "std"))]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn explicit_policy_wins_over_config() {
108        let _guard = tests_policy_lock();
109
110        set_policy_from_config(StreamPolicy::Single);
111        assert_eq!(policy(), StreamPolicy::Single);
112
113        set_policy(StreamPolicy::PerTask);
114        set_policy_from_config(StreamPolicy::Single);
115        assert_eq!(policy(), StreamPolicy::PerTask);
116
117        tests_reset_policy();
118    }
119}
cubecl_environment/stream/policy.rs

cubecl_environment/stream/
policy.rs
