cubecl_runtime/execution/
base.rs

1use super::{ProcessMode, STREAM_MODES, ServiceStream, StreamMode};
2
3/// What a server does with one launch: the verdict its stream's mode and the
4/// process mode resolve to.
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum LaunchMode {
7    /// Compile if needed, then run it. The normal case.
8    Execute,
9    /// Compile if needed, cache the artifact, and discard the launch.
10    ///
11    /// A server honoring this must still do everything a first launch does
12    /// short of dispatching — expand, compile, validate, populate its caches —
13    /// or the pass buys nothing.
14    Compile,
15    /// Queue the kernel to be compiled with others, and discard the launch.
16    ///
17    /// A server honoring this compiles the queue when it next loads a kernel
18    /// for a launch, or when asked to (`Server::compile_queued`): flushing or
19    /// syncing compiles nothing, so a pass that only queues gathers everything
20    /// it reaches into one batch.
21    /// A kernel that fails to compile there reports it when it is launched.
22    Queue,
23}
24
25impl LaunchMode {
26    /// What a launch issued on `stream` does now.
27    pub(crate) fn new(stream: ServiceStream) -> Self {
28        let mode = ProcessMode::current();
29        match (STREAM_MODES.mode(stream, mode.stream_mode()), mode) {
30            (StreamMode::Execute, _) => Self::Execute,
31            // A stream discarding its launches while the process executes
32            // queues them: they compile with the next launch that loads a
33            // kernel.
34            (StreamMode::Discard, ProcessMode::CompileOnly | ProcessMode::Execute) => Self::Queue,
35            (StreamMode::Discard, ProcessMode::CompileAndAutotune) => Self::Compile,
36        }
37    }
38
39    /// Whether the launch is discarded rather than run.
40    pub fn discards_launch(self) -> bool {
41        matches!(self, Self::Compile | Self::Queue)
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::super::{ProcessModeOverride, StatisticsCollector};
48    use super::*;
49    use cubecl_common::device::{DeviceId, ServiceId};
50    use cubecl_environment::stream::StreamId;
51    // `serial_test`'s macro expands to `vec!`, which a `no_std` crate has to
52    // bring in itself.
53    use alloc::vec;
54
55    /// The service of device `index`.
56    fn device(index: u16) -> ServiceId {
57        ServiceId::of::<()>(DeviceId {
58            type_id: 0,
59            index_id: index,
60        })
61    }
62
63    /// Stream `value`.
64    fn stream(value: u64) -> StreamId {
65        StreamId { value }
66    }
67
68    /// The verdict is a table of a stream's mode and the process mode.
69    #[test]
70    #[serial_test::serial]
71    fn a_launch_follows_its_stream_and_the_process_mode() {
72        let collector = StatisticsCollector::new();
73        let seventh = ServiceStream {
74            service: device(0),
75            stream: stream(7),
76        };
77        let mode = || LaunchMode::new(seventh);
78        assert_eq!(mode(), LaunchMode::Execute);
79        {
80            let compile = ProcessModeOverride::new(ProcessMode::CompileOnly, &collector);
81            assert_eq!(mode(), LaunchMode::Queue);
82            core::mem::drop(compile);
83        }
84        {
85            let tune = ProcessModeOverride::new(ProcessMode::CompileAndAutotune, &collector);
86            assert_eq!(mode(), LaunchMode::Compile);
87            core::mem::drop(tune);
88        }
89        {
90            // A stream discarding its launches while the process executes
91            // queues them.
92            let queuing = STREAM_MODES.set(StreamMode::Discard, seventh);
93            assert_eq!(mode(), LaunchMode::Queue);
94            core::mem::drop(queuing);
95        }
96        let execute = ProcessModeOverride::new(ProcessMode::Execute, &collector);
97        assert_eq!(mode(), LaunchMode::Execute);
98        core::mem::drop(execute);
99    }
100
101    /// A stream's mode overrides the process mode's default for that stream of
102    /// that device alone, the newest override deciding, and overrides drop
103    /// in any order.
104    #[test]
105    #[serial_test::serial]
106    fn a_stream_mode_holds_for_its_stream_on_its_device() {
107        let collector = StatisticsCollector::new();
108        let tune = ProcessModeOverride::new(ProcessMode::CompileAndAutotune, &collector);
109        let measured = ServiceStream {
110            service: device(0),
111            stream: stream(1),
112        };
113        let other_stream = ServiceStream {
114            stream: stream(2),
115            ..measured
116        };
117        let other_device = ServiceStream {
118            service: device(1),
119            ..measured
120        };
121
122        let measuring = STREAM_MODES.set(StreamMode::Execute, measured);
123        assert_eq!(LaunchMode::new(measured), LaunchMode::Execute);
124        assert_eq!(LaunchMode::new(other_stream), LaunchMode::Compile);
125        assert_eq!(LaunchMode::new(other_device), LaunchMode::Compile);
126
127        let nested = STREAM_MODES.set(StreamMode::Discard, measured);
128        assert_eq!(
129            LaunchMode::new(measured),
130            LaunchMode::Compile,
131            "the newest decides"
132        );
133        core::mem::drop(measuring);
134        assert_eq!(LaunchMode::new(measured), LaunchMode::Compile);
135        core::mem::drop(nested);
136        assert_eq!(
137            LaunchMode::new(measured),
138            LaunchMode::Compile,
139            "the mode's again"
140        );
141        core::mem::drop(tune);
142    }
143}
cubecl_runtime/execution/base.rs

cubecl_runtime/execution/
base.rs
