fidget_core/vm/
data.rs

1//! General-purpose tapes for use during evaluation or further compilation
2use crate::{
3    compiler::{RegOp, RegTape, RegisterAllocator, SsaOp, SsaTape},
4    context::{BadNode, Context, Node},
5    var::VarMap,
6    vm::Choice,
7};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10
11/// A flattened math expression, ready for evaluation or further compilation.
12///
13/// Under the hood, [`VmData`] stores two different representations:
14/// - A tape in [single static assignment form](https://en.wikipedia.org/wiki/Static_single-assignment_form)
15///   ([`SsaTape`]), which is suitable for use during tape simplification
16/// - A tape in register-allocated form ([`RegTape`]), which can be efficiently
17///   evaluated or lowered into machine assembly
18///
19/// # Example
20/// Consider the expression `x + y`.  The SSA tape will look something like
21/// this:
22/// ```text
23/// $0 = INPUT 0   // X
24/// $1 = INPUT 1   // Y
25/// $2 = ADD $0 $1 // (X + Y)
26/// ```
27///
28/// This will be lowered into a tape using real (or VM) registers:
29/// ```text
30/// r0 = INPUT 0 // X
31/// r1 = INPUT 1 // Y
32/// r0 = ADD r0 r1 // (X + Y)
33/// ```
34///
35/// Note that in this form, registers are reused (e.g. `r0` stores both `X` and
36/// `X + Y`).
37///
38/// We can peek at the internals and see this register-allocated tape:
39/// ```
40/// use fidget_core::{
41///     compiler::RegOp,
42///     context::{Context, Tree},
43///     vm::VmData,
44///     var::Var,
45/// };
46///
47/// let tree = Tree::x() + Tree::y();
48/// let mut ctx = Context::new();
49/// let sum = ctx.import(&tree);
50/// let data = VmData::<255>::new(&ctx, &[sum])?;
51/// assert_eq!(data.len(), 4); // X, Y, (X + Y), and output
52///
53/// let mut iter = data.iter_asm();
54/// let vars = &data.vars; // map from var to index
55/// assert_eq!(iter.next().unwrap(), RegOp::Input(0, vars[&Var::X] as u32));
56/// assert_eq!(iter.next().unwrap(), RegOp::Input(1, vars[&Var::Y] as u32));
57/// assert_eq!(iter.next().unwrap(), RegOp::AddRegReg(0, 0, 1));
58/// # Ok::<(), Box<dyn std::error::Error>>(())
59/// ```
60///
61/// Despite this peek at its internals, users are unlikely to touch `VmData`
62/// directly; a [`VmShape`](crate::vm::VmShape) wraps the `VmData` and
63/// implements our common traits.
64#[derive(Default, Serialize, Deserialize)]
65pub struct VmData<const N: usize = { u8::MAX as usize }> {
66    ssa: SsaTape,
67    asm: RegTape,
68
69    /// Mapping from variables to indices during evaluation
70    ///
71    /// This member is stored in a shared pointer because it's passed down to
72    /// children (constructed with [`VmData::simplify`]).
73    pub vars: Arc<VarMap>,
74}
75
76impl<const N: usize> VmData<N> {
77    /// Builds a new tape for the given node
78    pub fn new(context: &Context, nodes: &[Node]) -> Result<Self, BadNode> {
79        let (ssa, vars) = SsaTape::new(context, nodes)?;
80        let asm = RegTape::new::<N>(&ssa);
81        Ok(Self {
82            ssa,
83            asm,
84            vars: vars.into(),
85        })
86    }
87
88    /// Returns the length of the internal VM tape
89    pub fn len(&self) -> usize {
90        self.asm.len()
91    }
92
93    /// Returns true if the internal VM tape is empty
94    pub fn is_empty(&self) -> bool {
95        self.asm.is_empty()
96    }
97
98    /// Returns the number of choice (min/max) nodes in the tape.
99    ///
100    /// This is required because some evaluators pre-allocate spaces for the
101    /// choice array.
102    pub fn choice_count(&self) -> usize {
103        self.ssa.choice_count
104    }
105
106    /// Returns the number of output nodes in the tape.
107    ///
108    /// This is required because some evaluators pre-allocate spaces for the
109    /// output array.
110    pub fn output_count(&self) -> usize {
111        self.ssa.output_count
112    }
113
114    /// Returns the number of slots used by the inner VM tape
115    pub fn slot_count(&self) -> usize {
116        self.asm.slot_count()
117    }
118
119    /// Simplifies both inner tapes, using the provided choice array
120    ///
121    /// To minimize allocations, this function takes a [`VmWorkspace`] and
122    /// spare [`VmData`]; it will reuse those allocations.
123    pub fn simplify<const M: usize>(
124        &self,
125        choices: &[Choice],
126        workspace: &mut VmWorkspace<M>,
127        mut tape: VmData<M>,
128    ) -> Result<VmData<M>, BadChoiceSlice> {
129        if choices.len() != self.choice_count() {
130            return Err(BadChoiceSlice {
131                actual: choices.len(),
132                expected: self.choice_count(),
133            });
134        }
135        tape.ssa.reset();
136
137        // Steal `tape.asm` and hand it to the workspace for use in allocator
138        workspace.reset(self.ssa.tape.len(), tape.asm);
139
140        let mut choice_count = 0;
141        let mut output_count = 0;
142
143        // Other iterators to consume various arrays in order
144        let mut choice_iter = choices.iter().rev();
145
146        let mut ops_out = tape.ssa.tape;
147
148        for mut op in self.ssa.tape.iter().cloned() {
149            let index = match &mut op {
150                SsaOp::Output(reg, _i) => {
151                    *reg = workspace.get_or_insert_active(*reg);
152                    workspace.alloc.op(op);
153                    ops_out.push(op);
154                    output_count += 1;
155                    continue;
156                }
157                _ => op.output().unwrap(),
158            };
159
160            if workspace.active(index).is_none() {
161                if op.has_choice() {
162                    choice_iter.next().unwrap();
163                }
164                continue;
165            }
166
167            // Because we reassign nodes when they're used as an *input*
168            // (while walking the tape in reverse), this node must have been
169            // assigned already.
170            let new_index = workspace.active(index).unwrap();
171
172            match &mut op {
173                SsaOp::Output(..) => unreachable!(),
174                SsaOp::Input(index, ..) | SsaOp::CopyImm(index, ..) => {
175                    *index = new_index;
176                }
177                SsaOp::NegReg(index, arg)
178                | SsaOp::AbsReg(index, arg)
179                | SsaOp::RecipReg(index, arg)
180                | SsaOp::SqrtReg(index, arg)
181                | SsaOp::SquareReg(index, arg)
182                | SsaOp::FloorReg(index, arg)
183                | SsaOp::CeilReg(index, arg)
184                | SsaOp::RoundReg(index, arg)
185                | SsaOp::SinReg(index, arg)
186                | SsaOp::CosReg(index, arg)
187                | SsaOp::TanReg(index, arg)
188                | SsaOp::AsinReg(index, arg)
189                | SsaOp::AcosReg(index, arg)
190                | SsaOp::AtanReg(index, arg)
191                | SsaOp::ExpReg(index, arg)
192                | SsaOp::LnReg(index, arg)
193                | SsaOp::NotReg(index, arg)
194                | SsaOp::RandReg(index, arg) => {
195                    *index = new_index;
196                    *arg = workspace.get_or_insert_active(*arg);
197                }
198                SsaOp::CopyReg(index, src) => {
199                    // CopyReg effectively does
200                    //      dst <= src
201                    // If src has not yet been used (as we iterate backwards
202                    // through the tape), then we can replace it with dst
203                    // everywhere!
204                    match workspace.active(*src) {
205                        Some(new_src) => {
206                            *index = new_index;
207                            *src = new_src;
208                        }
209                        None => {
210                            workspace.set_active(*src, new_index);
211                            continue;
212                        }
213                    }
214                }
215                SsaOp::MinRegImm(index, arg, imm)
216                | SsaOp::MaxRegImm(index, arg, imm)
217                | SsaOp::AndRegImm(index, arg, imm)
218                | SsaOp::OrRegImm(index, arg, imm) => {
219                    match choice_iter.next().unwrap() {
220                        Choice::Left => match workspace.active(*arg) {
221                            Some(new_arg) => {
222                                op = SsaOp::CopyReg(new_index, new_arg);
223                            }
224                            None => {
225                                workspace.set_active(*arg, new_index);
226                                continue;
227                            }
228                        },
229                        Choice::Right => {
230                            op = SsaOp::CopyImm(new_index, *imm);
231                        }
232                        Choice::Both => {
233                            choice_count += 1;
234                            *index = new_index;
235                            *arg = workspace.get_or_insert_active(*arg);
236                        }
237                        Choice::Unknown => panic!("oh no"),
238                    }
239                }
240                SsaOp::MinRegReg(index, lhs, rhs)
241                | SsaOp::MaxRegReg(index, lhs, rhs)
242                | SsaOp::AndRegReg(index, lhs, rhs)
243                | SsaOp::OrRegReg(index, lhs, rhs) => {
244                    match choice_iter.next().unwrap() {
245                        Choice::Left => match workspace.active(*lhs) {
246                            Some(new_lhs) => {
247                                op = SsaOp::CopyReg(new_index, new_lhs);
248                            }
249                            None => {
250                                workspace.set_active(*lhs, new_index);
251                                continue;
252                            }
253                        },
254                        Choice::Right => match workspace.active(*rhs) {
255                            Some(new_rhs) => {
256                                op = SsaOp::CopyReg(new_index, new_rhs);
257                            }
258                            None => {
259                                workspace.set_active(*rhs, new_index);
260                                continue;
261                            }
262                        },
263                        Choice::Both => {
264                            choice_count += 1;
265                            *index = new_index;
266                            *lhs = workspace.get_or_insert_active(*lhs);
267                            *rhs = workspace.get_or_insert_active(*rhs);
268                        }
269                        Choice::Unknown => panic!("oh no"),
270                    }
271                }
272                SsaOp::AddRegReg(index, lhs, rhs)
273                | SsaOp::MulRegReg(index, lhs, rhs)
274                | SsaOp::SubRegReg(index, lhs, rhs)
275                | SsaOp::DivRegReg(index, lhs, rhs)
276                | SsaOp::AtanRegReg(index, lhs, rhs)
277                | SsaOp::CompareRegReg(index, lhs, rhs)
278                | SsaOp::MixRegReg(index, lhs, rhs)
279                | SsaOp::ModRegReg(index, lhs, rhs) => {
280                    *index = new_index;
281                    *lhs = workspace.get_or_insert_active(*lhs);
282                    *rhs = workspace.get_or_insert_active(*rhs);
283                }
284                SsaOp::AddRegImm(index, arg, _imm)
285                | SsaOp::MulRegImm(index, arg, _imm)
286                | SsaOp::SubRegImm(index, arg, _imm)
287                | SsaOp::SubImmReg(index, arg, _imm)
288                | SsaOp::DivRegImm(index, arg, _imm)
289                | SsaOp::DivImmReg(index, arg, _imm)
290                | SsaOp::AtanImmReg(index, arg, _imm)
291                | SsaOp::AtanRegImm(index, arg, _imm)
292                | SsaOp::CompareRegImm(index, arg, _imm)
293                | SsaOp::CompareImmReg(index, arg, _imm)
294                | SsaOp::MixRegImm(index, arg, _imm)
295                | SsaOp::MixImmReg(index, arg, _imm)
296                | SsaOp::ModRegImm(index, arg, _imm)
297                | SsaOp::ModImmReg(index, arg, _imm) => {
298                    *index = new_index;
299                    *arg = workspace.get_or_insert_active(*arg);
300                }
301            }
302            workspace.alloc.op(op);
303            ops_out.push(op);
304        }
305
306        assert_eq!(workspace.count as usize + 1, ops_out.len());
307        let asm_tape = workspace.alloc.finalize();
308
309        Ok(VmData {
310            ssa: SsaTape {
311                tape: ops_out,
312                choice_count,
313                output_count,
314            },
315            asm: asm_tape,
316            vars: self.vars.clone(),
317        })
318    }
319
320    /// Produces an iterator that visits [`RegOp`] values in evaluation order
321    pub fn iter_asm(&self) -> impl Iterator<Item = RegOp> + '_ {
322        self.asm.iter().cloned().rev()
323    }
324
325    /// Returns a reference to the inner [`RegTape`]
326    pub fn asm(&self) -> &RegTape {
327        &self.asm
328    }
329
330    /// Pretty-prints the inner SSA tape
331    pub fn pretty_print(&self) {
332        self.ssa.pretty_print();
333        for a in self.iter_asm() {
334            println!("{a:?}");
335        }
336    }
337}
338
339/// Error type for simplification
340#[derive(thiserror::Error, Debug)]
341#[error(
342    "choice slice length ({actual}) does not \
343     match choice count ({expected})"
344)]
345pub struct BadChoiceSlice {
346    actual: usize,
347    expected: usize,
348}
349
350////////////////////////////////////////////////////////////////////////////////
351
352/// Data structures used during [`VmData::simplify`]
353///
354/// This is exposed to minimize reallocations in hot loops.
355pub struct VmWorkspace<const N: usize> {
356    /// Register allocator
357    pub(crate) alloc: RegisterAllocator<N>,
358
359    /// Current bindings from SSA variables to registers
360    pub(crate) bind: Vec<u32>,
361
362    /// Number of active SSA bindings
363    ///
364    /// This value is monotonically increasing; each SSA variable gets the next
365    /// value if it is unassigned when encountered.
366    count: u32,
367}
368
369impl<const N: usize> Default for VmWorkspace<N> {
370    fn default() -> Self {
371        Self {
372            alloc: RegisterAllocator::empty(),
373            bind: vec![],
374            count: 0,
375        }
376    }
377}
378
379impl<const N: usize> VmWorkspace<N> {
380    fn active(&self, i: u32) -> Option<u32> {
381        if self.bind[i as usize] != u32::MAX {
382            Some(self.bind[i as usize])
383        } else {
384            None
385        }
386    }
387
388    fn get_or_insert_active(&mut self, i: u32) -> u32 {
389        if self.bind[i as usize] == u32::MAX {
390            self.bind[i as usize] = self.count;
391            self.count += 1;
392        }
393        self.bind[i as usize]
394    }
395
396    fn set_active(&mut self, i: u32, bind: u32) {
397        self.bind[i as usize] = bind;
398    }
399
400    /// Resets the workspace, preserving allocations and claiming the given
401    /// [`RegTape`].
402    pub fn reset(&mut self, tape_len: usize, tape: RegTape) {
403        self.alloc.reset(tape_len, tape);
404        self.bind.fill(u32::MAX);
405        self.bind.resize(tape_len, u32::MAX);
406        self.count = 0;
407    }
408}
409
410#[cfg(test)]
411mod test {
412    use super::*;
413
414    #[test]
415    fn simplify_reg_count_change() {
416        let mut ctx = Context::new();
417        let x = ctx.x();
418        let y = ctx.y();
419        let z = ctx.z();
420        let xy = ctx.add(x, y).unwrap();
421        let xyz = ctx.add(xy, z).unwrap();
422
423        let data = VmData::<3>::new(&ctx, &[xyz]).unwrap();
424        assert_eq!(data.len(), 6); // 3x input, 2x add, 1x output
425        let next = data
426            .simplify::<2>(&[], &mut Default::default(), Default::default())
427            .unwrap();
428        assert_eq!(next.len(), 8); // extra load + store
429
430        let data = VmData::<2>::new(&ctx, &[xyz]).unwrap();
431        assert_eq!(data.len(), 8);
432        let next = data
433            .simplify::<3>(&[], &mut Default::default(), Default::default())
434            .unwrap();
435        assert_eq!(next.len(), 6);
436    }
437}
fidget_core/vm/data.rs

fidget_core/vm/
data.rs
