cubecl_runtime/stream/capture.rs
1//! The stream-side graph-capture lifecycle, shared by every backend with
2//! graph support (see [`ComputeServer::graph_prepare`](crate::server::ComputeServer::graph_prepare)).
3
4use crate::metadata_cache::CacheMode;
5use crate::server::ServerError;
6
7/// Where a stream sits in the graph-capture lifecycle, and the only thing
8/// allowed to move it. Capture is a strict `NoCapture → Prepare → Capture →
9/// NoCapture` progression, driven by [`prepare`](Self::prepare),
10/// [`begin`](Self::begin) and [`end`](Self::end); each rejects an out-of-order
11/// call, so a capture can never start unprepared and two captures can never
12/// overlap on one stream.
13///
14/// The transitions live here rather than in each backend server because the
15/// rule is the same on every one of them — a backend supplies only the work a
16/// transition brackets (arming its pools, opening the driver's capture), never
17/// the ordering rule itself.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum StreamCaptureState {
20 /// No capture is prepared or recording.
21 NoCapture,
22 /// `graph_prepare` has armed the persistent pools for the warmup run;
23 /// `begin_capture` may now open the window. Slices the warmup run reserves
24 /// are retained by the memory manager's priming until `begin_capture` calls
25 /// [`capture_priming_end`](crate::memory_management::MemoryManagement::capture_priming_end),
26 /// so the pool ends up owning the capture run's full working set.
27 Prepare,
28 /// Launches are being recorded into a graph instead of executing. On a
29 /// hardware-graph backend (CUDA, HIP) a host sync issued now aborts the
30 /// driver capture, so the execution path defers fenced flushes until
31 /// `end_capture`. A software-graph backend (wgpu) has no driver capture to
32 /// abort and instead refuses the operations it cannot record: a read, sync
33 /// or profile fails on the spot, while a write is rejected lazily — queued
34 /// as an error that fails `end_capture`, since a graph missing an operation
35 /// is worse than a late diagnostic.
36 Capture,
37}
38
39impl StreamCaptureState {
40 /// Whether launches on the stream are being recorded into a graph right
41 /// now — the window during which a host sync would abort (or is rejected
42 /// by) the capture.
43 pub fn is_recording(&self) -> bool {
44 matches!(self, StreamCaptureState::Capture)
45 }
46
47 /// The [`CacheMode`] the metadata info cache should run in at this lifecycle
48 /// position. Both while a graph is being *prepared* (warmup, which primes
49 /// the cache) and while it is being *recorded* the cache runs in
50 /// [`CacheMode::Capture`] — caching every buffer and invalidating none — so
51 /// the capture window finds every info buffer warm and drops none out from
52 /// under a recorded launch. Normal operation uses [`CacheMode::Normal`].
53 pub fn cache_mode(&self) -> CacheMode {
54 match self {
55 StreamCaptureState::NoCapture => CacheMode::Normal,
56 StreamCaptureState::Prepare | StreamCaptureState::Capture => CacheMode::Capture,
57 }
58 }
59
60 /// `NoCapture → Prepare`, for `graph_prepare`. Call before arming the
61 /// pools; the caller owns the arming, this owns the rule that it happens
62 /// exactly once per capture.
63 ///
64 /// # Errors
65 ///
66 /// Fails when a capture is already prepared or already recording on this
67 /// stream, leaving the state untouched — two captures may never overlap on
68 /// one stream. The caller can retry after `end_capture`.
69 pub fn prepare(&mut self) -> Result<(), ServerError> {
70 match self {
71 StreamCaptureState::NoCapture => {
72 *self = StreamCaptureState::Prepare;
73 Ok(())
74 }
75 StreamCaptureState::Prepare => Err(ServerError::graph_state(
76 "graph_prepare: a graph capture is already prepared on this stream",
77 )),
78 StreamCaptureState::Capture => Err(ServerError::graph_state(
79 "graph_prepare: a graph capture is already recording on this stream",
80 )),
81 }
82 }
83
84 /// `Prepare → Capture`, for `begin_capture`. Call *before* the work that
85 /// opens the window (ending the priming phase, starting the driver's
86 /// capture) so a rejected call cannot run any of it: on a stream that is
87 /// already recording, a drop-queue flush issued on the way to the rejection
88 /// would abort the live capture.
89 ///
90 /// Since the state moves before that work, a backend whose window fails to
91 /// open must undo it with [`abort`](Self::abort).
92 ///
93 /// # Errors
94 ///
95 /// Fails when [`prepare`](Self::prepare) has not run — the persistent pools
96 /// have to be primed by a warmup run first — or when a capture is already
97 /// recording. The state is left untouched.
98 pub fn begin(&mut self) -> Result<(), ServerError> {
99 match self {
100 StreamCaptureState::Prepare => {
101 *self = StreamCaptureState::Capture;
102 Ok(())
103 }
104 StreamCaptureState::NoCapture => Err(ServerError::graph_state(
105 "begin_capture: call graph_prepare before starting a capture",
106 )),
107 StreamCaptureState::Capture => Err(ServerError::graph_state(
108 "begin_capture: a graph capture is already recording on this stream",
109 )),
110 }
111 }
112
113 /// `Capture → NoCapture`, for `end_capture`. Call before closing the
114 /// window, so the stream leaves capture state even if sealing the graph
115 /// then fails — a backend that returned an error with the state still set
116 /// would wedge the stream in capture mode forever.
117 ///
118 /// # Errors
119 ///
120 /// Fails when no capture is recording (nothing prepared or started, or the
121 /// capture already ended), leaving the state untouched — a stray
122 /// `end_capture` must not close a window that was never opened.
123 pub fn end(&mut self) -> Result<(), ServerError> {
124 match self {
125 StreamCaptureState::Capture => {
126 *self = StreamCaptureState::NoCapture;
127 Ok(())
128 }
129 StreamCaptureState::NoCapture | StreamCaptureState::Prepare => {
130 Err(ServerError::graph_state(
131 "end_capture: no graph capture is recording on this stream",
132 ))
133 }
134 }
135 }
136
137 /// Return to `NoCapture` from anywhere, for the failure path of a
138 /// transition's own work: the window never opened, so the stream must be
139 /// left fully usable and re-capturable rather than stuck arming its
140 /// persistent pools forever. Unlike [`end`](Self::end) this asserts
141 /// nothing, because the state it is recovering from is precisely the one
142 /// that could not be completed.
143 pub fn abort(&mut self) {
144 *self = StreamCaptureState::NoCapture;
145 }
146}
147
148#[cfg(test)]
149mod tests {
150 use super::*;
151
152 /// The ordering rule the three backends rely on: a capture cannot start
153 /// unprepared, and two cannot overlap on one stream. A backend that could
154 /// reach `Capture` without `Prepare` would record against pools no warmup
155 /// primed, and every allocation the window then makes is one the graph
156 /// replays against but nothing pins.
157 #[test]
158 fn transitions_follow_the_capture_order() {
159 let mut state = StreamCaptureState::NoCapture;
160
161 assert!(state.begin().is_err(), "a capture must be prepared first");
162 assert!(state.end().is_err(), "nothing is recording yet");
163 assert_eq!(state, StreamCaptureState::NoCapture);
164
165 state.prepare().unwrap();
166 assert_eq!(state, StreamCaptureState::Prepare);
167 assert!(state.prepare().is_err(), "one prepare per capture");
168 assert!(state.end().is_err(), "the window never opened");
169
170 state.begin().unwrap();
171 assert_eq!(state, StreamCaptureState::Capture);
172 assert!(state.begin().is_err(), "captures may not overlap");
173 assert!(state.prepare().is_err(), "captures may not overlap");
174
175 state.end().unwrap();
176 assert_eq!(state, StreamCaptureState::NoCapture);
177 }
178
179 /// A rejected transition leaves the stream exactly as it was, so a caller
180 /// that miss orders a call can recover by issuing the right one — the
181 /// property `wgpu_graph_lifecycle_state_errors` defends end to end.
182 #[test]
183 fn a_rejected_transition_changes_nothing() {
184 let mut state = StreamCaptureState::Prepare;
185 assert!(state.prepare().is_err());
186 assert_eq!(state, StreamCaptureState::Prepare);
187 state.begin().unwrap();
188 }
189
190 /// `abort` recovers from a window that failed to open, from either of the
191 /// states a backend can be holding when that happens.
192 #[test]
193 fn abort_recovers_a_window_that_never_opened() {
194 for state in [StreamCaptureState::Prepare, StreamCaptureState::Capture] {
195 let mut state = state;
196 state.abort();
197 assert_eq!(state, StreamCaptureState::NoCapture);
198 state.prepare().expect("the stream is re-capturable");
199 }
200 }
201
202 /// The cache runs in capture mode for the *whole* prepare → record window,
203 /// not just while recording: warmup is what makes the recorded launches hit
204 /// warm info buffers, and an entry evicted between the two would be one a
205 /// recorded launch dropped out from under itself.
206 #[test]
207 fn the_cache_captures_across_the_whole_window() {
208 assert_eq!(
209 StreamCaptureState::NoCapture.cache_mode(),
210 CacheMode::Normal
211 );
212 assert_eq!(StreamCaptureState::Prepare.cache_mode(), CacheMode::Capture);
213 assert_eq!(StreamCaptureState::Capture.cache_mode(), CacheMode::Capture);
214 }
215}