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                    *index = new_index;
195                    *arg = workspace.get_or_insert_active(*arg);
196                }
197                SsaOp::CopyReg(index, src) => {
198                    // CopyReg effectively does
199                    //      dst <= src
200                    // If src has not yet been used (as we iterate backwards
201                    // through the tape), then we can replace it with dst
202                    // everywhere!
203                    match workspace.active(*src) {
204                        Some(new_src) => {
205                            *index = new_index;
206                            *src = new_src;
207                        }
208                        None => {
209                            workspace.set_active(*src, new_index);
210                            continue;
211                        }
212                    }
213                }
214                SsaOp::MinRegImm(index, arg, imm)
215                | SsaOp::MaxRegImm(index, arg, imm)
216                | SsaOp::AndRegImm(index, arg, imm)
217                | SsaOp::OrRegImm(index, arg, imm) => {
218                    match choice_iter.next().unwrap() {
219                        Choice::Left => match workspace.active(*arg) {
220                            Some(new_arg) => {
221                                op = SsaOp::CopyReg(new_index, new_arg);
222                            }
223                            None => {
224                                workspace.set_active(*arg, new_index);
225                                continue;
226                            }
227                        },
228                        Choice::Right => {
229                            op = SsaOp::CopyImm(new_index, *imm);
230                        }
231                        Choice::Both => {
232                            choice_count += 1;
233                            *index = new_index;
234                            *arg = workspace.get_or_insert_active(*arg);
235                        }
236                        Choice::Unknown => panic!("oh no"),
237                    }
238                }
239                SsaOp::MinRegReg(index, lhs, rhs)
240                | SsaOp::MaxRegReg(index, lhs, rhs)
241                | SsaOp::AndRegReg(index, lhs, rhs)
242                | SsaOp::OrRegReg(index, lhs, rhs) => {
243                    match choice_iter.next().unwrap() {
244                        Choice::Left => match workspace.active(*lhs) {
245                            Some(new_lhs) => {
246                                op = SsaOp::CopyReg(new_index, new_lhs);
247                            }
248                            None => {
249                                workspace.set_active(*lhs, new_index);
250                                continue;
251                            }
252                        },
253                        Choice::Right => match workspace.active(*rhs) {
254                            Some(new_rhs) => {
255                                op = SsaOp::CopyReg(new_index, new_rhs);
256                            }
257                            None => {
258                                workspace.set_active(*rhs, new_index);
259                                continue;
260                            }
261                        },
262                        Choice::Both => {
263                            choice_count += 1;
264                            *index = new_index;
265                            *lhs = workspace.get_or_insert_active(*lhs);
266                            *rhs = workspace.get_or_insert_active(*rhs);
267                        }
268                        Choice::Unknown => panic!("oh no"),
269                    }
270                }
271                SsaOp::AddRegReg(index, lhs, rhs)
272                | SsaOp::MulRegReg(index, lhs, rhs)
273                | SsaOp::SubRegReg(index, lhs, rhs)
274                | SsaOp::DivRegReg(index, lhs, rhs)
275                | SsaOp::AtanRegReg(index, lhs, rhs)
276                | SsaOp::CompareRegReg(index, lhs, rhs)
277                | SsaOp::ModRegReg(index, lhs, rhs) => {
278                    *index = new_index;
279                    *lhs = workspace.get_or_insert_active(*lhs);
280                    *rhs = workspace.get_or_insert_active(*rhs);
281                }
282                SsaOp::AddRegImm(index, arg, _imm)
283                | SsaOp::MulRegImm(index, arg, _imm)
284                | SsaOp::SubRegImm(index, arg, _imm)
285                | SsaOp::SubImmReg(index, arg, _imm)
286                | SsaOp::DivRegImm(index, arg, _imm)
287                | SsaOp::DivImmReg(index, arg, _imm)
288                | SsaOp::AtanImmReg(index, arg, _imm)
289                | SsaOp::AtanRegImm(index, arg, _imm)
290                | SsaOp::CompareRegImm(index, arg, _imm)
291                | SsaOp::CompareImmReg(index, arg, _imm)
292                | SsaOp::ModRegImm(index, arg, _imm)
293                | SsaOp::ModImmReg(index, arg, _imm) => {
294                    *index = new_index;
295                    *arg = workspace.get_or_insert_active(*arg);
296                }
297            }
298            workspace.alloc.op(op);
299            ops_out.push(op);
300        }
301
302        assert_eq!(workspace.count as usize + 1, ops_out.len());
303        let asm_tape = workspace.alloc.finalize();
304
305        Ok(VmData {
306            ssa: SsaTape {
307                tape: ops_out,
308                choice_count,
309                output_count,
310            },
311            asm: asm_tape,
312            vars: self.vars.clone(),
313        })
314    }
315
316    /// Produces an iterator that visits [`RegOp`] values in evaluation order
317    pub fn iter_asm(&self) -> impl Iterator<Item = RegOp> + '_ {
318        self.asm.iter().cloned().rev()
319    }
320
321    /// Returns a reference to the inner [`RegTape`]
322    pub fn asm(&self) -> &RegTape {
323        &self.asm
324    }
325
326    /// Pretty-prints the inner SSA tape
327    pub fn pretty_print(&self) {
328        self.ssa.pretty_print();
329        for a in self.iter_asm() {
330            println!("{a:?}");
331        }
332    }
333}
334
335/// Error type for simplification
336#[derive(thiserror::Error, Debug)]
337#[error(
338    "choice slice length ({actual}) does not \
339     match choice count ({expected})"
340)]
341pub struct BadChoiceSlice {
342    actual: usize,
343    expected: usize,
344}
345
346////////////////////////////////////////////////////////////////////////////////
347
348/// Data structures used during [`VmData::simplify`]
349///
350/// This is exposed to minimize reallocations in hot loops.
351pub struct VmWorkspace<const N: usize> {
352    /// Register allocator
353    pub(crate) alloc: RegisterAllocator<N>,
354
355    /// Current bindings from SSA variables to registers
356    pub(crate) bind: Vec<u32>,
357
358    /// Number of active SSA bindings
359    ///
360    /// This value is monotonically increasing; each SSA variable gets the next
361    /// value if it is unassigned when encountered.
362    count: u32,
363}
364
365impl<const N: usize> Default for VmWorkspace<N> {
366    fn default() -> Self {
367        Self {
368            alloc: RegisterAllocator::empty(),
369            bind: vec![],
370            count: 0,
371        }
372    }
373}
374
375impl<const N: usize> VmWorkspace<N> {
376    fn active(&self, i: u32) -> Option<u32> {
377        if self.bind[i as usize] != u32::MAX {
378            Some(self.bind[i as usize])
379        } else {
380            None
381        }
382    }
383
384    fn get_or_insert_active(&mut self, i: u32) -> u32 {
385        if self.bind[i as usize] == u32::MAX {
386            self.bind[i as usize] = self.count;
387            self.count += 1;
388        }
389        self.bind[i as usize]
390    }
391
392    fn set_active(&mut self, i: u32, bind: u32) {
393        self.bind[i as usize] = bind;
394    }
395
396    /// Resets the workspace, preserving allocations and claiming the given
397    /// [`RegTape`].
398    pub fn reset(&mut self, tape_len: usize, tape: RegTape) {
399        self.alloc.reset(tape_len, tape);
400        self.bind.fill(u32::MAX);
401        self.bind.resize(tape_len, u32::MAX);
402        self.count = 0;
403    }
404}
405
406#[cfg(test)]
407mod test {
408    use super::*;
409
410    #[test]
411    fn simplify_reg_count_change() {
412        let mut ctx = Context::new();
413        let x = ctx.x();
414        let y = ctx.y();
415        let z = ctx.z();
416        let xy = ctx.add(x, y).unwrap();
417        let xyz = ctx.add(xy, z).unwrap();
418
419        let data = VmData::<3>::new(&ctx, &[xyz]).unwrap();
420        assert_eq!(data.len(), 6); // 3x input, 2x add, 1x output
421        let next = data
422            .simplify::<2>(&[], &mut Default::default(), Default::default())
423            .unwrap();
424        assert_eq!(next.len(), 8); // extra load + store
425
426        let data = VmData::<2>::new(&ctx, &[xyz]).unwrap();
427        assert_eq!(data.len(), 8);
428        let next = data
429            .simplify::<3>(&[], &mut Default::default(), Default::default())
430            .unwrap();
431        assert_eq!(next.len(), 6);
432    }
433}
fidget_core/vm/data.rs

fidget_core/vm/
data.rs
