tang_expr/int/
emit.rs

1//! Source lowerings: WGSL and Rust, from one schedule.
2//!
3//! Single-use nodes are inlined into their user (same scope only, so loop
4//! invariants stay hoisted); everything else is a `let`. Loop states are
5//! `var`s, snapshotted at the top of each iteration so the state update is
6//! parallel.
7
8use std::fmt::Write;
9
10use super::sched::{Item, Sched};
11use super::{BinOp, CmpOp, Elem, IId, INode, IType, IntGraph, UnOp};
12
13/// Helpers the WGSL lowering calls (exact high words of 32×32 products in
14/// 16-bit limbs). Include once per module.
15pub const WGSL_PRELUDE: &str = "\
16fn tx_mulhi_u32(a: u32, b: u32) -> u32 {
17    let al = a & 0xffffu; let ah = a >> 16u;
18    let bl = b & 0xffffu; let bh = b >> 16u;
19    let lh = al * bh; let hl = ah * bl;
20    let mid = ((al * bl) >> 16u) + (lh & 0xffffu) + (hl & 0xffffu);
21    return ah * bh + (lh >> 16u) + (hl >> 16u) + (mid >> 16u);
22}
23fn tx_mulhi_i32(a: i32, b: i32) -> i32 {
24    let ua = bitcast<u32>(a); let ub = bitcast<u32>(b);
25    return bitcast<i32>(tx_mulhi_u32(ua, ub) - select(0u, ub, a < 0i) - select(0u, ua, b < 0i));
26}
27";
28
29#[derive(Clone, Copy, PartialEq)]
30enum Lang {
31    Wgsl,
32    Rust,
33}
34
35struct Em<'a> {
36    g: &'a IntGraph,
37    lang: Lang,
38    sched: Sched,
39    params: Vec<String>,
40    out: String,
41    /// Bounds-checked loads (out of bounds reads 0, like the interpreter).
42    guard: bool,
43}
44
45fn ty_name(lang: Lang, t: IType) -> &'static str {
46    match (lang, t) {
47        (_, IType::I32) => "i32",
48        (_, IType::U32) => "u32",
49        (_, IType::Bool) => "bool",
50    }
51}
52
53impl Em<'_> {
54    fn lit(&self, t: IType, v: u32) -> String {
55        match (self.lang, t) {
56            (_, IType::Bool) => (v != 0).to_string(),
57            (Lang::Wgsl, IType::I32) if v as i32 == i32::MIN => "i32(-2147483648)".into(),
58            (Lang::Rust, IType::I32) if v as i32 == i32::MIN => "i32::MIN".into(),
59            (Lang::Wgsl, IType::I32) if (v as i32) < 0 => format!("({}i)", v as i32),
60            (Lang::Wgsl, IType::I32) => format!("{}i", v as i32),
61            (Lang::Rust, IType::I32) if (v as i32) < 0 => format!("({}i32)", v as i32),
62            (Lang::Rust, IType::I32) => format!("{}i32", v as i32),
63            (Lang::Wgsl, IType::U32) => format!("{v}u"),
64            (Lang::Rust, IType::U32) => format!("{v}u32"),
65        }
66    }
67
68    /// Reference to a node's value (atom, name, or inlined expression).
69    fn r(&self, id: IId) -> String {
70        match self.g.node(id) {
71            INode::Const(t, v) => self.lit(t, v),
72            INode::Param(p, _) => format!("({})", self.params[p as usize]),
73            INode::State(l, k) => format!("st{l}_{k}"),
74            INode::LoopOut(l, k) => format!("s{l}_{k}"),
75            INode::Tok => self.lit(IType::U32, 0),
76            _ if self.sched.inline[id.index()] => self.def(id),
77            _ => format!("t{}", id.0),
78        }
79    }
80
81    fn def(&self, id: IId) -> String {
82        let g = self.g;
83        let w = self.lang == Lang::Wgsl;
84        match g.node(id) {
85            INode::Bin(op, a, b) => {
86                let t = g.ty(a);
87                let (x, y) = (self.r(a), self.r(b));
88                let tn = ty_name(self.lang, t);
89                if w {
90                    let amt = |y: String| {
91                        if g.ty(b) == IType::U32 {
92                            y
93                        } else {
94                            format!("bitcast<u32>({y})")
95                        }
96                    };
97                    match op {
98                        BinOp::Add => format!("({x} + {y})"),
99                        BinOp::Sub => format!("({x} - {y})"),
100                        BinOp::Mul => format!("({x} * {y})"),
101                        BinOp::MulHi => format!("tx_mulhi_{tn}({x}, {y})"),
102                        BinOp::And => format!("({x} & {y})"),
103                        BinOp::Or => format!("({x} | {y})"),
104                        BinOp::Xor if t == IType::Bool => format!("({x} != {y})"),
105                        BinOp::Xor => format!("({x} ^ {y})"),
106                        BinOp::Shl => format!("({x} << {})", amt(y)),
107                        BinOp::Shr => format!("({x} >> {})", amt(y)),
108                        BinOp::Min => format!("min({x}, {y})"),
109                        BinOp::Max => format!("max({x}, {y})"),
110                        BinOp::Div => format!("({x} / {y})"),
111                        BinOp::Rem => format!("({x} % {y})"),
112                    }
113                } else {
114                    match op {
115                        BinOp::Add => format!("{x}.wrapping_add({y})"),
116                        BinOp::Sub => format!("{x}.wrapping_sub({y})"),
117                        BinOp::Mul => format!("{x}.wrapping_mul({y})"),
118                        BinOp::MulHi if t == IType::I32 => {
119                            format!("(((({x} as i64) * ({y} as i64)) >> 32) as i32)")
120                        }
121                        BinOp::MulHi => format!("(((({x} as u64) * ({y} as u64)) >> 32) as u32)"),
122                        BinOp::And => format!("({x} & {y})"),
123                        BinOp::Or => format!("({x} | {y})"),
124                        BinOp::Xor => format!("({x} ^ {y})"),
125                        BinOp::Shl => format!("{x}.wrapping_shl({y} as u32)"),
126                        BinOp::Shr => format!("{x}.wrapping_shr({y} as u32)"),
127                        BinOp::Min => format!("{x}.min({y})"),
128                        BinOp::Max => format!("{x}.max({y})"),
129                        BinOp::Div => {
130                            format!("(if {y} == 0 {{ {x} }} else {{ {x}.wrapping_div({y}) }})")
131                        }
132                        BinOp::Rem => {
133                            format!("(if {y} == 0 {{ 0 }} else {{ {x}.wrapping_rem({y}) }})")
134                        }
135                    }
136                }
137            }
138            INode::Un(op, a) => {
139                let x = self.r(a);
140                match (op, g.ty(a), self.lang) {
141                    (UnOp::Neg, IType::U32, Lang::Wgsl) => format!("(0u - {x})"),
142                    (UnOp::Neg, _, Lang::Wgsl) => format!("(-{x})"),
143                    (UnOp::Neg, _, Lang::Rust) => format!("{x}.wrapping_neg()"),
144                    (UnOp::Not, IType::Bool, _) | (UnOp::Not, _, Lang::Rust) => format!("(!{x})"),
145                    (UnOp::Not, _, Lang::Wgsl) => format!("(~{x})"),
146                    (UnOp::Clz, _, Lang::Wgsl) => format!("countLeadingZeros({x})"),
147                    (UnOp::Clz, t, Lang::Rust) => {
148                        format!("({x}.leading_zeros() as {})", ty_name(Lang::Rust, t))
149                    }
150                }
151            }
152            INode::Cmp(op, a, b) => {
153                let s = match op {
154                    CmpOp::Eq => "==",
155                    CmpOp::Ne => "!=",
156                    CmpOp::Lt => "<",
157                    CmpOp::Le => "<=",
158                };
159                format!("({} {s} {})", self.r(a), self.r(b))
160            }
161            INode::Select(c, a, b) => {
162                if w {
163                    format!("select({}, {}, {})", self.r(b), self.r(a), self.r(c))
164                } else {
165                    format!(
166                        "(if {} {{ {} }} else {{ {} }})",
167                        self.r(c),
168                        self.r(a),
169                        self.r(b)
170                    )
171                }
172            }
173            INode::Bitcast(t, a) => {
174                if w {
175                    format!("bitcast<{}>({})", ty_name(self.lang, t), self.r(a))
176                } else {
177                    format!("({} as {})", self.r(a), ty_name(self.lang, t))
178                }
179            }
180            INode::Load(b, i) => {
181                let buf = &g.bufs[b as usize];
182                let (n, i) = (&buf.name, self.r(i));
183                if w && self.guard {
184                    let (word, z) = match buf.elem {
185                        Elem::I32 => (i.to_string(), "0i"),
186                        Elem::U32 => (i.to_string(), "0u"),
187                        Elem::U8 => (format!("({i} >> 2u)"), "0u"),
188                        Elem::I8 => (format!("({i} >> 2u)"), "0i"),
189                    };
190                    let v = match buf.elem {
191                        Elem::I32 | Elem::U32 => format!("{n}[{i}]"),
192                        Elem::U8 => format!("extractBits({n}[{i} >> 2u], ({i} & 3u) * 8u, 8u)"),
193                        Elem::I8 => format!(
194                            "extractBits(bitcast<i32>({n}[{i} >> 2u]), ({i} & 3u) * 8u, 8u)"
195                        ),
196                    };
197                    return format!("select({z}, {v}, {word} < arrayLength(&{n}))");
198                }
199                if !w {
200                    let get = |j: String| format!("{n}.get({j} as usize).copied().unwrap_or(0)");
201                    return match buf.elem {
202                        Elem::U32 => get(i.clone()),
203                        Elem::I32 => format!("({} as i32)", get(i.clone())),
204                        Elem::U8 => format!(
205                            "(({} >> (({i} & 3) * 8)) & 255)",
206                            get(format!("({i} >> 2)"))
207                        ),
208                        Elem::I8 => format!(
209                            "(({} >> (({i} & 3) * 8)) as u8 as i8 as i32)",
210                            get(format!("({i} >> 2)"))
211                        ),
212                    };
213                }
214                match (buf.elem, self.lang) {
215                    (Elem::I32 | Elem::U32, Lang::Wgsl) => format!("{n}[{i}]"),
216                    (Elem::U8, Lang::Wgsl) => {
217                        format!("extractBits({n}[{i} >> 2u], ({i} & 3u) * 8u, 8u)")
218                    }
219                    (Elem::I8, Lang::Wgsl) => {
220                        format!("extractBits(bitcast<i32>({n}[{i} >> 2u]), ({i} & 3u) * 8u, 8u)")
221                    }
222                    (Elem::U32, Lang::Rust) => format!("{n}[{i} as usize]"),
223                    (Elem::I32, Lang::Rust) => format!("({n}[{i} as usize] as i32)"),
224                    (Elem::U8, Lang::Rust) => {
225                        format!("(({n}[({i} >> 2) as usize] >> (({i} & 3) * 8)) & 255)")
226                    }
227                    (Elem::I8, Lang::Rust) => format!(
228                        "(({n}[({i} >> 2) as usize] >> (({i} & 3) * 8)) as u8 as i8 as i32)"
229                    ),
230                }
231            }
232            INode::LoadRw(b, i, _) => {
233                let buf = &g.bufs[b as usize];
234                let (len, base) = buf.rw.as_ref().unwrap();
235                let (n, i) = (&buf.name, self.r(i));
236                if w {
237                    let z = if buf.elem == Elem::I32 { "0i" } else { "0u" };
238                    if self.guard {
239                        return format!(
240                            "select({z}, {n}[({base}) + min({i}, {}u)], {i} < {len}u)",
241                            len.saturating_sub(1)
242                        );
243                    }
244                    return format!("{n}[({base}) + {i}]");
245                }
246                let cast = if buf.elem == Elem::I32 { " as i32" } else { "" };
247                format!("((if {i} < {len}u32 {{ {n}.get({i} as usize).copied().unwrap_or(0) }} else {{ 0 }}){cast})")
248            }
249            INode::Store(..) => self.lit(IType::U32, 0),
250            INode::Const(..)
251            | INode::Param(..)
252            | INode::State(..)
253            | INode::LoopOut(..)
254            | INode::Tok => self.r(id),
255        }
256    }
257
258    /// The statement of a scratch store.
259    fn store(&self, b: u16, i: IId, v: IId) -> String {
260        let buf = &self.g.bufs[b as usize];
261        let (len, base) = buf.rw.as_ref().unwrap();
262        let (n, i, v) = (&buf.name, self.r(i), self.r(v));
263        match self.lang {
264            Lang::Wgsl if self.guard => format!("if ({i} < {len}u) {{ {n}[({base}) + {i}] = {v}; }}"),
265            Lang::Wgsl => format!("{n}[({base}) + {i}] = {v};"),
266            Lang::Rust => format!("if {i} < {len}u32 {{ if let Some(w) = {n}.get_mut({i} as usize) {{ *w = {v} as u32; }} }}"),
267        }
268    }
269
270    fn line(&mut self, depth: usize, s: &str) {
271        for _ in 0..depth {
272            self.out.push_str("    ");
273        }
274        self.out.push_str(s);
275        self.out.push('\n');
276    }
277
278    fn items(&mut self, items: &[Item], depth: usize) {
279        for it in items {
280            match it {
281                Item::Node(id) => {
282                    let id = *id;
283                    let skip = matches!(
284                        self.g.node(id),
285                        INode::Const(..)
286                            | INode::Param(..)
287                            | INode::State(..)
288                            | INode::LoopOut(..)
289                            | INode::Tok
290                    );
291                    if skip || self.sched.inline[id.index()] {
292                        continue;
293                    }
294                    if let INode::Store(b, i, v, _) = self.g.node(id) {
295                        let s = self.store(b, i, v);
296                        self.line(depth, &s);
297                    }
298                    let s = format!(
299                        "let t{}: {} = {};",
300                        id.0,
301                        ty_name(self.lang, self.g.ty(id)),
302                        self.def(id)
303                    );
304                    self.line(depth, &s);
305                }
306                Item::Loop { l, pre, post } => self.lp(*l, pre, post, depth),
307            }
308        }
309    }
310
311    fn lp(&mut self, l: u32, pre: &[Item], post: &[Item], depth: usize) {
312        let g = self.g;
313        let lp = &g.loops[l as usize];
314        for (k, &i) in lp.init.iter().enumerate() {
315            let t = ty_name(self.lang, g.ty(i));
316            let s = match self.lang {
317                Lang::Wgsl => format!("var s{l}_{k}: {t} = {};", self.r(i)),
318                Lang::Rust => format!("let mut s{l}_{k}: {t} = {};", self.r(i)),
319            };
320            self.line(depth, &s);
321        }
322        if lp.cap.is_some() {
323            let s = match self.lang {
324                Lang::Wgsl => format!("var c{l}: u32 = 0u;"),
325                Lang::Rust => format!("let mut c{l}: u32 = 0;"),
326            };
327            self.line(depth, &s);
328        }
329        self.line(depth, "loop {");
330        if let Some(cap) = lp.cap {
331            self.line(
332                depth + 1,
333                &format!("if c{l} >= {} {{ break; }}", self.lit(IType::U32, cap)),
334            );
335            self.line(
336                depth + 1,
337                &format!("c{l} = c{l} + {};", self.lit(IType::U32, 1)),
338            );
339        }
340        for k in 0..lp.init.len() {
341            self.line(depth + 1, &format!("let st{l}_{k} = s{l}_{k};"));
342        }
343        self.items(pre, depth + 1);
344        let c = self.r(lp.cond);
345        self.line(depth + 1, &format!("if !{} {{ break; }}", paren(&c)));
346        self.items(post, depth + 1);
347        for (k, &x) in lp.next.iter().enumerate() {
348            let s = format!("s{l}_{k} = {};", self.r(x));
349            self.line(depth + 1, &s);
350        }
351        self.line(depth, "}");
352    }
353}
354
355fn paren(s: &str) -> String {
356    if s.starts_with('(') && s.ends_with(')') {
357        s.to_string()
358    } else {
359        format!("({s})")
360    }
361}
362
363impl IntGraph {
364    fn emitter(&self, lang: Lang, outputs: &[IId], params: &[&str]) -> Em<'_> {
365        Em {
366            g: self,
367            lang,
368            sched: self.schedule(outputs),
369            params: params.iter().map(|s| s.to_string()).collect(),
370            out: String::new(),
371            guard: false,
372        }
373    }
374
375    pub(crate) fn nparams(&self) -> Vec<IType> {
376        let mut t: Vec<IType> = Vec::new();
377        for n in &self.nodes {
378            if let INode::Param(p, ty) = *n {
379                if t.len() <= p as usize {
380                    t.resize(p as usize + 1, IType::U32);
381                }
382                t[p as usize] = ty;
383            }
384        }
385        t
386    }
387
388    /// WGSL statements computing `outputs`: param `i` reads the expression
389    /// `params[i]`, output `k` is assigned to `dests[k]`. Needs
390    /// [`WGSL_PRELUDE`] in the module when the graph uses `MulHi`; buffers are
391    /// referenced by their declared names.
392    pub fn to_wgsl_body(
393        &self,
394        outputs: &[IId],
395        params: &[&str],
396        dests: &[&str],
397        depth: usize,
398    ) -> String {
399        let mut e = self.emitter(Lang::Wgsl, outputs, params);
400        let items = std::mem::take(&mut e.sched.items);
401        e.items(&items, depth);
402        // all outputs are computed before any destination is written (a
403        // destination may be a param's source)
404        let rs: Vec<String> = outputs.iter().map(|o| e.r(*o)).collect();
405        for (k, r) in rs.iter().enumerate() {
406            e.line(depth, &format!("let r{k} = {r};"));
407        }
408        for (k, d) in dests.iter().enumerate().take(outputs.len()) {
409            e.line(depth, &format!("{d} = r{k};"));
410        }
411        e.out
412    }
413
414    /// [`Self::to_wgsl_body`] with every load bounds-checked against
415    /// `arrayLength` (out of bounds reads 0, as in [`Program`](super::Program)).
416    /// For untrusted graphs; buffers must be runtime-sized storage arrays.
417    pub fn to_wgsl_body_guarded(
418        &self,
419        outputs: &[IId],
420        params: &[&str],
421        dests: &[&str],
422        depth: usize,
423    ) -> String {
424        let mut e = self.emitter(Lang::Wgsl, outputs, params);
425        e.guard = true;
426        let items = std::mem::take(&mut e.sched.items);
427        e.items(&items, depth);
428        let rs: Vec<String> = outputs.iter().map(|o| e.r(*o)).collect();
429        for (k, r) in rs.iter().enumerate() {
430            e.line(depth, &format!("let r{k} = {r};"));
431        }
432        for (k, d) in dests.iter().enumerate().take(outputs.len()) {
433            e.line(depth, &format!("{d} = r{k};"));
434        }
435        e.out
436    }
437
438    /// A WGSL function `fn name(p0: T0, …) -> name_out` returning a struct of
439    /// the outputs (`o0`, `o1`, …).
440    pub fn to_wgsl_fn(&self, name: &str, outputs: &[IId]) -> String {
441        let pt = self.nparams();
442        let pn: Vec<String> = (0..pt.len()).map(|i| format!("p{i}")).collect();
443        let pr: Vec<&str> = pn.iter().map(|s| s.as_str()).collect();
444        let mut s = format!("struct {name}_out {{\n");
445        for (k, o) in outputs.iter().enumerate() {
446            let _ = writeln!(s, "    o{k}: {},", ty_name(Lang::Wgsl, self.ty(*o)));
447        }
448        s.push_str("}\n");
449        let args: Vec<String> = pt
450            .iter()
451            .enumerate()
452            .map(|(i, t)| format!("p{i}: {}", ty_name(Lang::Wgsl, *t)))
453            .collect();
454        let _ = writeln!(s, "fn {name}({}) -> {name}_out {{", args.join(", "));
455        let dn: Vec<String> = (0..outputs.len()).map(|k| format!("r.o{k}")).collect();
456        let dr: Vec<&str> = dn.iter().map(|s| s.as_str()).collect();
457        s.push_str(&format!("    var r: {name}_out;\n"));
458        s.push_str(&self.to_wgsl_body(outputs, &pr, &dr, 1));
459        s.push_str("    return r;\n}\n");
460        s
461    }
462
463    /// A Rust function `fn name(p0: T0, …, buf: &[u32], …) -> [T; N]` (or a
464    /// tuple when the output types differ); buffers
465    /// take `u32` words (8-bit elements packed little endian).
466    pub fn to_rust(&self, name: &str, outputs: &[IId]) -> String {
467        let pt = self.nparams();
468        let pn: Vec<String> = (0..pt.len()).map(|i| format!("p{i}")).collect();
469        let pr: Vec<&str> = pn.iter().map(|s| s.as_str()).collect();
470        let mut args: Vec<String> = pt
471            .iter()
472            .enumerate()
473            .map(|(i, t)| format!("p{i}: {}", ty_name(Lang::Rust, *t)))
474            .collect();
475        args.extend(self.bufs.iter().map(|b| {
476            format!(
477                "{}: &{}[u32]",
478                b.name,
479                if b.rw.is_some() { "mut " } else { "" }
480            )
481        }));
482        let rt: Vec<&str> = outputs
483            .iter()
484            .map(|o| ty_name(Lang::Rust, self.ty(*o)))
485            .collect();
486        let mut e = self.emitter(Lang::Rust, outputs, &pr);
487        // one type: an array; mixed: a tuple
488        let same = rt.windows(2).all(|w| w[0] == w[1]) && !rt.is_empty();
489        let ret = if same {
490            format!("[{}; {}]", rt[0], rt.len())
491        } else {
492            format!("({},)", rt.join(", "))
493        };
494        e.out = format!(
495            "#[allow(unused, clippy::all)]\npub fn {name}({}) -> {ret} {{\n",
496            args.join(", ")
497        );
498        let items = std::mem::take(&mut e.sched.items);
499        e.items(&items, 1);
500        let rs: Vec<String> = outputs.iter().map(|o| e.r(*o)).collect();
501        let s = if same {
502            format!("[{}]", rs.join(", "))
503        } else {
504            format!("({},)", rs.join(", "))
505        };
506        e.line(1, &s);
507        e.out.push_str("}\n");
508        e.out
509    }
510}
tang_expr/int/emit.rs

tang_expr/int/
emit.rs
