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

fidget_core/compiler/
ssa_tape.rs
