Skip to main content

iridia_exec/
value.rs

1//! A node of a program's value graph ([`iridia::flow::Values`]) evaluated from the words of the leaves it reads, with the
2//! executor's own operations: what a run of the program computes for that lane where it is given the same leaves, bit
3//! for bit. The caller answers the leaves (an input lane, a supplied value, a global's word, a texel, a query) or leaves
4//! them unknown; a node that needs an unknown leaf is unknown, naming it. A select whose condition is unknown is known
5//! where both its values are and are the same bits.
6//!
7//! Refused with an error, never guessed: an operation the graph states no kind for, a derivative, a subgroup operation,
8//! a loop's value, a merge whose arms differ (the graph does not state which condition picks which arm), an element
9//! picked by a run-time index ([`Op::Index`]: the graph does not state which candidate an index picks), a math function
10//! that mixes lanes other than a dot product, and an operation the executor refuses.
11use std::collections::HashMap;
12
13use iridia::flow::{At, MathFunction, Node, NodeId, Op, ScalarKind, Values};
14
15use crate::compile::{Kind, binary_op, convert_op, kind_of, math_op, relational_op, unary_op};
16use crate::flat::MathOp;
17use crate::run::{bin, convert, math, unary};
18
19/// What a node evaluates to.
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub enum Evaluated {
22    /// The lane's word.
23    Word(u32),
24    /// The leaf whose word is not known (the caller's answer was `None`, or a value the leaf is read at is unknown).
25    Unknown(NodeId),
26}
27
28/// Evaluates nodes of one value graph, keeping what it has evaluated: the leaves are asked once each.
29pub struct Evaluator<'v, L> {
30    values: &'v Values,
31    flush_denormals: bool,
32    leaf: L,
33    known: HashMap<NodeId, Evaluated>,
34}
35
36impl<L> std::fmt::Debug for Evaluator<'_, L> {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        f.debug_struct("Evaluator")
39            .field("flush_denormals", &self.flush_denormals)
40            .field("known", &self.known.len())
41            .finish_non_exhaustive()
42    }
43}
44
45/// One step of the walk: the node's result, or the operands it needs first.
46enum Step {
47    Done(Evaluated),
48    Need(Vec<NodeId>),
49}
50
51impl<'v, L: FnMut(NodeId, &Node, &[u32]) -> Option<u32>> Evaluator<'v, L> {
52    /// An evaluator of `values`, whose program flushes subnormal floats to zero where `flush_denormals` (the lifted
53    /// program's modes), asking `leaf(id, node, at)` for each leaf's word: `at` the words it is read at, a run-time
54    /// indexed global's index values (outermost first) or a texel's coordinate then its controls, empty for the other
55    /// leaves.
56    pub fn new(values: &'v Values, flush_denormals: bool, leaf: L) -> Self {
57        Evaluator {
58            values,
59            flush_denormals,
60            leaf,
61            known: HashMap::new(),
62        }
63    }
64
65    /// The word node `id` evaluates to, or the leaf it needs that is unknown.
66    pub fn evaluate(&mut self, id: NodeId) -> Result<Evaluated, String> {
67        let mut stack = vec![id];
68        while let Some(&top) = stack.last() {
69            if self.known.contains_key(&top) {
70                stack.pop();
71                continue;
72            }
73            match self.step(top)? {
74                Step::Done(v) => {
75                    self.known.insert(top, v);
76                    stack.pop();
77                }
78                // a node may be pushed again while an earlier push of it waits; the later one is evaluated first and
79                // the earlier then found known. The graph's only cycles are loops' values, refused before their
80                // operands are asked for, so the walk ends
81                Step::Need(more) => stack.extend(more),
82            }
83        }
84        self.known
85            .get(&id)
86            .copied()
87            .ok_or_else(|| "the walk ended without the node's value".into())
88    }
89
90    /// The words of `ids`, or the first unknown leaf among them; `None` while some are not evaluated yet.
91    fn words(&self, ids: &[NodeId]) -> Option<Result<Vec<u32>, NodeId>> {
92        let mut out = Vec::with_capacity(ids.len());
93        for id in ids {
94            match self.known.get(id)? {
95                Evaluated::Word(w) => out.push(*w),
96                Evaluated::Unknown(leaf) => return Some(Err(*leaf)),
97            }
98        }
99        Some(Ok(out))
100    }
101
102    /// The operands of `ids` not yet evaluated.
103    fn pending(&self, ids: &[NodeId]) -> Vec<NodeId> {
104        ids.iter()
105            .filter(|id| !self.known.contains_key(id))
106            .copied()
107            .collect()
108    }
109
110    /// The kind operation `id` reads its operands in.
111    fn reads(&self, id: NodeId) -> Result<Kind, String> {
112        kind_of(
113            self.values
114                .reads(id)
115                .ok_or_else(|| format!("an operation (node {id:?}) the graph states no kind for"))?,
116        )
117    }
118
119    /// A leaf read at the words of `at`: unknown where one of them is (the leaf itself is what is not known).
120    fn leaf(&mut self, id: NodeId, node: &Node, at: &[NodeId]) -> Step {
121        match self.words(at) {
122            None => Step::Need(self.pending(at)),
123            Some(Err(_)) => Step::Done(Evaluated::Unknown(id)),
124            Some(Ok(words)) => Step::Done(match (self.leaf)(id, node, &words) {
125                Some(w) => Evaluated::Word(w),
126                None => Evaluated::Unknown(id),
127            }),
128        }
129    }
130
131    fn step(&mut self, id: NodeId) -> Result<Step, String> {
132        let node = self
133            .values
134            .get(id)
135            .ok_or_else(|| format!("node {id:?} is not of this graph"))?
136            .clone();
137        let ftz = self.flush_denormals;
138        Ok(match &node {
139            Node::Input { .. } | Node::Supplied { .. } | Node::Query { .. } => self.leaf(id, &node, &[]),
140            Node::Global { at, .. } => match at {
141                At::Component(_) => self.leaf(id, &node, &[]),
142                At::Indexed { index, .. } => self.leaf(id, &node, index),
143            },
144            Node::Texel {
145                coordinate, controls, ..
146            } => {
147                let at: Vec<NodeId> = coordinate.iter().chain(controls).copied().collect();
148                self.leaf(id, &node, &at)
149            }
150            Node::Literal(l) => Step::Done(Evaluated::Word(literal(l.0, self.values.kind(id))?)),
151            Node::Loop { .. } => return Err("a loop's value is not evaluated".into()),
152            Node::Merge { arms, .. } => match self.words(arms) {
153                None => Step::Need(self.pending(arms)),
154                Some(Err(leaf)) => Step::Done(Evaluated::Unknown(leaf)),
155                Some(Ok(words))
156                    if let Some(&first) = words.first()
157                        && words.iter().all(|w| *w == first) =>
158                {
159                    Step::Done(Evaluated::Word(first))
160                }
161                Some(Ok(_)) => {
162                    return Err(
163                        "a merge whose arms differ: the graph does not state which condition picks which"
164                            .into(),
165                    );
166                }
167            },
168            Node::Op { op, operands } => match op {
169                Op::Select => {
170                    let [condition, accept, reject] = operands[..] else {
171                        return Err(format!("a select of {} operands", operands.len()));
172                    };
173                    match self.known.get(&condition) {
174                        None => Step::Need(vec![condition]),
175                        Some(Evaluated::Word(c)) => {
176                            let chosen = if *c != 0 { accept } else { reject };
177                            match self.known.get(&chosen) {
178                                None => Step::Need(vec![chosen]),
179                                Some(v) => Step::Done(*v),
180                            }
181                        }
182                        // unknown which: known where both values are and are the same bits
183                        Some(Evaluated::Unknown(leaf)) => {
184                            let leaf = *leaf;
185                            match self.words(&[accept, reject]) {
186                                None => Step::Need(self.pending(&[accept, reject])),
187                                Some(Ok(w))
188                                    if let &[a, b] = w.as_slice()
189                                        && a == b =>
190                                {
191                                    Step::Done(Evaluated::Word(a))
192                                }
193                                Some(_) => Step::Done(Evaluated::Unknown(leaf)),
194                            }
195                        }
196                    }
197                }
198                _ => match self.words(operands) {
199                    None => Step::Need(self.pending(operands)),
200                    Some(Err(leaf)) => Step::Done(Evaluated::Unknown(leaf)),
201                    Some(Ok(w)) => Step::Done(Evaluated::Word(self.operation(id, *op, &w, ftz)?)),
202                },
203            },
204        })
205    }
206
207    /// Operation `op` of node `id` on the operand words `w`.
208    fn operation(&self, id: NodeId, op: Op, w: &[u32], ftz: bool) -> Result<u32, String> {
209        let lane = |i: usize| -> Result<u32, String> {
210            w.get(i)
211                .copied()
212                .ok_or_else(|| format!("{op:?} with {} operands", w.len()))
213        };
214        Ok(match op {
215            Op::Unary(u) => unary(unary_op(u, self.reads(id)?), lane(0)?),
216            Op::Binary(b) => {
217                if w.len() != 2 {
218                    return Err("a matrix product is not evaluated".into());
219                }
220                bin(binary_op(b, self.reads(id)?)?, lane(0)?, lane(1)?, ftz)?
221            }
222            Op::Convert(to) => match convert_op(self.reads(id)?, to)? {
223                Some(c) => convert(c, lane(0)?),
224                None => lane(0)?,
225            },
226            Op::Math(fun) => {
227                let m = math_op(fun, self.reads(id)?)?;
228                match m {
229                    MathOp::Dot => {
230                        let n = w.len() / 2;
231                        let dot = || format!("a dot product of {} operands", w.len());
232                        if !w.len().is_multiple_of(2) || !(1..=4).contains(&n) {
233                            return Err(dot());
234                        }
235                        let (a, b) = w.split_at_checked(n).ok_or_else(dot)?;
236                        math(
237                            m,
238                            reg(a),
239                            reg(b),
240                            [0; 4],
241                            u8::try_from(n).map_err(|_| dot())?,
242                            ftz,
243                        )[0]
244                    }
245                    MathOp::Normalize | MathOp::PackHalf | MathOp::UnpackHalf => {
246                        return Err(format!("{fun:?} mixes lanes: not evaluated"));
247                    }
248                    _ if lanes_mixed(fun) => return Err(format!("{fun:?} mixes lanes: not evaluated")),
249                    _ => {
250                        let at = |i: usize| [w.get(i).copied().unwrap_or(0), 0, 0, 0];
251                        math(m, at(0), at(1), at(2), 1, ftz)[0]
252                    }
253                }
254            }
255            Op::Relational(fun) => {
256                let m = relational_op(fun);
257                let lanes = u8::try_from(w.len())
258                    .ok()
259                    .filter(|n| (1..=4).contains(n))
260                    .ok_or_else(|| format!("{fun:?} of {} lanes", w.len()))?;
261                math(m, reg(w), [0; 4], [0; 4], lanes, ftz)[0]
262            }
263            Op::Select => {
264                if lane(0)? != 0 {
265                    lane(1)?
266                } else {
267                    lane(2)?
268                }
269            }
270            Op::Index => {
271                return Err("an element picked by a run-time index: the graph does not state which candidate an index picks".into());
272            }
273            Op::Derivative(..) => return Err("a derivative is not evaluated".into()),
274            Op::Subgroup => return Err("a subgroup operation is not evaluated".into()),
275        })
276    }
277}
278
279/// Up to four words as a register.
280fn reg(w: &[u32]) -> [u32; 4] {
281    let mut r = [0; 4];
282    for (out, v) in r.iter_mut().zip(w) {
283        *out = *v;
284    }
285    r
286}
287
288/// Whether math function `fun` makes each lane of its result from every lane of its arguments (the value graph keeps one
289/// node for all of them): the ones the executor runs that are not a dot product.
290fn lanes_mixed(fun: MathFunction) -> bool {
291    matches!(
292        fun,
293        MathFunction::Cross
294            | MathFunction::Length
295            | MathFunction::Distance
296            | MathFunction::Outer
297            | MathFunction::Determinant
298            | MathFunction::Transpose
299            | MathFunction::Inverse
300            | MathFunction::Reflect
301            | MathFunction::Refract
302            | MathFunction::FaceForward
303    )
304}
305
306/// A literal's word: its value as the scalar kind the graph states it of.
307#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] // each is checked to hold the value exactly
308fn literal(v: f64, kind: Option<ScalarKind>) -> Result<u32, String> {
309    let exact = |w: u32, back: f64| {
310        if back.to_bits() == v.to_bits() || (v.is_nan() && back.is_nan()) {
311            Ok(w)
312        } else {
313            Err(format!("literal {v} is no 32-bit {kind:?}"))
314        }
315    };
316    match kind {
317        Some(ScalarKind::Float) => {
318            let f = v as f32;
319            exact(f.to_bits(), f64::from(f))
320        }
321        Some(ScalarKind::Uint) => exact(v as u32, f64::from(v as u32)),
322        Some(ScalarKind::Sint) => exact((v as i32) as u32, f64::from(v as i32)),
323        Some(ScalarKind::Bool) => Ok(u32::from(v != 0.)),
324        other => Err(format!("a literal of kind {other:?}")),
325    }
326}