kime_tensor/
bucket.rs

1//! The bucket table: which padded shapes plans are built for.
2//!
3//! The table is data. The default is `buckets.txt` next to this file, compiled in, and an operator
4//! can load another with [`Buckets::parse`] without touching code.
5
6use std::fmt;
7
8use crate::plan::Rows;
9
10/// The shape a plan is built for. A batch fits a bucket when it has no more of each.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub struct Bucket {
13    /// Tokens.
14    pub tokens: usize,
15    /// Sequences.
16    pub seqs: usize,
17    /// Markers.
18    pub markers: usize,
19}
20
21impl Bucket {
22    /// The row count of `rows` in this bucket.
23    #[must_use]
24    pub fn rows(&self, rows: Rows) -> usize {
25        match rows {
26            Rows::Tokens => self.tokens,
27            Rows::Seqs => self.seqs,
28            Rows::Markers => self.markers,
29        }
30    }
31
32    /// Whether a batch of this size fits.
33    #[must_use]
34    pub fn holds(&self, tokens: usize, seqs: usize, markers: usize) -> bool {
35        tokens <= self.tokens && seqs <= self.seqs && markers <= self.markers
36    }
37}
38
39impl fmt::Display for Bucket {
40    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41        write!(f, "{} tokens, {} seqs, {} markers", self.tokens, self.seqs, self.markers)
42    }
43}
44
45/// A bucket table line that does not parse.
46#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct ParseError {
48    /// 1 based.
49    pub line: usize,
50    /// What is wrong with it.
51    pub reason: String,
52}
53
54impl fmt::Display for ParseError {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        write!(f, "bucket table line {}: {}", self.line, self.reason)
57    }
58}
59
60impl std::error::Error for ParseError {}
61
62/// Buckets per stage, each stage sorted smallest first.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct Buckets {
65    stages: Vec<(String, Vec<Bucket>)>,
66}
67
68/// The table kime ships with.
69pub const DEFAULT: &str = include_str!("buckets.txt");
70
71impl Default for Buckets {
72    fn default() -> Self {
73        Self::parse(DEFAULT).expect("the default bucket table parses")
74    }
75}
76
77impl Buckets {
78    /// Reads a table: one bucket per line as `stage tokens seqs markers`, with `#` comments.
79    ///
80    /// # Errors
81    ///
82    /// A line without four fields, a count that is not a number, or a bucket that repeats.
83    pub fn parse(text: &str) -> Result<Self, ParseError> {
84        let mut stages: Vec<(String, Vec<Bucket>)> = Vec::new();
85        for (i, line) in text.lines().enumerate() {
86            let err = |reason: String| ParseError { line: i + 1, reason };
87            let line = line.split('#').next().unwrap_or_default().trim();
88            if line.is_empty() {
89                continue;
90            }
91            let f: Vec<&str> = line.split_whitespace().collect();
92            let [stage, tokens, seqs, markers] = f[..] else {
93                return Err(err(format!("expected 4 fields, found {}", f.len())));
94            };
95            let num =
96                |s: &str| s.parse::<usize>().map_err(|_| err(format!("{s:?} is not a count")));
97            let b = Bucket { tokens: num(tokens)?, seqs: num(seqs)?, markers: num(markers)? };
98            let at = match stages.iter().position(|s| s.0 == stage) {
99                Some(at) => at,
100                None => {
101                    stages.push((stage.to_string(), Vec::new()));
102                    stages.len() - 1
103                }
104            };
105            if stages[at].1.contains(&b) {
106                return Err(err(format!("{stage} {b} is listed twice")));
107            }
108            stages[at].1.push(b);
109        }
110        for s in &mut stages {
111            s.1.sort_by_key(|b| (b.tokens, b.seqs, b.markers));
112        }
113        Ok(Self { stages })
114    }
115
116    /// The buckets of `stage`, smallest first, empty for an unknown stage.
117    #[must_use]
118    pub fn stage(&self, stage: &str) -> &[Bucket] {
119        self.stages.iter().find(|s| s.0 == stage).map_or(&[], |s| &s.1)
120    }
121
122    /// The smallest bucket of `stage` that holds the batch.
123    #[must_use]
124    pub fn pick(&self, stage: &str, tokens: usize, seqs: usize, markers: usize) -> Option<Bucket> {
125        self.stage(stage).iter().copied().find(|b| b.holds(tokens, seqs, markers))
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    #[test]
134    fn default_table() {
135        let t = Buckets::default();
136        assert_eq!(t.stage("state").len(), 11);
137        assert_eq!(t.stage("question").len(), 9);
138        let b = t.pick("compat", 119, 1, 4).unwrap();
139        assert_eq!(b, Bucket { tokens: 128, seqs: 32, markers: 128 });
140        assert_eq!(t.stage("compat").len(), 20);
141        // Many short sequences push past the sequence limit before the token limit.
142        assert_eq!(t.pick("compat", 100, 41, 80).unwrap().tokens, 192);
143        assert_eq!(t.pick("compat", 20000, 1, 1), None);
144        assert_eq!(t.pick("nothing", 1, 1, 1), None);
145    }
146
147    #[test]
148    fn another_table_changes_the_choice() {
149        let t = Buckets::parse("compat 100 1 8 # one\n\ncompat 50 2 8\n").unwrap();
150        assert_eq!(t.stage("compat")[0].tokens, 50);
151        assert_eq!(t.pick("compat", 60, 1, 1).unwrap().tokens, 100);
152        assert_eq!(t.pick("compat", 60, 2, 1), None);
153    }
154
155    #[test]
156    fn bad_tables() {
157        assert_eq!(Buckets::parse("compat 1 2").unwrap_err().line, 1);
158        assert!(Buckets::parse("# c\ncompat 1 x 3").unwrap_err().reason.contains("\"x\""));
159        assert!(Buckets::parse("a 1 1 1\na 1 1 1").unwrap_err().reason.contains("twice"));
160    }
161}
kime_tensor/bucket.rs

kime_tensor/
bucket.rs
