vyre_foundation/ir_inner/model/
expr.rs

1// Expression nodes  -  produce values.
2//
3// Every expression evaluates to a typed value. Expressions are pure:
4// they read state but do not modify it.
5
6use crate::ir_inner::model::types::{DataType, SubgroupReduceOp};
7use rustc_hash::FxHasher;
8use std::borrow::Borrow;
9use std::fmt;
10use std::hash::{Hash, Hasher};
11use std::ops::Deref;
12use std::sync::Arc;
13
14/// Reference to the generator/macro that produced an AST region.
15/// Used for source-mapping and DWARF-like debugging context.
16#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub struct GeneratorRef {
18    /// The name of the generator (e.g., `vyre-nn::flash_attention`).
19    pub name: String,
20}
21
22/// Interned identifier used by expression nodes.
23///
24/// `Ident` is cheap to clone and keeps expression trees from repeatedly
25/// allocating owned `String` values for the same variable or buffer names.
26#[derive(Clone, Eq, PartialEq)]
27pub struct Ident {
28    text: Arc<str>,
29    hash: u64,
30}
31
32impl Ident {
33    #[inline]
34    fn prehash(text: &str) -> u64 {
35        let mut hasher = FxHasher::default();
36        text.hash(&mut hasher);
37        hasher.finish()
38    }
39
40    #[must_use]
41    #[inline]
42    /// Construct an identifier from shared text while caching its hash once.
43    pub fn new(text: Arc<str>) -> Self {
44        let hash = Self::prehash(&text);
45        Self { text, hash }
46    }
47
48    /// Clone the underlying interned string handle without copying UTF-8 bytes.
49    #[must_use]
50    #[inline]
51    pub fn shared_text(&self) -> Arc<str> {
52        Arc::clone(&self.text)
53    }
54
55    /// Return another identifier handle to the same interned text without
56    /// reallocating text or recomputing the cached hash.
57    #[must_use]
58    #[inline]
59    pub fn duplicate_handle(&self) -> Self {
60        Self {
61            text: Arc::clone(&self.text),
62            hash: self.hash,
63        }
64    }
65
66    /// Return the identifier text.
67    #[must_use]
68    #[inline]
69    pub fn as_str(&self) -> &str {
70        &self.text
71    }
72
73    /// Return the cached hash used by hash-map/set lookups.
74    #[must_use]
75    #[inline]
76    pub fn cached_hash(&self) -> u64 {
77        self.hash
78    }
79}
80
81impl From<&str> for Ident {
82    #[inline]
83    fn from(value: &str) -> Self {
84        Self::new(Arc::from(value))
85    }
86}
87
88impl From<String> for Ident {
89    #[inline]
90    fn from(value: String) -> Self {
91        Self::new(Arc::from(value))
92    }
93}
94
95impl From<Arc<str>> for Ident {
96    #[inline]
97    fn from(value: Arc<str>) -> Self {
98        Self::new(value)
99    }
100}
101
102impl From<&String> for Ident {
103    #[inline]
104    fn from(value: &String) -> Self {
105        Self::from(value.as_str())
106    }
107}
108
109impl From<&Ident> for Ident {
110    #[inline]
111    fn from(value: &Ident) -> Self {
112        value.clone()
113    }
114}
115
116impl fmt::Debug for Ident {
117    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118        f.debug_tuple("Ident").field(&self.as_str()).finish()
119    }
120}
121
122impl Hash for Ident {
123    /// Audit P-IDENT-BORROW (2026-04-29): hash via the underlying str so the
124    /// `Hash` impl matches the `Borrow<str>` impl, preserving the
125    /// `HashMap::get<Q: Borrow<K> + Hash + Eq>` invariant. The
126    /// pre-fix `state.write_u64(self.hash)` produced a different u64 than
127    /// `<str as Hash>::hash` for the same hasher (which writes bytes + a
128    /// length terminator), so any `FxHashMap<Ident, V>::get(&str)` lookup
129    /// silently missed the inserted entry. Callers that want the cached
130    /// `FxHash` for a fast equality-check key call [`Ident::cached_hash`]
131    /// directly.
132    #[inline]
133    fn hash<H: Hasher>(&self, state: &mut H) {
134        self.text.hash(state);
135    }
136}
137
138impl Deref for Ident {
139    type Target = str;
140
141    #[inline]
142    fn deref(&self) -> &Self::Target {
143        self.as_str()
144    }
145}
146
147impl AsRef<str> for Ident {
148    #[inline]
149    fn as_ref(&self) -> &str {
150        self.as_str()
151    }
152}
153
154impl Borrow<str> for Ident {
155    #[inline]
156    fn borrow(&self) -> &str {
157        self.as_str()
158    }
159}
160
161impl fmt::Display for Ident {
162    #[inline]
163    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
164        f.write_str(self.as_str())
165    }
166}
167
168impl PartialEq<str> for Ident {
169    #[inline]
170    fn eq(&self, other: &str) -> bool {
171        self.as_str() == other
172    }
173}
174
175impl PartialEq<&str> for Ident {
176    #[inline]
177    fn eq(&self, other: &&str) -> bool {
178        self.as_str() == *other
179    }
180}
181
182impl PartialOrd for Ident {
183    #[inline]
184    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
185        Some(self.cmp(other))
186    }
187}
188
189impl Ord for Ident {
190    #[inline]
191    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
192        self.as_str().cmp(other.as_str())
193    }
194}
195
196/// An expression that produces a value.
197///
198/// # Examples
199///
200/// ```
201/// use vyre::ir::Expr;
202///
203/// let lit = Expr::u32(42);
204/// let var = Expr::var("x");
205/// let add = Expr::add(lit, var);
206/// ```
207pub use crate::ir_inner::model::generated::Expr;
208
209/// Public contract for downstream expression extension nodes.
210///
211/// Extension nodes are intentionally opaque to core. A downstream crate owns
212/// the semantic payload and provides the stable metadata core needs for
213/// validation, debug output, equality, and CSE identity. Backends that
214/// understand the extension can downcast through their own wrapper type before
215/// constructing target code; backends that do not understand it must reject it
216/// with an actionable error.
217pub trait ExprNode: fmt::Debug + Send + Sync + 'static {
218    /// Stable extension namespace, for example `my_backend.tensor.shuffle`.
219    fn extension_kind(&self) -> &'static str;
220
221    /// Human-readable identity used in diagnostics and debug logs.
222    fn debug_identity(&self) -> &str;
223
224    /// Static result type produced by this expression.
225    fn result_type(&self) -> Option<DataType>;
226
227    /// Whether CSE may treat this extension as a pure, repeatable expression.
228    fn cse_safe(&self) -> bool;
229
230    /// Stable, content-addressed identity for equality and optimizer keys.
231    fn stable_fingerprint(&self) -> [u8; 32];
232
233    /// Validate extension-local invariants.
234    ///
235    /// # Errors
236    ///
237    /// The returned error must explain the bad invariant and include `Fix:`.
238    fn validate_extension(&self) -> Result<(), String>;
239
240    /// Downcast to Any to allow backend-specific dispatch from opaque payloads.
241    fn as_any(&self) -> &dyn std::any::Any;
242
243    /// Serialize the extension payload into stable bytes used by the wire
244    /// encoder's `Expr::Opaque` path (tag `0x80`). Default: empty payload  -
245    /// suitable for extensions that carry no state beyond their type
246    /// identity. Extensions with state must override this to emit the exact
247    /// bytes `wire_payload`'s matching `OpaqueExprResolver` will consume.
248    ///
249    /// The payload contract is endian-fixed: any numeric field wider than
250    /// one byte MUST be written with `to_le_bytes`, and the matching decoder
251    /// MUST reconstruct it with `from_le_bytes`. Host-endian encodings such as
252    /// `to_ne_bytes` are forbidden because the wire format must stay
253    /// byte-identical across architectures.
254    ///
255    /// Extension authors are recommended (but not required, for API
256    /// compatibility) to use [`crate::opaque_payload::LeBytesWriter`] when
257    /// building payloads  -  it makes the right endianness the only choice at
258    /// the type level.
259    ///
260    /// Literal extensions that encode regex payloads must also canonicalize
261    /// inline flag prefixes before emitting bytes. For example, `(?mi)` and
262    /// `(?im)` are the same semantic payload and MUST serialize to the same
263    /// flag ordering.
264    fn wire_payload(&self) -> Vec<u8> {
265        Vec::new()
266    }
267}
268
269impl Expr {
270    /// Load from buffer at index.
271    ///
272    /// # Examples
273    ///
274    /// ```
275    /// use vyre::ir::Expr;
276    /// let _ = Expr::load("a", Expr::u32(0));
277    /// ```
278    #[must_use]
279    #[inline]
280    pub fn load(buffer: impl Into<Ident>, index: Self) -> Self {
281        Self::Load {
282            buffer: buffer.into(),
283            index: Box::new(index),
284        }
285    }
286
287    /// Buffer element count.
288    ///
289    /// # Examples
290    ///
291    /// ```
292    /// use vyre::ir::Expr;
293    /// let _ = Expr::buf_len("a");
294    /// ```
295    #[must_use]
296    #[inline]
297    pub fn buf_len(buffer: impl Into<Ident>) -> Self {
298        Self::BufLen {
299            buffer: buffer.into(),
300        }
301    }
302
303    /// Name a whole buffer as an argument to a composite op.
304    ///
305    /// This is not a value: it has no type and only call-argument position
306    /// accepts it. Inlining rebinds the callee's matching parameter onto
307    /// this buffer, so a callee that reads `table[i]` ends up reading the
308    /// caller's buffer at the same index. To read one element instead, use
309    /// [`Expr::load`].
310    ///
311    /// ```
312    /// use vyre::ir::Expr;
313    /// let _ = Expr::call("dialect::lookup", vec![Expr::buffer_ref("table"), Expr::u32(3)]);
314    /// ```
315    #[must_use]
316    #[inline]
317    pub fn buffer_ref(buffer: impl Into<Ident>) -> Self {
318        Self::BufferRef {
319            buffer: buffer.into(),
320        }
321    }
322
323    /// `global_invocation_id.x`
324    #[must_use]
325    #[inline]
326    pub fn gid_x() -> Self {
327        Self::InvocationId { axis: 0 }
328    }
329
330    /// `global_invocation_id.y`
331    #[must_use]
332    #[inline]
333    pub fn gid_y() -> Self {
334        Self::InvocationId { axis: 1 }
335    }
336
337    /// `global_invocation_id.z`
338    #[must_use]
339    #[inline]
340    pub fn gid_z() -> Self {
341        Self::InvocationId { axis: 2 }
342    }
343
344    /// `workgroup_id.x`
345    #[must_use]
346    #[inline]
347    pub fn workgroup_x() -> Self {
348        Self::WorkgroupId { axis: 0 }
349    }
350
351    /// `workgroup_id.y`
352    #[must_use]
353    #[inline]
354    pub fn workgroup_y() -> Self {
355        Self::WorkgroupId { axis: 1 }
356    }
357
358    /// `workgroup_id.z`
359    #[must_use]
360    #[inline]
361    pub fn workgroup_z() -> Self {
362        Self::WorkgroupId { axis: 2 }
363    }
364
365    /// Predicate `workgroup_id.x == 0`: the canonical "first parallel region only" guard.
366    ///
367    /// Single-workgroup kernels (scalar reductions, workgroup-local tree reductions, and any
368    /// kernel meant to run in exactly one parallel region) gate their body on this so a dispatch
369    /// of more than one workgroup leaves the extra regions as no-ops instead of double-counting or
370    /// racing on the shared output. Prefer this over re-spelling `eq(WorkgroupId{axis:0}, 0)`
371    /// inline so the "first workgroup" contract has one owner.
372    #[must_use]
373    #[inline]
374    pub fn is_first_workgroup() -> Self {
375        Self::eq(Self::WorkgroupId { axis: 0 }, Self::u32(0))
376    }
377
378    /// `local_invocation_id.x`
379    #[must_use]
380    #[inline]
381    pub fn local_x() -> Self {
382        Self::LocalId { axis: 0 }
383    }
384
385    /// `subgroup_invocation_id` (lane index within subgroup).
386    #[must_use]
387    #[inline]
388    pub fn subgroup_local_id() -> Self {
389        Self::SubgroupLocalId
390    }
391
392    /// `subgroup_size` (number of lanes per subgroup).
393    #[must_use]
394    #[inline]
395    pub fn subgroup_size() -> Self {
396        Self::SubgroupSize
397    }
398
399    /// `local_invocation_id.y`
400    #[must_use]
401    #[inline]
402    pub fn local_y() -> Self {
403        Self::LocalId { axis: 1 }
404    }
405
406    /// `local_invocation_id.z`
407    #[must_use]
408    #[inline]
409    pub fn local_z() -> Self {
410        Self::LocalId { axis: 2 }
411    }
412
413    /// Substrate-neutral alias for [`workgroup_x`](Self::workgroup_x).
414    ///
415    /// "Parallel region" is the vocabulary used in vyre-core's public
416    /// surface. Concrete drivers translate this concept into their own
417    /// target vocabulary at the boundary.
418    #[must_use]
419    #[inline]
420    pub fn parallel_region_x() -> Self {
421        Self::WorkgroupId { axis: 0 }
422    }
423
424    /// Substrate-neutral alias for [`workgroup_y`](Self::workgroup_y).
425    #[must_use]
426    #[inline]
427    pub fn parallel_region_y() -> Self {
428        Self::WorkgroupId { axis: 1 }
429    }
430
431    /// Substrate-neutral alias for [`workgroup_z`](Self::workgroup_z).
432    #[must_use]
433    #[inline]
434    pub fn parallel_region_z() -> Self {
435        Self::WorkgroupId { axis: 2 }
436    }
437
438    /// Substrate-neutral alias for [`local_x`](Self::local_x).
439    #[must_use]
440    #[inline]
441    pub fn invocation_local_x() -> Self {
442        Self::LocalId { axis: 0 }
443    }
444
445    /// Substrate-neutral alias for [`local_y`](Self::local_y).
446    #[must_use]
447    #[inline]
448    pub fn invocation_local_y() -> Self {
449        Self::LocalId { axis: 1 }
450    }
451
452    /// Substrate-neutral alias for [`local_z`](Self::local_z).
453    #[must_use]
454    #[inline]
455    pub fn invocation_local_z() -> Self {
456        Self::LocalId { axis: 2 }
457    }
458
459    /// Conditional select.
460    #[must_use]
461    #[inline]
462    pub fn select(cond: Self, true_val: Self, false_val: Self) -> Self {
463        Self::Select {
464            cond: Box::new(cond),
465            true_val: Box::new(true_val),
466            false_val: Box::new(false_val),
467        }
468    }
469
470    /// Subgroup reduction across the active subgroup with the given operator.
471    #[must_use]
472    #[inline]
473    pub fn subgroup_reduce(op: SubgroupReduceOp, value: Self) -> Self {
474        Self::SubgroupReduce {
475            op,
476            value: Box::new(value),
477        }
478    }
479
480    /// Subgroup sum reduction across the active subgroup.
481    #[must_use]
482    #[inline]
483    pub fn subgroup_add(value: Self) -> Self {
484        Self::subgroup_reduce(SubgroupReduceOp::Add, value)
485    }
486
487    /// Subgroup product reduction across the active subgroup.
488    #[must_use]
489    #[inline]
490    pub fn subgroup_mul(value: Self) -> Self {
491        Self::subgroup_reduce(SubgroupReduceOp::Mul, value)
492    }
493
494    /// Subgroup minimum reduction across the active subgroup.
495    #[must_use]
496    #[inline]
497    pub fn subgroup_min(value: Self) -> Self {
498        Self::subgroup_reduce(SubgroupReduceOp::Min, value)
499    }
500
501    /// Subgroup maximum reduction across the active subgroup.
502    #[must_use]
503    #[inline]
504    pub fn subgroup_max(value: Self) -> Self {
505        Self::subgroup_reduce(SubgroupReduceOp::Max, value)
506    }
507
508    /// Subgroup bitwise-AND reduction across the active subgroup.
509    #[must_use]
510    #[inline]
511    pub fn subgroup_and(value: Self) -> Self {
512        Self::subgroup_reduce(SubgroupReduceOp::And, value)
513    }
514
515    /// Subgroup bitwise-OR reduction across the active subgroup.
516    #[must_use]
517    #[inline]
518    pub fn subgroup_or(value: Self) -> Self {
519        Self::subgroup_reduce(SubgroupReduceOp::Or, value)
520    }
521
522    /// Subgroup bitwise-XOR reduction across the active subgroup.
523    #[must_use]
524    #[inline]
525    pub fn subgroup_xor(value: Self) -> Self {
526        Self::subgroup_reduce(SubgroupReduceOp::Xor, value)
527    }
528
529    /// Subgroup shuffle: broadcast `value` from the given lane id to
530    /// every active lane in the subgroup.
531    #[must_use]
532    #[inline]
533    pub fn subgroup_shuffle(value: Self, lane: Self) -> Self {
534        Self::SubgroupShuffle {
535            value: Box::new(value),
536            lane: Box::new(lane),
537        }
538    }
539
540    /// Subgroup ballot: gather the boolean predicate `cond` across
541    /// the active subgroup into a single bitmask.
542    #[must_use]
543    #[inline]
544    pub fn subgroup_ballot(cond: Self) -> Self {
545        Self::SubgroupBallot {
546            cond: Box::new(cond),
547        }
548    }
549
550    /// Named variable reference.
551    #[must_use]
552    #[inline]
553    pub fn var(name: impl Into<Ident>) -> Self {
554        Self::Var(name.into())
555    }
556
557    /// Unsigned 32-bit literal.
558    #[must_use]
559    #[inline]
560    pub fn u32(value: u32) -> Self {
561        Self::LitU32(value)
562    }
563
564    /// Signed 32-bit literal.
565    #[must_use]
566    #[inline]
567    pub fn i32(value: i32) -> Self {
568        Self::LitI32(value)
569    }
570
571    /// 32-bit floating-point literal.
572    #[must_use]
573    #[inline]
574    pub fn f32(value: f32) -> Self {
575        Self::LitF32(value)
576    }
577
578    /// Boolean literal.
579    #[must_use]
580    #[inline]
581    pub fn bool(value: bool) -> Self {
582        Self::LitBool(value)
583    }
584
585    /// Operation call by stable operation ID.
586    #[must_use]
587    #[inline]
588    pub fn call(op_id: impl Into<Ident>, args: Vec<Self>) -> Self {
589        Self::Call {
590            op_id: op_id.into(),
591            args,
592        }
593    }
594
595    /// Fused multiply-add `a * b + c` (f32).
596    #[must_use]
597    #[inline]
598    pub fn fma(a: Self, b: Self, c: Self) -> Self {
599        Self::Fma {
600            a: Box::new(a),
601            b: Box::new(b),
602            c: Box::new(c),
603        }
604    }
605
606    /// Cast a value to `target`.
607    #[must_use]
608    #[inline]
609    pub fn cast(target: DataType, value: Self) -> Self {
610        Self::Cast {
611            target,
612            value: Box::new(value),
613        }
614    }
615
616    /// Wrap a downstream extension expression node.
617    #[must_use]
618    #[inline]
619    pub fn opaque(node: impl ExprNode) -> Self {
620        Self::Opaque(Arc::new(node))
621    }
622
623    /// Wrap a shared downstream extension expression node.
624    #[must_use]
625    #[inline]
626    pub fn opaque_arc(node: Arc<dyn ExprNode>) -> Self {
627        Self::Opaque(node)
628    }
629}
630
631mod atomics;
632mod builders;
633
634#[cfg(test)]
635mod tests {
636    use super::Expr;
637
638    #[test]
639    fn expr_size_is_bounded() {
640        let size = std::mem::size_of::<Expr>();
641        assert!(
642            size <= 128,
643            "Expr grew to {size} bytes. Fix: box the largest variant before adding more fields."
644        );
645    }
646}
vyre_foundation/ir_inner/model/expr.rs

vyre_foundation/ir_inner/model/
expr.rs
