gpu_async_reactor/core/
bands.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! the banded plan: which [`WaitPlan`] a wait
4//! runs, chosen by the rows of the current decode step. The band edges are typed fields parsed
5//! once from the lever text (ROTI, PLEX); the only per-tick input is the row count, published
6//! once per step through [`StepRows`] (BPPT: O(1) per tick, nothing per token).
7//!
8//! Text: `<plan>` or `<plan>;rows>=<n>:<plan>;…`: the first band starts at 0 rows, every later
9//! band names its first row count, the edges strictly increase, at most [`BANDS_MAX`] bands.
10//! `<plan>` is [`WaitPlan`]'s text.
11//!
12//! Invariants: a parsed [`Bands`] is non-empty, starts at 0 rows and its edges strictly
13//! increase, so [`Bands::select`] always finds exactly one band.
14
15use alloc::format;
16use alloc::string::String;
17use alloc::sync::Arc;
18use alloc::vec::Vec;
19use core::sync::atomic::{AtomicU32, Ordering};
20
21use crate::core::error::Result;
22use crate::error;
23
24use crate::core::plan::{HostCaps, Resolved, WaitPlan};
25
26/// The most bands a plan may state.
27pub const BANDS_MAX: usize = 4;
28
29/// One band: `plan` runs from `from_rows` rows up to the next band's edge.
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct Band<P> {
32    /// The first row count of the band.
33    pub from_rows: u32,
34    /// The plan the band runs.
35    pub plan: P,
36}
37
38/// The bands, in edge order (module invariants).
39#[derive(Clone, Debug, PartialEq, Eq)]
40pub struct Bands<P> {
41    bands: Vec<Band<P>>,
42}
43
44/// The requested bands, as the lever states them.
45pub type WaitBands = Bands<WaitPlan>;
46
47impl<P> Bands<P> {
48    /// The plan of the band `rows` falls in.
49    #[must_use]
50    pub fn select(&self, rows: u32) -> &P {
51        let i = self.bands.partition_point(|b| b.from_rows <= rows);
52        // The first edge is 0, so `i >= 1` for every `rows`.
53        &self.bands[i - 1].plan
54    }
55
56    /// The bands in edge order.
57    pub fn iter(&self) -> impl Iterator<Item = &Band<P>> {
58        self.bands.iter()
59    }
60
61    /// The same edges over `f` of each plan.
62    #[must_use]
63    pub fn map<Q>(&self, mut f: impl FnMut(&P) -> Q) -> Bands<Q> {
64        Bands {
65            bands: self
66                .bands
67                .iter()
68                .map(|b| Band {
69                    from_rows: b.from_rows,
70                    plan: f(&b.plan),
71                })
72                .collect(),
73        }
74    }
75
76    fn label_with(&self, f: impl Fn(&P) -> String) -> String {
77        let mut out = String::new();
78        for (i, b) in self.bands.iter().enumerate() {
79            if i > 0 {
80                out.push_str(&format!(";rows>={}:", b.from_rows));
81            }
82            out.push_str(&f(&b.plan));
83        }
84        out
85    }
86}
87
88impl WaitBands {
89    /// Parse the lever text (module docs).
90    ///
91    /// # Errors
92    /// `InvalidArgument` naming the text and the first thing wrong with it, including every
93    /// [`WaitPlan::parse`] refusal of a band's plan.
94    pub fn parse(text: &str) -> Result<Self> {
95        let refused = |why: String| {
96            error!(InvalidPlan, "completion wait plan", format!("{text:?}: {why}"))
97        };
98        let mut bands: Vec<Band<WaitPlan>> = Vec::new();
99        for (i, part) in text.split(';').enumerate() {
100            if i == BANDS_MAX {
101                return Err(refused(format!("more than {BANDS_MAX} bands")));
102            }
103            let (from_rows, plan) = if i == 0 {
104                (0, part)
105            } else {
106                let (edge, plan) = part
107                    .strip_prefix("rows>=")
108                    .and_then(|p| p.split_once(':'))
109                    .ok_or_else(|| refused(format!("band {i}: expected rows>=<n>:<plan>")))?;
110                let edge: u32 = edge
111                    .parse()
112                    .map_err(|_| refused(format!("band {i}: {edge:?} is not a u32")))?;
113                (edge, plan)
114            };
115            if let Some(prev) = bands.last()
116                && from_rows <= prev.from_rows
117            {
118                return Err(refused(format!(
119                    "band {i}: rows>={from_rows} does not exceed the previous edge {}",
120                    prev.from_rows
121                )));
122            }
123            bands.push(Band {
124                from_rows,
125                plan: WaitPlan::parse(plan)?,
126            });
127        }
128        Ok(Self { bands })
129    }
130
131    /// The bands as `parse` reads them (round-trips exactly).
132    #[must_use]
133    pub fn label(&self) -> String {
134        self.label_with(WaitPlan::label)
135    }
136
137    /// Each band resolved against `caps` ([`WaitPlan::resolve`]); pure.
138    #[must_use]
139    pub fn resolve(&self, caps: &HostCaps) -> Bands<Resolved> {
140        self.map(|p| p.resolve(caps))
141    }
142}
143
144impl Bands<Resolved> {
145    /// The running bands' label, naming every fallback taken.
146    #[must_use]
147    pub fn label(&self) -> String {
148        self.label_with(Resolved::label)
149    }
150}
151
152/// The rows of the current decode step: written once per step by the router that
153/// launches it, read by every wait of the host that owns the bands. 0 until the first step.
154#[derive(Clone, Debug, Default)]
155pub struct StepRows(Arc<AtomicU32>);
156
157impl StepRows {
158    /// A fresh cell at 0 rows.
159    #[must_use]
160    pub fn new() -> Self {
161        Self::default()
162    }
163
164    /// The step about to launch has `rows` rows.
165    pub fn publish(&self, rows: u32) {
166        self.0.store(rows, Ordering::Relaxed);
167    }
168
169    /// The rows of the last step published.
170    #[must_use]
171    pub fn current(&self) -> u32 {
172        self.0.load(Ordering::Relaxed)
173    }
174}
175
176#[cfg(test)]
177#[path = "bands_tests.rs"]
178mod tests;
gpu_async_reactor/core/bands.rs

gpu_async_reactor/core/
bands.rs
