fidget_core/compiler/
ssa_tape.rs

1//use crate::vm::{RegisterAllocator, Tape as VmTape};
2use crate::{
3    Context,
4    compiler::SsaOp,
5    context::{BadNode, BinaryOpcode, Node, Op, UnaryOpcode},
6    var::VarMap,
7};
8use serde::{Deserialize, Serialize};
9
10use std::collections::{HashMap, HashSet};
11
12/// Instruction tape, storing [opcodes in SSA form](crate::compiler::SsaOp)
13///
14/// Each operation has the following parameters
15/// - 4-byte opcode (required)
16/// - 4-byte output register (required)
17/// - 4-byte LHS register
18/// - 4-byte RHS register (or immediate `f32`)
19///
20/// All register addressing is absolute.
21#[derive(Clone, Debug, Default, Serialize, Deserialize)]
22pub struct SsaTape {
23    /// The tape is stored in reverse order, such that the root of the tree is
24    /// the first item in the tape.
25    pub tape: Vec<SsaOp>,
26
27    /// Number of choice operations in the tape
28    pub choice_count: usize,
29
30    /// Number of output operations in the tape
31    pub output_count: usize,
32}
33
34impl SsaTape {
35    /// Flattens a subtree of the graph into straight-line code.
36    ///
37    /// This should always succeed unless the `root` is from a different
38    /// `Context`, in which case [`BadNode`] will be returned.
39    pub fn new(
40        ctx: &Context,
41        roots: &[Node],
42    ) -> Result<(Self, VarMap), BadNode> {
43        let mut mapping = HashMap::new();
44        let mut parent_count: HashMap<Node, usize> = HashMap::new();
45        let mut slot_count = 0;
46
47        // Get either a node or constant index
48        #[derive(Copy, Clone, Debug)]
49        enum Slot {
50            Reg(u32),
51            Immediate(f32),
52        }
53
54        // Accumulate parent counts and declare all nodes
55        let mut seen = HashSet::new();
56        let mut vars = VarMap::new();
57        let mut todo = roots.to_vec();
58        while let Some(node) = todo.pop() {
59            if !seen.insert(node) {
60                continue;
61            }
62            let op = ctx.get_op(node).ok_or(BadNode)?;
63            let prev = match op {
64                Op::Const(c) => {
65                    mapping.insert(node, Slot::Immediate(c.0 as f32))
66                }
67                _ => {
68                    if let Op::Input(v) = op {
69                        vars.insert(*v);
70                    }
71                    let i = slot_count;
72                    slot_count += 1;
73                    mapping.insert(node, Slot::Reg(i))
74                }
75            };
76            assert!(prev.is_none());
77            for child in op.iter_children() {
78                *parent_count.entry(child).or_default() += 1;
79                todo.push(child);
80            }
81        }
82
83        // Now that we've populated our parents, flatten the graph
84        let mut seen = HashSet::new();
85        let mut todo = roots.to_vec();
86        let mut choice_count = 0;
87
88        let mut tape = vec![];
89        for (i, r) in roots.iter().enumerate() {
90            let i = i as u32;
91            match mapping[r] {
92                Slot::Reg(out_reg) => tape.push(SsaOp::Output(out_reg, i)),
93                Slot::Immediate(imm) => {
94                    let o = slot_count;
95                    slot_count += 1;
96                    tape.push(SsaOp::Output(o, i));
97                    tape.push(SsaOp::CopyImm(o, imm));
98                }
99            }
100        }
101
102        while let Some(node) = todo.pop() {
103            if *parent_count.get(&node).unwrap_or(&0) > 0 || !seen.insert(node)
104            {
105                continue;
106            }
107
108            let op = ctx.get_op(node).unwrap();
109            for child in op.iter_children() {
110                todo.push(child);
111                *parent_count.get_mut(&child).unwrap() -= 1;
112            }
113
114            let Slot::Reg(i) = mapping[&node] else {
115                // Constants are skipped, because they become immediates
116                continue;
117            };
118            let op = match op {
119                Op::Input(v) => {
120                    let arg = vars[v];
121                    SsaOp::Input(i, arg.try_into().unwrap())
122                }
123                Op::Const(..) => {
124                    unreachable!("skipped above")
125                }
126                Op::Binary(op, lhs, rhs) => {
127                    let lhs = mapping[lhs];
128                    let rhs = mapping[rhs];
129
130                    type RegFn = fn(u32, u32, u32) -> SsaOp;
131                    type ImmFn = fn(u32, u32, f32) -> SsaOp;
132                    let f: (RegFn, ImmFn, ImmFn) = match op {
133                        BinaryOpcode::Add => (
134                            SsaOp::AddRegReg,
135                            SsaOp::AddRegImm,
136                            SsaOp::AddRegImm,
137                        ),
138                        BinaryOpcode::Sub => (
139                            SsaOp::SubRegReg,
140                            SsaOp::SubRegImm,
141                            SsaOp::SubImmReg,
142                        ),
143                        BinaryOpcode::Mul => (
144                            SsaOp::MulRegReg,
145                            SsaOp::MulRegImm,
146                            SsaOp::MulRegImm,
147                        ),
148                        BinaryOpcode::Div => (
149                            SsaOp::DivRegReg,
150                            SsaOp::DivRegImm,
151                            SsaOp::DivImmReg,
152                        ),
153                        BinaryOpcode::Atan => (
154                            SsaOp::AtanRegReg,
155                            SsaOp::AtanRegImm,
156                            SsaOp::AtanImmReg,
157                        ),
158                        BinaryOpcode::Min => (
159                            SsaOp::MinRegReg,
160                            SsaOp::MinRegImm,
161                            SsaOp::MinRegImm,
162                        ),
163                        BinaryOpcode::Max => (
164                            SsaOp::MaxRegReg,
165                            SsaOp::MaxRegImm,
166                            SsaOp::MaxRegImm,
167                        ),
168                        BinaryOpcode::And => (
169                            SsaOp::AndRegReg,
170                            SsaOp::AndRegImm,
171                            |_out, _lhs, _rhs| {
172                                panic!("AndImmReg must be collapsed")
173                            },
174                        ),
175                        BinaryOpcode::Or => (
176                            SsaOp::OrRegReg,
177                            SsaOp::OrRegImm,
178                            |_out, _lhs, _rhs| {
179                                panic!("OrImmReg must be collapsed")
180                            },
181                        ),
182                        BinaryOpcode::Compare => (
183                            SsaOp::CompareRegReg,
184                            SsaOp::CompareRegImm,
185                            SsaOp::CompareImmReg,
186                        ),
187                        BinaryOpcode::Mod => (
188                            SsaOp::ModRegReg,
189                            SsaOp::ModRegImm,
190                            SsaOp::ModImmReg,
191                        ),
192                    };
193
194                    if matches!(
195                        op,
196                        BinaryOpcode::Min
197                            | BinaryOpcode::Max
198                            | BinaryOpcode::And
199                            | BinaryOpcode::Or
200                    ) {
201                        choice_count += 1;
202                    }
203
204                    match (lhs, rhs) {
205                        (Slot::Reg(lhs), Slot::Reg(rhs)) => f.0(i, lhs, rhs),
206                        (Slot::Reg(arg), Slot::Immediate(imm)) => {
207                            f.1(i, arg, imm)
208                        }
209                        (Slot::Immediate(imm), Slot::Reg(arg)) => {
210                            f.2(i, arg, imm)
211                        }
212                        (Slot::Immediate(..), Slot::Immediate(..)) => {
213                            panic!("Cannot handle f(imm, imm)")
214                        }
215                    }
216                }
217                Op::Unary(op, lhs) => {
218                    let lhs = match mapping[lhs] {
219                        Slot::Reg(r) => r,
220                        Slot::Immediate(..) => {
221                            panic!("Cannot handle f(imm)")
222                        }
223                    };
224                    let op = match op {
225                        UnaryOpcode::Neg => SsaOp::NegReg,
226                        UnaryOpcode::Abs => SsaOp::AbsReg,
227                        UnaryOpcode::Recip => SsaOp::RecipReg,
228                        UnaryOpcode::Sqrt => SsaOp::SqrtReg,
229                        UnaryOpcode::Square => SsaOp::SquareReg,
230                        UnaryOpcode::Floor => SsaOp::FloorReg,
231                        UnaryOpcode::Ceil => SsaOp::CeilReg,
232                        UnaryOpcode::Round => SsaOp::RoundReg,
233                        UnaryOpcode::Sin => SsaOp::SinReg,
234                        UnaryOpcode::Cos => SsaOp::CosReg,
235                        UnaryOpcode::Tan => SsaOp::TanReg,
236                        UnaryOpcode::Asin => SsaOp::AsinReg,
237                        UnaryOpcode::Acos => SsaOp::AcosReg,
238                        UnaryOpcode::Atan => SsaOp::AtanReg,
239                        UnaryOpcode::Exp => SsaOp::ExpReg,
240                        UnaryOpcode::Ln => SsaOp::LnReg,
241                        UnaryOpcode::Not => SsaOp::NotReg,
242                    };
243                    op(i, lhs)
244                }
245            };
246            tape.push(op);
247        }
248
249        Ok((
250            SsaTape {
251                tape,
252                choice_count,
253                output_count: roots.len(),
254            },
255            vars,
256        ))
257    }
258
259    /// Checks whether the tape is empty
260    pub fn is_empty(&self) -> bool {
261        self.tape.is_empty()
262    }
263
264    /// Returns the length of the tape
265    pub fn len(&self) -> usize {
266        self.tape.len()
267    }
268
269    /// Iterates over clauses in the tape in reverse-evaluation order
270    ///
271    /// The root (output) of the tape will be first in the iterator
272    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &SsaOp> {
273        self.tape.iter()
274    }
275
276    /// Resets to an empty tape, preserving allocations
277    pub fn reset(&mut self) {
278        self.tape.clear();
279        self.choice_count = 0;
280    }
281    /// Pretty-prints the given tape to `stdout`
282    pub fn pretty_print(&self) {
283        for &op in self.tape.iter().rev() {
284            match op {
285                SsaOp::Output(arg, i) => {
286                    println!("OUTPUT[{i}] = ${arg}");
287                }
288                SsaOp::Input(out, i) => {
289                    println!("${out} = INPUT[{i}]");
290                }
291                SsaOp::NegReg(out, arg)
292                | SsaOp::AbsReg(out, arg)
293                | SsaOp::RecipReg(out, arg)
294                | SsaOp::SqrtReg(out, arg)
295                | SsaOp::CopyReg(out, arg)
296                | SsaOp::SquareReg(out, arg)
297                | SsaOp::FloorReg(out, arg)
298                | SsaOp::CeilReg(out, arg)
299                | SsaOp::RoundReg(out, arg)
300                | SsaOp::SinReg(out, arg)
301                | SsaOp::CosReg(out, arg)
302                | SsaOp::TanReg(out, arg)
303                | SsaOp::AsinReg(out, arg)
304                | SsaOp::AcosReg(out, arg)
305                | SsaOp::AtanReg(out, arg)
306                | SsaOp::ExpReg(out, arg)
307                | SsaOp::LnReg(out, arg)
308                | SsaOp::NotReg(out, arg) => {
309                    let op = match op {
310                        SsaOp::NegReg(..) => "NEG",
311                        SsaOp::AbsReg(..) => "ABS",
312                        SsaOp::RecipReg(..) => "RECIP",
313                        SsaOp::SqrtReg(..) => "SQRT",
314                        SsaOp::SquareReg(..) => "SQUARE",
315                        SsaOp::FloorReg(..) => "FLOOR",
316                        SsaOp::CeilReg(..) => "CEIL",
317                        SsaOp::RoundReg(..) => "ROUND",
318                        SsaOp::SinReg(..) => "SIN",
319                        SsaOp::CosReg(..) => "COS",
320                        SsaOp::TanReg(..) => "TAN",
321                        SsaOp::AsinReg(..) => "ASIN",
322                        SsaOp::AcosReg(..) => "ACOS",
323                        SsaOp::AtanReg(..) => "ATAN",
324                        SsaOp::ExpReg(..) => "EXP",
325                        SsaOp::LnReg(..) => "LN",
326                        SsaOp::NotReg(..) => "NOT",
327                        SsaOp::CopyReg(..) => "COPY",
328                        _ => unreachable!(),
329                    };
330                    println!("${out} = {op} ${arg}");
331                }
332
333                SsaOp::AddRegReg(out, lhs, rhs)
334                | SsaOp::MulRegReg(out, lhs, rhs)
335                | SsaOp::DivRegReg(out, lhs, rhs)
336                | SsaOp::SubRegReg(out, lhs, rhs)
337                | SsaOp::MinRegReg(out, lhs, rhs)
338                | SsaOp::MaxRegReg(out, lhs, rhs)
339                | SsaOp::ModRegReg(out, lhs, rhs)
340                | SsaOp::AndRegReg(out, lhs, rhs)
341                | SsaOp::AtanRegReg(out, lhs, rhs)
342                | SsaOp::OrRegReg(out, lhs, rhs)
343                | SsaOp::CompareRegReg(out, lhs, rhs) => {
344                    let op = match op {
345                        SsaOp::AddRegReg(..) => "ADD",
346                        SsaOp::MulRegReg(..) => "MUL",
347                        SsaOp::DivRegReg(..) => "DIV",
348                        SsaOp::AtanRegReg(..) => "ATAN",
349                        SsaOp::SubRegReg(..) => "SUB",
350                        SsaOp::MinRegReg(..) => "MIN",
351                        SsaOp::MaxRegReg(..) => "MAX",
352                        SsaOp::ModRegReg(..) => "MAX",
353                        SsaOp::AndRegReg(..) => "AND",
354                        SsaOp::OrRegReg(..) => "OR",
355                        SsaOp::CompareRegReg(..) => "COMPARE",
356                        _ => unreachable!(),
357                    };
358                    println!("${out} = {op} ${lhs} ${rhs}");
359                }
360
361                SsaOp::AddRegImm(out, arg, imm)
362                | SsaOp::MulRegImm(out, arg, imm)
363                | SsaOp::DivRegImm(out, arg, imm)
364                | SsaOp::DivImmReg(out, arg, imm)
365                | SsaOp::SubImmReg(out, arg, imm)
366                | SsaOp::SubRegImm(out, arg, imm)
367                | SsaOp::AtanRegImm(out, arg, imm)
368                | SsaOp::AtanImmReg(out, arg, imm)
369                | SsaOp::MinRegImm(out, arg, imm)
370                | SsaOp::MaxRegImm(out, arg, imm)
371                | SsaOp::ModRegImm(out, arg, imm)
372                | SsaOp::ModImmReg(out, arg, imm)
373                | SsaOp::AndRegImm(out, arg, imm)
374                | SsaOp::OrRegImm(out, arg, imm)
375                | SsaOp::CompareRegImm(out, arg, imm)
376                | SsaOp::CompareImmReg(out, arg, imm) => {
377                    let (op, swap) = match op {
378                        SsaOp::AddRegImm(..) => ("ADD", false),
379                        SsaOp::MulRegImm(..) => ("MUL", false),
380                        SsaOp::DivImmReg(..) => ("DIV", true),
381                        SsaOp::DivRegImm(..) => ("DIV", false),
382                        SsaOp::SubImmReg(..) => ("SUB", true),
383                        SsaOp::SubRegImm(..) => ("SUB", false),
384                        SsaOp::AtanImmReg(..) => ("ATAN", true),
385                        SsaOp::AtanRegImm(..) => ("ATAN", false),
386                        SsaOp::MinRegImm(..) => ("MIN", false),
387                        SsaOp::MaxRegImm(..) => ("MAX", false),
388                        SsaOp::ModRegImm(..) => ("MOD", false),
389                        SsaOp::ModImmReg(..) => ("MOD", true),
390                        SsaOp::AndRegImm(..) => ("AND", false),
391                        SsaOp::OrRegImm(..) => ("OR", false),
392                        SsaOp::CompareRegImm(..) => ("COMPARE", false),
393                        SsaOp::CompareImmReg(..) => ("COMPARE", true),
394                        _ => unreachable!(),
395                    };
396                    if swap {
397                        println!("${out} = {op} {imm} ${arg}");
398                    } else {
399                        println!("${out} = {op} ${arg} {imm}");
400                    }
401                }
402                SsaOp::CopyImm(out, imm) => {
403                    println!("${out} = COPY {imm}");
404                }
405            }
406        }
407    }
408}
409
410#[cfg(test)]
411mod test {
412    use super::*;
413
414    #[test]
415    fn test_ring() {
416        let mut ctx = Context::new();
417        let c0 = ctx.constant(0.5);
418        let x = ctx.x();
419        let y = ctx.y();
420        let x2 = ctx.square(x).unwrap();
421        let y2 = ctx.square(y).unwrap();
422        let r = ctx.add(x2, y2).unwrap();
423        let c6 = ctx.sub(r, c0).unwrap();
424        let c7 = ctx.constant(0.25);
425        let c8 = ctx.sub(c7, r).unwrap();
426        let c9 = ctx.max(c8, c6).unwrap();
427
428        let (tape, vs) = SsaTape::new(&ctx, &[c9]).unwrap();
429        assert_eq!(tape.len(), 9);
430        assert_eq!(vs.len(), 2);
431    }
432
433    #[test]
434    fn test_dupe() {
435        let mut ctx = Context::new();
436        let x = ctx.x();
437        let x_squared = ctx.mul(x, x).unwrap();
438
439        let (tape, vs) = SsaTape::new(&ctx, &[x_squared]).unwrap();
440        assert_eq!(tape.len(), 3); // x, square, output
441        assert_eq!(vs.len(), 1);
442    }
443
444    #[test]
445    fn test_constant() {
446        let mut ctx = Context::new();
447        let p = ctx.constant(1.5);
448        let (tape, vs) = SsaTape::new(&ctx, &[p]).unwrap();
449        assert_eq!(tape.len(), 2); // CopyImm, output
450        assert_eq!(vs.len(), 0);
451    }
452}
fidget_core/compiler/ssa_tape.rs

fidget_core/compiler/
ssa_tape.rs
