nxpu_ir/
expr.rs

1//! Expressions — pure SSA values with no side effects.
2
3use crate::arena::Handle;
4use crate::types::{Bytes, Scalar, ScalarKind, Type, VectorSize};
5
6/// A vector swizzle component.
7#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
8pub enum SwizzleComponent {
9    X = 0,
10    Y = 1,
11    Z = 2,
12    W = 3,
13}
14
15/// A literal constant value.
16#[derive(Clone, Copy, Debug)]
17pub enum Literal {
18    Bool(bool),
19    I32(i32),
20    U32(u32),
21    F32(f32),
22    F64(f64),
23    AbstractInt(i64),
24    AbstractFloat(f64),
25}
26
27impl Literal {
28    /// Returns the scalar type of this literal.
29    pub fn scalar(&self) -> Scalar {
30        match *self {
31            Self::Bool(_) => Scalar::BOOL,
32            Self::I32(_) => Scalar::I32,
33            Self::U32(_) => Scalar::U32,
34            Self::F32(_) => Scalar::F32,
35            Self::F64(_) => Scalar {
36                kind: ScalarKind::Float,
37                width: 8,
38            },
39            Self::AbstractInt(_) => Scalar {
40                kind: ScalarKind::Sint,
41                width: 8,
42            },
43            Self::AbstractFloat(_) => Scalar {
44                kind: ScalarKind::Float,
45                width: 8,
46            },
47        }
48    }
49}
50
51/// A unary operator.
52#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
53pub enum UnaryOp {
54    Negate,
55    LogicalNot,
56    BitwiseNot,
57}
58
59/// A binary operator.
60#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
61pub enum BinaryOp {
62    Add,
63    Subtract,
64    Multiply,
65    Divide,
66    Modulo,
67    Equal,
68    NotEqual,
69    Less,
70    LessEqual,
71    Greater,
72    GreaterEqual,
73    LogicalAnd,
74    LogicalOr,
75    BitwiseAnd,
76    BitwiseOr,
77    BitwiseXor,
78    ShiftLeft,
79    ShiftRight,
80}
81
82/// A built-in math function.
83#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
84pub enum MathFunction {
85    // Component-wise
86    Abs,
87    Min,
88    Max,
89    Clamp,
90    Saturate,
91    // Rounding
92    Floor,
93    Ceil,
94    Round,
95    Fract,
96    Trunc,
97    // Trigonometric
98    Sin,
99    Cos,
100    Tan,
101    Asin,
102    Acos,
103    Atan,
104    Atan2,
105    Sinh,
106    Cosh,
107    Tanh,
108    // Exponential
109    Sqrt,
110    InverseSqrt,
111    Log,
112    Log2,
113    Exp,
114    Exp2,
115    Pow,
116    // Linear algebra
117    Dot,
118    Cross,
119    Normalize,
120    Length,
121    Distance,
122    // Interpolation
123    Mix,
124    Step,
125    SmoothStep,
126    // Fused multiply-add
127    Fma,
128    // Bit manipulation
129    //
130    // Quantized weights arrive packed several codes to a u32, so unpacking one
131    // is the first thing every int4/int8 matmul kernel does.
132    ExtractBits,
133    InsertBits,
134}
135
136/// An atomic operation.
137#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
138pub enum AtomicFunction {
139    Add,
140    Subtract,
141    And,
142    ExclusiveOr,
143    InclusiveOr,
144    Min,
145    Max,
146    Exchange { compare: Option<Handle<Expression>> },
147}
148
149/// An expression in the IR — a pure SSA value with no side effects.
150///
151/// Expressions are stored in per-function or module-level arenas.
152/// They are referenced by [`Handle<Expression>`].
153#[derive(Clone, Debug)]
154pub enum Expression {
155    /// A literal constant.
156    Literal(Literal),
157    /// Construct a composite type from components.
158    Compose {
159        ty: Handle<Type>,
160        components: Vec<Handle<Expression>>,
161    },
162    /// Reference to a function argument by index.
163    FunctionArgument(u32),
164    /// Reference to a global variable (produces a pointer).
165    GlobalVariable(Handle<crate::GlobalVariable>),
166    /// Reference to a local variable (produces a pointer).
167    LocalVariable(Handle<crate::LocalVariable>),
168    /// Load a value through a pointer.
169    Load { pointer: Handle<Expression> },
170    /// Dynamic index into a composite (array, vector, matrix).
171    Access {
172        base: Handle<Expression>,
173        index: Handle<Expression>,
174    },
175    /// Static index into a composite.
176    AccessIndex {
177        base: Handle<Expression>,
178        index: u32,
179    },
180    /// Swizzle vector components.
181    Swizzle {
182        size: VectorSize,
183        vector: Handle<Expression>,
184        pattern: [SwizzleComponent; 4],
185    },
186    /// Broadcast a scalar to a vector.
187    Splat {
188        size: VectorSize,
189        value: Handle<Expression>,
190    },
191    /// Apply a unary operator.
192    Unary {
193        op: UnaryOp,
194        expr: Handle<Expression>,
195    },
196    /// Apply a binary operator.
197    Binary {
198        op: BinaryOp,
199        left: Handle<Expression>,
200        right: Handle<Expression>,
201    },
202    /// Select between two values based on a condition.
203    Select {
204        condition: Handle<Expression>,
205        accept: Handle<Expression>,
206        reject: Handle<Expression>,
207    },
208    /// Call a built-in math function.
209    Math {
210        fun: MathFunction,
211        arg: Handle<Expression>,
212        arg1: Option<Handle<Expression>>,
213        arg2: Option<Handle<Expression>>,
214        arg3: Option<Handle<Expression>>,
215    },
216    /// Type cast / bitcast.
217    As {
218        expr: Handle<Expression>,
219        kind: ScalarKind,
220        convert: Option<Bytes>,
221    },
222    /// Get the length of a runtime-sized array.
223    ArrayLength(Handle<Expression>),
224    /// The result of a function call (paired with a `Call` statement).
225    CallResult(Handle<crate::Function>),
226    /// The result of an atomic operation (paired with an `Atomic` statement).
227    AtomicResult { ty: Handle<Type>, comparison: bool },
228    /// A zero-initialized value of the given type.
229    ///
230    /// Used for vector, matrix, struct, and array zero-initialization where a
231    /// simple scalar literal would be type-incorrect.
232    ZeroValue(Handle<Type>),
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use crate::arena::Arena;
239
240    #[test]
241    fn literal_scalars() {
242        assert_eq!(Literal::F32(1.0).scalar(), Scalar::F32);
243        assert_eq!(Literal::I32(-1).scalar(), Scalar::I32);
244        assert_eq!(Literal::U32(42).scalar(), Scalar::U32);
245        assert_eq!(Literal::Bool(true).scalar(), Scalar::BOOL);
246    }
247
248    #[test]
249    fn expression_arena() {
250        let mut exprs = Arena::new();
251        let lit = exprs.append(Expression::Literal(Literal::F32(3.125)));
252        let neg = exprs.append(Expression::Unary {
253            op: UnaryOp::Negate,
254            expr: lit,
255        });
256        assert_eq!(lit.index(), 0);
257        assert_eq!(neg.index(), 1);
258        assert_eq!(exprs.len(), 2);
259    }
260
261    #[test]
262    fn binary_expression() {
263        let mut exprs = Arena::new();
264        let left = exprs.append(Expression::Literal(Literal::F32(1.0)));
265        let right = exprs.append(Expression::Literal(Literal::F32(2.0)));
266        let add = exprs.append(Expression::Binary {
267            op: BinaryOp::Add,
268            left,
269            right,
270        });
271        if let Expression::Binary { op, .. } = &exprs[add] {
272            assert_eq!(*op, BinaryOp::Add);
273        } else {
274            panic!("expected Binary");
275        }
276    }
277}
nxpu_ir/expr.rs

nxpu_ir/
expr.rs
