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::spec_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 should use [`crate::opaque_payload::endian::LeBytesWriter`] when
256    /// building payloads because it makes the required endianness explicit in the type.
257    ///
258    /// Literal extensions that encode regex payloads must also canonicalize
259    /// inline flag prefixes before emitting bytes. For example, `(?mi)` and
260    /// `(?im)` are the same semantic payload and MUST serialize to the same
261    /// flag ordering.
262    fn wire_payload(&self) -> Vec<u8> {
263        Vec::new()
264    }
265}
266
267impl Expr {
268    /// Load from buffer at index.
269    ///
270    /// # Examples
271    ///
272    /// ```
273    /// use vyre::ir::Expr;
274    /// let _ = Expr::load("a", Expr::u32(0));
275    /// ```
276    #[must_use]
277    #[inline]
278    pub fn load(buffer: impl Into<Ident>, index: Self) -> Self {
279        Self::Load {
280            buffer: buffer.into(),
281            index: Box::new(index),
282        }
283    }
284
285    /// Buffer element count.
286    ///
287    /// # Examples
288    ///
289    /// ```
290    /// use vyre::ir::Expr;
291    /// let _ = Expr::buf_len("a");
292    /// ```
293    #[must_use]
294    #[inline]
295    pub fn buf_len(buffer: impl Into<Ident>) -> Self {
296        Self::BufLen {
297            buffer: buffer.into(),
298        }
299    }
300
301    /// Name a whole buffer as an argument to a composite op.
302    ///
303    /// This is not a value: it has no type and only call-argument position
304    /// accepts it. Inlining rebinds the callee's matching parameter onto
305    /// this buffer, so a callee that reads `table[i]` ends up reading the
306    /// caller's buffer at the same index. To read one element instead, use
307    /// [`Expr::load`].
308    ///
309    /// ```
310    /// use vyre::ir::Expr;
311    /// let _ = Expr::call("dialect::lookup", vec![Expr::buffer_ref("table"), Expr::u32(3)]);
312    /// ```
313    #[must_use]
314    #[inline]
315    pub fn buffer_ref(buffer: impl Into<Ident>) -> Self {
316        Self::BufferRef {
317            buffer: buffer.into(),
318        }
319    }
320
321    /// `global_invocation_id.x`
322    #[must_use]
323    #[inline]
324    pub fn gid_x() -> Self {
325        Self::InvocationId { axis: 0 }
326    }
327
328    /// `global_invocation_id.y`
329    #[must_use]
330    #[inline]
331    pub fn gid_y() -> Self {
332        Self::InvocationId { axis: 1 }
333    }
334
335    /// `global_invocation_id.z`
336    #[must_use]
337    #[inline]
338    pub fn gid_z() -> Self {
339        Self::InvocationId { axis: 2 }
340    }
341
342    /// `workgroup_id.x`
343    #[must_use]
344    #[inline]
345    pub fn workgroup_x() -> Self {
346        Self::WorkgroupId { axis: 0 }
347    }
348
349    /// `workgroup_id.y`
350    #[must_use]
351    #[inline]
352    pub fn workgroup_y() -> Self {
353        Self::WorkgroupId { axis: 1 }
354    }
355
356    /// `workgroup_id.z`
357    #[must_use]
358    #[inline]
359    pub fn workgroup_z() -> Self {
360        Self::WorkgroupId { axis: 2 }
361    }
362
363    /// Predicate `workgroup_id.x == 0`: the canonical "first parallel region only" guard.
364    ///
365    /// Single-workgroup kernels (scalar reductions, workgroup-local tree reductions, and any
366    /// kernel meant to run in exactly one parallel region) gate their body on this so a dispatch
367    /// of more than one workgroup leaves the extra regions as no-ops instead of double-counting or
368    /// racing on the shared output. Prefer this over re-spelling `eq(WorkgroupId{axis:0}, 0)`
369    /// inline so the "first workgroup" contract has one owner.
370    #[must_use]
371    #[inline]
372    pub fn is_first_workgroup() -> Self {
373        Self::eq(Self::WorkgroupId { axis: 0 }, Self::u32(0))
374    }
375
376    /// `local_invocation_id.x`
377    #[must_use]
378    #[inline]
379    pub fn local_x() -> Self {
380        Self::LocalId { axis: 0 }
381    }
382
383    /// `subgroup_invocation_id` (lane index within subgroup).
384    #[must_use]
385    #[inline]
386    pub fn subgroup_local_id() -> Self {
387        Self::SubgroupLocalId
388    }
389
390    /// `subgroup_size` (number of lanes per subgroup).
391    #[must_use]
392    #[inline]
393    pub fn subgroup_size() -> Self {
394        Self::SubgroupSize
395    }
396
397    /// `local_invocation_id.y`
398    #[must_use]
399    #[inline]
400    pub fn local_y() -> Self {
401        Self::LocalId { axis: 1 }
402    }
403
404    /// `local_invocation_id.z`
405    #[must_use]
406    #[inline]
407    pub fn local_z() -> Self {
408        Self::LocalId { axis: 2 }
409    }
410
411    /// Substrate-neutral alias for [`workgroup_x`](Self::workgroup_x).
412    ///
413    /// "Parallel region" is the vocabulary used in the public `vyre` facade.
414    /// Concrete drivers translate this concept into target vocabulary at the
415    /// lowering boundary.
416    #[must_use]
417    #[inline]
418    pub fn parallel_region_x() -> Self {
419        Self::WorkgroupId { axis: 0 }
420    }
421
422    /// Substrate-neutral alias for [`workgroup_y`](Self::workgroup_y).
423    #[must_use]
424    #[inline]
425    pub fn parallel_region_y() -> Self {
426        Self::WorkgroupId { axis: 1 }
427    }
428
429    /// Substrate-neutral alias for [`workgroup_z`](Self::workgroup_z).
430    #[must_use]
431    #[inline]
432    pub fn parallel_region_z() -> Self {
433        Self::WorkgroupId { axis: 2 }
434    }
435
436    /// Substrate-neutral alias for [`local_x`](Self::local_x).
437    #[must_use]
438    #[inline]
439    pub fn invocation_local_x() -> Self {
440        Self::LocalId { axis: 0 }
441    }
442
443    /// Substrate-neutral alias for [`local_y`](Self::local_y).
444    #[must_use]
445    #[inline]
446    pub fn invocation_local_y() -> Self {
447        Self::LocalId { axis: 1 }
448    }
449
450    /// Substrate-neutral alias for [`local_z`](Self::local_z).
451    #[must_use]
452    #[inline]
453    pub fn invocation_local_z() -> Self {
454        Self::LocalId { axis: 2 }
455    }
456
457    /// Conditional select.
458    #[must_use]
459    #[inline]
460    pub fn select(cond: Self, true_val: Self, false_val: Self) -> Self {
461        Self::Select {
462            cond: Box::new(cond),
463            true_val: Box::new(true_val),
464            false_val: Box::new(false_val),
465        }
466    }
467
468    /// Subgroup reduction across the active subgroup with the given operator.
469    #[must_use]
470    #[inline]
471    pub fn subgroup_reduce(op: SubgroupReduceOp, value: Self) -> Self {
472        Self::SubgroupReduce {
473            op,
474            value: Box::new(value),
475        }
476    }
477
478    /// Subgroup sum reduction across the active subgroup.
479    #[must_use]
480    #[inline]
481    pub fn subgroup_add(value: Self) -> Self {
482        Self::subgroup_reduce(SubgroupReduceOp::Add, value)
483    }
484
485    /// Subgroup product reduction across the active subgroup.
486    #[must_use]
487    #[inline]
488    pub fn subgroup_mul(value: Self) -> Self {
489        Self::subgroup_reduce(SubgroupReduceOp::Mul, value)
490    }
491
492    /// Subgroup minimum reduction across the active subgroup.
493    #[must_use]
494    #[inline]
495    pub fn subgroup_min(value: Self) -> Self {
496        Self::subgroup_reduce(SubgroupReduceOp::Min, value)
497    }
498
499    /// Subgroup maximum reduction across the active subgroup.
500    #[must_use]
501    #[inline]
502    pub fn subgroup_max(value: Self) -> Self {
503        Self::subgroup_reduce(SubgroupReduceOp::Max, value)
504    }
505
506    /// Subgroup bitwise-AND reduction across the active subgroup.
507    #[must_use]
508    #[inline]
509    pub fn subgroup_and(value: Self) -> Self {
510        Self::subgroup_reduce(SubgroupReduceOp::And, value)
511    }
512
513    /// Subgroup bitwise-OR reduction across the active subgroup.
514    #[must_use]
515    #[inline]
516    pub fn subgroup_or(value: Self) -> Self {
517        Self::subgroup_reduce(SubgroupReduceOp::Or, value)
518    }
519
520    /// Subgroup bitwise-XOR reduction across the active subgroup.
521    #[must_use]
522    #[inline]
523    pub fn subgroup_xor(value: Self) -> Self {
524        Self::subgroup_reduce(SubgroupReduceOp::Xor, value)
525    }
526
527    /// Subgroup shuffle: broadcast `value` from the given lane id to
528    /// every active lane in the subgroup.
529    #[must_use]
530    #[inline]
531    pub fn subgroup_shuffle(value: Self, lane: Self) -> Self {
532        Self::SubgroupShuffle {
533            value: Box::new(value),
534            lane: Box::new(lane),
535        }
536    }
537
538    /// Subgroup ballot: gather the boolean predicate `cond` across
539    /// the active subgroup into a single bitmask.
540    #[must_use]
541    #[inline]
542    pub fn subgroup_ballot(cond: Self) -> Self {
543        Self::SubgroupBallot {
544            cond: Box::new(cond),
545        }
546    }
547
548    /// Named variable reference.
549    #[must_use]
550    #[inline]
551    pub fn var(name: impl Into<Ident>) -> Self {
552        Self::Var(name.into())
553    }
554
555    /// Unsigned 32-bit literal.
556    #[must_use]
557    #[inline]
558    pub fn u32(value: u32) -> Self {
559        Self::LitU32(value)
560    }
561
562    /// Signed 32-bit literal.
563    #[must_use]
564    #[inline]
565    pub fn i32(value: i32) -> Self {
566        Self::LitI32(value)
567    }
568
569    /// 32-bit floating-point literal.
570    #[must_use]
571    #[inline]
572    pub fn f32(value: f32) -> Self {
573        Self::LitF32(value)
574    }
575
576    /// Boolean literal.
577    #[must_use]
578    #[inline]
579    pub fn bool(value: bool) -> Self {
580        Self::LitBool(value)
581    }
582
583    /// Operation call by stable operation ID.
584    #[must_use]
585    #[inline]
586    pub fn call(op_id: impl Into<Ident>, args: Vec<Self>) -> Self {
587        Self::Call {
588            op_id: op_id.into(),
589            args,
590        }
591    }
592
593    /// Fused multiply-add `a * b + c` (f32).
594    #[must_use]
595    #[inline]
596    pub fn fma(a: Self, b: Self, c: Self) -> Self {
597        Self::Fma {
598            a: Box::new(a),
599            b: Box::new(b),
600            c: Box::new(c),
601        }
602    }
603
604    /// Cast a value to `target`.
605    #[must_use]
606    #[inline]
607    pub fn cast(target: DataType, value: Self) -> Self {
608        Self::Cast {
609            target,
610            value: Box::new(value),
611        }
612    }
613
614    /// Wrap a downstream extension expression node.
615    #[must_use]
616    #[inline]
617    pub fn opaque(node: impl ExprNode) -> Self {
618        Self::Opaque(Arc::new(node))
619    }
620
621    /// Wrap a shared downstream extension expression node.
622    #[must_use]
623    #[inline]
624    pub fn opaque_arc(node: Arc<dyn ExprNode>) -> Self {
625        Self::Opaque(node)
626    }
627}
628
629mod atomics;
630mod builders;
631
632#[cfg(test)]
633mod tests {
634    use super::Expr;
635
636    #[test]
637    fn expr_size_is_bounded() {
638        let size = std::mem::size_of::<Expr>();
639        assert!(
640            size <= 128,
641            "Expr grew to {size} bytes. Fix: box the largest variant before adding more fields."
642        );
643    }
644}
vyre_foundation/ir_inner/model/expr.rs

vyre_foundation/ir_inner/model/
expr.rs
