Skip to main content

gpu_async_reactor/core/
wait_plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! The wait plan: how a host waits at a readback point, as typed fields parsed once from text
4//! (resolve once, then immutable). Every field is stated in the text; nothing is defaulted.
5//!
6//! Text: `<park>` or `<park>+predict:<deadline fields>`, where `<park>` is `block`, `poll` or
7//! `flag:park_us=<n>,backstop_us=<n>` and the deadline fields are
8//! `margin_us=<n>,min_sleep_us=<n>,max_sleep_us=<n>,window=<n>`, each key exactly once.
9//! `block` takes no prediction: the vendor's blocking wait has nothing to sleep in front of.
10//!
11//! Invariants: pure; a plan that parses is in range (see the bounds below).
12
13use alloc::format;
14use alloc::string::String;
15
16use crate::core::error::{Error, Result};
17use crate::error;
18
19/// The largest `window` a deadline plan may state.
20pub const WINDOW_MAX: usize = 16;
21/// The longest single park on the device word a plan may state (10 ms).
22const PARK_US_MAX: u64 = 10_000;
23/// The longest gap between two status queries in a flag wait a plan may state (1 s).
24const BACKSTOP_US_MAX: u64 = 1_000_000;
25
26/// How a host waits at a readback point.
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub struct WaitPlan {
29    /// Sleep to the point's predicted completion minus a margin before parking.
30    pub predict: Option<DeadlinePlan>,
31    /// How the wait is spent once it is not (or no longer) sleeping.
32    pub park: Park,
33}
34
35/// How the rest of a wait is spent.
36#[derive(Clone, Copy, Debug, PartialEq, Eq)]
37pub enum Park {
38    /// The vendor's blocking wait (CUDA: `cuStreamSynchronize`).
39    Block,
40    /// Poll the status until it reports done.
41    QueryPoll,
42    /// Park on a word the device writes behind the work; query the status every `backstop_ns`.
43    FlagWait {
44        /// Longest single park before the waiter looks at the clock again.
45        park_ns: u64,
46        /// Longest time between two status queries (the error authority).
47        backstop_ns: u64,
48    },
49}
50
51/// The deadline fields.
52#[derive(Clone, Copy, Debug, PartialEq, Eq)]
53pub struct DeadlinePlan {
54    /// Wake this long before the predicted completion.
55    pub margin_ns: u64,
56    /// A computed sleep shorter than this is skipped.
57    pub min_sleep_ns: u64,
58    /// No single sleep is longer than this.
59    pub max_sleep_ns: u64,
60    /// Exact residuals kept per readback point (`1..=WINDOW_MAX`); the prediction is their min.
61    pub window: u8,
62}
63
64/// The plan text is refused for `why`.
65fn refused(text: &str, why: &str) -> Error {
66    error!(InvalidPlan, "completion wait plan", format!("{text:?}: {why}"))
67}
68
69/// The values of the `key=<u64>` fields of `body`, in `keys` order, each once.
70fn fields<const N: usize>(text: &str, body: &str, keys: [&str; N]) -> Result<[u64; N]> {
71    let mut got: [Option<u64>; N] = [None; N];
72    for kv in body.split(',') {
73        let (k, v) = kv
74            .split_once('=')
75            .ok_or_else(|| refused(text, "a field is not key=<n>"))?;
76        let i = keys
77            .iter()
78            .position(|x| *x == k)
79            .ok_or_else(|| refused(text, &format!("unknown field {k:?}; expected {keys:?}")))?;
80        if got[i].is_some() {
81            return Err(refused(text, &format!("field {k:?} stated twice")));
82        }
83        got[i] = Some(
84            v.parse()
85                .map_err(|_| refused(text, &format!("field {k:?}: {v:?} is not a u64")))?,
86        );
87    }
88    let mut out = [0u64; N];
89    for ((o, g), k) in out.iter_mut().zip(got).zip(keys) {
90        *o = g.ok_or_else(|| {
91            refused(
92                text,
93                &format!("field {k:?} missing (every field is required)"),
94            )
95        })?;
96    }
97    Ok(out)
98}
99
100/// `us` in nanoseconds, refused on overflow.
101fn ns(text: &str, us: u64) -> Result<u64> {
102    us.checked_mul(1_000)
103        .ok_or_else(|| refused(text, "a duration overflows u64 nanoseconds"))
104}
105
106impl DeadlinePlan {
107    fn parse(text: &str, body: &str) -> Result<Self> {
108        let [m, lo, hi, w] = fields(
109            text,
110            body,
111            ["margin_us", "min_sleep_us", "max_sleep_us", "window"],
112        )?;
113        let window = u8::try_from(w)
114            .ok()
115            .filter(|w| (1..=WINDOW_MAX).contains(&usize::from(*w)));
116        let window = window
117            .ok_or_else(|| refused(text, &format!("window {w}: expected 1..={WINDOW_MAX}")))?;
118        if lo > hi {
119            return Err(refused(text, "min_sleep_us is above max_sleep_us"));
120        }
121        Ok(Self {
122            margin_ns: ns(text, m)?,
123            min_sleep_ns: ns(text, lo)?,
124            max_sleep_ns: ns(text, hi)?,
125            window,
126        })
127    }
128}
129
130impl Park {
131    fn parse(text: &str, park: &str) -> Result<Self> {
132        match park {
133            "block" => Ok(Self::Block),
134            "poll" => Ok(Self::QueryPoll),
135            p => {
136                let body = p
137                    .strip_prefix("flag:")
138                    .ok_or_else(|| refused(text, "park: expected block, poll or flag:<fields>"))?;
139                let [park_us, backstop_us] = fields(text, body, ["park_us", "backstop_us"])?;
140                if !(1..=PARK_US_MAX).contains(&park_us) {
141                    return Err(refused(
142                        text,
143                        &format!("park_us {park_us}: expected 1..={PARK_US_MAX}"),
144                    ));
145                }
146                if !(1..=BACKSTOP_US_MAX).contains(&backstop_us) {
147                    return Err(refused(
148                        text,
149                        &format!("backstop_us {backstop_us}: expected 1..={BACKSTOP_US_MAX}"),
150                    ));
151                }
152                Ok(Self::FlagWait {
153                    park_ns: ns(text, park_us)?,
154                    backstop_ns: ns(text, backstop_us)?,
155                })
156            }
157        }
158    }
159}
160
161impl WaitPlan {
162    /// Parse the lever text (module docs).
163    ///
164    /// # Errors
165    /// `InvalidArgument` naming the text and the first thing wrong with it.
166    pub fn parse(text: &str) -> Result<Self> {
167        let (park, predict) = match text.split_once('+') {
168            None => (text, None),
169            Some((park, rest)) => {
170                let body = rest
171                    .strip_prefix("predict:")
172                    .ok_or_else(|| refused(text, "after '+': expected predict:<fields>"))?;
173                (park, Some(DeadlinePlan::parse(text, body)?))
174            }
175        };
176        let park = Park::parse(text, park)?;
177        if park == Park::Block && predict.is_some() {
178            return Err(refused(text, "block takes no predict"));
179        }
180        Ok(Self { predict, park })
181    }
182
183    /// The plan as `parse` reads it (round-trips exactly).
184    #[must_use]
185    pub fn label(&self) -> String {
186        let park = match self.park {
187            Park::Block => String::from("block"),
188            Park::QueryPoll => String::from("poll"),
189            Park::FlagWait {
190                park_ns,
191                backstop_ns,
192            } => {
193                format!(
194                    "flag:park_us={},backstop_us={}",
195                    park_ns / 1_000,
196                    backstop_ns / 1_000
197                )
198            }
199        };
200        match self.predict {
201            None => park,
202            Some(d) => format!(
203                "{park}+predict:margin_us={},min_sleep_us={},max_sleep_us={},window={}",
204                d.margin_ns / 1_000,
205                d.min_sleep_ns / 1_000,
206                d.max_sleep_ns / 1_000,
207                d.window
208            ),
209        }
210    }
211}
212
213/// A host capability a park needs.
214#[derive(Clone, Copy, Debug, PartialEq, Eq)]
215pub enum Capability {
216    /// A device-written word the host can wait on in user mode (aarch64 `ldxr` + `wfe`/`wfet`;
217    /// later x86 MONITORX/UMWAIT), and a driver that stores into it behind stream work.
218    DeviceWordWait,
219}
220
221/// One link of the declared fallback chain.
222#[derive(Clone, Copy, Debug, PartialEq, Eq)]
223pub struct Fallback {
224    /// The park this link applies to (its fields kept only for matching).
225    pub from: ParkKind,
226    /// What the host must offer for `from` to run.
227    pub needs: Capability,
228    /// The park the plan resolves to without it (the plan's `predict` is kept unchanged).
229    pub to: ParkKind,
230}
231
232/// A park's kind, without its fields.
233#[derive(Clone, Copy, Debug, PartialEq, Eq)]
234pub enum ParkKind {
235    /// [`Park::Block`].
236    Block,
237    /// [`Park::QueryPoll`].
238    QueryPoll,
239    /// [`Park::FlagWait`].
240    FlagWait,
241}
242
243/// the whole chain, as data. The executor never chooses;
244/// [`WaitPlan::resolve`] walks this table once at boot.
245/// The chain ends at `QueryPoll`, which needs nothing and runs on every host; `Block`, the
246/// vendor's blocking wait, is the conservative arm a plan may name (declared, not
247/// a default).
248pub const FALLBACKS: [Fallback; 1] = [Fallback {
249    from: ParkKind::FlagWait,
250    needs: Capability::DeviceWordWait,
251    to: ParkKind::QueryPoll,
252}];
253
254/// What a host offers, probed once at boot by the host (never by the core).
255#[derive(Clone, Debug, PartialEq, Eq)]
256pub struct HostCaps {
257    /// `Ok` when [`Capability::DeviceWordWait`] is present, else why not.
258    pub device_word_wait: core::result::Result<(), String>,
259}
260
261impl HostCaps {
262    fn lacks(&self, c: Capability) -> Option<&str> {
263        match c {
264            Capability::DeviceWordWait => self.device_word_wait.as_ref().err().map(String::as_str),
265        }
266    }
267}
268
269/// The plan a host runs, and the coded warning when a fallback was taken.
270#[derive(Clone, Debug)]
271pub struct Resolved {
272    /// The plan as requested.
273    pub requested: WaitPlan,
274    /// The plan that runs.
275    pub plan: WaitPlan,
276    /// [`ErrorCode::ParkFallback`](crate::ErrorCode::ParkFallback) when `plan != requested`.
277    pub warning: Option<Error>,
278}
279
280impl Resolved {
281    /// The running plan's label, naming the fallback when one was taken (boot log,
282    /// `/metrics`, bench record).
283    #[must_use]
284    pub fn label(&self) -> String {
285        if self.warning.is_some() {
286            format!(
287                "{} (fallback from {})",
288                self.plan.label(),
289                self.requested.label()
290            )
291        } else {
292            self.plan.label()
293        }
294    }
295}
296
297impl Park {
298    /// This park's kind.
299    #[must_use]
300    pub fn kind(&self) -> ParkKind {
301        match self {
302            Self::Block => ParkKind::Block,
303            Self::QueryPoll => ParkKind::QueryPoll,
304            Self::FlagWait { .. } => ParkKind::FlagWait,
305        }
306    }
307}
308
309impl WaitPlan {
310    /// Resolve against `caps` by walking [`FALLBACKS`]; pure.
311    #[must_use]
312    pub fn resolve(&self, caps: &HostCaps) -> Resolved {
313        let mut plan = *self;
314        let mut warning = None;
315        let kind = plan.park.kind();
316        for link in FALLBACKS.iter().filter(|l| l.from == kind) {
317            if let Some(why) = caps.lacks(link.needs) {
318                let from = plan.park;
319                plan.park = match link.to {
320                    ParkKind::Block => Park::Block,
321                    ParkKind::QueryPoll => Park::QueryPoll,
322                    ParkKind::FlagWait => from,
323                };
324                warning = Some(error!(
325                    ParkFallback,
326                    self.label(),
327                    format!("{:?} ({:?})", link.from, link.needs),
328                    String::from(why),
329                    plan.label()
330                ));
331            }
332        }
333        Resolved {
334            requested: *self,
335            plan,
336            warning,
337        }
338    }
339}
340
341#[cfg(test)]
342#[path = "wait_plan_tests.rs"]
343mod tests;